optimize timeout dictionary

j4ckzh0u · Dec 29, 2022 · ebeb075 · ebeb075
1 parent 5dcf5cd
commit ebeb075
Show file tree

Hide file tree

Showing 5 changed files with 54 additions and 37 deletions.
diff --git a/VpnHood.Common/Collections/TimeoutDictionary.cs b/VpnHood.Common/Collections/TimeoutDictionary.cs
@@ -29,16 +29,22 @@ public int Count
 
     public TValue GetOrAdd(TKey key, Func<TKey, TValue> valueFactory)
     {
-        lock(_items)
-        {
-            if (TryGetValue(key, out var value))
-                return value;
+        Cleanup();
 
-            value = valueFactory(key);
-            if (!TryAdd(key, value, false))
-                throw new Exception($"Could not add an item to {GetType().Name}");
-            return value;
-        }
+        // update old item if expired or return the old item
+        var res = _items.AddOrUpdate(key, valueFactory,
+            (k, oldValue) =>
+            {
+                if (IsExpired(oldValue))
+                {
+                    oldValue.Dispose();
+                    return valueFactory(k);
+                }
+                return oldValue;
+            });
+
+        res.AccessedTime = FastDateTime.Now;
+        return res;
     }
 
     public bool TryGetValue(TKey key, out TValue value)
@@ -62,31 +68,31 @@ public bool TryGetValue(TKey key, out TValue value)
         return true;
     }
 
-    public bool TryAdd(TKey key, TValue value, bool overwrite)
+    public TValue AddOrUpdate(TKey key, TValue value)
     {
         Cleanup();
-        value.AccessedTime = FastDateTime.Now;
 
-        // return true if added
-        if (_items.TryAdd(key, value))
-            return true;
-
-        // remove and retry if overwrite is on
-        if (overwrite)
+        var res = _items.AddOrUpdate(key, value, (_, oldValue) =>
         {
-            TryRemove(key, out _);
-            return _items.TryAdd(key, value);
-        }
+            oldValue.Dispose();
+            return value;
+        });
 
-        // remove & retry of an item that has been expired
-        if (_items.TryGetValue(key, out var prevValue) && IsExpired(prevValue))
-        {
-            TryRemove(key, out _);
-            return _items.TryAdd(key, value);
-        }
+        res.AccessedTime = FastDateTime.Now;
+        return res;
+    }
+
+    public bool TryAdd(TKey key, TValue value)
+    {
+        Cleanup();
+
+        // return true if added
+        if (!_items.TryAdd(key, value))
+            return false;
+
+        value.AccessedTime = FastDateTime.Now;
+        return true;
 
-        // couldn't add
-        return false;
     }
 
     public bool TryRemove(TKey key, out TValue value)

diff --git a/VpnHood.Server/Session.cs b/VpnHood.Server/Session.cs
@@ -8,6 +8,7 @@
 using Microsoft.Extensions.Logging;
 using PacketDotNet;
 using VpnHood.Common.Client;
+using VpnHood.Common.Collections;
 using VpnHood.Common.Logging;
 using VpnHood.Common.Messaging;
 using VpnHood.Common.Utils;

diff --git a/VpnHood.Server/TcpHost.cs b/VpnHood.Server/TcpHost.cs
@@ -445,9 +445,10 @@ private async Task ProcessTcpProxyChannel(TcpClientStream tcpClientStream, Cance
             VhLogger.Instance.LogTrace(GeneralEventId.StreamChannel,
                 $"Connecting to the requested endpoint. RequestedEP: {VhLogger.Format(request.DestinationEndPoint)}");
 
-            // Check tcp wait limit
+            // Apply limitation
             lock (session)
             {
+                // Check tcp wait limit
                 if (session.TcpConnectWaitCount >= MaxTcpConnectWaitCount)
                     throw new MaxTcpConnectException(session.SessionId);
 
@@ -456,6 +457,8 @@ private async Task ProcessTcpProxyChannel(TcpClientStream tcpClientStream, Cance
 
                 Interlocked.Increment(ref session.TcpConnectWaitCount);
                 isTcpConnectIncreased = true;
+
+                // Anti TcpScan
             }
 
             // prepare client

diff --git a/VpnHood.ZTest/Tests/AccessTest.cs b/VpnHood.ZTest/Tests/AccessTest.cs
@@ -3,14 +3,29 @@
 using System.Threading.Tasks;
 using Microsoft.VisualStudio.TestTools.UnitTesting;
 using VpnHood.Client;
+using VpnHood.Common.Collections;
 using VpnHood.Common.Logging;
 using VpnHood.Common.Messaging;
+using VpnHood.Common.Utils;
 
 namespace VpnHood.Test.Tests;
 
 [TestClass]
 public class AccessTest
 {
+    [TestMethod]
+    public void Foo()
+    {
+        var dic = new TimeoutDictionary<int, TimeoutItem<string>>(TimeSpan.FromSeconds(1));
+        var a1 = dic.GetOrAdd(1, k => new TimeoutItem<string>("aa", false));
+        Console.WriteLine(FastDateTime.Now);
+        Thread.Sleep(2000);
+        a1 = dic.GetOrAdd(1, k => new TimeoutItem<string>("bb", false));
+
+        //var a2 = dic.GetOrAdd(1, k => new TimeoutItem<string>("222", false));
+        Console.WriteLine($@"{a1.AccessedTime}, {a1.Value}");
+    }
+
     [TestInitialize]
     public void Initialize()
     {

diff --git a/VpnHood.ZTest/Tests/ClientServerTest.cs b/VpnHood.ZTest/Tests/ClientServerTest.cs
@@ -28,14 +28,6 @@ public void Initialize()
         VhLogger.Instance = VhLogger.CreateConsoleLogger(true);
     }
 
-    [TestMethod]
-    public void Foo()
-    {
-        var s = new TimeSpan();
-        var b = s == TimeSpan.Zero;
-        Console.WriteLine(b);
-    }
-
     [TestMethod]
     public void Redirect_Server()
     {