Returns the owned token if the source is not yet completed

Author: sakno

src/DotNext.Threading/Threading/CancellationTokenMultiplexer.Scope.cs

@@ -6,6 +6,7 @@
 namespace DotNext.Threading;
 
 using static Timeout;
+using MultiplexerOrToken = ValueTuple<object>;
 
 partial class CancellationTokenMultiplexer
 {
@@ -23,7 +24,7 @@ private interface IMultiplexedTokenScope : IMultiplexedCancellationTokenSource,
         // CancellationToken is just a wrapper over CancellationTokenSource.
         // For optimization purposes, if only one token is passed to the scope, we can inline the underlying CTS
         // to this structure.
-        private readonly ValueTuple<object> multiplexerOrToken;
+        private readonly MultiplexerOrToken multiplexerOrToken;
         private readonly PooledCancellationTokenSource? source;
 
         internal Scope(CancellationTokenMultiplexer multiplexer, ReadOnlySpan<CancellationToken> tokens, bool timeoutSupport = false)
@@ -46,14 +47,14 @@ internal Scope(CancellationTokenMultiplexer multiplexer, TimeSpan timeout, ReadO
         internal Scope(CancellationToken token)
             => multiplexerOrToken = InlineToken(token);
 
-        private static ValueTuple<object> InlineToken(CancellationToken token)
+        private static MultiplexerOrToken InlineToken(CancellationToken token)
             => LinkedCancellationTokenSource.CanInlineToken
-                ? Unsafe.BitCast<CancellationToken, ValueTuple<object>>(token)
+                ? Unsafe.BitCast<CancellationToken, MultiplexerOrToken>(token)
                 : new(token);
 
-        private static CancellationToken GetToken(ValueTuple<object> value)
+        private static CancellationToken GetToken(MultiplexerOrToken value)
             => LinkedCancellationTokenSource.CanInlineToken
-                ? Unsafe.BitCast<ValueTuple<object>, CancellationToken>(value)
+                ? Unsafe.BitCast<MultiplexerOrToken, CancellationToken>(value)
                 : (CancellationToken)value.Item1;
 
         /// <inheritdoc cref="IMultiplexedCancellationTokenSource.Token"/>

src/DotNext.Threading/Threading/LinkedCancellationTokenSource.cs

@@ -6,6 +6,7 @@
 namespace DotNext.Threading;
 
 using Runtime;
+using InlinedToken = ValueTuple<object?>;
 
 /// <summary>
 /// Gets cancellation token source that allows to obtain the token that causes
@@ -18,7 +19,7 @@ namespace DotNext.Threading;
 public abstract class LinkedCancellationTokenSource : CancellationTokenSource, IMultiplexedCancellationTokenSource
 {
     // represents inlined CancellationToken
-    private ValueTuple<object?> cancellationOrigin;
+    private InlinedToken cancellationOrigin;
 
     private protected LinkedCancellationTokenSource()
     {
@@ -36,8 +37,8 @@ static void OnCanceled(object? source, CancellationToken token)
         }
     }
 
-    private static ValueTuple<object?> InlineToken(CancellationToken token) => CanInlineToken
-        ? Unsafe.BitCast<CancellationToken, ValueTuple<object?>>(token)
+    private static InlinedToken InlineToken(CancellationToken token) => CanInlineToken
+        ? Unsafe.BitCast<CancellationToken, InlinedToken>(token)
         : new(token);
 
     internal void RegisterTimeoutHandler()
@@ -81,11 +82,11 @@ private bool TrySetCancellationOrigin(CancellationToken token)
     /// </remarks>
     public CancellationToken CancellationOrigin
     {
-        get => CanInlineToken
-            ? Unsafe.BitCast<ValueTuple<object?>, CancellationToken>(cancellationOrigin)
-            : cancellationOrigin.Item1 is null
-                ? CancellationToken.None
-                : Unsafe.Unbox<CancellationToken>(cancellationOrigin.Item1);
+        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;
 
         private protected set => cancellationOrigin = InlineToken(value);
     }
@@ -99,7 +100,7 @@ public CancellationToken CancellationOrigin
 
     // This property checks whether the reinterpret cast CancellationToken => CancellationTokenSource
     // is safe. If not, just box the token.
-    internal static bool CanInlineToken => Intrinsics.AreCompatible<CancellationToken, ValueTuple<object>>()
+    internal static bool CanInlineToken => Intrinsics.AreCompatible<CancellationToken, InlinedToken>()
                                            && RuntimeHelpers.IsReferenceOrContainsReferences<CancellationToken>();
 
     internal static LinkedCancellationTokenSource? Combine(ref CancellationToken first, CancellationToken second)