From 96bd7473caa048cb3d8d076c7d18312720b49efe Mon Sep 17 00:00:00 2001
From: ChsBuffer <33744752+chsbuffer@users.noreply.github.com>
Date: Sun, 12 Sep 2021 03:21:34 +0800
Subject: [PATCH] Refactor AsyncLock

---
 Netch/Controllers/MainController.cs | 58 ++++++++++++++++-------------
 Netch/Utils/DelayTestHelper.cs      | 26 ++++---------
 2 files changed, 39 insertions(+), 45 deletions(-)

diff --git a/Netch/Controllers/MainController.cs b/Netch/Controllers/MainController.cs
index fa332191..320e35a9 100644
--- a/Netch/Controllers/MainController.cs
+++ b/Netch/Controllers/MainController.cs
@@ -28,8 +28,12 @@ namespace Netch.Controllers
 
         public static ModeFeature ModeFeatures { get; private set; }
 
+        private static readonly AsyncSemaphore Lock = new(1);
+
         public static async Task StartAsync(Server server, Mode mode)
         {
+            using var _ = await Lock.EnterAsync();
+
             Log.Information("Start MainController: {Server} {Mode}", $"{server.Type}", $"[{(int)mode.Type}]{mode.Remark}");
 
             if (await DnsUtils.LookupAsync(server.Hostname) == null)
@@ -59,34 +63,28 @@ namespace Netch.Controllers
                 if (modePort != null)
                     TryReleaseTcpPort((ushort)modePort, portName);
 
-
-                switch (Server)
+                if (Server is Socks5Server socks5 && (!socks5.Auth() || ModeFeatures.HasFlag(ModeFeature.SupportSocks5Auth)))
                 {
-                    case Socks5Server socks5 when !socks5.Auth():
-                    case Socks5Server socks5B when socks5B.Auth() && ModeFeatures.HasFlag(ModeFeature.SupportSocks5Auth):
-                        // Directly Start ModeController
-                        Socks5Server = (Socks5Server)Server;
-                        Global.MainForm.StatusText(i18N.TranslateFormat("Starting {0}", ModeController.Name));
-                        await ModeController.StartAsync(Socks5Server, mode);
-                        break;
-                    default:
-                        // Start Server Controller to get a local socks5 server
-                        Log.Debug("Server Information: {Data}", $"{server.Type} {server.MaskedData()}");
-
-                        ServerController = ServerHelper.GetUtilByTypeName(server.Type).GetController();
-                        Global.MainForm.StatusText(i18N.TranslateFormat("Starting {0}", ServerController.Name));
-
-                        TryReleaseTcpPort(ServerController.Socks5LocalPort(), "Socks5");
-                        Socks5Server = await ServerController.StartAsync(server);
-
-                        StatusPortInfoText.Socks5Port = Socks5Server.Port;
-                        StatusPortInfoText.UpdateShareLan();
-
-                        // Start Mode Controller
-                        Global.MainForm.StatusText(i18N.TranslateFormat("Starting {0}", ModeController.Name));
-                        await ModeController.StartAsync(Socks5Server, mode);
-                        break;
+                    Socks5Server = socks5;
                 }
+                else
+                {
+                    // Start Server Controller to get a local socks5 server
+                    Log.Debug("Server Information: {Data}", $"{server.Type} {server.MaskedData()}");
+
+                    ServerController = ServerHelper.GetUtilByTypeName(server.Type).GetController();
+                    Global.MainForm.StatusText(i18N.TranslateFormat("Starting {0}", ServerController.Name));
+
+                    TryReleaseTcpPort(ServerController.Socks5LocalPort(), "Socks5");
+                    Socks5Server = await ServerController.StartAsync(server);
+
+                    StatusPortInfoText.Socks5Port = Socks5Server.Port;
+                    StatusPortInfoText.UpdateShareLan();
+                }
+
+                // Start Mode Controller
+                Global.MainForm.StatusText(i18N.TranslateFormat("Starting {0}", ModeController.Name));
+                await ModeController.StartAsync(Socks5Server, mode);
             }
             catch (Exception e)
             {
@@ -109,6 +107,14 @@ namespace Netch.Controllers
 
         public static async Task StopAsync()
         {
+            if (Lock.CurrentCount == 0)
+            {
+                (await Lock.EnterAsync()).Dispose();
+                return;
+            }
+
+            using var _ = await Lock.EnterAsync();
+
             if (ServerController == null && ModeController == null)
                 return;
 
diff --git a/Netch/Utils/DelayTestHelper.cs b/Netch/Utils/DelayTestHelper.cs
index 383393f7..1be4f8c5 100644
--- a/Netch/Utils/DelayTestHelper.cs
+++ b/Netch/Utils/DelayTestHelper.cs
@@ -1,10 +1,9 @@
 using System;
 using System.Linq;
-using System.Threading;
 using System.Threading.Tasks;
+using System.Timers;
 using Microsoft.VisualStudio.Threading;
 using Netch.Models;
-using Timer = System.Timers.Timer;
 
 namespace Netch.Utils
 {
@@ -12,9 +11,9 @@ namespace Netch.Utils
     {
         private static readonly Timer Timer;
 
-        private static readonly SemaphoreSlim Lock = new(1, 1);
+        private static readonly AsyncSemaphore Lock = new(1);
 
-        private static readonly SemaphoreSlim PoolLock = new(16, 16);
+        private static readonly AsyncSemaphore PoolLock = new(16);
 
         public static readonly NumberRange Range = new(0, int.MaxValue / 1000);
 
@@ -47,28 +46,21 @@ namespace Netch.Utils
             if (Lock.CurrentCount == 0)
             {
                 if (waitFinish)
-                {
-                    await Lock.WaitAsync();
-                    Lock.Release();
-                }
+                    (await Lock.EnterAsync()).Dispose();
 
                 return;
             }
 
-            await Lock.WaitAsync();
+            using var _ = await Lock.EnterAsync();
+
             try
             {
                 var tasks = Global.Settings.Server.Select(async s =>
                 {
-                    await PoolLock.WaitAsync();
-                    try
+                    using (await PoolLock.EnterAsync())
                     {
                         await s.PingAsync();
                     }
-                    finally
-                    {
-                        PoolLock.Release();
-                    }
                 });
 
                 await Task.WhenAll(tasks);
@@ -77,10 +69,6 @@ namespace Netch.Utils
             {
                 // ignored
             }
-            finally
-            {
-                Lock.Release();
-            }
         }
 
         public static void UpdateTick(bool performTestAtOnce = false)