From 1558d47a317baf3a4ed7a0abeca64a5a621592d6 Mon Sep 17 00:00:00 2001
From: Patrick Ellis <319655+pje@users.noreply.github.com>
Date: Thu, 18 May 2023 13:21:32 -0700
Subject: [PATCH] WIP

---
 src/Runner.Common/BrokerServer.cs            | 15 ++++++++++++++-
 src/Runner.Listener/BrokerMessageListener.cs | 13 +++++++++++--
 src/Runner.Listener/Runner.cs                |  8 ++------
 src/Sdk/WebApi/WebApi/BrokerHttpClient.cs    |  9 ++-------
 4 files changed, 29 insertions(+), 16 deletions(-)

diff --git a/src/Runner.Common/BrokerServer.cs b/src/Runner.Common/BrokerServer.cs
index bb0691771..34a977813 100644
--- a/src/Runner.Common/BrokerServer.cs
+++ b/src/Runner.Common/BrokerServer.cs
@@ -5,6 +5,7 @@ using System.Threading.Tasks;
 using GitHub.Actions.RunService.WebApi;
 using GitHub.DistributedTask.Pipelines;
 using GitHub.DistributedTask.WebApi;
+using GitHub.Runner.Common.Util;
 using GitHub.Runner.Sdk;
 using GitHub.Services.Common;
 using Sdk.RSWebApi.Contracts;
@@ -18,6 +19,7 @@ namespace GitHub.Runner.Common
         Task ConnectAsync(Uri serverUrl, VssCredentials credentials);
 
         Task<TaskAgentMessage> GetRunnerMessageAsync(CancellationToken token, TaskAgentStatus status, string version);
+        Task<bool> DeleteRunnerMessageAsync(string messageID, CancellationToken token);
     }
 
     public sealed class BrokerServer : RunnerService, IBrokerServer
@@ -48,9 +50,20 @@ namespace GitHub.Runner.Common
         {
             CheckConnection();
             var jobMessage = RetryRequest<TaskAgentMessage>(
-                async () => await _brokerHttpClient.GetRunnerMessageAsync(version, status, cancellationToken), cancellationToken);
+                async () => await _brokerHttpClient.GetRunnerMessageAsync(version, status, cancellationToken), cancellationToken
+                );
 
             return jobMessage;
         }
+
+        public Task<bool> DeleteRunnerMessageAsync(string messageID, CancellationToken cancellationToken)
+        {
+            CheckConnection();
+
+            return RetryRequest(
+                async () => await _brokerHttpClient.DeleteRunnerMessageAsync(messageID, cancellationToken),
+                cancellationToken
+            );
+        }
     }
 }
diff --git a/src/Runner.Listener/BrokerMessageListener.cs b/src/Runner.Listener/BrokerMessageListener.cs
index 1ea4e977b..4ae8aca7b 100644
--- a/src/Runner.Listener/BrokerMessageListener.cs
+++ b/src/Runner.Listener/BrokerMessageListener.cs
@@ -168,10 +168,19 @@ namespace GitHub.Runner.Listener
             }
         }
 
-        public async Task DeleteMessageAsync(string messageID, CancellationToken token)
+        public async Task DeleteMessageAsync(TaskAgentMessage message)
         {
             Trace.Entering();
-            await _brokerServer.DeleteRunnerMessageAsync(messageID, token);
+
+            if (MessageUtil.IsRunServiceJob(message.MessageType))
+            {
+                var messageRef = StringUtil.ConvertFromJson<RunnerJobRequestRef>(message.Body);
+
+                using (var cs = new CancellationTokenSource(TimeSpan.FromSeconds(30)))
+                {
+                    await _brokerServer.DeleteRunnerMessageAsync(messageRef.RunnerRequestId, cs.Token);
+                }
+            }
         }
 
         private bool IsGetNextMessageExceptionRetriable(Exception ex)
diff --git a/src/Runner.Listener/Runner.cs b/src/Runner.Listener/Runner.cs
index 10edd83aa..56e9500bf 100644
--- a/src/Runner.Listener/Runner.cs
+++ b/src/Runner.Listener/Runner.cs
@@ -393,8 +393,6 @@ namespace GitHub.Runner.Listener
                         TaskAgentMessage message = null;
                         bool skipMessageDeletion = false;
                         bool useBrokerDeletion = false;
-                        string brokerDeletionParamsMessageID = null;
-                        CancellationToken brokerDeletionParamsToken = null;
                         try
                         {
                             Task<TaskAgentMessage> getNextMessage = _listener.GetNextMessageAsync(messageQueueLoopTokenSource.Token);
@@ -557,9 +555,6 @@ namespace GitHub.Runner.Listener
                                             jobRequestMessage =
                                             await runServer.GetJobMessageAsync(messageRef.RunnerRequestId,
                                             messageQueueLoopTokenSource.Token);
-                                            useBrokerDeletion = true;
-                                            brokerDeletionParamsMessageID = messageRef.RunnerRequestId;
-                                            brokerDeletionParamsToken = messageQueueLoopTokenSource.Token;
                                         }
                                         catch (TaskOrchestrationJobAlreadyAcquiredException)
                                         {
@@ -608,7 +603,8 @@ namespace GitHub.Runner.Listener
                                 {
                                     if (useBrokerDeletion)
                                     {
-                                        await _listener.DeleteMessageAsync(brokerDeletionParamsMessageID, brokerDeletionParamsToken);
+                                        // await _listener.DeleteMessageAsync(brokerDeletionParamsMessageID, brokerDeletionParamsToken);
+                                        await _listener.DeleteMessageAsync(message);
                                     }
                                     else
                                     {
diff --git a/src/Sdk/WebApi/WebApi/BrokerHttpClient.cs b/src/Sdk/WebApi/WebApi/BrokerHttpClient.cs
index eb35f0542..592ff3e74 100644
--- a/src/Sdk/WebApi/WebApi/BrokerHttpClient.cs
+++ b/src/Sdk/WebApi/WebApi/BrokerHttpClient.cs
@@ -89,7 +89,7 @@ namespace GitHub.Actions.RunService.WebApi
             throw new Exception($"Failed to get job message: {result.Error}");
         }
 
-        public async Task<TaskAgentMessage> DeleteRunnerMessageAsync(
+        public async Task<bool> DeleteRunnerMessageAsync(
             string messageID,
             CancellationToken cancellationToken = default
         )
@@ -109,12 +109,7 @@ namespace GitHub.Actions.RunService.WebApi
                 queryParameters: queryParams,
                 cancellationToken: cancellationToken);
 
-            if (result.IsSuccess)
-            {
-                return result.Value;
-            }
-
-            throw new Exception($"Failed to delete job message: {result.Error}");
+            return (result.IsSuccess);
         }
     }
 }