Composite Code Cleanup (#1232)

* composite polish

* Cleanup Condition Handling

* Refactor ConditionTraceWriter

* pr feedback

* cleanup

src/Runner.Worker/ActionManager.cs

@@ -247,7 +247,10 @@ namespace GitHub.Runner.Worker
                                 {
                                     _cachedEmbeddedPreSteps[parentStepId] = new List<Pipelines.ActionStep>();
                                 }
-                                _cachedEmbeddedPreSteps[parentStepId].Add(action);
+                                // Clone action so we can modify the condition without affecting the original
+                                var clonedAction = action.Clone() as Pipelines.ActionStep;
+                                clonedAction.Condition = definition.Data.Execution.InitCondition;
+                                _cachedEmbeddedPreSteps[parentStepId].Add(clonedAction);
                             }
                         }
 
@@ -258,7 +261,10 @@ namespace GitHub.Runner.Worker
                                 // If we haven't done so already, add the parent to the post steps
                                 _cachedEmbeddedPostSteps[parentStepId] = new Stack<Pipelines.ActionStep>();
                             }
-                            _cachedEmbeddedPostSteps[parentStepId].Push(action);
+                            // Clone action so we can modify the condition without affecting the original
+                            var clonedAction = action.Clone() as Pipelines.ActionStep;
+                            clonedAction.Condition = definition.Data.Execution.CleanupCondition;
+                            _cachedEmbeddedPostSteps[parentStepId].Push(clonedAction);
                         }
                     }
                 }

src/Runner.Worker/ActionRunner.cs

@@ -82,12 +82,6 @@ namespace GitHub.Runner.Worker
             ActionExecutionData handlerData = definition.Data?.Execution;
             ArgUtil.NotNull(handlerData, nameof(handlerData));
 
-            if (handlerData.HasPre &&
-                Action.Reference is Pipelines.RepositoryPathReference repoAction &&
-                string.Equals(repoAction.RepositoryType, Pipelines.PipelineConstants.SelfAlias, StringComparison.OrdinalIgnoreCase))
-            {
-                ExecutionContext.Warning($"`pre` execution is not supported for local action from '{repoAction.Path}'");
-            }
             List<JobExtensionRunner> localActionContainerSetupSteps = null;
             // Handle Composite Local Actions
             // Need to download and expand the tree of referenced actions
@@ -110,6 +104,13 @@ namespace GitHub.Runner.Worker
                 localActionContainerSetupSteps = prepareResult.ContainerSetupSteps;
             }
 
+            if (handlerData.HasPre &&
+                Action.Reference is Pipelines.RepositoryPathReference repoAction &&
+                string.Equals(repoAction.RepositoryType, Pipelines.PipelineConstants.SelfAlias, StringComparison.OrdinalIgnoreCase))
+            {
+                ExecutionContext.Warning($"`pre` execution is not supported for local action from '{repoAction.Path}'");
+            }
+
             // The action has post cleanup defined.
             // we need to create timeline record for them and add them to the step list that StepRunner is using
             if (handlerData.HasPost && (Stage == ActionRunStage.Pre || Stage == ActionRunStage.Main))

src/Runner.Worker/ConditionTraceWriter.cs

@@ -0,0 +1,47 @@
+using System;
+using System.Text;
+using GitHub.DistributedTask.Pipelines;
+using GitHub.Runner.Common;
+using GitHub.Runner.Sdk;
+using ObjectTemplating = GitHub.DistributedTask.ObjectTemplating;
+
+namespace GitHub.Runner.Worker
+{
+    public sealed class ConditionTraceWriter : ObjectTemplating::ITraceWriter
+    {
+        private readonly IExecutionContext _executionContext;
+        private readonly Tracing _trace;
+        private readonly StringBuilder _traceBuilder = new StringBuilder();
+
+        public string Trace => _traceBuilder.ToString();
+
+        public ConditionTraceWriter(Tracing trace, IExecutionContext executionContext)
+        {
+            ArgUtil.NotNull(trace, nameof(trace));
+            _trace = trace;
+            _executionContext = executionContext;
+        }
+
+        public void Error(string format, params Object[] args)
+        {
+            var message = StringUtil.Format(format, args);
+            _trace.Error(message);
+            _executionContext?.Debug(message);
+        }
+
+        public void Info(string format, params Object[] args)
+        {
+            var message = StringUtil.Format(format, args);
+            _trace.Info(message);
+            _executionContext?.Debug(message);
+            _traceBuilder.AppendLine(message);
+        }
+
+        public void Verbose(string format, params Object[] args)
+        {
+            var message = StringUtil.Format(format, args);
+            _trace.Verbose(message);
+            _executionContext?.Debug(message);
+        }
+    }
+}

