From 43bf73a36f0bd1b091c70575bda94f68d2756270 Mon Sep 17 00:00:00 2001 From: Radek Zikmund Date: Tue, 19 Sep 2023 13:21:29 +0200 Subject: [PATCH] Fix data race between cancellation and dispose --- .../Internal/ResettableValueTaskSource.cs | 29 +++++++++++++++++-- .../src/System/Net/Quic/QuicStream.cs | 20 ++++++++++--- 2 files changed, 42 insertions(+), 7 deletions(-) diff --git a/src/libraries/System.Net.Quic/src/System/Net/Quic/Internal/ResettableValueTaskSource.cs b/src/libraries/System.Net.Quic/src/System/Net/Quic/Internal/ResettableValueTaskSource.cs index e7c0cf87bfd5d8..42020ea9711a3a 100644 --- a/src/libraries/System.Net.Quic/src/System/Net/Quic/Internal/ResettableValueTaskSource.cs +++ b/src/libraries/System.Net.Quic/src/System/Net/Quic/Internal/ResettableValueTaskSource.cs @@ -29,6 +29,7 @@ private enum State private ManualResetValueTaskSourceCore _valueTaskSource; private CancellationTokenRegistration _cancellationRegistration; private Action? _cancellationAction; + private object? _cancellationActionLock; private GCHandle _keepAlive; private readonly TaskCompletionSource _finalTaskSource; @@ -50,6 +51,11 @@ public ResettableValueTaskSource(bool runContinuationsAsynchronously = true) /// public Action CancellationAction { init { _cancellationAction = value; } } + /// + /// Optional lock object to be held during the callback. If no lock is provided, the callback will be invoked without holding any lock. + /// + public object? CancellationActionLock { init { _cancellationActionLock = value; } } + /// /// Returns true is this task source has entered its final state, i.e. or /// was called with final set to true and the result was propagated. @@ -91,10 +97,27 @@ public bool TryGetValueTask(out ValueTask valueTask, object? keepAlive = null, C _cancellationRegistration = cancellationToken.UnsafeRegister(static (obj, cancellationToken) => { (ResettableValueTaskSource thisRef, object? target) = ((ResettableValueTaskSource, object?))obj!; - // This will transition the state to Ready. - if (thisRef.TrySetException(new OperationCanceledException(cancellationToken))) + + object? lockObject = thisRef._cancellationActionLock; + try + { + if (lockObject != null) + { + Monitor.Enter(lockObject); + } + + // This will transition the state to Ready. + if (thisRef.TrySetException(new OperationCanceledException(cancellationToken))) + { + thisRef._cancellationAction?.Invoke(target); + } + } + finally { - thisRef._cancellationAction?.Invoke(target); + if (lockObject != null) + { + Monitor.Exit(lockObject); + } } }, (this, keepAlive)); } diff --git a/src/libraries/System.Net.Quic/src/System/Net/Quic/QuicStream.cs b/src/libraries/System.Net.Quic/src/System/Net/Quic/QuicStream.cs index 6165f2085cb5f0..34f7bc076955a7 100644 --- a/src/libraries/System.Net.Quic/src/System/Net/Quic/QuicStream.cs +++ b/src/libraries/System.Net.Quic/src/System/Net/Quic/QuicStream.cs @@ -89,7 +89,9 @@ public sealed partial class QuicStream private ReceiveBuffers _receiveBuffers = new ReceiveBuffers(); private int _receivedNeedsEnable; - private readonly ResettableValueTaskSource _sendTcs = new ResettableValueTaskSource() + private readonly ResettableValueTaskSource _sendTcs; + // helper function since we need to pass the _shutdownTcs to the _sendTcs + private ResettableValueTaskSource CreateSendTaskCompletionSource() => new ResettableValueTaskSource() { CancellationAction = target => { @@ -106,7 +108,12 @@ public sealed partial class QuicStream // when using CancellationTokenSource.CancelAfter. // Ignore the exception } - } + }, + + // since we're using _sendTcs (Abortive) and _shutdownTcs (Graceful) for controlling write direction shutdown, + // CancellationAction above can race with disposing the stream. Holding _shutdownTcs will prevent Dispose from + // performing graceful shutdown while we are Aborting it. + CancellationActionLock = _shutdownTcs }; private MsQuicBuffers _sendBuffers = new MsQuicBuffers(); private readonly object _sendBuffersLock = new object(); @@ -185,6 +192,7 @@ internal unsafe QuicStream(MsQuicContextSafeHandle connectionHandle, QuicStreamT _receiveTcs.TrySetResult(final: true); } _type = type; + _sendTcs = CreateSendTaskCompletionSource(); } /// @@ -211,6 +219,7 @@ internal unsafe QuicStream(MsQuicContextSafeHandle connectionHandle, QUIC_HANDLE context.Free(); throw; } + _sendTcs = CreateSendTaskCompletionSource(); _defaultErrorCode = defaultErrorCode; @@ -355,15 +364,18 @@ public ValueTask WriteAsync(ReadOnlyMemory buffer, bool completeWrites, Ca return ValueTask.FromCanceled(cancellationToken); } - // Concurrent call, this one lost the race. + // attempt to reserve the task. If the token was already canceled, this will also invoke Abort(Write, DefaultErrorCode) and transition to final state. if (!_sendTcs.TryGetValueTask(out ValueTask valueTask, this, cancellationToken)) { + // Concurrent call, this one lost the race. throw new InvalidOperationException(SR.Format(SR.net_io_invalidnestedcall, "write")); } - // No need to call anything since we already have a result, most likely an exception. + // Task.Delay(TimeSpan.FromSeconds(1), CancellationToken.None).Wait(CancellationToken.None); + if (valueTask.IsCompleted) { + // No need to continue since we already have a result, most likely an exception. return valueTask; }