From 3b149bd8db7a8c7ac38a9c2508a81e55c8e974e3 Mon Sep 17 00:00:00 2001
From: Patrick Ellis <319655+pje@users.noreply.github.com>
Date: Thu, 18 May 2023 12:18:24 -0700
Subject: [PATCH] WIP WIP WIP

---
 src/Runner.Listener/BrokerMessageListener.cs |  5 ++--
 src/Runner.Listener/Runner.cs                | 15 ++++++++++-
 src/Sdk/WebApi/WebApi/BrokerHttpClient.cs    | 28 ++++++++++++++++++++
 3 files changed, 45 insertions(+), 3 deletions(-)

diff --git a/src/Runner.Listener/BrokerMessageListener.cs b/src/Runner.Listener/BrokerMessageListener.cs
index aa3e81886..1ea4e977b 100644
--- a/src/Runner.Listener/BrokerMessageListener.cs
+++ b/src/Runner.Listener/BrokerMessageListener.cs
@@ -168,9 +168,10 @@ namespace GitHub.Runner.Listener
             }
         }
 
-        public async Task DeleteMessageAsync(TaskAgentMessage message)
+        public async Task DeleteMessageAsync(string messageID, CancellationToken token)
         {
-            await Task.CompletedTask;
+            Trace.Entering();
+            await _brokerServer.DeleteRunnerMessageAsync(messageID, token);
         }
 
         private bool IsGetNextMessageExceptionRetriable(Exception ex)
diff --git a/src/Runner.Listener/Runner.cs b/src/Runner.Listener/Runner.cs
index c367f8852..10edd83aa 100644
--- a/src/Runner.Listener/Runner.cs
+++ b/src/Runner.Listener/Runner.cs
@@ -392,6 +392,9 @@ 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);
@@ -554,6 +557,9 @@ namespace GitHub.Runner.Listener
                                             jobRequestMessage =
                                             await runServer.GetJobMessageAsync(messageRef.RunnerRequestId,
                                             messageQueueLoopTokenSource.Token);
+                                            useBrokerDeletion = true;
+                                            brokerDeletionParamsMessageID = messageRef.RunnerRequestId;
+                                            brokerDeletionParamsToken = messageQueueLoopTokenSource.Token;
                                         }
                                         catch (TaskOrchestrationJobAlreadyAcquiredException)
                                         {
@@ -600,7 +606,14 @@ namespace GitHub.Runner.Listener
                             {
                                 try
                                 {
-                                    await _listener.DeleteMessageAsync(message);
+                                    if (useBrokerDeletion)
+                                    {
+                                        await _listener.DeleteMessageAsync(brokerDeletionParamsMessageID, brokerDeletionParamsToken);
+                                    }
+                                    else
+                                    {
+                                        await _listener.DeleteMessageAsync(message);
+                                    }
                                 }
                                 catch (Exception ex)
                                 {
diff --git a/src/Sdk/WebApi/WebApi/BrokerHttpClient.cs b/src/Sdk/WebApi/WebApi/BrokerHttpClient.cs
index e02541866..eb35f0542 100644
--- a/src/Sdk/WebApi/WebApi/BrokerHttpClient.cs
+++ b/src/Sdk/WebApi/WebApi/BrokerHttpClient.cs
@@ -88,5 +88,33 @@ namespace GitHub.Actions.RunService.WebApi
 
             throw new Exception($"Failed to get job message: {result.Error}");
         }
+
+        public async Task<TaskAgentMessage> DeleteRunnerMessageAsync(
+            string messageID,
+            CancellationToken cancellationToken = default
+        )
+        {
+            var requestUri = new Uri(Client.BaseAddress, "message");
+
+            List<KeyValuePair<string, string>> queryParams = new List<KeyValuePair<string, string>>();
+
+            if (messageID != null)
+            {
+                queryParams.Add("messageID", messageID);
+            }
+
+            var result = await SendAsync<TaskAgentMessage>(
+                new HttpMethod("DELETE"),
+                requestUri: requestUri,
+                queryParameters: queryParams,
+                cancellationToken: cancellationToken);
+
+            if (result.IsSuccess)
+            {
+                return result.Value;
+            }
+
+            throw new Exception($"Failed to delete job message: {result.Error}");
+        }
     }
 }