src/Runner.Worker/Handlers/CompositeActionHandler.cs

@@ -11,6 +11,7 @@ using GitHub.DistributedTask.Pipelines.ContextData;
 using GitHub.DistributedTask.Pipelines.ObjectTemplating;
 using GitHub.DistributedTask.WebApi;
 using GitHub.Runner.Common;
+using GitHub.Runner.Common.Util;
 using GitHub.Runner.Sdk;
 using GitHub.Runner.Worker;
 using GitHub.Runner.Worker.Expressions;
@@ -45,13 +46,17 @@ namespace GitHub.Runner.Worker.Handlers
             else if (stage == ActionRunStage.Post)
             {
                 ArgUtil.NotNull(Data.PostSteps, nameof(Data.PostSteps));
-                steps = Data.PostSteps.ToList();
+                steps = new List<Pipelines.ActionStep>();
                 // Only register post steps for steps that actually ran
-                foreach (var step in steps)
+                foreach (var step in Data.PostSteps.ToList())
                 {
-                    if (!ExecutionContext.Root.EmbeddedStepsWithPostRegistered.Contains(step.Id))
+                    if (ExecutionContext.Root.EmbeddedStepsWithPostRegistered.Contains(step.Id))
                     {
-                        steps.Remove(step);
+                        steps.Add(step);
+                    }
+                    else
+                    {
+                        Trace.Info($"Skipping executing post step id: {step.Id}, name: ${step.DisplayName}");
                     }
                 }
             }  
@@ -217,6 +222,10 @@ namespace GitHub.Runner.Worker.Handlers
 
                 // Add Expression Functions
                 step.ExecutionContext.ExpressionFunctions.Add(new FunctionInfo<HashFilesFunction>(PipelineTemplateConstants.HashFiles, 1, byte.MaxValue));
+                step.ExecutionContext.ExpressionFunctions.Add(new FunctionInfo<AlwaysFunction>(PipelineTemplateConstants.Always, 0, 0));
+                step.ExecutionContext.ExpressionFunctions.Add(new FunctionInfo<CancelledFunction>(PipelineTemplateConstants.Cancelled, 0, 0));
+                step.ExecutionContext.ExpressionFunctions.Add(new FunctionInfo<FailureFunction>(PipelineTemplateConstants.Failure, 0, 0));
+                step.ExecutionContext.ExpressionFunctions.Add(new FunctionInfo<SuccessFunction>(PipelineTemplateConstants.Success, 0, 0));
 
                 // Initialize env context
                 Trace.Info("Initialize Env context for embedded step");
@@ -268,13 +277,120 @@ namespace GitHub.Runner.Worker.Handlers
                     step.ExecutionContext.Complete(TaskResult.Failed);
                 }
 
