diff --git a/src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs b/src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs index dc0fafd0a0e4cb..d39d92867891d6 100644 --- a/src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs +++ b/src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs @@ -217,7 +217,7 @@ private struct RuntimeAsyncAwaitState // to one of these notifiers. public ICriticalNotifyCompletion? CriticalNotifier; public INotifyCompletion? Notifier; - public IValueTaskSourceNotifier? ValueTaskSourceNotifier; + public ValueTaskSourceNotifier? ValueTaskSourceNotifier; public Task? TaskNotifier; public ExecutionContext? ExecutionContext; @@ -310,7 +310,7 @@ private static void TransparentAwait(object o) } else { - state.ValueTaskSourceNotifier = (IValueTaskSourceNotifier)o; + state.ValueTaskSourceNotifier = (ValueTaskSourceNotifier)o; } state.CaptureContexts(); @@ -366,7 +366,7 @@ internal void HandleSuspended() ICriticalNotifyCompletion? critNotifier = state.CriticalNotifier; INotifyCompletion? notifier = state.Notifier; - IValueTaskSourceNotifier? vtsNotifier = state.ValueTaskSourceNotifier; + ValueTaskSourceNotifier? vtsNotifier = state.ValueTaskSourceNotifier; Task? taskNotifier = state.TaskNotifier; state.CriticalNotifier = null; diff --git a/src/libraries/System.Private.CoreLib/src/System/Threading/Tasks/ValueTask.cs b/src/libraries/System.Private.CoreLib/src/System/Threading/Tasks/ValueTask.cs index 692cfbb0a32cc3..c20a2e8b7cb7b9 100644 --- a/src/libraries/System.Private.CoreLib/src/System/Threading/Tasks/ValueTask.cs +++ b/src/libraries/System.Private.CoreLib/src/System/Threading/Tasks/ValueTask.cs @@ -179,28 +179,19 @@ obj as Task ?? GetTaskForValueTaskSource(Unsafe.As(obj)); } - internal object AsTaskOrNotifier() + /// + /// Helper to invoke IValueTaskSource.OnCompleted from a caller that does not know the actual type of the source. + /// + private static void OnCompleted(object obj, Action continuation, object? state, short token, ValueTaskSourceOnCompletedFlags flags) => + Unsafe.As(obj).OnCompleted(continuation, state, token, flags); + + internal unsafe object AsTaskOrNotifier() { object? obj = _obj; Debug.Assert(obj is Task || obj is IValueTaskSource); return obj as Task ?? - (object)new ValueTaskSourceNotifier(Unsafe.As(obj), _token); - } - - private sealed class ValueTaskSourceNotifier : IValueTaskSourceNotifier - { - private IValueTaskSource _valueTaskSource; - private short _token; - - public ValueTaskSourceNotifier(IValueTaskSource valueTaskSource, short token) - { - _valueTaskSource = valueTaskSource; - _token = token; - } - - public void OnCompleted(Action continuation, object? state, ValueTaskSourceOnCompletedFlags flags) => - _valueTaskSource.OnCompleted(continuation, state, _token, flags); + (object)ValueTaskSourceNotifier.GetInstance(obj, &OnCompleted, _token); } /// Gets a that may be used at any point in the future. @@ -612,28 +603,19 @@ public Task AsTask() return GetTaskForValueTaskSource(Unsafe.As>(obj)); } - internal object AsTaskOrNotifier() + /// + /// Helper to invoke IValueTaskSource.OnCompleted from a caller that does not know the actual type of the source. + /// + private static void OnCompleted(object obj, Action continuation, object? state, short token, ValueTaskSourceOnCompletedFlags flags) => + Unsafe.As>(obj).OnCompleted(continuation, state, token, flags); + + internal unsafe object AsTaskOrNotifier() { object? obj = _obj; Debug.Assert(obj is Task || obj is IValueTaskSource); return obj as Task ?? - (object)new ValueTaskSourceNotifier(Unsafe.As>(obj), _token); - } - - private sealed class ValueTaskSourceNotifier : IValueTaskSourceNotifier - { - private IValueTaskSource _valueTaskSource; - private short _token; - - public ValueTaskSourceNotifier(IValueTaskSource valueTaskSource, short token) - { - _valueTaskSource = valueTaskSource; - _token = token; - } - - public void OnCompleted(Action continuation, object? state, ValueTaskSourceOnCompletedFlags flags) => - _valueTaskSource.OnCompleted(continuation, state, _token, flags); + (object)ValueTaskSourceNotifier.GetInstance(obj, &OnCompleted, _token); } /// Gets a that may be used at any point in the future. @@ -897,8 +879,53 @@ public ConfiguredValueTaskAwaitable ConfigureAwait(bool continueOnCaptu } } - internal interface IValueTaskSourceNotifier + internal unsafe sealed class ValueTaskSourceNotifier { - void OnCompleted(Action continuation, object? state, ValueTaskSourceOnCompletedFlags flags); + // ValueTaskSourceNotifier is used only during suspension sequence, thus + // a given thread will never need more than one instance. + // We will just reuse the same instance when needed. + [ThreadStatic] + private static ValueTaskSourceNotifier? t_instance; + + private object _source; + private delegate* managed, object?, short, ValueTaskSourceOnCompletedFlags, void> _onCompleted; + private short _token; + + private ValueTaskSourceNotifier( + object source, + delegate* managed, object?, short, ValueTaskSourceOnCompletedFlags, void> onCompleted, + short token) + { + _source = source; + _onCompleted = onCompleted; + _token = token; + } + + public static ValueTaskSourceNotifier GetInstance( + object source, + delegate* managed, object?, short, ValueTaskSourceOnCompletedFlags, void> onCompleted, + short token) + { + ValueTaskSourceNotifier? instance = t_instance; + if (instance == null) + { + return t_instance = new ValueTaskSourceNotifier(source, onCompleted, token); + } + + instance._source = source; + instance._onCompleted = onCompleted; + instance._token = token; + + return instance; + } + + public void OnCompleted(Action continuation, object? state, ValueTaskSourceOnCompletedFlags flags) + { + _onCompleted(_source, continuation, state, _token, flags); + + // The data that we store is effectively single-use. + // Once used, clear the _source to not retain unknown data. + _source = null!; + } } }