Switch back to an implementation inspired by the System/Net/Quic/Internal/ValueTaskSource.cs

Author: danielmarbach

projects/RabbitMQ.Client/Impl/AsyncManualResetEvent.cs

@@ -29,44 +29,129 @@
 //  Copyright (c) 2007-2025 Broadcom. All Rights Reserved.
 //---------------------------------------------------------------------------
 
+using System;
+using System.Runtime.CompilerServices;
 using System.Threading;
 using System.Threading.Tasks;
+using System.Threading.Tasks.Sources;
 
 namespace RabbitMQ.Client.Impl
 {
-    sealed class AsyncManualResetEvent
+    sealed class AsyncManualResetEvent(bool initialState = false) : IValueTaskSource
     {
-        volatile TaskCompletionSource<bool> _taskCompletionSource = new(TaskCreationOptions.RunContinuationsAsynchronously);
+        private readonly object _lock = new();
 
-        public AsyncManualResetEvent(bool initialState = false)
+        private State _state = initialState ? State.Set : State.Reset;
+
+        // Do not make this field readonly
+        private ManualResetValueTaskSourceCore<bool> _valueTaskSource = new() { RunContinuationsAsynchronously = true };
+
+        public bool IsSet => (State)Volatile.Read(ref Unsafe.As<State, byte>(ref _state)) == State.Set;
+
+        public async ValueTask WaitAsync(CancellationToken cancellationToken = default)
         {
-            if (initialState)
+            ValueTask valueTask;
+            CancellationTokenRegistration tokenRegistration = default;
+
+            lock (_lock)
+            {
+                // If already set, return immediately
+                if (_state == State.Set)
+                {
+                    return;
+                }
+
+                // Create the ValueTask with current version
+                valueTask = new ValueTask(this, _valueTaskSource.Version);
+
+                // Only transition to Awaiting if we're in Reset state
+                if (_state == State.Reset)
+                {
+                    _state = State.Awaiting;
+
+                    // Register cancellation if token can be cancelled
+                    if (cancellationToken.CanBeCanceled)
+                    {
+                        tokenRegistration = cancellationToken.UnsafeRegister(
+                            static state =>
+                            {
+                                (AsyncManualResetEvent amre, CancellationToken token) =
+                                    (Tuple<AsyncManualResetEvent, CancellationToken>)state!;
+                                amre.SetCancelled(token);
+                            }, Tuple.Create(this, cancellationToken));
+                    }
+                }
+            }
+
+            try
             {
-                _taskCompletionSource.SetResult(true);
+                await valueTask.ConfigureAwait(false);
+            }
+            finally
+            {
+                // Dispose cancellation registration outside of lock to avoid deadlock
+                if (tokenRegistration != default)
+                {
+#if NET
+                    await tokenRegistration.DisposeAsync().ConfigureAwait(false);
+#else
+                tokenRegistration.Dispose();
+#endif
+                }
             }
         }
 
-        public bool IsSet => _taskCompletionSource.Task.IsCompleted;
-
-        public Task WaitAsync(CancellationToken cancellationToken = default)
+        public void Set()
         {
-            Task<bool> task = _taskCompletionSource.Task;
-            return task.IsCompleted ? task : task.WaitAsync(cancellationToken);
-        }
+            lock (_lock)
+            {
+                State previousState = _state;
+                _state = State.Set;
 
-        public void Set() => _taskCompletionSource.TrySetResult(true);
+                // Only set result if we were in Awaiting state
+                if (previousState == State.Awaiting)
+                {
+                    _valueTaskSource.SetResult(true);
+                }
+            }
+        }
 
         public void Reset()
         {
-            while (true)
+            lock (_lock)
             {
-                TaskCompletionSource<bool> currentTcs = _taskCompletionSource;
-                if (!currentTcs.Task.IsCompleted ||
-                    Interlocked.CompareExchange(ref _taskCompletionSource, new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously), currentTcs) == currentTcs)
+                if (_state == State.Set)
                 {
-                    return;
+                    _state = State.Reset;
+                    _valueTaskSource.Reset();
+                }
+            }
+        }
+
+        void SetCancelled(CancellationToken cancellationToken)
+        {
+            lock (_lock)
+            {
+                if (_state == State.Awaiting)
+                {
+                    _state = State.Reset; // Reset to allow future waits
+                    _valueTaskSource.SetException(new OperationCanceledException(cancellationToken));
                 }
             }
         }
+
+        void IValueTaskSource.GetResult(short token) => _valueTaskSource.GetResult(token);
+
+        ValueTaskSourceStatus IValueTaskSource.GetStatus(short token) => _valueTaskSource.GetStatus(token);
+
+        void IValueTaskSource.OnCompleted(Action<object?> continuation, object? state, short token,
+            ValueTaskSourceOnCompletedFlags flags) => _valueTaskSource.OnCompleted(continuation, state, token, flags);
+
+        enum State : byte
+        {
+            Reset,
+            Set,
+            Awaiting
+        }
     }
 }

projects/RabbitMQ.Client/Impl/Channel.BasicPublish.cs

@@ -229,11 +229,11 @@ void MaybeAddPublishSequenceNumberToHeaders(IDictionary<string, object?> headers
         }
 
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
-        private Task MaybeEnforceFlowControlAsync(CancellationToken cancellationToken)
+        private ValueTask MaybeEnforceFlowControlAsync(CancellationToken cancellationToken)
         {
             if (_flowControlBlock.IsSet)
             {
-                return Task.CompletedTask;
+                return default;
             }
 
             return _flowControlBlock.WaitAsync(cancellationToken);