Moved logic to drain queues after each task under knob "AGENT_DRAIN_QUEUES_AFTER_TASK"

src/Agent.Sdk/Knob/AgentKnobs.cs

@@ -410,5 +410,12 @@ public class AgentKnobs
             new RuntimeKnobSource("AGENT_USE_NODE"),
             new EnvironmentKnobSource("AGENT_USE_NODE"),
             new BuiltInDefaultKnobSource(string.Empty));
+
+        public static readonly Knob DrainQueuesAfterTask = new Knob(
+            nameof(DrainQueuesAfterTask),
+            "Forces the agent to drain queues after each task",
+            new RuntimeKnobSource("AGENT_DRAIN_QUEUES_AFTER_TASK"),
+            new EnvironmentKnobSource("AGENT_DRAIN_QUEUES_AFTER_TASK"),
+            new BuiltInDefaultKnobSource("false"));
     }
 }

src/Agent.Worker/StepsRunner.cs

@@ -11,6 +11,7 @@
 using Microsoft.TeamFoundation.DistributedTask.Expressions;
 using Pipelines = Microsoft.TeamFoundation.DistributedTask.Pipelines;
 using Microsoft.VisualStudio.Services.CircuitBreaker;
+using Agent.Sdk.Knob;
 
 namespace Microsoft.VisualStudio.Services.Agent.Worker
 {
@@ -309,16 +310,19 @@ private async Task RunStepAsync(IStep step, CancellationToken jobCancellationTok
             step.ExecutionContext.Section(StringUtil.Loc("StepFinishing", step.DisplayName));
             step.ExecutionContext.Complete();
 
-            try
-            {
-                // We need to drain the queues after a task just in case if
-                // there are a lot of items since it can cause some UI hangs.
-                await JobServerQueue.DrainQueues();
-            }
-            catch (Exception ex)
+            if (AgentKnobs.DrainQueuesAfterTask.GetValue(step.ExecutionContext).AsBoolean() == true)
             {
-                Trace.Error($"Error has occurred while draining queues, it can cause some UI glitches but it doesn't affect a pipeline execution itself: {ex}");
-                step.ExecutionContext.Error(ex);
+                try
+                {
+                    // We need to drain the queues after a task just in case if
+                    // there are a lot of items since it can cause some UI hangs.
+                    await JobServerQueue.DrainQueues();
+                }
+                catch (Exception ex)
+                {
+                    Trace.Error($"Error has occurred while draining queues, it can cause some UI glitches but it doesn't affect a pipeline execution itself: {ex}");
+                    step.ExecutionContext.Error(ex);
+                }
             }
         }
 

src/Test/L0/Worker/StepsRunnerL0.cs

@@ -13,6 +13,7 @@
 using Xunit;
 using Microsoft.TeamFoundation.DistributedTask.Expressions;
 using Pipelines = Microsoft.TeamFoundation.DistributedTask.Pipelines;
+using Agent.Sdk;
 
 namespace Microsoft.VisualStudio.Services.Agent.Tests.Worker
 {
@@ -426,6 +427,7 @@ public async Task SetStepTarget()
                             stepContext.Object.Result = r;
                         }
                     });
+                stepContext.Setup(x => x.GetScopedEnvironment()).Returns(new SystemEnvironment());
                 step.Setup(x => x.ExecutionContext).Returns(stepContext.Object);
 
                 // Act.
@@ -460,6 +462,7 @@ private Mock<IStep> CreateStep(TaskResult result, IExpressionNode condition, Boo
                     }
                 });
             stepContext.Object.Result = result;
+            stepContext.Setup(x => x.GetScopedEnvironment()).Returns(new SystemEnvironment());
             step.Setup(x => x.ExecutionContext).Returns(stepContext.Object);
 
             return step;