Improved reliability of the shutdown process.

SharpChat/ClientCommands/ShutdownRestartClientCommand.cs

@@ -2,10 +2,7 @@ using SharpChat.SockChat.S2CPackets;
 
 namespace SharpChat.ClientCommands;
 
-public class ShutdownRestartClientCommand(ManualResetEvent waitHandle, Func<bool> shutdownCheck) : ClientCommand {
+public class ShutdownRestartClientCommand(CancellationTokenSource cancellationTokenSource) : ClientCommand {
-    private readonly ManualResetEvent WaitHandle = waitHandle ?? throw new ArgumentNullException(nameof(waitHandle));
-    private readonly Func<bool> ShutdownCheck = shutdownCheck ?? throw new ArgumentNullException(nameof(shutdownCheck));
-
     public bool IsMatch(ClientCommandContext ctx) {
         return ctx.NameEquals("shutdown")
             || ctx.NameEquals("restart");
@@ -18,14 +15,16 @@ public class ShutdownRestartClientCommand(ManualResetEvent waitHandle, Func<bool
             return;
         }
 
-        if(!ShutdownCheck())
+        if(cancellationTokenSource.IsCancellationRequested)
             return;
 
+        Logger.Write("Shutdown requested through Sock Chat command...");
+
         if(ctx.NameEquals("restart"))
             foreach(Connection conn in ctx.Chat.Connections)
                 conn.PrepareForRestart();
 
         await ctx.Chat.Update();
-        WaitHandle?.Set();
+        await cancellationTokenSource.CancelAsync();
     }
 }

SharpChat/Program.cs

@@ -2,7 +2,6 @@ using SharpChat;
 using SharpChat.Configuration;
 using SharpChat.Flashii;
 using SharpChat.MariaDB;
-using SharpChat.Messages;
 using SharpChat.SQLite;
 using System.Text;
 
@@ -22,24 +21,27 @@ if(SharpInfo.IsDebugBuild) {
 } else
     Console.WriteLine(SharpInfo.VersionStringShort.PadLeft(28, ' '));
 
-using ManualResetEvent mre = new(false);
+using CancellationTokenSource cts = new();
-bool hasCancelled = false;
 
-void cancelKeyPressHandler(object? sender, ConsoleCancelEventArgs ev) {
+void exitHandler() {
-    Console.CancelKeyPress -= cancelKeyPressHandler;
+    if(cts.IsCancellationRequested)
-    hasCancelled = true;
+        return;
-    ev.Cancel = true;
+
-    mre.Set();
+    cts.Cancel();
+    Logger.Write("Shutdown requested through console...");
 }
-;
-Console.CancelKeyPress += cancelKeyPressHandler;
 
-if(hasCancelled) return;
+AppDomain.CurrentDomain.ProcessExit += (sender, ev) => { exitHandler(); };
+Console.CancelKeyPress += (sender, ev) => { ev.Cancel = true; exitHandler(); };
+
+if(cts.IsCancellationRequested) return;
 
 string configFile = CONFIG;
 
 // If the config file doesn't exist and we're using the default path, run the converter
 if(!File.Exists(configFile) && configFile == CONFIG) {
+    Logger.Write("Creating example configuration file...");
+
     using Stream s = new FileStream(CONFIG, FileMode.OpenOrCreate, FileAccess.ReadWrite, FileShare.None);
     s.SetLength(0);
     s.Flush();
@@ -110,38 +112,44 @@ if(!File.Exists(configFile) && configFile == CONFIG) {
     sw.Flush();
 }
 
+Logger.Write("Initialising configuration...");
 using StreamConfig config = StreamConfig.FromPath(configFile);
 
-if(hasCancelled) return;
+if(cts.IsCancellationRequested) return;
 
+Logger.Write("Initialising HTTP client...");
 using HttpClient httpClient = new(new HttpClientHandler() {
     UseProxy = false,
 });
 httpClient.DefaultRequestHeaders.Add("User-Agent", SharpInfo.ProgramName);
 
-if(hasCancelled) return;
+if(cts.IsCancellationRequested) return;
 
+Logger.Write("Initialising Flashii client...");
 FlashiiClient flashii = new(httpClient, config.ScopeTo("msz"));
 
-if(hasCancelled) return;
+if(cts.IsCancellationRequested) return;
 
+Logger.Write("Initialising storage...");
 Storage storage = string.IsNullOrWhiteSpace(config.SafeReadValue("mariadb:host", string.Empty))
     ? new SQLiteStorage(SQLiteStorage.BuildConnectionString(config.ScopeTo("sqlite")))
     : new MariaDBStorage(MariaDBStorage.BuildConnectionString(config.ScopeTo("mariadb")));
 
 try {
-    if(hasCancelled) return;
+    if(cts.IsCancellationRequested) return;
 
+    Logger.Write("Upgrading storage...");
     await storage.UpgradeStorage();
 
-    if(hasCancelled) return;
+    if(cts.IsCancellationRequested) return;
 
-    using SockChatServer scs = new(flashii, flashii, storage.CreateMessageStorage(), config.ScopeTo("chat"));
+    Logger.Write("Preparing server...");
-    scs.Listen(mre);
+    await new SockChatServer(cts, flashii, flashii, storage.CreateMessageStorage(), config.ScopeTo("chat")).Listen(cts.Token);
-
-    mre.WaitOne();
 } finally {
-    if(storage is IDisposable disp)
+    if(storage is IDisposable disp) {
+        Logger.Write("Cleaning storage...");
         disp.Dispose();
+    }
 }
 
+Logger.Write("Exiting...");

SharpChat/SockChatServer.cs

@@ -10,18 +10,18 @@ using System.Net;
 
 namespace SharpChat;
 
-public class SockChatServer : IDisposable {
+public class SockChatServer {
     public const ushort DEFAULT_PORT = 6770;
     public const int DEFAULT_MSG_LENGTH_MAX = 5000;
     public const int DEFAULT_MAX_CONNECTIONS = 5;
     public const int DEFAULT_FLOOD_KICK_LENGTH = 30;
     public const int DEFAULT_FLOOD_KICK_EXEMPT_RANK = 9;
 
-    public IWebSocketServer Server { get; }
     public Context Context { get; }
 
     private readonly BansClient BansClient;
 
+    private readonly CachedValue<ushort> Port;
     private readonly CachedValue<int> MaxMessageLength;
     private readonly CachedValue<int> MaxConnections;
     private readonly CachedValue<int> FloodKickLength;
@@ -31,11 +31,10 @@ public class SockChatServer : IDisposable {
     private readonly List<C2SPacketHandler> AuthedHandlers = [];
     private readonly SendMessageC2SPacketHandler SendMessageHandler;
 
-    private bool IsShuttingDown = false;
-
     private static readonly string[] DEFAULT_CHANNELS = ["lounge"];
 
     public SockChatServer(
+        CancellationTokenSource cancellationTokenSource,
         AuthClient authClient,
         BansClient bansClient,
         MessageStorage msgStorage,
@@ -45,6 +44,7 @@ public class SockChatServer : IDisposable {
 
         BansClient = bansClient;
 
+        Port = config.ReadCached("port", DEFAULT_PORT);
         MaxMessageLength = config.ReadCached("msgMaxLength", DEFAULT_MSG_LENGTH_MAX);
         MaxConnections = config.ReadCached("connMaxCount", DEFAULT_MAX_CONNECTIONS);
         FloodKickLength = config.ReadCached("floodKickLength", DEFAULT_FLOOD_KICK_LENGTH);
@@ -100,18 +100,14 @@ public class SockChatServer : IDisposable {
             new PardonAddressClientCommand(bansClient),
             new BanListClientCommand(bansClient),
             new RemoteAddressClientCommand(),
+            new ShutdownRestartClientCommand(cancellationTokenSource)
         ]);
-
-        ushort port = config.SafeReadValue("port", DEFAULT_PORT);
-        Server = new SharpChatWebSocketServer($"ws://0.0.0.0:{port}");
     }
 
-    public void Listen(ManualResetEvent waitHandle) {
+    public async Task Listen(CancellationToken cancellationToken) {
-        if(waitHandle != null)
+        using IWebSocketServer server = new SharpChatWebSocketServer($"ws://0.0.0.0:{Port}");
-            SendMessageHandler.AddCommand(new ShutdownRestartClientCommand(waitHandle, () => !IsShuttingDown && (IsShuttingDown = true)));
+        server.Start(sock => {
-
+            if(cancellationToken.IsCancellationRequested) {
-        Server.Start(sock => {
-            if(IsShuttingDown) {
                 sock.Close(1013);
                 return;
             }
@@ -126,6 +122,10 @@ public class SockChatServer : IDisposable {
         });
 
         Logger.Write("Listening...");
+        await Task.Delay(Timeout.Infinite, cancellationToken).ConfigureAwait(ConfigureAwaitOptions.SuppressThrowing);
+
+        foreach(Connection conn in Context.Connections)
+            conn.Dispose();
     }
 
     private async Task OnOpen(Connection conn) {
@@ -214,27 +214,4 @@ public class SockChatServer : IDisposable {
         if(handler is not null)
             await handler.Handle(context);
     }
-
-    private bool IsDisposed;
-
-    ~SockChatServer() {
-        DoDispose();
-    }
-
-    public void Dispose() {
-        DoDispose();
-        GC.SuppressFinalize(this);
-    }
-
-    private void DoDispose() {
-        if(IsDisposed)
-            return;
-        IsDisposed = true;
-        IsShuttingDown = true;
-
-        foreach(Connection conn in Context.Connections)
-            conn.Dispose();
-
-        Server?.Dispose();
-    }
 }