Handle HostedRunnerShutdownMessage from service to shutdown hosted runner faster. (#1922)

src/Runner.Listener/Runner.cs

@@ -1,18 +1,18 @@
-using GitHub.DistributedTask.WebApi;
-using GitHub.Runner.Listener.Configuration;
 using System;
-using System.Threading;
+using System.Collections.Generic;
-using System.Threading.Tasks;
-using GitHub.Services.WebApi;
-using Pipelines = GitHub.DistributedTask.Pipelines;
 using System.IO;
+using System.Linq;
 using System.Reflection;
 using System.Runtime.CompilerServices;
+using System.Threading;
+using System.Threading.Tasks;
+using GitHub.DistributedTask.WebApi;
 using GitHub.Runner.Common;
-using GitHub.Runner.Sdk;
-using System.Linq;
 using GitHub.Runner.Listener.Check;
-using System.Collections.Generic;
+using GitHub.Runner.Listener.Configuration;
+using GitHub.Runner.Sdk;
+using GitHub.Services.WebApi;
+using Pipelines = GitHub.DistributedTask.Pipelines;
 
 namespace GitHub.Runner.Listener
 {
@@ -322,6 +322,7 @@ namespace GitHub.Runner.Listener
 
                 // Should we try to cleanup ephemeral runners
                 bool runOnceJobCompleted = false;
+                bool skipSessionDeletion = false;
                 try
                 {
                     var notification = HostContext.GetService<IJobNotification>();
@@ -468,6 +469,14 @@ namespace GitHub.Runner.Listener
                                     Trace.Info($"Skip message deletion for cancellation message '{message.MessageId}'.");
                                 }
                             }
+                            else if (string.Equals(message.MessageType, Pipelines.HostedRunnerShutdownMessage.MessageType, StringComparison.OrdinalIgnoreCase))
+                            {
+                                var HostedRunnerShutdownMessage = JsonUtility.FromString<Pipelines.HostedRunnerShutdownMessage>(message.Body);
+                                skipMessageDeletion = true;
+                                skipSessionDeletion = true;
+                                Trace.Info($"Service requests the hosted runner to shutdown. Reason: '{HostedRunnerShutdownMessage.Reason}'.");
+                                return Constants.Runner.ReturnCode.Success;
+                            }
                             else
                             {
                                 Trace.Error($"Received message {message.MessageId} with unsupported message type {message.MessageType}.");
@@ -501,6 +510,8 @@ namespace GitHub.Runner.Listener
                         await jobDispatcher.ShutdownAsync();
                     }
 
+                    if (!skipSessionDeletion)
+                    {
                         try
                         {
                             await _listener.DeleteSessionAsync();
@@ -511,6 +522,7 @@ namespace GitHub.Runner.Listener
                             // and the delete session call will ends up with 401.
                             Trace.Info($"Ignore any exception during DeleteSession for an ephemeral runner. {ex}");
                         }
+                    }
 
                     messageQueueLoopTokenSource.Dispose();
 

src/Sdk/DTPipelines/Pipelines/RunnerShutdownMessage.cs

@@ -0,0 +1,39 @@
+using System;
+using System.Runtime.Serialization;
+using GitHub.Services.WebApi;
+using Newtonsoft.Json;
+
+namespace GitHub.DistributedTask.Pipelines
+{
+    [DataContract]
+    public sealed class HostedRunnerShutdownMessage
+    {
+        public static readonly String MessageType = "RunnerShutdown";
+
+        [JsonConstructor]
+        internal HostedRunnerShutdownMessage()
+        {
+        }
+
+        public HostedRunnerShutdownMessage(String reason)
+        {
+            this.Reason = reason;
+        }
+
+        [DataMember]
+        public String Reason
+        {
+            get;
+            private set;
+        }
+
+        public WebApi.TaskAgentMessage GetAgentMessage()
+        {
+            return new WebApi.TaskAgentMessage
+            {
+                Body = JsonUtility.ToString(this),
+                MessageType = HostedRunnerShutdownMessage.MessageType,
+            };
+        }
+    }
+}