diff --git a/src/libraries/System.IO.Pipelines/src/System/IO/Pipelines/PipeAwaitable.cs b/src/libraries/System.IO.Pipelines/src/System/IO/Pipelines/PipeAwaitable.cs index fb25780a3fab4..aba52f13e9646 100644 --- a/src/libraries/System.IO.Pipelines/src/System/IO/Pipelines/PipeAwaitable.cs +++ b/src/libraries/System.IO.Pipelines/src/System/IO/Pipelines/PipeAwaitable.cs @@ -47,16 +47,33 @@ public void BeginOperation(CancellationToken cancellationToken, Action { cancellationToken.ThrowIfCancellationRequested(); - _awaitableState |= AwaitableState.Running; - // Don't register if already completed, we would immediately unregistered in ObserveCancellation if (cancellationToken.CanBeCanceled && !IsCompleted) { +#if DEBUG + var previousAwaitableState = _awaitableState; +#endif + + _cancellationTokenRegistration = cancellationToken.UnsafeRegister(callback, state); + + // If we get back the default CancellationTokenRegistration then it means the + // callback synchronously and we can throw inline. This is safe because we haven't changed + // the state of the awaitable as yet. + if (_cancellationTokenRegistration == default(CancellationTokenRegistration)) + { +#if DEBUG + Debug.Assert(previousAwaitableState == _awaitableState, "The awaitable state changed!"); +#endif + + cancellationToken.ThrowIfCancellationRequested(); + } + #if (NETSTANDARD2_0 || NETFRAMEWORK) _cancellationToken = cancellationToken; #endif - _cancellationTokenRegistration = cancellationToken.UnsafeRegister(callback, state); } + + _awaitableState |= AwaitableState.Running; } [MethodImpl(MethodImplOptions.AggressiveInlining)]