From 496064f72db93ea8596c87cfd7c4d2ed44f3519f Mon Sep 17 00:00:00 2001
From: Ethan Chiu <ethanchewy@github.com>
Date: Thu, 18 Jun 2020 17:34:17 -0400
Subject: [PATCH] Optimize runtime of code

---
 src/Runner.Worker/ExecutionContext.cs         | 12 ++++++-----
 .../Handlers/CompositeActionHandler.cs        | 20 +++++++++++++++++--
 2 files changed, 25 insertions(+), 7 deletions(-)

diff --git a/src/Runner.Worker/ExecutionContext.cs b/src/Runner.Worker/ExecutionContext.cs
index 11116e96e..526f9af56 100644
--- a/src/Runner.Worker/ExecutionContext.cs
+++ b/src/Runner.Worker/ExecutionContext.cs
@@ -107,6 +107,7 @@ namespace GitHub.Runner.Worker
         void RegisterPostJobStep(IStep step);
         IStep RegisterCompositeStep(IStep step, DictionaryContextData inputsData, Dictionary<string, string> envData);
         void EnqueueAllCompositeSteps(Queue<IStep> steps);
+        public void SetEnvironmentVariables(Dictionary<string, string> dict);
     }
 
     public sealed class ExecutionContext : RunnerService, IExecutionContext
@@ -298,14 +299,16 @@ namespace GitHub.Runner.Worker
             // Add the composite action environment variables to each step.
             // If the key already exists, we override it since the composite action env variables will have higher precedence
             // Note that for each composite action step, it's environment variables will be set in the StepRunner automatically
-            foreach (var e in envData)
-            {
-                step.ExecutionContext.EnvironmentVariables[e.Key] = e.Value;
-            }
+            step.ExecutionContext.SetEnvironmentVariables(envData);
 
             return step;
         }
 
+        public void SetEnvironmentVariables(Dictionary<string, string> dict)
+        {
+            this.EnvironmentVariables = dict;
+        }
+
         // Add Composite Steps first and then requeue the rest of the job steps. 
         public void EnqueueAllCompositeSteps(Queue<IStep> steps)
         {
@@ -333,7 +336,6 @@ namespace GitHub.Runner.Worker
                 }
             }
         }
-
         public IExecutionContext CreateChild(Guid recordId, string displayName, string refName, string scopeName, string contextName, Dictionary<string, string> intraActionState = null, int? recordOrder = null)
         {
             Trace.Entering();
diff --git a/src/Runner.Worker/Handlers/CompositeActionHandler.cs b/src/Runner.Worker/Handlers/CompositeActionHandler.cs
index 0902fafae..db3ff2cdd 100644
--- a/src/Runner.Worker/Handlers/CompositeActionHandler.cs
+++ b/src/Runner.Worker/Handlers/CompositeActionHandler.cs
@@ -44,12 +44,28 @@ namespace GitHub.Runner.Worker.Handlers
                 inputsData[i.Key] = new StringContextData(i.Value);
             }
 
-            
+
             // Get Environment Data for Composite Action
             var extraExpressionValues = new Dictionary<string, PipelineContextData>(StringComparer.OrdinalIgnoreCase);
             extraExpressionValues["inputs"] = inputsData;
             var manifestManager = HostContext.GetService<IActionManifestManager>();
-            var envData = manifestManager.EvaluateCompositeActionEnvironment(ExecutionContext, Data.Environment, extraExpressionValues);
+
+            // Add the composite action environment variables to each step.
+            // If the key already exists, we override it since the composite action env variables will have higher precedence
+            // Note that for each composite action step, it's environment variables will be set in the StepRunner automatically
+            var compositeEnvData = manifestManager.EvaluateCompositeActionEnvironment(ExecutionContext, Data.Environment, extraExpressionValues);
+            var envData = new Dictionary<string, string>();
+
+            // Copy over parent environment
+            foreach (var env in ExecutionContext.EnvironmentVariables)
+            {
+                envData[env.Key] = env.Value;
+            }
+            // Overwrite with current env
+            foreach (var env in compositeEnvData)
+            {
+                envData[env.Key] = env.Value;
+            }
 
             // Add each composite action step to the front of the queue
             var compositeActionSteps = new Queue<IStep>();