Reduce MultiplayerServer.initData API surface and simplify it

Author: mibac138

Source/Client/Networking/HostUtil.cs

@@ -74,10 +74,8 @@ private static void PrepareLocalServer(ServerSettings settings, bool fromReplay)
                 localServer.startingTimer = TickPatch.Timer;
             }
 
-            localServer.initDataSource = new TaskCompletionSource<ServerInitData>();
-            localServer.CompleteInitData(
-                ServerInitData.Deserialize(new ByteReader(ClientJoiningState.PackInitData(settings.syncConfigs)))
-            );
+            localServer.initData =
+                ServerInitData.Deserialize(new ByteReader(ClientJoiningState.PackInitData(settings.syncConfigs)));
         }
 
         private static void PrepareGame()

Source/Common/MultiplayerServer.cs

@@ -48,8 +48,13 @@ static MultiplayerServer()
         public ServerSettings settings;
 
         public ServerInitData? initData;
-        public TaskCompletionSource<ServerInitData?> initDataSource = new();
-        public InitDataState initDataState = InitDataState.Waiting;
+        private TaskCompletionSource<ServerInitData?> initDataSource = new();
+        public InitDataState InitDataState =>
+            initData != null ? InitDataState.Complete :
+            // started init data must've completed with null, meaning the client disconnected while waiting for the data,
+            // so we are waiting again
+            initDataSource.Task.IsCompleted ? InitDataState.Waiting :
+            InitDataState.Requested;
 
         public volatile bool running;
         public event Action<MultiplayerServer>? TickEvent;
@@ -245,16 +250,22 @@ public void RegisterChatCmd(string cmdName, ChatCmdHandler handler) =>
 
         public void HandleChatCmd(IChatSource source, string cmd) => chatCmdManager.Handle(source, cmd);
 
-        public Task<ServerInitData?> InitData()
-        {
-            return initDataSource.Task;
-        }
+        public Task<ServerInitData?> InitData() => initDataSource.Task;
 
-        public void CompleteInitData(ServerInitData data)
+        /// Can only start one init data at a time. A StartInitData is considered complete once
+        /// TaskCompletionResult.SetResult is called. Until that time no new calls to StartInitData will succeed.
+        public TaskCompletionSource<ServerInitData?> StartInitData()
         {
-            initData = data;
-            initDataState = InitDataState.Complete;
-            initDataSource.SetResult(data);
+            if (InitDataState != InitDataState.Waiting)
+                throw new InvalidOperationException($"Can't start init data in state {InitDataState}");
+            var currInitDataSource = initDataSource = new TaskCompletionSource<ServerInitData?>();
+            currInitDataSource.Task.ContinueWith(task =>
+            {
+                if (currInitDataSource != initDataSource)
+                    ServerLog.Error("InitDataSource changed during StartInitData");
+                initData = task.Result;
+            }, TaskContinuationOptions.ExecuteSynchronously);
+            return currInitDataSource;
         }
     }
 

Source/Common/Networking/State/ServerJoiningState.cs

@@ -14,7 +14,7 @@ protected override async Task RunState()
         HandleUsername(await Packet(Packets.Client_Username));
 
         while (await Server.InitData() is null && await EndIfDead())
-            if (Server.initDataState == InitDataState.Waiting)
+            if (Server.InitDataState == InitDataState.Waiting)
                 await RequestInitData();
 
         connection.Send(Packets.Server_UsernameOk);
@@ -83,23 +83,20 @@ private void HandleUsername(ByteReader data)
 
     private async Task RequestInitData()
     {
-        Server.initDataState = InitDataState.Requested;
-        Server.initDataSource = new TaskCompletionSource<ServerInitData?>();
-
-        Player.SendPacket(Packets.Server_InitDataRequest, ByteWriter.GetBytes(Server.settings.syncConfigs));
-
-        ServerLog.Verbose("Sent initial data request");
-
-        var initData = await PacketOrNull(Packets.Client_InitData).Fragmented();
-
-        if (initData != null)
+        ByteReader? initData = null;
+        var completionSource = Server.StartInitData();
+        try
         {
-            Server.CompleteInitData(ServerInitData.Deserialize(initData));
+            Player.SendPacket(Packets.Server_InitDataRequest, ByteWriter.GetBytes(Server.settings.syncConfigs));
+
+            ServerLog.Verbose("Sent initial data request");
+            initData = await PacketOrNull(Packets.Client_InitData).Fragmented();
         }
-        else
+        finally
         {
-            Server.initDataState = InitDataState.Waiting;
-            Server.initDataSource.SetResult(null);
+            // Invoking StartInitData and abandoning the completion source in case of an exception would mean a server
+            // restart is necessary to try and set init data again. Make sure the server is more graceful in that case.
+            completionSource.SetResult(initData != null ? ServerInitData.Deserialize(initData) : null);
         }
     }
 

Source/Server/Server.cs

@@ -26,7 +26,6 @@
 var server = MultiplayerServer.instance = new MultiplayerServer(settings)
 {
     running = true,
-    initDataState = InitDataState.Waiting
 };
 
 var consoleSource = new ConsoleSource();

Source/Tests/ServerTest.cs

@@ -38,7 +38,7 @@ public void Test()
         var timeoutWatch = Stopwatch.StartNew();
         while (true)
         {
-            if (server.initDataState == InitDataState.Complete && server.playerManager.Players.Count == 0)
+            if (server.InitDataState == InitDataState.Complete && server.playerManager.Players.Count == 0)
                 break; // Success
 
             if (timeoutWatch.ElapsedMilliseconds > 2000)
@@ -83,7 +83,6 @@ private MultiplayerServer MakeServer(int port)
         })
         {
             running = true,
-            initDataState = InitDataState.Waiting
         };
 
         server.worldData.savedGame = Array.Empty<byte>();