From 9499f477a2688e46dc35c1fe4dc7955a9fcce67a Mon Sep 17 00:00:00 2001
From: Tatyana Kostromskaya <32135588+takost@users.noreply.github.com>
Date: Tue, 21 Jun 2022 16:12:07 +0200
Subject: [PATCH] Add retry logic around getting job messages from broker
 (#1939)

* Jsut simple solution without additional funcs

* Delete old comment

* resolve

* Refactor retry function, make it more common

* Make retry function generic, get rid of extra params

* delete extra using

* Add cancellation token and limit of attempts

* Add some additional logging

* Rework condition

* replace to do..while

* return `while (true)` to simplify code structure

* Add other cancelling token, add TODO comment
---
 src/Runner.Common/RunServer.cs | 36 +++++++++++++++++++++++++++++++---
 src/Runner.Listener/Runner.cs  |  3 +--
 2 files changed, 34 insertions(+), 5 deletions(-)

diff --git a/src/Runner.Common/RunServer.cs b/src/Runner.Common/RunServer.cs
index 582054133..7df8da7a8 100644
--- a/src/Runner.Common/RunServer.cs
+++ b/src/Runner.Common/RunServer.cs
@@ -17,7 +17,7 @@ namespace GitHub.Runner.Common
     {
         Task ConnectAsync(Uri serverUrl, VssCredentials credentials);
 
-        Task<AgentJobRequestMessage> GetJobMessageAsync(string id);
+        Task<AgentJobRequestMessage> GetJobMessageAsync(string id, CancellationToken token);
     }
 
     public sealed class RunServer : RunnerService, IRunServer
@@ -67,10 +67,40 @@ namespace GitHub.Runner.Common
             }
         }
 
-        public Task<AgentJobRequestMessage> GetJobMessageAsync(string id)
+        public Task<AgentJobRequestMessage> GetJobMessageAsync(string id, CancellationToken cancellationToken)
         {
             CheckConnection();
-            return _taskAgentClient.GetJobMessageAsync(id);
+            var jobMessage = RetryRequest<AgentJobRequestMessage>(async () =>
+                                                    {
+                                                        return await _taskAgentClient.GetJobMessageAsync(id, cancellationToken);
+                                                    }, cancellationToken);
+            return jobMessage;
+        }
+
+        private async Task<T> RetryRequest<T>(Func<Task<T>> func,
+                                              CancellationToken cancellationToken,
+                                              int maxRetryAttemptsCount = 5
+                                             )
+        {
+            var retryCount = 0;
+            while (true)
+            {
+                retryCount++;
+                cancellationToken.ThrowIfCancellationRequested();
+                try
+                {
+                    return await func();
+                }
+                // TODO: Add handling of non-retriable exceptions: https://github.com/github/actions-broker/issues/122
+                catch (Exception ex) when (retryCount < maxRetryAttemptsCount)
+                {
+                    Trace.Error("Catch exception during get full job message");
+                    Trace.Error(ex);
+                    var backOff = BackoffTimerHelper.GetRandomBackoff(TimeSpan.FromSeconds(5), TimeSpan.FromSeconds(15));
+                    Trace.Warning($"Back off {backOff.TotalSeconds} seconds before next retry. {maxRetryAttemptsCount - retryCount} attempt left.");
+                    await Task.Delay(backOff, cancellationToken);
+                }
+            }
         }
     }
 }
diff --git a/src/Runner.Listener/Runner.cs b/src/Runner.Listener/Runner.cs
index 8455ebef2..f7505bc08 100644
--- a/src/Runner.Listener/Runner.cs
+++ b/src/Runner.Listener/Runner.cs
@@ -474,10 +474,9 @@ namespace GitHub.Runner.Listener
                                     var credMgr = HostContext.GetService<ICredentialManager>();
                                     var creds = credMgr.LoadCredentials();
 
-                                    // todo: add retries https://github.com/github/actions-broker/issues/49
                                     var runServer = HostContext.CreateService<IRunServer>();
                                     await runServer.ConnectAsync(new Uri(settings.ServerUrl), creds);
-                                    var jobMessage = await runServer.GetJobMessageAsync(messageRef.RunnerRequestId);
+                                    var jobMessage = await runServer.GetJobMessageAsync(messageRef.RunnerRequestId, messageQueueLoopTokenSource.Token);
 
                                     jobDispatcher.Run(jobMessage, runOnce);
                                     if (runOnce)