From 2cb1f9431acb65faaf64c481a809a370f4b435f9 Mon Sep 17 00:00:00 2001 From: Tingluo Huang Date: Fri, 28 Mar 2025 16:25:12 -0400 Subject: [PATCH] Small runner code cleanup. (#3773) --- src/Runner.Common/BrokerServer.cs | 7 ++++++- src/Runner.Listener/BrokerMessageListener.cs | 22 +++++++------------- src/Runner.Listener/MessageListener.cs | 9 ++++---- src/Runner.Listener/Runner.cs | 2 +- 4 files changed, 20 insertions(+), 20 deletions(-) diff --git a/src/Runner.Common/BrokerServer.cs b/src/Runner.Common/BrokerServer.cs index b0774178f..04d48aa0c 100644 --- a/src/Runner.Common/BrokerServer.cs +++ b/src/Runner.Common/BrokerServer.cs @@ -88,7 +88,12 @@ namespace GitHub.Runner.Common public Task ForceRefreshConnection(VssCredentials credentials) { - return ConnectAsync(_brokerUri, credentials); + if (!string.IsNullOrEmpty(_brokerUri?.AbsoluteUri)) + { + return ConnectAsync(_brokerUri, credentials); + } + + return Task.CompletedTask; } public bool ShouldRetryException(Exception ex) diff --git a/src/Runner.Listener/BrokerMessageListener.cs b/src/Runner.Listener/BrokerMessageListener.cs index 2a994631e..6bab66e14 100644 --- a/src/Runner.Listener/BrokerMessageListener.cs +++ b/src/Runner.Listener/BrokerMessageListener.cs @@ -29,6 +29,7 @@ namespace GitHub.Runner.Listener private TaskAgentSession _session; private IRunnerServer _runnerServer; private IBrokerServer _brokerServer; + private ICredentialManager _credMgr; private readonly Dictionary _sessionCreationExceptionTracker = new(); private bool _accessTokenRevoked = false; private readonly TimeSpan _sessionCreationRetryInterval = TimeSpan.FromSeconds(30); @@ -43,6 +44,7 @@ namespace GitHub.Runner.Listener _term = HostContext.GetService(); _runnerServer = HostContext.GetService(); _brokerServer = HostContext.GetService(); + _credMgr = HostContext.GetService(); } public async Task CreateSessionAsync(CancellationToken token) @@ -63,8 +65,7 @@ namespace GitHub.Runner.Listener // Create connection. Trace.Info("Loading Credentials"); - var credMgr = HostContext.GetService(); - _creds = credMgr.LoadCredentials(); + _creds = _credMgr.LoadCredentials(); var agent = new TaskAgentReference { @@ -329,7 +330,7 @@ namespace GitHub.Runner.Listener } } - public async Task RefreshListenerTokenAsync(CancellationToken cancellationToken) + public async Task RefreshListenerTokenAsync() { await RefreshBrokerConnectionAsync(); } @@ -432,17 +433,10 @@ namespace GitHub.Runner.Listener private async Task RefreshBrokerConnectionAsync() { - var configManager = HostContext.GetService(); - _settings = configManager.LoadSettings(); - - if (string.IsNullOrEmpty(_settings.ServerUrlV2)) - { - throw new InvalidOperationException("ServerUrlV2 is not set"); - } - - var credMgr = HostContext.GetService(); - VssCredentials creds = credMgr.LoadCredentials(); - await _brokerServer.ConnectAsync(new Uri(_settings.ServerUrlV2), creds); + Trace.Info("Reload credentials."); + _creds = _credMgr.LoadCredentials(); + await _brokerServer.ConnectAsync(new Uri(_settings.ServerUrlV2), _creds); + Trace.Info("Connection to Broker Server recreated."); } } } diff --git a/src/Runner.Listener/MessageListener.cs b/src/Runner.Listener/MessageListener.cs index fcff3a507..f2c7873fd 100644 --- a/src/Runner.Listener/MessageListener.cs +++ b/src/Runner.Listener/MessageListener.cs @@ -33,7 +33,7 @@ namespace GitHub.Runner.Listener Task GetNextMessageAsync(CancellationToken token); Task DeleteMessageAsync(TaskAgentMessage message); - Task RefreshListenerTokenAsync(CancellationToken token); + Task RefreshListenerTokenAsync(); void OnJobStatus(object sender, JobStatusEventArgs e); } @@ -44,6 +44,7 @@ namespace GitHub.Runner.Listener private ITerminal _term; private IRunnerServer _runnerServer; private IBrokerServer _brokerServer; + private ICredentialManager _credMgr; private TaskAgentSession _session; private TimeSpan _getNextMessageRetryInterval; private bool _accessTokenRevoked = false; @@ -64,6 +65,7 @@ namespace GitHub.Runner.Listener _term = HostContext.GetService(); _runnerServer = HostContext.GetService(); _brokerServer = hostContext.GetService(); + _credMgr = hostContext.GetService(); } public async Task CreateSessionAsync(CancellationToken token) @@ -78,8 +80,7 @@ namespace GitHub.Runner.Listener // Create connection. Trace.Info("Loading Credentials"); - var credMgr = HostContext.GetService(); - _creds = credMgr.LoadCredentials(); + _creds = _credMgr.LoadCredentials(); var agent = new TaskAgentReference { @@ -411,7 +412,7 @@ namespace GitHub.Runner.Listener } } - public async Task RefreshListenerTokenAsync(CancellationToken cancellationToken) + public async Task RefreshListenerTokenAsync() { await _runnerServer.RefreshConnectionAsync(RunnerConnectionType.MessageQueue, TimeSpan.FromSeconds(60)); await _brokerServer.ForceRefreshConnection(_creds); diff --git a/src/Runner.Listener/Runner.cs b/src/Runner.Listener/Runner.cs index 28b65d877..7aae6bc24 100644 --- a/src/Runner.Listener/Runner.cs +++ b/src/Runner.Listener/Runner.cs @@ -633,7 +633,7 @@ namespace GitHub.Runner.Listener else if (string.Equals(message.MessageType, TaskAgentMessageTypes.ForceTokenRefresh)) { Trace.Info("Received ForceTokenRefreshMessage"); - await _listener.RefreshListenerTokenAsync(messageQueueLoopTokenSource.Token); + await _listener.RefreshListenerTokenAsync(); } else if (string.Equals(message.MessageType, RunnerRefreshConfigMessage.MessageType)) {