Shrink Task.Delay when used without Cancellation (2)

src/System.Private.CoreLib/shared/System/Threading/Tasks/Task.cs

@@ -5370,12 +5370,18 @@ public static Task Delay(int millisecondsDelay, CancellationToken cancellationTo
             }
 
             // Construct a promise-style Task to encapsulate our return value
-            var promise = new DelayPromise(cancellationToken);
+            DelayPromise promise;
 
-            // Register our cancellation token, if necessary.
             if (cancellationToken.CanBeCanceled)
             {
-                promise.Registration = cancellationToken.UnsafeRegister(state => ((DelayPromise)state).Complete(), promise);
+                var promiseWithCancellation = new DelayPromiseWithCancellation(cancellationToken);
+                // Register our cancellation token, if necessary.
+                promiseWithCancellation.Registration = cancellationToken.UnsafeRegister(state => ((DelayPromise)state).Complete(), promiseWithCancellation);
+                promise = promiseWithCancellation;
+            }
+            else
+            {
+                promise = new DelayPromise();
             }
 
             // ... and create our timer and make sure that it stays rooted.
@@ -5389,49 +5395,74 @@ public static Task Delay(int millisecondsDelay, CancellationToken cancellationTo
         }
 
         /// <summary>Task that also stores the completion closure and logic for Task.Delay implementation.</summary>
-        private sealed class DelayPromise : Task<VoidTaskResult>
+        private class DelayPromise : Task<VoidTaskResult>
         {
-            internal DelayPromise(CancellationToken token)
-                : base()
+            internal DelayPromise() : base()
             {
-                this.Token = token;
                 if (AsyncCausalityTracer.LoggingOn)
                     AsyncCausalityTracer.TraceOperationCreation(this, "Task.Delay");
 
                 if (s_asyncDebuggingEnabled)
                     AddToActiveTasks(this);
             }
 
-            internal readonly CancellationToken Token;
-            internal CancellationTokenRegistration Registration;
             internal TimerQueueTimer Timer;
 
-            internal void Complete()
+            internal virtual bool Complete()
             {
                 // Transition the task to completed.
                 bool setSucceeded;
 
+                if (AsyncCausalityTracer.LoggingOn)
+                    AsyncCausalityTracer.TraceOperationCompletion(this, AsyncCausalityStatus.Completed);
+
+                if (s_asyncDebuggingEnabled)
+                    RemoveFromActiveTasks(this);
+
+                setSucceeded = TrySetResult(default);
+
+                // If we set the value, also clean up.
+                if (setSucceeded)
+                {
+                    Timer?.Close();
+                }
+
+                return setSucceeded;
+            }
+        }
+
+        private sealed class DelayPromiseWithCancellation : DelayPromise
+        {
+            internal DelayPromiseWithCancellation(CancellationToken token) : base()
+                => Token = token;
+
+            internal readonly CancellationToken Token;
+            internal CancellationTokenRegistration Registration;
+
+            internal override bool Complete()
+            {
+                // Transition the task to completed.
+                bool setSucceeded;
                 if (Token.IsCancellationRequested)
                 {
                     setSucceeded = TrySetCanceled(Token);
+                    if (setSucceeded)
+                    {
+                        Timer?.Close();
+                    }
                 }
                 else
                 {
-                    if (AsyncCausalityTracer.LoggingOn)
-                        AsyncCausalityTracer.TraceOperationCompletion(this, AsyncCausalityStatus.Completed);
-
-                    if (s_asyncDebuggingEnabled)
-                        RemoveFromActiveTasks(this);
-
-                    setSucceeded = TrySetResult(default);
+                    setSucceeded = base.Complete();
                 }
 
                 // If we set the value, also clean up.
                 if (setSucceeded)
                 {
-                    Timer?.Close();
                     Registration.Dispose();
                 }
+
+                return setSucceeded;
             }
         }
         #endregion