Skip to content

Commit b30b777

Browse files
lokesh755lokesh755
committed
Adds the core background step execution engine to the runner
1 parent fb78489 commit b30b777

8 files changed

Lines changed: 1180 additions & 11 deletions

File tree

‎src/Runner.Worker/BackgroundStepCoordinator.cs‎

Lines changed: 330 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,330 @@
1+
using System;
2+
using System.Collections.Generic;
3+
using System.Linq;
4+
using System.Threading;
5+
using System.Threading.Tasks;
6+
using GitHub.DistributedTask.Pipelines.ContextData;
7+
using GitHub.DistributedTask.WebApi;
8+
using GitHub.Runner.Common;
9+
using GitHub.Runner.Common.Util;
10+
using GitHub.Runner.Sdk;
11+
12+
namespace GitHub.Runner.Worker
13+
{
14+
[ServiceLocator(Default = typeof(BackgroundStepCoordinator))]
15+
public interface IBackgroundStepCoordinator : IRunnerService
16+
{
17+
bool HasActiveBackgroundSteps { get; }
18+
bool HasUnwaitedSteps { get; }
19+
IReadOnlyList<string> UnwaitedStepIds { get; }
20+
void InitializeCoordinator(int maxConcurrent);
21+
void StartBackgroundStep(IStep step, CancellationToken jobCancellationToken);
22+
Task<TaskResult> WaitForStepsAsync(string[] stepIds, CancellationToken cancellationToken);
23+
Task<TaskResult> WaitForAllAsync(CancellationToken cancellationToken);
24+
Task CancelStepsAsync(string[] cancelStepIds);
25+
void PropagateFailures(IExecutionContext jobContext);
26+
}
27+
28+
/// <summary>
29+
/// Coordinates background step execution, waiting, cancellation, and deferred state.
30+
/// Extracted from StepsRunner so the main step loop stays clean.
31+
/// </summary>
32+
public sealed class BackgroundStepCoordinator : RunnerService, IBackgroundStepCoordinator
33+
{
34+
private const int DefaultMaxBackgroundSteps = 10;
35+
private readonly Dictionary<string, (IStep Step, Task Task, CancellationTokenSource Cts)> _backgroundSteps = new();
36+
private readonly HashSet<string> _waitedStepIds = new();
37+
private SemaphoreSlim _backgroundSlotSemaphore;
38+
39+
/// <summary>
40+
/// Reset per-job state. Call at the start of each job.
41+
/// </summary>
42+
public void InitializeCoordinator(int maxConcurrent)
43+
{
44+
_backgroundSteps.Clear();
45+
_waitedStepIds.Clear();
46+
var max = maxConcurrent > 0 ? maxConcurrent : DefaultMaxBackgroundSteps;
47+
_backgroundSlotSemaphore = new SemaphoreSlim(max);
48+
}
49+
50+
public bool HasActiveBackgroundSteps => _backgroundSteps.Count > 0;
51+
52+
public bool HasUnwaitedSteps => _backgroundSteps.Keys.Any(id => !_waitedStepIds.Contains(id));
53+
54+
public IReadOnlyList<string> UnwaitedStepIds =>
55+
_backgroundSteps.Keys.Where(id => !_waitedStepIds.Contains(id)).ToList();
56+
57+
// -----------------------------------------------------------------
58+
// Starting background steps
59+
// -----------------------------------------------------------------
60+
61+
/// <summary>
62+
/// Prepare and launch a background step. Does not block the caller.
63+
/// </summary>
64+
public void StartBackgroundStep(IStep step, CancellationToken jobCancellationToken)
65+
{
66+
var stepId = step.ExecutionContext?.ContextName ?? step.DisplayName;
67+
68+
// Isolate GitHubContext so concurrent steps don't overwrite each other's GITHUB_OUTPUT paths
69+
if (step.ExecutionContext.ExpressionValues.TryGetValue("github", out var ghCtx) && ghCtx is GitHubContext sharedGitHub)
70+
{
71+
step.ExecutionContext.ExpressionValues["github"] = sharedGitHub.ShallowCopy();
72+
}
73+
74+
var bgCts = CancellationTokenSource.CreateLinkedTokenSource(jobCancellationToken);
75+
76+
// Evaluate timeout on the main thread (needs expression context)
77+
var timeoutMinutes = 0;
78+
try
79+
{
80+
var templateEvaluator = step.ExecutionContext.ToPipelineTemplateEvaluator();
81+
timeoutMinutes = templateEvaluator.EvaluateStepTimeout(step.Timeout, step.ExecutionContext.ExpressionValues, step.ExecutionContext.ExpressionFunctions);
82+
}
83+
catch (Exception ex)
84+
{
85+
Trace.Info($"Error determining timeout for background step '{stepId}': {ex.Message}");
86+
}
87+
88+
var task = ExecuteBackgroundStepCoreAsync(step, bgCts, stepId, timeoutMinutes);
89+
_backgroundSteps[stepId] = (step, task, bgCts);
90+
Trace.Info($"Background step '{stepId}' queued (slot will be acquired asynchronously).");
91+
}
92+
93+
private async Task ExecuteBackgroundStepCoreAsync(
94+
IStep step, CancellationTokenSource bgCts,
95+
string stepId, int timeoutMinutes)
96+
{
97+
Trace.Info($"Background step '{stepId}' waiting for slot.");
98+
await _backgroundSlotSemaphore.WaitAsync(bgCts.Token);
99+
Trace.Info($"Background step '{stepId}' acquired slot.");
100+
101+
step.ExecutionContext.Start();
102+
103+
if (timeoutMinutes > 0)
104+
{
105+
step.ExecutionContext.SetTimeout(TimeSpan.FromMinutes(timeoutMinutes));
106+
}
107+
108+
using var cancelReg = bgCts.Token.Register(() =>
109+
{
110+
Trace.Info($"Background step '{stepId}': cancellation signalled, sending CancelToken to process.");
111+
step.ExecutionContext.CancelToken();
112+
});
113+
114+
TaskResult? result = null;
115+
try
116+
{
117+
await step.RunAsync();
118+
result = step.ExecutionContext.Result ?? TaskResult.Succeeded;
119+
}
120+
catch (OperationCanceledException) when (bgCts.Token.IsCancellationRequested)
121+
{
122+
result = TaskResult.Canceled;
123+
}
124+
catch (OperationCanceledException) when (step.ExecutionContext.CancellationToken.IsCancellationRequested)
125+
{
126+
Trace.Info($"Background step '{stepId}' timed out after {timeoutMinutes} minutes.");
127+
step.ExecutionContext.Error($"The background step '{step.DisplayName}' has timed out after {timeoutMinutes} minutes.");
128+
result = TaskResult.Failed;
129+
}
130+
catch (Exception ex)
131+
{
132+
Trace.Info($"Background step '{stepId}' failed: {ex.Message}");
133+
step.ExecutionContext.Error(ex);
134+
result = TaskResult.Failed;
135+
}
136+
finally
137+
{
138+
_backgroundSlotSemaphore.Release();
139+
140+
if (step.ExecutionContext.CommandResult != null)
141+
{
142+
result = TaskResultUtil.MergeTaskResults(result, step.ExecutionContext.CommandResult.Value);
143+
}
144+
145+
step.ExecutionContext.Result = result;
146+
step.ExecutionContext.ApplyContinueOnError(step.ContinueOnError);
147+
148+
step.ExecutionContext.Complete(step.ExecutionContext.Result);
149+
Trace.Info($"Background step '{stepId}' completed with result: {step.ExecutionContext.Result}");
150+
}
151+
}
152+
153+
// -----------------------------------------------------------------
154+
// Wait
155+
// -----------------------------------------------------------------
156+
157+
/// <summary>
158+
/// Wait for specific background steps by ID. Flushes deferred state and
159+
/// returns the aggregate result.
160+
/// </summary>
161+
public async Task<TaskResult> WaitForStepsAsync(string[] stepIds, CancellationToken cancellationToken)
162+
{
163+
var ids = stepIds ?? Array.Empty<string>();
164+
165+
await WaitForStepTasksAsync(ids, cancellationToken);
166+
167+
foreach (var id in ids)
168+
{
169+
FlushDeferredState(id);
170+
_waitedStepIds.Add(id);
171+
}
172+
173+
return GetWaitResult(ids);
174+
}
175+
176+
/// <summary>
177+
/// Wait for all not-yet-waited background steps. Returns the aggregate result.
178+
/// </summary>
179+
public async Task<TaskResult> WaitForAllAsync(CancellationToken cancellationToken)
180+
{
181+
var remaining = _backgroundSteps.Keys.Where(id => !_waitedStepIds.Contains(id)).ToList();
182+
183+
await WaitForStepTasksAsync(remaining, cancellationToken);
184+
185+
foreach (var id in remaining)
186+
{
187+
FlushDeferredState(id);
188+
_waitedStepIds.Add(id);
189+
}
190+
191+
return GetWaitResult(remaining);
192+
}
193+
194+
// -----------------------------------------------------------------
195+
// Cancel
196+
// -----------------------------------------------------------------
197+
198+
/// <summary>
199+
/// Cancel specific background steps by ID. Flushes deferred state after cancellation.
200+
/// </summary>
201+
public async Task CancelStepsAsync(string[] cancelStepIds)
202+
{
203+
if (cancelStepIds == null || cancelStepIds.Length == 0) return;
204+
205+
var idsToCancel = cancelStepIds
206+
.Where(id => _backgroundSteps.ContainsKey(id) && !_backgroundSteps[id].Task.IsCompleted)
207+
.ToArray();
208+
209+
if (idsToCancel.Length > 0)
210+
{
211+
Trace.Info($"Cancelling {idsToCancel.Length} background step(s): {string.Join(", ", idsToCancel)}");
212+
await CancelWithGracePeriodAsync(idsToCancel);
213+
}
214+
215+
foreach (var id in cancelStepIds)
216+
{
217+
FlushDeferredState(id);
218+
}
219+
}
220+
221+
// -----------------------------------------------------------------
222+
// Failure propagation
223+
// -----------------------------------------------------------------
224+
225+
/// <summary>
226+
/// Propagate any unhandled background step failures to the job result.
227+
/// </summary>
228+
public void PropagateFailures(IExecutionContext jobContext)
229+
{
230+
foreach (var (_, (step, _, _)) in _backgroundSteps)
231+
{
232+
if (step.ExecutionContext.Result == TaskResult.Failed)
233+
{
234+
Trace.Info($"Propagating failure from background step '{step.ExecutionContext.ContextName}' to job result.");
235+
jobContext.Result = TaskResultUtil.MergeTaskResults(jobContext.Result, TaskResult.Failed);
236+
jobContext.JobContext.Status = jobContext.Result?.ToActionResult();
237+
break;
238+
}
239+
}
240+
}
241+
242+
// -----------------------------------------------------------------
243+
// Private helpers
244+
// -----------------------------------------------------------------
245+
246+
private async Task WaitForStepTasksAsync(IEnumerable<string> stepIds, CancellationToken cancellationToken)
247+
{
248+
var ids = stepIds.ToList();
249+
var tasks = new List<Task>();
250+
251+
foreach (var stepId in ids)
252+
{
253+
if (_backgroundSteps.TryGetValue(stepId, out var entry) && !entry.Task.IsCompleted)
254+
{
255+
tasks.Add(entry.Task);
256+
}
257+
else if (!_backgroundSteps.ContainsKey(stepId))
258+
{
259+
Trace.Info($"Wait references unknown background step: {stepId}");
260+
}
261+
}
262+
263+
if (tasks.Count > 0)
264+
{
265+
Trace.Info($"Waiting for {tasks.Count} background step(s)...");
266+
try
267+
{
268+
await Task.WhenAll(tasks).WaitAsync(cancellationToken);
269+
}
270+
catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
271+
{
272+
Trace.Info("Wait interrupted by job cancellation — cancelling background steps.");
273+
await CancelWithGracePeriodAsync(ids);
274+
}
275+
}
276+
}
277+
278+
private async Task CancelWithGracePeriodAsync(IEnumerable<string> stepIds, double graceSeconds = 7.5)
279+
{
280+
var tasks = new List<Task>();
281+
foreach (var stepId in stepIds)
282+
{
283+
if (_backgroundSteps.TryGetValue(stepId, out var entry) && !entry.Task.IsCompleted)
284+
{
285+
entry.Step.ExecutionContext.CancelToken();
286+
entry.Cts.Cancel();
287+
tasks.Add(entry.Task);
288+
}
289+
}
290+
291+
if (tasks.Count > 0)
292+
{
293+
try
294+
{
295+
await Task.WhenAll(tasks).WaitAsync(TimeSpan.FromSeconds(graceSeconds));
296+
}
297+
catch (TimeoutException)
298+
{
299+
Trace.Info($"Some background steps did not terminate within {graceSeconds}s grace period.");
300+
}
301+
}
302+
}
303+
304+
private void FlushDeferredState(string stepId)
305+
{
306+
if (_backgroundSteps.TryGetValue(stepId, out var entry))
307+
{
308+
entry.Step.ExecutionContext.FlushDeferredOutputs();
309+
entry.Step.ExecutionContext.FlushDeferredEnvironment();
310+
entry.Step.ExecutionContext.FlushDeferredOutcomeConclusion();
311+
Trace.Info($"Flushed deferred state for background step '{stepId}'.");
312+
}
313+
}
314+
315+
private TaskResult GetWaitResult(IEnumerable<string> stepIds)
316+
{
317+
if (stepIds == null) return TaskResult.Succeeded;
318+
319+
foreach (var stepId in stepIds)
320+
{
321+
if (_backgroundSteps.TryGetValue(stepId, out var entry) && entry.Step.ExecutionContext.Result == TaskResult.Failed)
322+
{
323+
Trace.Info($"Background step '{stepId}' failed.");
324+
return TaskResult.Failed;
325+
}
326+
}
327+
return TaskResult.Succeeded;
328+
}
329+
}
330+
}

0 commit comments

Comments
 (0)