Added async overload of ChangeToken.OnChange

src/libraries/Microsoft.Extensions.Primitives/ref/Microsoft.Extensions.Primitives.cs

@@ -16,7 +16,9 @@ public CancellationChangeToken(System.Threading.CancellationToken cancellationTo
     public static partial class ChangeToken
     {
         public static System.IDisposable OnChange(System.Func<Microsoft.Extensions.Primitives.IChangeToken?> changeTokenProducer, System.Action changeTokenConsumer) { throw null; }
+        public static System.IDisposable OnChange(System.Func<Microsoft.Extensions.Primitives.IChangeToken?> changeTokenProducer, System.Func<System.Threading.Tasks.Task> changeTokenConsumer) { throw null; }
         public static System.IDisposable OnChange<TState>(System.Func<Microsoft.Extensions.Primitives.IChangeToken?> changeTokenProducer, System.Action<TState> changeTokenConsumer, TState state) { throw null; }
+        public static System.IDisposable OnChange<TState>(System.Func<Microsoft.Extensions.Primitives.IChangeToken?> changeTokenProducer, System.Func<TState, System.Threading.Tasks.Task> changeTokenConsumer, TState state) { throw null; }
     }
     public partial class CompositeChangeToken : Microsoft.Extensions.Primitives.IChangeToken
     {

src/libraries/Microsoft.Extensions.Primitives/src/ChangeToken.cs

@@ -3,6 +3,7 @@
 
 using System;
 using System.Threading;
+using System.Threading.Tasks;
 
 namespace Microsoft.Extensions.Primitives
 {
@@ -15,8 +16,12 @@ public static class ChangeToken
         /// Registers the <paramref name="changeTokenConsumer"/> action to be called whenever the token produced changes.
         /// </summary>
         /// <param name="changeTokenProducer">Produces the change token.</param>
-        /// <param name="changeTokenConsumer">Action called when the token changes.</param>
-        /// <returns></returns>
+        /// <param name="changeTokenConsumer">Action called when the token changes. The token is re-registered once the action returns.</param>
+        /// <returns>An <see cref="IDisposable"/> that, when disposed, unregisters the consumer.</returns>
+        /// <remarks>
+        /// Exceptions from <paramref name="changeTokenProducer"/> are propagated to the caller of this method or to the code that triggers the change token.
+        /// Exceptions from <paramref name="changeTokenConsumer"/> are propagated to the code that triggers the change token.
+        /// </remarks>
         public static IDisposable OnChange(Func<IChangeToken?> changeTokenProducer, Action changeTokenConsumer)
         {
             if (changeTokenProducer is null)
@@ -28,16 +33,20 @@ public static IDisposable OnChange(Func<IChangeToken?> changeTokenProducer, Acti
                 ThrowHelper.ThrowArgumentNullException(ExceptionArgument.changeTokenConsumer);
             }
 
-            return new ChangeTokenRegistration<Action>(changeTokenProducer, callback => callback(), changeTokenConsumer);
+            return new SyncChangeTokenRegistration<Action>(changeTokenProducer, static callback => callback(), changeTokenConsumer);
         }
 
         /// <summary>
         /// Registers the <paramref name="changeTokenConsumer"/> action to be called whenever the token produced changes.
         /// </summary>
         /// <param name="changeTokenProducer">Produces the change token.</param>
-        /// <param name="changeTokenConsumer">Action called when the token changes.</param>
+        /// <param name="changeTokenConsumer">Action called when the token changes. The token is re-registered once the action returns.</param>
         /// <param name="state">state for the consumer.</param>
-        /// <returns></returns>
+        /// <returns>An <see cref="IDisposable"/> that, when disposed, unregisters the consumer.</returns>
+        /// <remarks>
+        /// Exceptions from <paramref name="changeTokenProducer"/> are propagated to the caller of this method or to the code that triggers the change token.
+        /// Exceptions from <paramref name="changeTokenConsumer"/> are propagated to the code that triggers the change token.
+        /// </remarks>
         public static IDisposable OnChange<TState>(Func<IChangeToken?> changeTokenProducer, Action<TState> changeTokenConsumer, TState state)
         {
             if (changeTokenProducer is null)
@@ -49,65 +58,97 @@ public static IDisposable OnChange<TState>(Func<IChangeToken?> changeTokenProduc
                 ThrowHelper.ThrowArgumentNullException(ExceptionArgument.changeTokenConsumer);
             }
 
-            return new ChangeTokenRegistration<TState>(changeTokenProducer, changeTokenConsumer, state);
+            return new SyncChangeTokenRegistration<TState>(changeTokenProducer, changeTokenConsumer, state);
         }
 
-        private sealed class ChangeTokenRegistration<TState> : IDisposable
+        /// <summary>
+        /// Registers the <paramref name="changeTokenConsumer"/> function to be called whenever the token produced changes.
+        /// </summary>
+        /// <param name="changeTokenProducer">Produces the change token.</param>
+        /// <param name="changeTokenConsumer">Function called when the token changes. The token is only re-registered once the returned <see cref="Task"/> completes.</param>
+        /// <returns>An <see cref="IDisposable"/> that, when disposed, unregisters the consumer.</returns>
+        /// <remarks>
+        /// Exceptions from <paramref name="changeTokenProducer"/> are propagated to the caller of this method or to the code that triggers the change token.
+        /// Synchronous exceptions from <paramref name="changeTokenConsumer"/> are propagated to the code that triggers the change token.
+        /// Asynchronous exceptions from <paramref name="changeTokenConsumer"/> are left unobserved.
+        /// </remarks>
+        public static IDisposable OnChange(Func<IChangeToken?> changeTokenProducer, Func<Task> changeTokenConsumer)
+        {
+            if (changeTokenProducer is null)
+            {
+                ThrowHelper.ThrowArgumentNullException(ExceptionArgument.changeTokenProducer);
+            }
+            if (changeTokenConsumer is null)
+            {
+                ThrowHelper.ThrowArgumentNullException(ExceptionArgument.changeTokenConsumer);
+            }
+
+            return new AsyncChangeTokenRegistration<Func<Task>>(changeTokenProducer, static callback => callback(), changeTokenConsumer);
+        }
+
+        /// <summary>
+        /// Registers the <paramref name="changeTokenConsumer"/> function to be called whenever the token produced changes.
+        /// </summary>
+        /// <param name="changeTokenProducer">Produces the change token.</param>
+        /// <param name="changeTokenConsumer">Function called when the token changes. The token is only re-registered once the returned <see cref="Task"/> completes.</param>
+        /// <param name="state">state for the consumer.</param>
+        /// <returns>An <see cref="IDisposable"/> that, when disposed, unregisters the consumer.</returns>
+        /// <remarks>
+        /// Exceptions from <paramref name="changeTokenProducer"/> are propagated to the caller of this method or to the code that triggers the change token.
+        /// Synchronous exceptions from <paramref name="changeTokenConsumer"/> are propagated to the code that triggers the change token.
+        /// Asynchronous exceptions from <paramref name="changeTokenConsumer"/> are left unobserved.
+        /// </remarks>
+        public static IDisposable OnChange<TState>(Func<IChangeToken?> changeTokenProducer, Func<TState, Task> changeTokenConsumer, TState state)
+        {
+            if (changeTokenProducer is null)
+            {
+                ThrowHelper.ThrowArgumentNullException(ExceptionArgument.changeTokenProducer);
+            }
+            if (changeTokenConsumer is null)
+            {
+                ThrowHelper.ThrowArgumentNullException(ExceptionArgument.changeTokenConsumer);
+            }
+
+            return new AsyncChangeTokenRegistration<TState>(changeTokenProducer, changeTokenConsumer, state);
+        }
+
+        private abstract class ChangeTokenRegistration<TState>(Func<IChangeToken?> changeTokenProducer, TState state) : IDisposable
         {
-            private readonly Func<IChangeToken?> _changeTokenProducer;
-            private readonly Action<TState> _changeTokenConsumer;
-            private readonly TState _state;
             private IDisposable? _disposable;
 
             private static readonly NoopDisposable _disposedSentinel = new NoopDisposable();
 
-            public ChangeTokenRegistration(Func<IChangeToken?> changeTokenProducer, Action<TState> changeTokenConsumer, TState state)
-            {
-                _changeTokenProducer = changeTokenProducer;
-                _changeTokenConsumer = changeTokenConsumer;
-                _state = state;
+            protected TState State { get; } = state;
 
-                IChangeToken? token = changeTokenProducer();
+            protected Func<IChangeToken?> ChangeTokenProducer { get; } = changeTokenProducer;
 
-                RegisterChangeTokenCallback(token);
-            }
+            protected abstract void OnChangeTokenFired();
 
-            private void OnChangeTokenFired()
+            protected void RegisterChangeTokenCallback(IChangeToken? token)
             {
-                // The order here is important. We need to take the token and then apply our changes BEFORE
-                // registering. This prevents us from possible having two change updates to process concurrently.
-                //
-                // If the token changes after we take the token, then we'll process the update immediately upon
-                // registering the callback.
-                IChangeToken? token = _changeTokenProducer();
-
-                try
-                {
-                    _changeTokenConsumer(_state);
-                }
-                finally
+                if (token is null)
                 {
-                    // We always want to ensure the callback is registered
-                    RegisterChangeTokenCallback(token);
+                    return;
                 }
-            }
 
-            private void RegisterChangeTokenCallback(IChangeToken? token)
-            {
-                if (token is null)
+                // If the registration has already been disposed, don't register again. This guards re-registration
+                // after disposal: registering on a token that has already changed invokes the callback synchronously, which
+                // would re-run the consumer after disposal.
+                if (Volatile.Read(ref _disposable) == _disposedSentinel)
                 {
                     return;
                 }
-                IDisposable registraton = token.RegisterChangeCallback(s => ((ChangeTokenRegistration<TState>?)s)!.OnChangeTokenFired(), this);
+
+                IDisposable? registration = token.RegisterChangeCallback(static s => ((ChangeTokenRegistration<TState>?)s)!.OnChangeTokenFired(), this);
                 if (token.HasChanged && token.ActiveChangeCallbacks)
                 {
-                    registraton?.Dispose();
+                    registration?.Dispose();
                     return;
                 }
-                SetDisposable(registraton);
+                SetDisposable(registration);
             }
 
-            private void SetDisposable(IDisposable disposable)
+            private void SetDisposable(IDisposable? disposable)
             {
                 // We don't want to transition from _disposedSentinel => anything since it's terminal
                 // but we want to allow going from previously assigned disposable, to another
@@ -117,7 +158,7 @@ private void SetDisposable(IDisposable disposable)
                 // If Dispose was called, then immediately dispose the disposable
                 if (current == _disposedSentinel)
                 {
-                    disposable.Dispose();
+                    disposable?.Dispose();
                     return;
                 }
 
@@ -127,7 +168,7 @@ private void SetDisposable(IDisposable disposable)
                 if (previous == _disposedSentinel)
                 {
                     // The subscription was disposed so we dispose immediately and return
-                    disposable.Dispose();
+                    disposable?.Dispose();
                 }
                 else if (previous == current)
                 {
@@ -136,7 +177,7 @@ private void SetDisposable(IDisposable disposable)
                 else
                 {
                     // Sets can never overlap with other SetDisposable calls so we should never get into this situation
-                    throw new InvalidOperationException("Somebody else set the _disposable field");
+                    ThrowHelper.ThrowInvalidOperationException(ExceptionResource.InvalidOperation_ConcurrentDisposableSet);
                 }
             }
 
@@ -154,5 +195,106 @@ public void Dispose()
                 }
             }
         }
+
+        private sealed class SyncChangeTokenRegistration<TState> : ChangeTokenRegistration<TState>
+        {
+            private readonly Action<TState> _changeTokenConsumer;
+
+            public SyncChangeTokenRegistration(Func<IChangeToken?> changeTokenProducer, Action<TState> changeTokenConsumer, TState state)
+                : base(changeTokenProducer, state)
+            {
+                _changeTokenConsumer = changeTokenConsumer;
+
+                RegisterChangeTokenCallback(changeTokenProducer());
+            }
+
+            protected override void OnChangeTokenFired()
+            {
+                // The order here is important. We need to take the token and then apply our changes BEFORE
+                // registering. This prevents us from possible having two change updates to process concurrently.
+                //
+                // If the token changes after we take the token, then we'll process the update immediately upon
+                // registering the callback.
+                IChangeToken? token = ChangeTokenProducer();
+
+                try
+                {
+                    _changeTokenConsumer(State);
+                }
+                finally
+                {
+                    // We always want to ensure the callback is registered
+                    RegisterChangeTokenCallback(token);
+                }
+            }
+        }
+
+        private sealed class AsyncChangeTokenRegistration<TState> : ChangeTokenRegistration<TState>
+        {
+            private readonly Func<TState, Task> _changeTokenConsumer;
+
+            public AsyncChangeTokenRegistration(Func<IChangeToken?> changeTokenProducer, Func<TState, Task> changeTokenConsumer, TState state)
+                : base(changeTokenProducer, state)
+            {
+                _changeTokenConsumer = changeTokenConsumer;
+
+                RegisterChangeTokenCallback(changeTokenProducer());
+            }
+
+            protected override void OnChangeTokenFired()
+            {
+                // The order here is important. We need to take the token and then apply our changes BEFORE
+                // registering. This prevents us from possible having two change updates to process concurrently.
+                //
+                // If the token changes after we take the token, then we'll process the update immediately upon
+                // registering the callback once the consumer's task completes.
+                IChangeToken? token = ChangeTokenProducer();
+
+                Task consumerTask;
+                try
+                {
+                    // The consumer is invoked synchronously here, so that synchronous exceptions from it are propagated
+                    // to the code that triggers the change token, just like the sync overload does.
+                    consumerTask = _changeTokenConsumer(State);
+                }
+                catch
+                {
+                    // We always want to ensure the callback is registered, even when the consumer throws synchronously.
+                    RegisterChangeTokenCallback(token);
+                    throw;
+                }
+
+                if (consumerTask is null)
+                {
+                    RegisterChangeTokenCallback(token);
+                    ThrowHelper.ThrowInvalidOperationException(ExceptionResource.InvalidOperation_NullConsumerTask);
+                }
+
+                if (consumerTask.Status == TaskStatus.RanToCompletion)
+                {
+                    // The common case where the consumer completes synchronously: re-register without allocations.
+                    RegisterChangeTokenCallback(token);
+                }
+                else
+                {
+                    // Asynchronous exceptions can't be propagated without blocking, so they are left unobserved
+                    // (meaning they can be observed only through TaskScheduler.UnobservedTaskException).
+                    _ = AwaitConsumerAndRegisterCallback(consumerTask, token);
+                }
+            }
+
+            private async Task AwaitConsumerAndRegisterCallback(Task consumerTask, IChangeToken? token)
+            {
+                try
+                {
+                    await consumerTask.ConfigureAwait(false);
+                }
+                finally
+                {
+                    // We always want to ensure the callback is registered
+                    RegisterChangeTokenCallback(token);
+                }
+            }
+        }
     }
 }

