| | | 1 | | // Licensed to the .NET Foundation under one or more agreements. |
| | | 2 | | // The .NET Foundation licenses this file to you under the MIT license. |
| | | 3 | | |
| | | 4 | | using System.Runtime.ExceptionServices; |
| | | 5 | | using System.Threading; |
| | | 6 | | using System.Threading.Tasks; |
| | | 7 | | using System.Threading.Tasks.Sources; |
| | | 8 | | |
| | | 9 | | namespace System.Net.Http |
| | | 10 | | { |
| | | 11 | | /// <summary>Represents a waiter for credit.</summary> |
| | | 12 | | internal sealed class CreditWaiter : IValueTaskSource<int> |
| | | 13 | | { |
| | | 14 | | // State for the implementation of the CreditWaiter. Note that neither _cancellationToken nor |
| | | 15 | | // _registration are zero'd out upon completion, because they're used for synchronization |
| | | 16 | | // between successful completion and cancellation. This means an instance may end up |
| | | 17 | | // referencing the underlying CancellationTokenSource even after the await operation has completed. |
| | | 18 | | |
| | | 19 | | /// <summary>Cancellation token for the current wait operation.</summary> |
| | | 20 | | private CancellationToken _cancellationToken; |
| | | 21 | | /// <summary>Cancellation registration for the current wait operation.</summary> |
| | | 22 | | private CancellationTokenRegistration _registration; |
| | | 23 | | /// <summary><see cref="IValueTaskSource"/> implementation.</summary> |
| | | 24 | | private ManualResetValueTaskSourceCore<int> _source; |
| | | 25 | | |
| | | 26 | | // State carried with the waiter for the consumer to use; these aren't used at all in the implementation. |
| | | 27 | | |
| | | 28 | | /// <summary>Amount of credit desired by this waiter.</summary> |
| | | 29 | | public int Amount; |
| | | 30 | | /// <summary>Next waiter in a list of waiters.</summary> |
| | | 31 | | public CreditWaiter? Next; |
| | | 32 | | |
| | | 33 | | /// <summary>Initializes a waiter for a credit wait operation.</summary> |
| | | 34 | | /// <param name="cancellationToken">The cancellation token for this wait operation.</param> |
| | 0 | 35 | | public CreditWaiter(CancellationToken cancellationToken) |
| | 0 | 36 | | { |
| | 0 | 37 | | _source.RunContinuationsAsynchronously = true; |
| | 0 | 38 | | RegisterCancellation(cancellationToken); |
| | 0 | 39 | | } |
| | | 40 | | |
| | | 41 | | /// <summary>Re-initializes a waiter for a credit wait operation.</summary> |
| | | 42 | | /// <param name="cancellationToken">The cancellation token for this wait operation.</param> |
| | | 43 | | public void ResetForAwait(CancellationToken cancellationToken) |
| | 0 | 44 | | { |
| | 0 | 45 | | _source.Reset(); |
| | 0 | 46 | | RegisterCancellation(cancellationToken); |
| | 0 | 47 | | } |
| | | 48 | | |
| | | 49 | | /// <summary>Registers with the cancellation token to transition the source to a canceled state.</summary> |
| | | 50 | | /// <param name="cancellationToken">The cancellation token with which to register.</param> |
| | | 51 | | private void RegisterCancellation(CancellationToken cancellationToken) |
| | 0 | 52 | | { |
| | 0 | 53 | | _cancellationToken = cancellationToken; |
| | 0 | 54 | | _registration = cancellationToken.UnsafeRegister(static (s, cancellationToken) => |
| | 0 | 55 | | { |
| | 0 | 56 | | // The callback will only fire if cancellation owns the right to complete the instance. |
| | 0 | 57 | | ((CreditWaiter)s!)._source.SetException(ExceptionDispatchInfo.SetCurrentStackTrace(new OperationCanceled |
| | 0 | 58 | | }, this); |
| | 0 | 59 | | } |
| | | 60 | | |
| | | 61 | | /// <summary>Wraps the instance as a <see cref="ValueTask{TResult}"/> to make it awaitable.</summary> |
| | 0 | 62 | | public ValueTask<int> AsValueTask() => new ValueTask<int>(this, _source.Version); |
| | | 63 | | |
| | | 64 | | /// <summary>Completes the instance with the specified result.</summary> |
| | | 65 | | /// <param name="result">The result value.</param> |
| | | 66 | | /// <returns>true if the instance was successfully completed; false if it was or is being canceled.</returns> |
| | | 67 | | public bool TrySetResult(int result) |
| | 0 | 68 | | { |
| | 0 | 69 | | if (UnregisterAndOwnCompletion()) |
| | 0 | 70 | | { |
| | 0 | 71 | | _source.SetResult(result); |
| | 0 | 72 | | return true; |
| | | 73 | | } |
| | | 74 | | |
| | 0 | 75 | | return false; |
| | 0 | 76 | | } |
| | | 77 | | |
| | | 78 | | /// <summary>Disposes the instance, failing any outstanding wait.</summary> |
| | | 79 | | public void Dispose() |
| | 0 | 80 | | { |
| | 0 | 81 | | if (UnregisterAndOwnCompletion()) |
| | 0 | 82 | | { |
| | 0 | 83 | | _source.SetException(ExceptionDispatchInfo.SetCurrentStackTrace(new ObjectDisposedException(nameof(Credi |
| | 0 | 84 | | } |
| | 0 | 85 | | } |
| | | 86 | | |
| | | 87 | | /// <summary>Unregisters the cancellation callback.</summary> |
| | | 88 | | /// <returns>true if the non-cancellation caller has the right to complete the instance; false if the instance w |
| | | 89 | | private bool UnregisterAndOwnCompletion() => |
| | | 90 | | // Unregister the cancellation callback. If Unregister returns true, then the cancellation callback was suc |
| | | 91 | | // meaning it hasn't run and won't ever run. If it returns false, a) cancellation already occurred or is oc |
| | | 92 | | // the callback couldn't be removed, b) cancellation occurred prior to the UnsafeRegister call such that _re |
| | | 93 | | // set to a default value (or hasn't been set yet), or c) a default CancellationToken was used. (a) and (b) |
| | | 94 | | // the same, and (c) can be checked via CanBeCanceled. |
| | 0 | 95 | | _registration.Unregister() || !_cancellationToken.CanBeCanceled; |
| | | 96 | | |
| | | 97 | | int IValueTaskSource<int>.GetResult(short token) => |
| | 0 | 98 | | _source.GetResult(token); |
| | | 99 | | ValueTaskSourceStatus IValueTaskSource<int>.GetStatus(short token) => |
| | 0 | 100 | | _source.GetStatus(token); |
| | | 101 | | void IValueTaskSource<int>.OnCompleted(Action<object?> continuation, object? state, short token, ValueTaskSource |
| | 0 | 102 | | _source.OnCompleted(continuation, state, token, flags); |
| | | 103 | | } |
| | | 104 | | } |