Acknowledge runner request (#3996)

This commit is contained in:
eric sciple
2025-08-22 13:52:32 -05:00
committed by GitHub
parent 20d82ad357
commit 5f1efec208
7 changed files with 155 additions and 19 deletions

View File

@@ -23,7 +23,7 @@ namespace GitHub.Runner.Listener
private RunnerSettings _settings;
private ITerminal _term;
private TimeSpan _getNextMessageRetryInterval;
private TaskAgentStatus runnerStatus = TaskAgentStatus.Online;
private TaskAgentStatus _runnerStatus = TaskAgentStatus.Online;
private CancellationTokenSource _getMessagesTokenSource;
private VssCredentials _creds;
private VssCredentials _credsV2;
@@ -258,7 +258,7 @@ namespace GitHub.Runner.Listener
public void OnJobStatus(object sender, JobStatusEventArgs e)
{
Trace.Info("Received job status event. JobState: {0}", e.Status);
runnerStatus = e.Status;
_runnerStatus = e.Status;
try
{
_getMessagesTokenSource?.Cancel();
@@ -291,7 +291,7 @@ namespace GitHub.Runner.Listener
}
message = await _brokerServer.GetRunnerMessageAsync(_session.SessionId,
runnerStatus,
_runnerStatus,
BuildConstants.RunnerPackage.Version,
VarUtil.OS,
VarUtil.OSArchitecture,
@@ -417,6 +417,21 @@ namespace GitHub.Runner.Listener
await Task.CompletedTask;
}
public async Task AcknowledgeMessageAsync(string runnerRequestId, CancellationToken cancellationToken)
{
using var timeoutCts = new CancellationTokenSource(TimeSpan.FromSeconds(5)); // Short timeout
using var linkedCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, timeoutCts.Token);
Trace.Info($"Acknowledging runner request '{runnerRequestId}'.");
await _brokerServer.AcknowledgeRunnerRequestAsync(
runnerRequestId,
_session.SessionId,
_runnerStatus,
BuildConstants.RunnerPackage.Version,
VarUtil.OS,
VarUtil.OSArchitecture,
linkedCts.Token);
}
private bool IsGetNextMessageExceptionRetriable(Exception ex)
{
if (ex is TaskAgentNotFoundException ||

View File

@@ -32,6 +32,7 @@ namespace GitHub.Runner.Listener
Task DeleteSessionAsync();
Task<TaskAgentMessage> GetNextMessageAsync(CancellationToken token);
Task DeleteMessageAsync(TaskAgentMessage message);
Task AcknowledgeMessageAsync(string runnerRequestId, CancellationToken cancellationToken);
Task RefreshListenerTokenAsync();
void OnJobStatus(object sender, JobStatusEventArgs e);
@@ -52,7 +53,7 @@ namespace GitHub.Runner.Listener
private readonly TimeSpan _sessionConflictRetryLimit = TimeSpan.FromMinutes(4);
private readonly TimeSpan _clockSkewRetryLimit = TimeSpan.FromMinutes(30);
private readonly Dictionary<string, int> _sessionCreationExceptionTracker = new();
private TaskAgentStatus runnerStatus = TaskAgentStatus.Online;
private TaskAgentStatus _runnerStatus = TaskAgentStatus.Online;
private CancellationTokenSource _getMessagesTokenSource;
private VssCredentials _creds;
private VssCredentials _credsV2;
@@ -217,7 +218,7 @@ namespace GitHub.Runner.Listener
public void OnJobStatus(object sender, JobStatusEventArgs e)
{
Trace.Info("Received job status event. JobState: {0}", e.Status);
runnerStatus = e.Status;
_runnerStatus = e.Status;
try
{
_getMessagesTokenSource?.Cancel();
@@ -250,7 +251,7 @@ namespace GitHub.Runner.Listener
message = await _runnerServer.GetAgentMessageAsync(_settings.PoolId,
_session.SessionId,
_lastMessageId,
runnerStatus,
_runnerStatus,
BuildConstants.RunnerPackage.Version,
VarUtil.OS,
VarUtil.OSArchitecture,
@@ -274,7 +275,7 @@ namespace GitHub.Runner.Listener
}
message = await _brokerServer.GetRunnerMessageAsync(_session.SessionId,
runnerStatus,
_runnerStatus,
BuildConstants.RunnerPackage.Version,
VarUtil.OS,
VarUtil.OSArchitecture,
@@ -437,6 +438,21 @@ namespace GitHub.Runner.Listener
await _brokerServer.ForceRefreshConnection(_credsV2);
}
public async Task AcknowledgeMessageAsync(string runnerRequestId, CancellationToken cancellationToken)
{
using var timeoutCts = new CancellationTokenSource(TimeSpan.FromSeconds(5)); // Short timeout
using var linkedCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, timeoutCts.Token);
Trace.Info($"Acknowledging runner request '{runnerRequestId}'.");
await _brokerServer.AcknowledgeRunnerRequestAsync(
runnerRequestId,
_session.SessionId,
_runnerStatus,
BuildConstants.RunnerPackage.Version,
VarUtil.OS,
VarUtil.OSArchitecture,
linkedCts.Token);
}
private TaskAgentMessage DecryptMessage(TaskAgentMessage message)
{
if (_session.EncryptionKey == null ||

View File

@@ -654,22 +654,42 @@ namespace GitHub.Runner.Listener
else
{
var messageRef = StringUtil.ConvertFromJson<RunnerJobRequestRef>(message.Body);
Pipelines.AgentJobRequestMessage jobRequestMessage = null;
// Create connection
var credMgr = HostContext.GetService<ICredentialManager>();
// Acknowledge (best-effort)
if (messageRef.ShouldAcknowledge) // Temporary feature flag
{
try
{
await _listener.AcknowledgeMessageAsync(messageRef.RunnerRequestId, messageQueueLoopTokenSource.Token);
}
catch (Exception ex)
{
Trace.Error($"Best-effort acknowledge failed for request '{messageRef.RunnerRequestId}'");
Trace.Error(ex);
}
}
Pipelines.AgentJobRequestMessage jobRequestMessage = null;
if (string.IsNullOrEmpty(messageRef.RunServiceUrl))
{
// Connect
var credMgr = HostContext.GetService<ICredentialManager>();
var creds = credMgr.LoadCredentials(allowAuthUrlV2: false);
var actionsRunServer = HostContext.CreateService<IActionsRunServer>();
await actionsRunServer.ConnectAsync(new Uri(settings.ServerUrl), creds);
// Get job message
jobRequestMessage = await actionsRunServer.GetJobMessageAsync(messageRef.RunnerRequestId, messageQueueLoopTokenSource.Token);
}
else
{
// Connect
var credMgr = HostContext.GetService<ICredentialManager>();
var credsV2 = credMgr.LoadCredentials(allowAuthUrlV2: true);
var runServer = HostContext.CreateService<IRunServer>();
await runServer.ConnectAsync(new Uri(messageRef.RunServiceUrl), credsV2);
// Get job message
try
{
jobRequestMessage = await runServer.GetJobMessageAsync(messageRef.RunnerRequestId, messageRef.BillingOwnerId, messageQueueLoopTokenSource.Token);
@@ -698,7 +718,10 @@ namespace GitHub.Runner.Listener
}
}
// Dispatch
jobDispatcher.Run(jobRequestMessage, runOnce);
// Run once?
if (runOnce)
{
Trace.Info("One time used runner received job message.");

View File

@@ -10,6 +10,9 @@ namespace GitHub.Runner.Listener
[DataMember(Name = "runner_request_id")]
public string RunnerRequestId { get; set; }
[DataMember(Name = "should_acknowledge")]
public bool ShouldAcknowledge { get; set; }
[DataMember(Name = "run_service_url")]
public string RunServiceUrl { get; set; }