Fix: 2 memory leaks in Tokens websocket endpoint

Piotrekol · May 14, 2023 · 6f30d4e · 6f30d4e
1 parent de5486e
commit 6f30d4e
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 4 deletions.
diff --git a/plugins/WebSocketDataSender/LockingQueue.cs b/plugins/WebSocketDataSender/LockingQueue.cs
@@ -0,0 +1,30 @@
+using System.Collections.Generic;
+
+namespace WebSocketDataSender
+{
+    internal class LockingQueue<T>
+    {
+        protected Queue<T> queue = new Queue<T>();
+        private readonly object syncRoot = new();
+
+        public int Count { get { lock (syncRoot) return queue.Count; } }
+
+        public bool TryPeek(out T value)
+        {
+            lock (syncRoot)
+                return queue.TryPeek(out value);
+        }
+
+        public bool TryDequeue(out T value)
+        {
+            lock (syncRoot)
+                return queue.TryDequeue(out value);
+        }
+
+        public void Enqueue(T value)
+        {
+            lock (syncRoot)
+                queue.Enqueue(value);
+        }
+    }
+}
diff --git a/plugins/WebSocketDataSender/WebSocketTokenEndpoint.cs b/plugins/WebSocketDataSender/WebSocketTokenEndpoint.cs
@@ -84,6 +84,9 @@ public async Task SendLoop(IWebSocketContext context)
             }
             finally
             {
+                foreach (var token in state.WatchedTokens)
+                    token.ValueUpdated -= state.TokenValueUpdated;
+
                 state.Dispose();
                 contextStates.Remove(context.Id);
             }
@@ -106,9 +109,13 @@ protected override Task OnMessageReceivedAsync(IWebSocketContext context, byte[]
                 return Task.CompletedTask;
 
             var settings = contextStates[context.Id];
-            settings.RequestedTokenNames.Clear();
-            settings.RequestedTokenNames.AddRange(kvNames);
-            UpdateListenedTokens(settings);
+
+            lock (settings)
+            {
+                settings.RequestedTokenNames.Clear();
+                settings.RequestedTokenNames.AddRange(kvNames);
+                UpdateListenedTokens(settings);
+            }
 
             return Task.CompletedTask;
         }
@@ -142,7 +149,7 @@ private class ContextTokensState : IDisposable
             public List<IToken> WatchedTokens { get; set; } = new();
             public List<string> RequestedTokenNames { get; set; } = new();
             public ManualResetEventSlim ManualResetEventSlim { get; set; } = new();
-            public ConcurrentQueue<IToken> TokensPendingUpdate { get; private set; } = new();
+            public LockingQueue<IToken> TokensPendingUpdate { get; private set; } = new();
 
             public void TokenValueUpdated(object _, IToken token)
             {