Fix race between Attach and Detach Tokens

aspnet · Dec 1, 2015 · f9035d5 · f9035d5
1 parent 834385e
commit f9035d5
Show file tree

Hide file tree

Showing 3 changed files with 91 additions and 20 deletions.
diff --git a/src/Microsoft.Extensions.Caching.Memory/CacheEntry.cs b/src/Microsoft.Extensions.Caching.Memory/CacheEntry.cs
@@ -16,6 +16,8 @@ internal class CacheEntry
 
         private readonly DateTimeOffset? _absoluteExpiration;
 
+        internal readonly object _lock = new object();
+
         internal CacheEntry(
             object key,
             object value,
@@ -100,47 +102,54 @@ internal bool CheckForExpiredTokens()
             return false;
         }
 
-        // TODO: There's a possible race between AttachTokens and DetachTokens if a token fires almost immediately.
-        // This may result in some registrations not getting disposed.
         internal void AttachTokens()
         {
             var expirationTokens = Options.ExpirationTokens;
             if (expirationTokens != null)
             {
-                for (int i = 0; i < expirationTokens.Count; i++)
+                lock (_lock)
                 {
-                    var expirationToken = expirationTokens[i];
-                    if (expirationToken.ActiveChangeCallbacks)
+                    for (int i = 0; i < expirationTokens.Count; i++)
                     {
-                        if (ExpirationTokenRegistrations == null)
+                        var expirationToken = expirationTokens[i];
+                        if (expirationToken.ActiveChangeCallbacks)
                         {
-                            ExpirationTokenRegistrations = new List<IDisposable>(1);
+                            if (ExpirationTokenRegistrations == null)
+                            {
+                                ExpirationTokenRegistrations = new List<IDisposable>(1);
+                            }
+                            var registration = expirationToken.RegisterChangeCallback(ExpirationCallback, this);
+                            ExpirationTokenRegistrations.Add(registration);
                         }
-                        var registration = expirationToken.RegisterChangeCallback(ExpirationCallback, this);
-                        ExpirationTokenRegistrations.Add(registration);
                     }
                 }
             }
         }
 
         private static void ExpirationTokensExpired(object obj)
         {
-            var entry = (CacheEntry)obj;
-            entry.SetExpired(EvictionReason.TokenExpired);
-            entry._notifyCacheOfExpiration(entry);
+            // start a new thread to avoid issues with callbacks called from RegisterChangeCallback
+            Task.Factory.StartNew(state =>
+            {
+                var entry = (CacheEntry)state;
+                entry.SetExpired(EvictionReason.TokenExpired);
+                entry._notifyCacheOfExpiration(entry);
+            }, obj, CancellationToken.None, TaskCreationOptions.DenyChildAttach, TaskScheduler.Default);
         }
 
-        // TODO: Thread safety
         private void DetachTokens()
         {
-            var registrations = ExpirationTokenRegistrations;
-            if (registrations != null)
+            lock(_lock)
             {
-                ExpirationTokenRegistrations = null;
-                for (int i = 0; i < registrations.Count; i++)
+                var registrations = ExpirationTokenRegistrations;
+                if (registrations != null)
                 {
-                    var registration = registrations[i];
-                    registration.Dispose();
+                    ExpirationTokenRegistrations = null;
+                    for (int i = 0; i < registrations.Count; i++)
+                    {
+                        var registration = registrations[i];
+                        registration.Dispose();
+                    }
                 }
             }
         }

diff --git a/test/Microsoft.Extensions.Caching.Memory.Tests/MemoryCacheSetAndRemoveTests.cs b/test/Microsoft.Extensions.Caching.Memory.Tests/MemoryCacheSetAndRemoveTests.cs
@@ -246,7 +246,7 @@ public void SetGetAndRemoveWorksWithObjectKeysWhenDifferentReferences()
             result = cache.Get(new TestKey());
             Assert.Same(obj, result);
 
-            var key = new TestKey(); 
+            var key = new TestKey();
             cache.Remove(key);
             result = cache.Get(key);
             Assert.Null(result);

diff --git a/test/Microsoft.Extensions.Caching.Memory.Tests/TokenExpirationTests.cs b/test/Microsoft.Extensions.Caching.Memory.Tests/TokenExpirationTests.cs
@@ -3,8 +3,10 @@
 
 using System;
 using System.Threading;
+using System.Threading.Tasks;
 using Microsoft.Extensions.Caching.Memory.Infrastructure;
 using Microsoft.Extensions.Internal;
+using Microsoft.Extensions.Primitives;
 using Xunit;
 
 namespace Microsoft.Extensions.Caching.Memory
@@ -187,5 +189,65 @@ public void AddExpiredTokenPreventsCaching()
             result = cache.Get(key);
             Assert.Null(result); // It wasn't cached
         }
+
+        [Fact]
+        public void TokenExpiresOnRegister()
+        {
+            var cache = CreateCache();
+            var key = "myKey";
+            var value = new object();
+            var callbackInvoked = new ManualResetEvent(false);
+            var expirationToken = new TestToken(callbackInvoked);
+            var task = Task.Run(() => cache.Set(key, value, new MemoryCacheEntryOptions()
+                .AddExpirationToken(expirationToken)));
+            callbackInvoked.WaitOne(TimeSpan.FromSeconds(30));
+            var result = task.Result;
+
+            Assert.Same(value, result);
+            result = cache.Get(key);
+            Assert.Null(result);
+        }
+
+        internal class TestToken : IChangeToken
+        {
+            private bool _hasChanged;
+            private ManualResetEvent _event;
+
+            public TestToken(ManualResetEvent mre)
+            {
+                _event = mre;
+            }
+
+            public bool ActiveChangeCallbacks
+            {
+                get
+                {
+                    return true;
+                }
+            }
+
+            public bool HasChanged
+            {
+                get
+                {
+                    return _hasChanged;
+                }
+            }
+
+            public IDisposable RegisterChangeCallback(Action<object> callback, object state)
+            {
+                _hasChanged = true;
+                callback(state);
+                _event.Set();
+                return new TestDisposable();
+            }
+        }
+
+        internal class TestDisposable : IDisposable
+        {
+            public void Dispose()
+            {
+            }
+        }
     }
 }