netcore/System.Private.CoreLib/shared/System/Threading/CancellationTokenRegistration.cs

   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 // See the LICENSE file in the project root for more information.
   4
   5 using System.Threading.Tasks;
   6
   7 namespace System.Threading
   8 {
   9     /// <summary>
  10     /// Represents a callback delegate that has been registered with a <see cref="System.Threading.CancellationToken">CancellationToken</see>.
  11     /// </summary>
  12     /// <remarks>
  13     /// To unregister a callback, dispose the corresponding Registration instance.
  14     /// </remarks>
  15     public readonly struct CancellationTokenRegistration : IEquatable<CancellationTokenRegistration>, IDisposable, IAsyncDisposable
  16     {
  17         private readonly long _id;
  18         private readonly CancellationTokenSource.CallbackNode _node;
  19
  20         internal CancellationTokenRegistration(long id, CancellationTokenSource.CallbackNode node)
  21         {
  22             _id = id;
  23             _node = node;
  24         }
  25
  26         /// <summary>
  27         /// Disposes of the registration and unregisters the target callback from the associated
  28         /// <see cref="System.Threading.CancellationToken">CancellationToken</see>.
  29         /// If the target callback is currently executing, this method will wait until it completes, except
  30         /// in the degenerate cases where a callback method unregisters itself.
  31         /// </summary>
  32         public void Dispose()
  33         {
  34             CancellationTokenSource.CallbackNode node = _node;
  35             if (node != null && !node.Partition.Unregister(_id, node))
  36             {
  37                 WaitForCallbackIfNecessary();
  38             }
  39         }
  40
  41         /// <summary>
  42         /// Disposes of the registration and unregisters the target callback from the associated
  43         /// <see cref="System.Threading.CancellationToken">CancellationToken</see>.
  44         /// The returned <see cref="ValueTask"/> will complete once the associated callback
  45         /// is unregistered without having executed or once it's finished executing, except
  46         /// in the degenerate case where the callback itself is unregistering itself.
  47         /// </summary>
  48         public ValueTask DisposeAsync()
  49         {
  50             CancellationTokenSource.CallbackNode node = _node;
  51             return node != null && !node.Partition.Unregister(_id, node) ?
  52                 WaitForCallbackIfNecessaryAsync() :
  53                 default;
  54         }
  55
  56         /// <summary>
  57         /// Gets the <see cref="CancellationToken"/> with which this registration is associated.  If the
  58         /// registration isn't associated with a token (such as after the registration has been disposed),
  59         /// this will return a default token.
  60         /// </summary>
  61         public CancellationToken Token
  62         {
  63             get
  64             {
  65                 CancellationTokenSource.CallbackNode node = _node;
  66                 return node != null ?
  67                     new CancellationToken(node.Partition.Source) : // avoid CTS.Token, which throws after disposal
  68                     default;
  69             }
  70         }
  71
  72         /// <summary>
  73         /// Disposes of the registration and unregisters the target callback from the associated
  74         /// <see cref="System.Threading.CancellationToken">CancellationToken</see>.
  75         /// </summary>
  76         public bool Unregister()
  77         {
  78             CancellationTokenSource.CallbackNode node = _node;
  79             return node != null && node.Partition.Unregister(_id, node);
  80         }
  81
  82         private void WaitForCallbackIfNecessary()
  83         {
  84             // We're a valid registration but we were unable to unregister, which means the callback wasn't in the list,
  85             // which means either it already executed or it's currently executing. We guarantee that we will not return
  86             // if the callback is being executed (assuming we are not currently called by the callback itself)
  87             // We achieve this by the following rules:
  88             //    1. If we are called in the context of an executing callback, no need to wait (determined by tracking callback-executor threadID)
  89             //       - if the currently executing callback is this CTR, then waiting would deadlock. (We choose to return rather than deadlock)
  90             //       - if not, then this CTR cannot be the one executing, hence no need to wait
  91             //    2. If unregistration failed, and we are on a different thread, then the callback may be running under control of cts.Cancel()
  92             //       => poll until cts.ExecutingCallback is not the one we are trying to unregister.
  93             CancellationTokenSource source = _node.Partition.Source;
  94             if (source.IsCancellationRequested && // Running callbacks has commenced.
  95                 !source.IsCancellationCompleted && // Running callbacks hasn't finished.
  96                 source.ThreadIDExecutingCallbacks != Environment.CurrentManagedThreadId) // The executing thread ID is not this thread's ID.
  97             {
  98                 // Callback execution is in progress, the executing thread is different from this thread and has taken the callback for execution
  99                 // so observe and wait until this target callback is no longer the executing callback.
 100                 source.WaitForCallbackToComplete(_id);
 101             }
 102         }
 103
 104         private ValueTask WaitForCallbackIfNecessaryAsync()
 105         {
 106             // Same as WaitForCallbackIfNecessary, except returning a task that'll be completed when callbacks complete.
 107
 108             CancellationTokenSource source = _node.Partition.Source;
 109             if (source.IsCancellationRequested && // Running callbacks has commenced.
 110                 !source.IsCancellationCompleted && // Running callbacks hasn't finished.
 111                 source.ThreadIDExecutingCallbacks != Environment.CurrentManagedThreadId) // The executing thread ID is not this thread's ID.
 112             {
 113                 // Callback execution is in progress, the executing thread is different from this thread and has taken the callback for execution
 114                 // so get a task that'll complete when this target callback is no longer the executing callback.
 115                 return source.WaitForCallbackToCompleteAsync(_id);
 116             }
 117
 118             // Callback is either already completed, won't execute, or the callback itself is calling this.
 119             return default;
 120         }
 121
 122         /// <summary>
 123         /// Determines whether two <see
 124         /// cref="System.Threading.CancellationTokenRegistration">CancellationTokenRegistration</see>
 125         /// instances are equal.
 126         /// </summary>
 127         /// <param name="left">The first instance.</param>
 128         /// <param name="right">The second instance.</param>
 129         /// <returns>True if the instances are equal; otherwise, false.</returns>
 130         public static bool operator ==(CancellationTokenRegistration left, CancellationTokenRegistration right) => left.Equals(right);
 131
 132         /// <summary>
 133         /// Determines whether two <see cref="System.Threading.CancellationTokenRegistration">CancellationTokenRegistration</see> instances are not equal.
 134         /// </summary>
 135         /// <param name="left">The first instance.</param>
 136         /// <param name="right">The second instance.</param>
 137         /// <returns>True if the instances are not equal; otherwise, false.</returns>
 138         public static bool operator !=(CancellationTokenRegistration left, CancellationTokenRegistration right) => !left.Equals(right);
 139
 140         /// <summary>
 141         /// Determines whether the current <see cref="System.Threading.CancellationTokenRegistration">CancellationTokenRegistration</see> instance is equal to the
 142         /// specified <see cref="object"/>.
 143         /// </summary>
 144         /// <param name="obj">The other object to which to compare this instance.</param>
 145         /// <returns>True, if both this and <paramref name="obj"/> are equal. False, otherwise.
 146         /// Two <see cref="System.Threading.CancellationTokenRegistration">CancellationTokenRegistration</see> instances are equal if
 147         /// they both refer to the output of a single call to the same Register method of a
 148         /// <see cref="System.Threading.CancellationToken">CancellationToken</see>.
 149         /// </returns>
 150         public override bool Equals(object? obj) => obj is CancellationTokenRegistration && Equals((CancellationTokenRegistration)obj);
 151
 152         /// <summary>
 153         /// Determines whether the current <see cref="System.Threading.CancellationToken">CancellationToken</see> instance is equal to the
 154         /// specified <see cref="object"/>.
 155         /// </summary>
 156         /// <param name="other">The other <see cref="System.Threading.CancellationTokenRegistration">CancellationTokenRegistration</see> to which to compare this instance.</param>
 157         /// <returns>True, if both this and <paramref name="other"/> are equal. False, otherwise.
 158         /// Two <see cref="System.Threading.CancellationTokenRegistration">CancellationTokenRegistration</see> instances are equal if
 159         /// they both refer to the output of a single call to the same Register method of a
 160         /// <see cref="System.Threading.CancellationToken">CancellationToken</see>.
 161         /// </returns>
 162         public bool Equals(CancellationTokenRegistration other) => _node == other._node && _id == other._id;
 163
 164         /// <summary>
 165         /// Serves as a hash function for a <see cref="System.Threading.CancellationTokenRegistration">CancellationTokenRegistration.</see>.
 166         /// </summary>
 167         /// <returns>A hash code for the current <see cref="System.Threading.CancellationTokenRegistration">CancellationTokenRegistration</see> instance.</returns>
 168         public override int GetHashCode() => _node != null ? _node.GetHashCode() ^ _id.GetHashCode()  : _id.GetHashCode();
 169     }
 170 }