From 720de064a027bac108f50553bdafb11448a78434 Mon Sep 17 00:00:00 2001
From: pg9182 <96569817+pg9182@users.noreply.github.com>
Date: Thu, 14 Sep 2023 18:21:43 -0400
Subject: [PATCH] pkg/api/api0: Implement isDraining flag for servers

* For backwards compatibility and doing what makes sense, the max
  players in metrics and the server list is set to the player count
  while a server is draining. Existing NorthstarMods versions will
  sort these servers to the end of the server list.
* NorthstarLauncher will enforce rejecting connections using the
  existing rejection mechanism while a server is draining.
---
 pkg/api/api0/server.go     |  9 +++++++++
 pkg/api/api0/serverlist.go | 36 +++++++++++++++++++++++++++++++-----
 2 files changed, 40 insertions(+), 5 deletions(-)

diff --git a/pkg/api/api0/server.go b/pkg/api/api0/server.go
index 98e2804..d16553c 100644
--- a/pkg/api/api0/server.go
+++ b/pkg/api/api0/server.go
@@ -284,6 +284,15 @@ func (h *Handler) handleServerUpsert(w http.ResponseWriter, r *http.Request) {
 				u.MaxPlayers = &x
 			}
 		}
+
+		if v, err := strconv.ParseBool(q.Get("isDraining")); err == nil {
+			if canCreate {
+				s.IsDraining = v
+			}
+			if canUpdate {
+				u.IsDraining = &v
+			}
+		}
 	}
 
 	if canCreate {
diff --git a/pkg/api/api0/serverlist.go b/pkg/api/api0/serverlist.go
index a1b3258..a17b5b5 100644
--- a/pkg/api/api0/serverlist.go
+++ b/pkg/api/api0/serverlist.go
@@ -99,6 +99,8 @@ type Server struct {
 	Map         string
 	Playlist    string
 
+	IsDraining bool
+
 	ServerAuthToken string // used for authenticating the masterserver to the gameserver authserver
 
 	ModInfo []ServerModInfo
@@ -118,6 +120,15 @@ func (s Server) AuthAddr() netip.AddrPort {
 	return netip.AddrPortFrom(s.Addr.Addr(), s.AuthPort)
 }
 
+// EffectiveMaxPlayers returns the max players which can actually join, taking
+// IsDraining into account.
+func (s Server) EffectiveMaxPlayers() int {
+	if s.IsDraining {
+		return s.PlayerCount
+	}
+	return s.MaxPlayers
+}
+
 // clone returns a deep copy of s.
 func (s Server) clone() Server {
 	m := make([]ServerModInfo, len(s.ModInfo))
@@ -140,6 +151,7 @@ type ServerUpdate struct {
 	MaxPlayers  *int
 	Map         *string
 	Playlist    *string
+	IsDraining  *bool
 }
 
 type ServerListLimit struct {
@@ -338,7 +350,7 @@ func csJSON(ss []*Server, est int, cfg ServerListConfig) ([]byte, int) {
 		b = append(b, `,"playerCount":`...)
 		b = strconv.AppendInt(b, int64(srv.PlayerCount), 10)
 		b = append(b, `,"maxPlayers":`...)
-		b = strconv.AppendInt(b, int64(srv.MaxPlayers), 10)
+		b = strconv.AppendInt(b, int64(srv.EffectiveMaxPlayers()), 10)
 		b = append(b, `,"map":`...)
 		b = appendJSONString(b, srv.Map)
 		b = append(b, `,"playlist":`...)
@@ -348,6 +360,11 @@ func csJSON(ss []*Server, est int, cfg ServerListConfig) ([]byte, int) {
 		} else {
 			b = append(b, `,"hasPassword":false`...)
 		}
+		if srv.IsDraining {
+			b = append(b, `,"isDraining":true`...)
+		} else {
+			b = append(b, `,"isDraining":false`...)
+		}
 		b = append(b, `,"modInfo":{"Mods":[`...)
 		for j, mi := range srv.ModInfo {
 			if j != 0 {
@@ -504,7 +521,7 @@ func (s *ServerList) GetMetrics() []byte {
 		mpls = append(mpls, mpl{m, nstypes.Playlist("")})
 	}
 
-	var players, maxPlayers, servers, serversWithPlayers, fullServers int
+	var players, maxPlayers, servers, serversWithPlayers, fullServers, drainingServers int
 	mplPlayers := make(map[mpl]int, len(mpls))
 	mplMaxPlayers := make(map[mpl]int, len(mpls))
 	mplServers := make(map[mpl]int, len(mpls))
@@ -523,16 +540,19 @@ func (s *ServerList) GetMetrics() []byte {
 					mplv.Playlist = pl
 				}
 				players += srv.PlayerCount
-				maxPlayers += srv.MaxPlayers
+				maxPlayers += srv.EffectiveMaxPlayers()
 				servers++
 				if srv.PlayerCount > 0 {
 					serversWithPlayers++
 				}
-				if srv.PlayerCount == srv.MaxPlayers {
+				if srv.PlayerCount == srv.MaxPlayers { // not EffectiveMaxPlayers
 					fullServers++
 				}
+				if srv.IsDraining {
+					drainingServers++
+				}
 				mplPlayers[mplv] += srv.PlayerCount
-				mplMaxPlayers[mplv] += srv.MaxPlayers
+				mplMaxPlayers[mplv] += srv.EffectiveMaxPlayers()
 				mplServers[mplv]++
 				verServers[srv.LauncherVersion]++
 				for _, mi := range srv.ModInfo {
@@ -690,6 +710,9 @@ func (s *ServerList) GetMetrics() []byte {
 	b.WriteString(`atlas_api0sl_fullservers `)
 	b.WriteString(strconv.Itoa(fullServers))
 	b.WriteByte('\n')
+	b.WriteString(`atlas_api0sl_drainingservers `)
+	b.WriteString(strconv.Itoa(drainingServers))
+	b.WriteByte('\n')
 
 	return b.Bytes()
 }
@@ -881,6 +904,9 @@ func (s *ServerList) ServerHybridUpdatePut(u *ServerUpdate, c *Server, l ServerL
 				if u.MaxPlayers != nil {
 					esrv.MaxPlayers, changed = *u.MaxPlayers, true
 				}
+				if u.IsDraining != nil {
+					esrv.IsDraining, changed = *u.IsDraining, true
+				}
 				if changed {
 					s.csForceUpdate()
 				}