-                await RunStepAsync(step);
+                // Register Callback
+                CancellationTokenRegistration? jobCancelRegister = null;
+                try
+                {
+                    // Register job cancellation call back only if job cancellation token not been fire before each step run
+                    if (!ExecutionContext.Root.CancellationToken.IsCancellationRequested)
+                    {
+                        // Test the condition again. The job was canceled after the condition was originally evaluated.
+                        jobCancelRegister = ExecutionContext.Root.CancellationToken.Register(() =>
+                        {
+                            // Mark job as cancelled
+                            ExecutionContext.Root.Result = TaskResult.Canceled;
+                            ExecutionContext.Root.JobContext.Status = ExecutionContext.Root.Result?.ToActionResult();
+
+                            step.ExecutionContext.Debug($"Re-evaluate condition on job cancellation for step: '{step.DisplayName}'.");
+                            var conditionReTestTraceWriter = new ConditionTraceWriter(Trace, null); // host tracing only
+                            var conditionReTestResult = false;
+                            if (HostContext.RunnerShutdownToken.IsCancellationRequested)
+                            {
+                                step.ExecutionContext.Debug($"Skip Re-evaluate condition on runner shutdown.");
+                            }
+                            else
+                            {
+                                try
+                                {
+                                    var templateEvaluator = step.ExecutionContext.ToPipelineTemplateEvaluator(conditionReTestTraceWriter);
+                                    var condition = new BasicExpressionToken(null, null, null, step.Condition);
+                                    conditionReTestResult = templateEvaluator.EvaluateStepIf(condition, step.ExecutionContext.ExpressionValues, step.ExecutionContext.ExpressionFunctions, step.ExecutionContext.ToExpressionState());
+                                }
+                                catch (Exception ex)
+                                {
+                                    // Cancel the step since we get exception while re-evaluate step condition
+                                    Trace.Info("Caught exception from expression when re-test condition on job cancellation.");
+                                    step.ExecutionContext.Error(ex);
+                                }
+                            }
+
+                            if (!conditionReTestResult)
+                            {
+                                // Cancel the step
+                                Trace.Info("Cancel current running step.");
+                                step.ExecutionContext.CancelToken();
+                            }
+                        });
+                    }
+                    else
+                    {
+                        if (ExecutionContext.Root.Result != TaskResult.Canceled)
+                        {
+                            // Mark job as cancelled
+                            ExecutionContext.Root.Result = TaskResult.Canceled;
+                            ExecutionContext.Root.JobContext.Status = ExecutionContext.Root.Result?.ToActionResult();
+                        }
+                    }
+                    // Evaluate condition
+                    step.ExecutionContext.Debug($"Evaluating condition for step: '{step.DisplayName}'");
+                    var conditionTraceWriter = new ConditionTraceWriter(Trace, step.ExecutionContext);
+                    var conditionResult = false;
+                    var conditionEvaluateError = default(Exception);
+                    if (HostContext.RunnerShutdownToken.IsCancellationRequested)
+                    {
+                        step.ExecutionContext.Debug($"Skip evaluate condition on runner shutdown.");
+                    }
+                    else
+                    {
+                        try
+                        {
+                            var templateEvaluator = step.ExecutionContext.ToPipelineTemplateEvaluator(conditionTraceWriter);
+                            var condition = new BasicExpressionToken(null, null, null, step.Condition);
+                            conditionResult = templateEvaluator.EvaluateStepIf(condition, step.ExecutionContext.ExpressionValues, step.ExecutionContext.ExpressionFunctions, step.ExecutionContext.ToExpressionState());
+                        }
+                        catch (Exception ex)
+                        {
+                            Trace.Info("Caught exception from expression.");
+                            Trace.Error(ex);
+                            conditionEvaluateError = ex;
+                        }
+                    }
+                    if (!conditionResult && conditionEvaluateError == null)
+                    {
+                        // Condition is false
+                        Trace.Info("Skipping step due to condition evaluation.");
+                        step.ExecutionContext.Result = TaskResult.Skipped;
+                        continue;
+                    }
+                    else if (conditionEvaluateError != null)
+                    {
+                        // Condition error
+                        step.ExecutionContext.Error(conditionEvaluateError);
+                        step.ExecutionContext.Result = TaskResult.Failed;
+                        ExecutionContext.Result = TaskResult.Failed;
+                        break;
+                    }
+                    else
+                    {
+                        await RunStepAsync(step);
+                    }
+                }
+                finally
+                {
+                    if (jobCancelRegister != null)
+                    {
+                        jobCancelRegister?.Dispose();
+                        jobCancelRegister = null;
+                    }
+                }
 
                 // Check failed or canceled
                 if (step.ExecutionContext.Result == TaskResult.Failed || step.ExecutionContext.Result == TaskResult.Canceled)
                 {
+                    Trace.Info($"Update job result with current composite step result '{step.ExecutionContext.Result}'.");
                     ExecutionContext.Result = step.ExecutionContext.Result;
-                    break;
+                    ExecutionContext.Root.Result = TaskResultUtil.MergeTaskResults(ExecutionContext.Root.Result, step.ExecutionContext.Result.Value);
+                    ExecutionContext.Root.JobContext.Status = ExecutionContext.Root.Result?.ToActionResult();
                 }
             }
         }

src/Runner.Worker/StepsRunner.cs

@@ -354,43 +354,5 @@ namespace GitHub.Runner.Worker
 
             executionContext.Complete(result, resultCode: resultCode);
         }
-
-        private sealed class ConditionTraceWriter : ObjectTemplating::ITraceWriter
-        {
-            private readonly IExecutionContext _executionContext;
-            private readonly Tracing _trace;
-            private readonly StringBuilder _traceBuilder = new StringBuilder();
-
-            public string Trace => _traceBuilder.ToString();
-
-            public ConditionTraceWriter(Tracing trace, IExecutionContext executionContext)
-            {
-                ArgUtil.NotNull(trace, nameof(trace));
-                _trace = trace;
-                _executionContext = executionContext;
-            }
-
-            public void Error(string format, params Object[] args)
-            {
-                var message = StringUtil.Format(format, args);
-                _trace.Error(message);
-                _executionContext?.Debug(message);
-            }
-
-            public void Info(string format, params Object[] args)
-            {
-                var message = StringUtil.Format(format, args);
-                _trace.Info(message);
-                _executionContext?.Debug(message);
-                _traceBuilder.AppendLine(message);
-            }
-
-            public void Verbose(string format, params Object[] args)
-            {
-                var message = StringUtil.Format(format, args);
-                _trace.Verbose(message);
-                _executionContext?.Debug(message);
-            }
-        }
     }
 }