Support auth migration using authUrlV2 in Runner/MessageListener. (#3787)

This commit is contained in:
Tingluo Huang
2025-04-10 12:58:33 -04:00
committed by GitHub
parent f1b5b5bd5c
commit d5ccbd10d1
6 changed files with 1204 additions and 12 deletions

View File

@@ -1,4 +1,5 @@
using System;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.IO;
using System.Linq;
@@ -31,7 +32,11 @@ namespace GitHub.Runner.Listener
private ITerminal _term;
private bool _inConfigStage;
private ManualResetEvent _completedCommand = new(false);
private readonly ConcurrentQueue<string> _authMigrationTelemetries = new();
private Task _authMigrationTelemetryTask;
private readonly object _authMigrationTelemetryLock = new();
private IRunnerServer _runnerServer;
private CancellationTokenSource _authMigrationTelemetryTokenSource = new();
// <summary>
// Helps avoid excessive calls to Run Service when encountering non-retriable errors from /acquirejob.
@@ -68,6 +73,8 @@ namespace GitHub.Runner.Listener
//register a SIGTERM handler
HostContext.Unloading += Runner_Unloading;
HostContext.AuthMigrationChanged += HandleAuthMigrationChanged;
// TODO Unit test to cover this logic
Trace.Info(nameof(ExecuteCommand));
var configManager = HostContext.GetService<IConfigurationManager>();
@@ -313,6 +320,8 @@ namespace GitHub.Runner.Listener
}
finally
{
_authMigrationTelemetryTokenSource?.Cancel();
HostContext.AuthMigrationChanged -= HandleAuthMigrationChanged;
_term.CancelKeyPress -= CtrlCHandler;
HostContext.Unloading -= Runner_Unloading;
_completedCommand.Set();
@@ -572,18 +581,18 @@ namespace GitHub.Runner.Listener
// Create connection
var credMgr = HostContext.GetService<ICredentialManager>();
var creds = credMgr.LoadCredentials(allowAuthUrlV2: false);
if (string.IsNullOrEmpty(messageRef.RunServiceUrl))
{
var creds = credMgr.LoadCredentials(allowAuthUrlV2: false);
var actionsRunServer = HostContext.CreateService<IActionsRunServer>();
await actionsRunServer.ConnectAsync(new Uri(settings.ServerUrl), creds);
jobRequestMessage = await actionsRunServer.GetJobMessageAsync(messageRef.RunnerRequestId, messageQueueLoopTokenSource.Token);
}
else
{
var credsV2 = credMgr.LoadCredentials(allowAuthUrlV2: true);
var runServer = HostContext.CreateService<IRunServer>();
await runServer.ConnectAsync(new Uri(messageRef.RunServiceUrl), creds);
await runServer.ConnectAsync(new Uri(messageRef.RunServiceUrl), credsV2);
try
{
jobRequestMessage = await runServer.GetJobMessageAsync(messageRef.RunnerRequestId, messageRef.BillingOwnerId, messageQueueLoopTokenSource.Token);
@@ -601,6 +610,13 @@ namespace GitHub.Runner.Listener
catch (Exception ex)
{
Trace.Error($"Caught exception from acquiring job message: {ex}");
if (HostContext.AllowAuthMigration)
{
Trace.Info("Disable migration mode for 60 minutes.");
HostContext.DeferAuthMigration(TimeSpan.FromMinutes(60), $"Acquire job failed with exception: {ex}");
}
continue;
}
}
@@ -718,6 +734,73 @@ namespace GitHub.Runner.Listener
return Constants.Runner.ReturnCode.Success;
}
private void HandleAuthMigrationChanged(object sender, AuthMigrationEventArgs e)
{
Trace.Verbose("Handle AuthMigrationChanged in Runner");
_authMigrationTelemetries.Enqueue($"{DateTime.UtcNow.ToString("O")}: {e.Trace}");
// only start the telemetry reporting task once auth migration is changed (enabled or disabled)
lock (_authMigrationTelemetryLock)
{
if (_authMigrationTelemetryTask == null)
{
_authMigrationTelemetryTask = ReportAuthMigrationTelemetryAsync(_authMigrationTelemetryTokenSource.Token);
}
}
}
private async Task ReportAuthMigrationTelemetryAsync(CancellationToken token)
{
var configManager = HostContext.GetService<IConfigurationManager>();
var runnerSettings = configManager.LoadSettings();
while (!token.IsCancellationRequested)
{
try
{
await HostContext.Delay(TimeSpan.FromSeconds(60), token);
}
catch (TaskCanceledException)
{
// Ignore cancellation
}
Trace.Verbose("Checking for auth migration telemetry to report");
while (_authMigrationTelemetries.TryDequeue(out var telemetry))
{
Trace.Verbose($"Reporting auth migration telemetry: {telemetry}");
if (runnerSettings != null)
{
try
{
using (var tokenSource = new CancellationTokenSource(TimeSpan.FromSeconds(30)))
{
await _runnerServer.UpdateAgentUpdateStateAsync(runnerSettings.PoolId, runnerSettings.AgentId, "RefreshConfig", telemetry, tokenSource.Token);
}
}
catch (Exception ex)
{
Trace.Error("Failed to report auth migration telemetry.");
Trace.Error(ex);
_authMigrationTelemetries.Enqueue(telemetry);
}
}
if (!token.IsCancellationRequested)
{
try
{
await HostContext.Delay(TimeSpan.FromSeconds(10), token);
}
catch (TaskCanceledException)
{
// Ignore cancellation
}
}
}
}
}
private void PrintUsage(CommandSettings command)
{
string separator;