Limit parallelism of forced key ring refreshes

src/DataProtection/DataProtection/src/Error.cs

@@ -97,4 +97,9 @@ public static InvalidOperationException KeyRingProvider_DefaultKeyRevoked(Guid i
         var message = string.Format(CultureInfo.CurrentCulture, Resources.KeyRingProvider_DefaultKeyRevoked, id);
         return new InvalidOperationException(message);
     }
+
+    public static InvalidOperationException KeyRingProvider_RefreshFailedOnOtherThread()
+    {
+        return new InvalidOperationException(Resources.KeyRingProvider_RefreshFailedOnOtherThread);
+    }
 }

src/DataProtection/DataProtection/src/KeyManagement/KeyRingProvider.cs

@@ -215,7 +215,7 @@ internal IKeyRing GetCurrentKeyRingCore(DateTime utcNow, bool forceRefresh = fal
         var acquiredLock = false;
         try
         {
-            Monitor.TryEnter(_cacheableKeyRingLockObj, (existingCacheableKeyRing != null) ? 0 : Timeout.Infinite, ref acquiredLock);
+            Monitor.TryEnter(_cacheableKeyRingLockObj, (forceRefresh || existingCacheableKeyRing != null) ? 0 : Timeout.Infinite, ref acquiredLock);
             if (acquiredLock)
             {
                 if (!forceRefresh)
@@ -274,9 +274,23 @@ internal IKeyRing GetCurrentKeyRingCore(DateTime utcNow, bool forceRefresh = fal
                 Volatile.Write(ref _cacheableKeyRing, newCacheableKeyRing);
                 return newCacheableKeyRing.KeyRing;
             }
+            else if (forceRefresh)
+            {
+                // If this is a forced refresh, we should not just reuse the existing key ring (if any).
+                // Instead, we'll wait for the thread that *did* acquire the critical section
+                // to release it, at which point it will have updated _cacheableKeyRing.
+                Monitor.Enter(_cacheableKeyRingLockObj, ref acquiredLock);
+                existingCacheableKeyRing = Volatile.Read(ref _cacheableKeyRing);
+                if (existingCacheableKeyRing == null)
+                {
+                    // There will have been a better exception from the thread that acquired the critical section.
+                    throw Error.KeyRingProvider_RefreshFailedOnOtherThread();
+                }
+                return existingCacheableKeyRing.KeyRing;
+            }
             else
             {
-                // We didn't acquire the critical section. This should only occur if we passed
+                // We didn't acquire the critical section for a non-forced refresh. This should only occur if we passed
                 // zero for the Monitor.TryEnter timeout, which implies that we had an existing
                 // (but outdated) keyring that we can use as a fallback.
                 Debug.Assert(existingCacheableKeyRing != null);

src/DataProtection/DataProtection/src/Resources.resx

@@ -190,6 +190,9 @@
   <data name="KeyRingProvider_DefaultKeyRevoked" xml:space="preserve">
     <value>The key ring's default data protection key {0:B} has been revoked.</value>
   </data>
+  <data name="KeyRingProvider_RefreshFailedOnOtherThread" xml:space="preserve">
+    <value>Another thread failed to refresh the key ring.</value>
+  </data>
   <data name="LifetimeMustNotBeNegative" xml:space="preserve">
     <value>{0} must not be negative.</value>
   </data>

src/DataProtection/DataProtection/test/Microsoft.AspNetCore.DataProtection.Tests/KeyManagement/KeyRingProviderTests.cs

@@ -844,6 +844,72 @@ private static ICacheableKeyRingProvider SetupCreateCacheableKeyRingTestAndCreat
         return CreateKeyRingProvider(mockKeyManager.Object, mockDefaultKeyResolver.Object, keyManagementOptions);
     }
 
+    [Theory]
+    [InlineData(true)]
+    [InlineData(false)]
+    public async Task MultipleThreadsForceRefresh(bool failsToReadKeyRing)
+    {
+        const int taskCount = 10;
+        var now = StringToDateTime("2015-03-01 00:00:00Z");
+
+        var expectedKeyRing = new Mock<IKeyRing>();
+        var expectedException = new InvalidOperationException(nameof(MultipleThreadsForceRefresh));
+
+        var mockCacheableKeyRingProvider = new Mock<ICacheableKeyRingProvider>();
+        if (failsToReadKeyRing)
+        {
+            mockCacheableKeyRingProvider
+                .Setup(o => o.GetCacheableKeyRing(now))
+                .Throws(expectedException);
+        }
+        else
+        {
+            mockCacheableKeyRingProvider
+                .Setup(o => o.GetCacheableKeyRing(now))
+                .Returns(new CacheableKeyRing(
+                    expirationToken: CancellationToken.None,
+                    expirationTime: now.AddDays(1),
+                    keyRing: expectedKeyRing.Object));
+        }
+
+        var keyRingProvider = CreateKeyRingProvider(mockCacheableKeyRingProvider.Object);
+
+        var tasks = new Task<IKeyRing>[taskCount];
+        for (var i = 0; i < taskCount; i++)
+        {
+            tasks[i] = Task.Run(() =>
+            {
+                var keyRing = keyRingProvider.GetCurrentKeyRingCore(now, forceRefresh: true);
+                return keyRing;
+            });
+        }
+
+        if (failsToReadKeyRing)
+        {
+            await Task.WhenAll(tasks).ContinueWith(t => { }, TaskContinuationOptions.OnlyOnFaulted); // Don't throw an exception
+            Assert.All(tasks, task => Assert.NotNull(task.Exception));
+
+            // We expect only one task to have thrown expectedException, but it's possible that multiple
+            // threads made it into the critical section (in sequence, obviously) and each saw expectedException.
+            // This check is descriptive, rather than normative - it would probably be preferable to have all of
+            // them see expectedException, but it's presently not propagated to threads that simply wait.
+            Assert.InRange(tasks.Count(task => ReferenceEquals(expectedException, task.Exception.InnerException)), 1, taskCount);
+        }
+        else
+        {
+            var actualKeyRings = await Task.WhenAll(tasks);
+            Assert.All(actualKeyRings, actualKeyRing => ReferenceEquals(expectedKeyRing, actualKeyRing));
+        }
+
+        // We'd like there to be exactly one call, but it's possible that the first thread will actually
+        // release the critical section before the last thread attempts to acquire it and the work will
+        // be redone (by design, since the refresh is forced).
+        // Even asserting < taskCount is probabilistic - it's possible, though very unlikely, that each
+        // thread could finish before the next even attempts to enter the criticial section.  If this
+        // proves to be flaky, we could increase taskCount or intentionally slow down GetCacheableKeyRing.
+        mockCacheableKeyRingProvider.Verify(o => o.GetCacheableKeyRing(It.IsAny<DateTimeOffset>()), Times.AtMost(taskCount - 1));
+    }
+
     private static KeyRingProvider CreateKeyRingProvider(ICacheableKeyRingProvider cacheableKeyRingProvider)
     {
         var mockEncryptorFactory = new Mock<IAuthenticatedEncryptorFactory>();