Deterministic behavior of CancellationOrigin in the case of timeout

Author: sakno

src/DotNext.Tests/Threading/CancellationTokenMultiplexerTests.cs

@@ -137,4 +137,22 @@ private static void CheckDefaultScope<TScope>(TScope scope)
         False(scope.Token.IsCancellationRequested);
         Equal(scope.Token, scope.CancellationOrigin);
     }
+
+    [Fact]
+    public static async Task CompleteAndCheckConcurrently()
+    {
+        var multiplexer = new CancellationTokenMultiplexer();
+
+        var scope = multiplexer.CombineAndSetTimeoutLater([]);
+        var source = new TaskCompletionSource<bool>();
+        await using var registration = scope.Token.Register(() =>
+        {
+            source.SetResult(scope.IsTimedOut);
+        });
+
+        scope.Timeout = TimeSpan.Zero;
+        True(await source.Task);
+
+        await scope.DisposeAsync();
+    }
 }

src/DotNext.Threading/Threading/LinkedCancellationTokenSource.cs

@@ -74,6 +74,43 @@ private bool TrySetCancellationOrigin(CancellationToken token)
         return Interlocked.CompareExchange(ref cancellationOrigin.Item1, inlinedToken.Item1, comparand: null) is null;
     }
 
+    [DebuggerBrowsable(DebuggerBrowsableState.Never)]
+    private object RawToken
+    {
+        get
+        {
+            // There is rare race condition that could lead to incorrect detection of the cancellation root:
+            // the linked token gets canceled in the same time as the timeout happens for this source.
+            // In this case, the timeout thread resumes the attached callbacks. However, the order
+            // of these callbacks is not guaranteed. Thus, OnTimeout is not yet called, and cancellationOrigin
+            // is still null. The resumed thread observes CancellationOrigin == Token. But then,
+            // the linked token calls Cancel callback that sets cancellationOrigin to the real token.
+            // In that case, CancellationOrigin != Token. It means that the consumer of CancellationOrigin
+            // can see two different values when the property getter is called sequentially. This
+            // is non-deterministic behavior. Currently, nothing we can do with incorrect detection
+            // of the cancellation root due to absence of the necessary methods in CTS. But we can
+            // achieve deterministic behavior for CancellationOrigin and IsRootCause properties:
+            // if cancellation is requested for this source by timeout, switch cancellationOrigin to not-null
+            // value in getter to prevent concurrent overwrite by the linked token cancellation callback.
+            object? tokenCopy;
+            if (CanInlineToken)
+            {
+                tokenCopy = IsCancellationRequested
+                    ? Interlocked.CompareExchange(ref cancellationOrigin.Item1, this, comparand: null)
+                    : Volatile.Read(in cancellationOrigin.Item1);
+
+                tokenCopy ??= this;
+            }
+            else if ((tokenCopy = Volatile.Read(in cancellationOrigin.Item1)) is null)
+            {
+                object boxedToken = Token;
+                tokenCopy = Interlocked.CompareExchange(ref cancellationOrigin.Item1, boxedToken, comparand: null) ?? boxedToken;
+            }
+
+            return tokenCopy;
+        }
+    }
+
     /// <summary>
     /// Gets the token caused cancellation.
     /// </summary>
@@ -82,11 +119,13 @@ private bool TrySetCancellationOrigin(CancellationToken token)
     /// </remarks>
     public CancellationToken CancellationOrigin
     {
-        get => new InlinedToken(Volatile.Read(in cancellationOrigin.Item1)) is { Item1: not null } tokenCopy
-            ? CanInlineToken
-                ? Unsafe.BitCast<InlinedToken, CancellationToken>(tokenCopy)
-                : Unsafe.Unbox<CancellationToken>(tokenCopy.Item1)
-            : Token;
+        get
+        {
+            var rawToken = RawToken;
+            return CanInlineToken
+                ? Unsafe.BitCast<InlinedToken, CancellationToken>(new(rawToken))
+                : Unsafe.Unbox<CancellationToken>(rawToken);
+        }
 
         private protected set
         {
@@ -99,11 +138,17 @@ private protected set
     /// Gets a value indicating that this token source is cancelled by the timeout associated with this source,
     /// or by calling <see cref="CancellationTokenSource.Cancel()"/> manually.
     /// </summary>
-    [DebuggerBrowsable(DebuggerBrowsableState.Never)]
-    internal bool IsRootCause => CanInlineToken
-        ? Volatile.Read(in cancellationOrigin.Item1) is not { } tokenCopy || ReferenceEquals(this, tokenCopy)
-        : Token == CancellationOrigin;
-    
+    internal bool IsRootCause
+    {
+        get
+        {
+            var rawToken = RawToken;
+            return CanInlineToken
+                ? ReferenceEquals(rawToken, this)
+                : Unsafe.Unbox<CancellationToken>(rawToken) == Token;
+        }
+    }
+
     // This property checks whether the reinterpret cast CancellationToken => CancellationTokenSource
     // is safe. If not, just box the token.
     internal static bool CanInlineToken => Intrinsics.AreCompatible<CancellationToken, InlinedToken>()