Skip to content

Commit 40c6e91

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

8 files changed

Lines changed: 1193 additions & 11 deletions

File tree

‎src/Runner.Worker/BackgroundStepCoordinator.cs‎

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

0 commit comments

Comments
 (0)