Fix CallbackManager to be thread-safe (#1428)

SteamRE · Sep 15, 2024 · 223468c · 223468c
1 parent 08e2df9
commit 223468c
Show file tree

Hide file tree

Showing 2 changed files with 50 additions and 15 deletions.
diff --git a/SteamKit2/SteamKit2/Steam/SteamClient/CallbackMgr/CallbackMgr.cs b/SteamKit2/SteamKit2/Steam/SteamClient/CallbackMgr/CallbackMgr.cs
@@ -6,7 +6,7 @@
 
 
 using System;
-using System.Collections.Generic;
+using System.Collections.Immutable;
 using System.Threading;
 using System.Threading.Tasks;
 using SteamKit2.Internal;
@@ -22,7 +22,7 @@ public sealed class CallbackManager : ICallbackMgrInternals
     {
         SteamClient client;
 
-        List<CallbackBase> registeredCallbacks;
+        ImmutableList<CallbackBase> registeredCallbacks = [];
 
 
 
@@ -34,8 +34,6 @@ public CallbackManager( SteamClient client )
         {
             ArgumentNullException.ThrowIfNull( client );
 
-            registeredCallbacks = [];
-
             this.client = client;
         }
 
@@ -141,19 +139,18 @@ public IDisposable Subscribe<TCallback>( Action<TCallback> callbackFunc )
         }
 
         void ICallbackMgrInternals.Register( CallbackBase call )
-        {
-            if ( registeredCallbacks.Contains( call ) )
-                return;
+            => ImmutableInterlocked.Update( ref registeredCallbacks, static ( list, item ) => list.Add( item ), call );
 
-            registeredCallbacks.Add( call );
-        }
+        void ICallbackMgrInternals.Unregister( CallbackBase call )
+            => ImmutableInterlocked.Update( ref registeredCallbacks, static ( list, item ) => list.Remove( item ), call );
 
         void Handle( ICallbackMsg call )
         {
+            var callbacks = registeredCallbacks;
             var type = call.GetType();
 
             // find handlers interested in this callback
-            foreach ( var callback in registeredCallbacks )
+            foreach ( var callback in callbacks )
             {
                 if ( callback.CallbackType.IsAssignableFrom( type ) )
                 {
@@ -162,11 +159,6 @@ void Handle( ICallbackMsg call )
             }
         }
 
-        void ICallbackMgrInternals.Unregister( CallbackBase call )
-        {
-            registeredCallbacks.Remove( call );
-        }
-
         sealed class Subscription : IDisposable
         {
             public Subscription( CallbackBase call, ICallbackMgrInternals manager )

diff --git a/SteamKit2/Tests/CallbackManagerFacts.cs b/SteamKit2/Tests/CallbackManagerFacts.cs
@@ -232,6 +232,49 @@ void action( CallbackForTest cb )
             }
         }
 
+        [Fact]
+        public void CorrectlyUnsubscribesFromInsideOfCallback()
+        {
+            static void nothing( CallbackForTest cb )
+            {
+                //
+            }
+
+            using var s1 = mgr.Subscribe<CallbackForTest>( nothing );
+
+            IDisposable subscription = null;
+
+            void unsubscribe( CallbackForTest cb )
+            {
+                Assert.NotNull( subscription );
+                subscription.Dispose();
+                subscription = null;
+            }
+
+            subscription = mgr.Subscribe<CallbackForTest>( unsubscribe );
+
+            PostAndRunCallback( new CallbackForTest { UniqueID = Guid.NewGuid() } );
+        }
+
+        [Fact]
+        public void CorrectlysubscribesFromInsideOfCallback()
+        {
+            static void nothing( CallbackForTest cb )
+            {
+                //
+            }
+
+            void subscribe( CallbackForTest cb )
+            {
+                using var s2 = mgr.Subscribe<CallbackForTest>( nothing );
+            }
+
+            using var s1 = mgr.Subscribe<CallbackForTest>( nothing );
+            using var se = mgr.Subscribe<CallbackForTest>( subscribe );
+
+            PostAndRunCallback( new CallbackForTest { UniqueID = Guid.NewGuid() } );
+        }
+
         void PostAndRunCallback(CallbackMsg callback)
         {
             client.PostCallback(callback);