src/libraries/Microsoft.Extensions.Primitives/src/Resources/Strings.resx

@@ -132,4 +132,10 @@
   <data name="Capacity_NotUsedEntirely" xml:space="preserve">
     <value>Entire reserved capacity was not used. Capacity: '{0}', written '{1}'.</value>
   </data>
+  <data name="InvalidOperation_ConcurrentDisposableSet" xml:space="preserve">
+    <value>The change token callback registration was set concurrently, which is not supported.</value>
+  </data>
+  <data name="InvalidOperation_NullConsumerTask" xml:space="preserve">
+    <value>The change token consumer returned a null task.</value>
+  </data>
 </root>

src/libraries/Microsoft.Extensions.Primitives/src/ThrowHelper.cs

@@ -68,6 +68,8 @@ private static string GetResourceText(ExceptionResource resource)
                 case ExceptionResource.Capacity_CannotChangeAfterWriteStarted: return SR.Capacity_CannotChangeAfterWriteStarted;
                 case ExceptionResource.Capacity_NotEnough: return SR.Capacity_NotEnough;
                 case ExceptionResource.Capacity_NotUsedEntirely: return SR.Capacity_NotUsedEntirely;
+                case ExceptionResource.InvalidOperation_ConcurrentDisposableSet: return SR.InvalidOperation_ConcurrentDisposableSet;
+                case ExceptionResource.InvalidOperation_NullConsumerTask: return SR.InvalidOperation_NullConsumerTask;
                 default:
                     Debug.Fail($"Unexpected resource {resource}");
                     return "";
@@ -108,6 +110,8 @@ internal enum ExceptionResource
         Argument_InvalidOffsetLengthStringSegment,
         Capacity_CannotChangeAfterWriteStarted,
         Capacity_NotEnough,
-        Capacity_NotUsedEntirely
+        Capacity_NotUsedEntirely,
+        InvalidOperation_ConcurrentDisposableSet,
+        InvalidOperation_NullConsumerTask
     }
 }