From ee7545b6ee9dcc657119aec7717c0700ae4032e3 Mon Sep 17 00:00:00 2001
From: zznty <94796179+zznty@users.noreply.github.com>
Date: Mon, 1 Jan 2024 23:41:18 +0700
Subject: [PATCH] [Maintenance] fix unload exceptions

---
 .../Managers/MaintenanceScheduleManager.cs    | 21 ++++++++++-
 Maintenance/Patches/SteamQueryPatch.cs        | 35 ++++++++++++++++---
 Maintenance/manifest.xml                      |  2 +-
 3 files changed, 52 insertions(+), 6 deletions(-)

diff --git a/Maintenance/Managers/MaintenanceScheduleManager.cs b/Maintenance/Managers/MaintenanceScheduleManager.cs
index ec8ca37..d374ab0 100644
--- a/Maintenance/Managers/MaintenanceScheduleManager.cs
+++ b/Maintenance/Managers/MaintenanceScheduleManager.cs
@@ -2,6 +2,7 @@
 using System.Text.Json;
 using System.Text.Json.Serialization;
 using Microsoft.Extensions.Configuration;
+using NLog;
 using Torch.API.Managers;
 using Torch.Managers;
 
@@ -9,6 +10,8 @@ namespace Maintenance.Managers;
 
 public class MaintenanceScheduleManager(string storagePath) : IManager
 {
+    private static readonly ILogger Log = LogManager.GetCurrentClassLogger();
+    
     [Manager.Dependency]
     private readonly MaintenanceManager _maintenanceManager = null!;
     
@@ -64,7 +67,23 @@ public class MaintenanceScheduleManager(string storagePath) : IManager
     private async void Scheduler()
     {
         var token = _cancellationTokenSource.Token;
-        
+
+        try
+        {
+            await SchedulerLoop(token);
+        }
+        catch (OperationCanceledException)
+        {
+            Log.Info("Maintenance scheduler loop was cancelled");
+        }
+        catch (Exception e)
+        {
+            Log.Fatal(e, "Exception in maintenance scheduler loop. Shutting down the scheduler.");
+        }
+    }
+
+    private async Task SchedulerLoop(CancellationToken token)
+    {
         var timerSecondsSection = _configManager.Configuration.GetSection(ConfigKeys.TimerBroadcastForSeconds);
         var timerSeconds = timerSecondsSection.GetChildren().Select(b => b.Get<int>()).ToArray();
 
diff --git a/Maintenance/Patches/SteamQueryPatch.cs b/Maintenance/Patches/SteamQueryPatch.cs
index 2f9d3a4..8f49f26 100644
--- a/Maintenance/Patches/SteamQueryPatch.cs
+++ b/Maintenance/Patches/SteamQueryPatch.cs
@@ -30,26 +30,48 @@ namespace Maintenance.Patches;
 public static class SteamQueryPatch
 {
     private static readonly ILogger Log = LogManager.GetCurrentClassLogger();
+
+    [ReflectedSetter(Name = "Running", TypeName = "VRage.Steam.MySteamGameServer, VRage.Steam")]
+    private static readonly Action<IMyGameServer, bool> RunningSetter = null!;
     
     [ReflectedMethodInfo(null, "SteamServerEntryPoint", TypeName = "VRage.Steam.MySteamGameServer, VRage.Steam")]
     private static readonly MethodInfo EntryPointMethod = null!;
 
-    [ReflectedMethodInfo(typeof(SteamQueryPatch), nameof(Prefix))]
-    private static readonly MethodInfo PrefixMethod = null!;
+    [ReflectedMethodInfo(typeof(SteamQueryPatch), nameof(EntryPointPrefix))]
+    private static readonly MethodInfo EntryPointPrefixMethod = null!;
+    
+    [ReflectedMethodInfo(null, "Shutdown", TypeName = "VRage.Steam.MySteamGameServer, VRage.Steam")]
+    private static readonly MethodInfo ShutdownMethod = null!;
+
+    [ReflectedMethodInfo(typeof(SteamQueryPatch), nameof(ShutdownPrefix))]
+    private static readonly MethodInfo ShutdownPrefixMethod = null!;
 
 #pragma warning disable CS0618 // Type or member is obsolete
     private static ITorchBase Torch => TorchBase.Instance;
 #pragma warning restore CS0618 // Type or member is obsolete
     
+    private static readonly ManualResetEventSlim ExitEvent = new(false);
+    
     private static bool MaintenanceEnabled =>
         Torch.CurrentSession?.Managers.GetManager<MaintenanceManager>().MaintenanceEnabled is true;
 
     public static void Patch(PatchContext context)
     {
-        context.GetPattern(EntryPointMethod).Prefixes.Add(PrefixMethod);
+        context.GetPattern(EntryPointMethod).Prefixes.Add(EntryPointPrefixMethod);
+        context.GetPattern(ShutdownMethod).Prefixes.Add(ShutdownPrefixMethod);
+    }
+
+    private static bool ShutdownPrefix(IMyGameServer __instance, Thread __field_m_serverThread)
+    {
+        if (!__instance.Running) return false;
+        
+        RunningSetter(__instance, false);
+        ExitEvent.Wait();
+
+        return false;
     }
     
-    private static bool Prefix(IMyGameServer __instance, object argument)
+    private static bool EntryPointPrefix(IMyGameServer __instance, object argument)
     {
         var socket = (Socket)argument;
         RunServerAsync(__instance, socket);
@@ -142,6 +164,11 @@ public static class SteamQueryPatch
                 await socket.SendToAsync(new(buffer, 0, length), SocketFlags.None, remoteEndPoint);
             }
         }
+        
+        socket.Dispose();
+        GameServer.Shutdown();
+        
+        ExitEvent.Set();
     }
 
     [MethodImpl(MethodImplOptions.AggressiveInlining)]
diff --git a/Maintenance/manifest.xml b/Maintenance/manifest.xml
index e59b4ac..430e4b4 100644
--- a/Maintenance/manifest.xml
+++ b/Maintenance/manifest.xml
@@ -2,5 +2,5 @@
 <PluginManifest xmlns:xsd="http://www.w3.org/2001/XMLSchema" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance">
     <Name>Maintenance</Name>
     <Guid>42AF9955-AAA7-442F-9BF4-AAC4FA4A923F</Guid>
-    <Version>v1.0.2</Version>
+    <Version>v1.0.3</Version>
 </PluginManifest>
\ No newline at end of file