Fix cancellation token race during parser comparison (#4280)

Author: ericsciple

src/Runner.Worker/ExecutionContext.cs

@@ -77,8 +77,7 @@ public interface IExecutionContext : IRunnerService
 
         List<string> StepEnvironmentOverrides { get; }
 
-        ExecutionContext Root { get; }
-        ExecutionContext Parent { get; }
+        IExecutionContext Root { get; }
 
         // Initialize
         void InitializeJob(Pipelines.AgentJobRequestMessage message, CancellationToken token);
@@ -251,7 +250,9 @@ public string ResultCode
             }
         }
 
-        public ExecutionContext Root
+        IExecutionContext IExecutionContext.Root => Root;
+
+        private ExecutionContext Root
         {
             get
             {
@@ -266,13 +267,7 @@ public ExecutionContext Root
             }
         }
 
-        public ExecutionContext Parent
-        {
-            get
-            {
-                return _parentExecutionContext;
-            }
-        }
+
 
         public JobContext JobContext
         {

src/Runner.Worker/PipelineTemplateEvaluatorWrapper.cs

@@ -1,5 +1,6 @@
 using System;
 using System.Collections.Generic;
+using System.Threading;
 using GitHub.Actions.WorkflowParser;
 using GitHub.DistributedTask.Expressions2;
 using GitHub.DistributedTask.ObjectTemplating.Tokens;
@@ -226,8 +227,12 @@ internal TLegacy EvaluateAndCompare<TLegacy, TNew>(
             Func<TNew> newEvaluator,
             Func<TLegacy, TNew, bool> resultComparer)
         {
-            // Capture cancellation state before evaluation
-            var cancellationRequestedBefore = _context.CancellationToken.IsCancellationRequested;
+            // Use the root (job-level) cancellation token to detect cancellation race conditions.
+            // The step-level token only fires on step timeout, not on job cancellation.
+            // Job cancellation mutates JobContext.Status which expression functions read,
+            // so we need the root token to properly detect cancellation between evaluator runs.
+            var rootCancellationToken = _context.Root?.CancellationToken ?? CancellationToken.None;
+            var cancellationRequestedBefore = rootCancellationToken.IsCancellationRequested;
 
             // Legacy evaluator
             var legacyException = default(Exception);
@@ -261,7 +266,7 @@ internal TLegacy EvaluateAndCompare<TLegacy, TNew>(
                 }
 
                 // Capture cancellation state after evaluation
-                var cancellationRequestedAfter = _context.CancellationToken.IsCancellationRequested;
+                var cancellationRequestedAfter = rootCancellationToken.IsCancellationRequested;
 
                 // Compare results or exceptions
                 bool hasMismatch = false;

src/Test/L0/Worker/PipelineTemplateEvaluatorWrapperL0.cs

@@ -19,6 +19,7 @@ namespace GitHub.Runner.Common.Tests.Worker
     public sealed class PipelineTemplateEvaluatorWrapperL0
     {
         private CancellationTokenSource _ecTokenSource;
+        private CancellationTokenSource _rootTokenSource;
         private Mock<IExecutionContext> _ec;
         private TestHostContext _hc;
 
@@ -65,7 +66,7 @@ public void EvaluateAndCompare_SkipsMismatchRecording_WhenCancellationOccursDuri
 
                 var wrapper = new PipelineTemplateEvaluatorWrapper(_hc, _ec.Object, allowServiceContainerCommand: false);
 
-                // Call EvaluateAndCompare directly: the new evaluator cancels the token
+                // Call EvaluateAndCompare directly: the new evaluator cancels the root token
                 // and returns a different value, forcing hasMismatch = true.
                 // Because cancellation flipped during the evaluation window, the
                 // mismatch should be skipped.
@@ -74,7 +75,7 @@ public void EvaluateAndCompare_SkipsMismatchRecording_WhenCancellationOccursDuri
                     () => "legacy-value",
                     () =>
                     {
-                        _ecTokenSource.Cancel();
+                        _rootTokenSource.Cancel();
                         return "different-value";
                     },
                     (legacy, @new) => string.Equals(legacy, @new, StringComparison.Ordinal));
@@ -88,6 +89,43 @@ public void EvaluateAndCompare_SkipsMismatchRecording_WhenCancellationOccursDuri
             }
         }
 
+        [Fact]
+        [Trait("Level", "L0")]
+        [Trait("Category", "Worker")]
+        public void EvaluateAndCompare_SkipsMismatchRecording_WhenRootCancellationOccursBetweenEvaluators()
+        {
+            // Simulates job-level cancellation firing between legacy and new evaluator runs.
+            // Root is mocked with a separate CancellationTokenSource to exercise the
+            // _context.Root?.CancellationToken path (the job-level token).
+            try
+            {
+                Setup();
+                _ec.Object.Global.Variables.Set(Constants.Runner.Features.CompareWorkflowParser, "true");
+
+                var wrapper = new PipelineTemplateEvaluatorWrapper(_hc, _ec.Object, allowServiceContainerCommand: false);
+
+                // Legacy evaluator cancels the root token (simulating job cancel) and returns a value.
+                // The new evaluator returns a different value. The mismatch should be skipped.
+                var result = wrapper.EvaluateAndCompare<string, string>(
+                    "TestRootCancellationSkip",
+                    () =>
+                    {
+                        var legacyValue = "legacy-value";
+                        _rootTokenSource.Cancel();
+                        return legacyValue;
+                    },
+                    () => "different-value",
+                    (legacy, @new) => string.Equals(legacy, @new, StringComparison.Ordinal));
+
+                Assert.Equal("legacy-value", result);
+                Assert.False(_ec.Object.Global.HasTemplateEvaluatorMismatch);
+            }
+            finally
+            {
+                Teardown();
+            }
+        }
+
         [Fact]
         [Trait("Level", "L0")]
         [Trait("Category", "Worker")]
@@ -862,6 +900,8 @@ private void Setup([CallerMemberName] string name = "")
         {
             _ecTokenSource?.Dispose();
             _ecTokenSource = new CancellationTokenSource();
+            _rootTokenSource?.Dispose();
+            _rootTokenSource = new CancellationTokenSource();
 
             _hc = new TestHostContext(this, name);
 
@@ -877,6 +917,9 @@ private void Setup([CallerMemberName] string name = "")
                     WriteDebug = true,
                 });
             _ec.Setup(x => x.CancellationToken).Returns(_ecTokenSource.Token);
+            var rootEc = new Mock<IExecutionContext>();
+            rootEc.Setup(x => x.CancellationToken).Returns(_rootTokenSource.Token);
+            _ec.Setup(x => x.Root).Returns(rootEc.Object);
             _ec.Setup(x => x.ExpressionValues).Returns(expressionValues);
             _ec.Setup(x => x.ExpressionFunctions).Returns(expressionFunctions);
             _ec.Setup(x => x.Write(It.IsAny<string>(), It.IsAny<string>())).Callback((string tag, string message) => { _hc.GetTrace().Info($"{tag}{message}"); });
@@ -885,6 +928,8 @@ private void Setup([CallerMemberName] string name = "")
 
         private void Teardown()
         {
+            _ecTokenSource?.Dispose();
+            _rootTokenSource?.Dispose();
             _hc?.Dispose();
         }
     }