Handle ForceTokenRefresh message (#3133)

* Handle ForceTokenRefresh message * move to constants * format
actions · Feb 7, 2024 · 651ea42 · 651ea42
1 parent bcc665a
commit 651ea42
Show file tree

Hide file tree

Showing 4 changed files with 28 additions and 2 deletions.
diff --git a/src/Runner.Listener/BrokerMessageListener.cs b/src/Runner.Listener/BrokerMessageListener.cs
@@ -273,7 +273,7 @@ public async Task<TaskAgentMessage> GetNextMessageAsync(CancellationToken token)
                         }
 
                         // re-create VssConnection before next retry
-                        await RefreshBrokerConnection();
+                        await RefreshBrokerConnectionAsync();
 
                         Trace.Info("Sleeping for {0} seconds before retrying.", _getNextMessageRetryInterval.TotalSeconds);
                         await HostContext.Delay(_getNextMessageRetryInterval, token);
@@ -303,6 +303,11 @@ public async Task<TaskAgentMessage> GetNextMessageAsync(CancellationToken token)
             }
         }
 
+        public async Task RefreshListenerTokenAsync(CancellationToken cancellationToken)
+        {
+            await RefreshBrokerConnectionAsync();
+        }
+
         public async Task DeleteMessageAsync(TaskAgentMessage message)
         {
             await Task.CompletedTask;
@@ -398,7 +403,7 @@ ex is AccessDeniedException ||
             }
         }
 
-        private async Task RefreshBrokerConnection()
+        private async Task RefreshBrokerConnectionAsync()
         {
             var configManager = HostContext.GetService<IConfigurationManager>();
             _settings = configManager.LoadSettings();

diff --git a/src/Runner.Listener/MessageListener.cs b/src/Runner.Listener/MessageListener.cs
@@ -25,6 +25,8 @@ public interface IMessageListener : IRunnerService
         Task DeleteSessionAsync();
         Task<TaskAgentMessage> GetNextMessageAsync(CancellationToken token);
         Task DeleteMessageAsync(TaskAgentMessage message);
+
+        Task RefreshListenerTokenAsync(CancellationToken token);
         void OnJobStatus(object sender, JobStatusEventArgs e);
     }
 
@@ -379,6 +381,11 @@ public async Task DeleteMessageAsync(TaskAgentMessage message)
             }
         }
 
+        public async Task RefreshListenerTokenAsync(CancellationToken cancellationToken)
+        {
+            await _runnerServer.RefreshConnectionAsync(RunnerConnectionType.MessageQueue, TimeSpan.FromSeconds(60));
+        }
+
         private TaskAgentMessage DecryptMessage(TaskAgentMessage message)
         {
             if (_session.EncryptionKey == null ||

diff --git a/src/Runner.Listener/Runner.cs b/src/Runner.Listener/Runner.cs
@@ -596,6 +596,10 @@ await runServer.GetJobMessageAsync(messageRef.RunnerRequestId,
                                 Trace.Info($"Service requests the hosted runner to shutdown. Reason: '{HostedRunnerShutdownMessage.Reason}'.");
                                 return Constants.Runner.ReturnCode.Success;
                             }
+                            else if (string.Equals(message.MessageType, TaskAgentMessageTypes.ForceTokenRefresh))
+                            {
+                                await _listener.RefreshListenerTokenAsync(messageQueueLoopTokenSource.Token);
+                            }
                             else
                             {
                                 Trace.Error($"Received message {message.MessageId} with unsupported message type {message.MessageType}.");

diff --git a/src/Sdk/DTWebApi/WebApi/TaskAgentMessageTypes.cs b/src/Sdk/DTWebApi/WebApi/TaskAgentMessageTypes.cs
@@ -0,0 +1,10 @@
+using System;
+using System.Runtime.Serialization;
+
+namespace GitHub.DistributedTask.WebApi
+{
+    public sealed class TaskAgentMessageTypes
+    {
+        public static readonly string ForceTokenRefresh = "ForceTokenRefresh";
+    }
+}