call run service renewjob (#2461)

* call run service renewjob

* format

* formatting

* make it private and expose internals

* lint

* fix exception class

* lint

* fix test as well

src/Runner.Common/RunServer.cs

@@ -1,4 +1,4 @@
-using System;
+using System;
 using System.Collections.Generic;
 using System.Threading;
 using System.Threading.Tasks;
@@ -7,6 +7,7 @@ using GitHub.DistributedTask.Pipelines;
 using GitHub.DistributedTask.WebApi;
 using GitHub.Runner.Sdk;
 using GitHub.Services.Common;
+using Sdk.RSWebApi.Contracts;
 using Sdk.WebApi.WebApi.RawClient;
 
 namespace GitHub.Runner.Common
@@ -19,6 +20,8 @@ namespace GitHub.Runner.Common
         Task<AgentJobRequestMessage> GetJobMessageAsync(string id, CancellationToken token);
 
         Task CompleteJobAsync(Guid planId, Guid jobId, TaskResult result, Dictionary<String, VariableValue> outputs, IList<StepResult> stepResults, CancellationToken token);
+
+        Task<RenewJobResponse> RenewJobAsync(Guid planId, Guid jobId, CancellationToken token);
     }
 
     public sealed class RunServer : RunnerService, IRunServer
@@ -64,5 +67,18 @@ namespace GitHub.Runner.Common
             return RetryRequest(
                 async () => await _runServiceHttpClient.CompleteJobAsync(requestUri, planId, jobId, result, outputs, stepResults, cancellationToken), cancellationToken);
         }
+
+        public Task<RenewJobResponse> RenewJobAsync(Guid planId, Guid jobId, CancellationToken cancellationToken)
+        {
+            CheckConnection();
+            var renewJobResponse = RetryRequest<RenewJobResponse>(
+                async () => await _runServiceHttpClient.RenewJobAsync(requestUri, planId, jobId, cancellationToken), cancellationToken);
+            if (renewJobResponse == null)
+            {
+                throw new TaskOrchestrationJobNotFoundException(jobId.ToString());
+            }
+
+            return renewJobResponse;
+        }
     }
 }

src/Runner.Common/Util/MessageUtil.cs

@@ -0,0 +1,14 @@
+namespace GitHub.Runner.Common.Util
+{
+    using System;
+    using GitHub.DistributedTask.WebApi;
+
+    public static class MessageUtil
+    {
+        public static bool IsRunServiceJob(string messageType)
+        {
+            return string.Equals(messageType, JobRequestMessageTypes.RunnerJobRequest, StringComparison.OrdinalIgnoreCase);
+        }
+    }
+}
+

src/Runner.Listener/InternalsVisibleTo.cs

@@ -0,0 +1,3 @@
+using System.Runtime.CompilerServices;
+
+[assembly: InternalsVisibleTo("Test")]

src/Runner.Listener/JobDispatcher.cs

@@ -7,6 +7,7 @@ using System.Text;
 using System.Text.RegularExpressions;
 using System.Threading;
 using System.Threading.Tasks;
+using GitHub.DistributedTask.Pipelines;
 using GitHub.DistributedTask.WebApi;
 using GitHub.Runner.Common;
 using GitHub.Runner.Common.Util;
@@ -58,6 +59,8 @@ namespace GitHub.Runner.Listener
 
         public event EventHandler<JobStatusEventArgs> JobStatus;
 
+        private bool _isRunServiceJob;
+
         public override void Initialize(IHostContext hostContext)
         {
             base.Initialize(hostContext);
@@ -86,6 +89,8 @@ namespace GitHub.Runner.Listener
         {
             Trace.Info($"Job request {jobRequestMessage.RequestId} for plan {jobRequestMessage.Plan.PlanId} job {jobRequestMessage.JobId} received.");
 
+            _isRunServiceJob = MessageUtil.IsRunServiceJob(jobRequestMessage.MessageType);
+
             WorkerDispatcher currentDispatch = null;
             if (_jobDispatchedQueue.Count > 0)
             {
@@ -239,6 +244,13 @@ namespace GitHub.Runner.Listener
                     return;
                 }
 
+                if (this._isRunServiceJob)
+                {
+                    Trace.Error($"We are not yet checking the state of jobrequest {jobDispatch.JobId} status. Cancel running worker right away.");
+                    jobDispatch.WorkerCancellationTokenSource.Cancel();
+                    return;
+                }
+
                 // based on the current design, server will only send one job for a given runner at a time.
                 // if the runner received a new job request while a previous job request is still running, this typically indicates two situations
                 // 1. a runner bug caused a server and runner mismatch on the state of the job request, e.g. the runner didn't renew the jobrequest
@@ -367,9 +379,11 @@ namespace GitHub.Runner.Listener
                     long requestId = message.RequestId;
                     Guid lockToken = Guid.Empty; // lockToken has never been used, keep this here of compat
 
+                    var systemConnection = message.Resources.Endpoints.SingleOrDefault(x => string.Equals(x.Name, WellKnownServiceEndpointNames.SystemVssConnection, StringComparison.OrdinalIgnoreCase));
+
                     // start renew job request
                     Trace.Info($"Start renew job request {requestId} for job {message.JobId}.");
-                    Task renewJobRequest = RenewJobRequestAsync(_poolId, requestId, lockToken, orchestrationId, firstJobRequestRenewed, lockRenewalTokenSource.Token);
+                    Task renewJobRequest = RenewJobRequestAsync(message, systemConnection, _poolId, requestId, lockToken, orchestrationId, firstJobRequestRenewed, lockRenewalTokenSource.Token);
 
                     // wait till first renew succeed or job request is cancelled
                     // not even start worker if the first renew fail
@@ -508,7 +522,6 @@ namespace GitHub.Runner.Listener
 
                         // we get first jobrequest renew succeed and start the worker process with the job message.
                         // send notification to machine provisioner.
-                        var systemConnection = message.Resources.Endpoints.SingleOrDefault(x => string.Equals(x.Name, WellKnownServiceEndpointNames.SystemVssConnection, StringComparison.OrdinalIgnoreCase));
                         var accessToken = systemConnection?.Authorization?.Parameters["AccessToken"];
                         notification.JobStarted(message.JobId, accessToken, systemConnection.Url);
 
@@ -531,11 +544,8 @@ namespace GitHub.Runner.Listener
                                     detailInfo = string.Join(Environment.NewLine, workerOutput);
                                     Trace.Info($"Return code {returnCode} indicate worker encounter an unhandled exception or app crash, attach worker stdout/stderr to JobRequest result.");
 
-                                    var jobServer = HostContext.GetService<IJobServer>();
-                                    VssCredentials jobServerCredential = VssUtil.GetVssCredential(systemConnection);
-                                    VssConnection jobConnection = VssUtil.CreateConnection(systemConnection.Url, jobServerCredential);
-                                    await jobServer.ConnectAsync(jobConnection);
 
+                                    var jobServer = await InitializeJobServerAsync(systemConnection);
                                     await LogWorkerProcessUnhandledException(jobServer, message, detailInfo);
 
                                     // Go ahead to finish the job with result 'Failed' if the STDERR from worker is System.IO.IOException, since it typically means we are running out of disk space.
@@ -675,9 +685,128 @@ namespace GitHub.Runner.Listener
             }
         }
 
-        public async Task RenewJobRequestAsync(int poolId, long requestId, Guid lockToken, string orchestrationId, TaskCompletionSource<int> firstJobRequestRenewed, CancellationToken token)
+        internal async Task RenewJobRequestAsync(Pipelines.AgentJobRequestMessage message, ServiceEndpoint systemConnection, int poolId, long requestId, Guid lockToken, string orchestrationId, TaskCompletionSource<int> firstJobRequestRenewed, CancellationToken token)
+        {
+            if (this._isRunServiceJob)
+            {
+                var runServer = await GetRunServerAsync(systemConnection);
+                await RenewJobRequestAsync(runServer, message.Plan.PlanId, message.JobId, firstJobRequestRenewed, token);
+            }
+            else
             {
                 var runnerServer = HostContext.GetService<IRunnerServer>();
+                await RenewJobRequestAsync(runnerServer, poolId, requestId, lockToken, orchestrationId, firstJobRequestRenewed, token);
+            }
+        }
+
+        private async Task RenewJobRequestAsync(IRunServer runServer, Guid planId, Guid jobId, TaskCompletionSource<int> firstJobRequestRenewed, CancellationToken token)
+        {
+            TaskAgentJobRequest request = null;
+            int firstRenewRetryLimit = 5;
+            int encounteringError = 0;
+
+            // renew lock during job running.
+            // stop renew only if cancellation token for lock renew task been signal or exception still happen after retry.
+            while (!token.IsCancellationRequested)
+            {
+                try
+                {
+                    var renewResponse = await runServer.RenewJobAsync(planId, jobId, token);
+                    Trace.Info($"Successfully renew job {jobId}, job is valid till {renewResponse.LockedUntil}");
+
+                    if (!firstJobRequestRenewed.Task.IsCompleted)
+                    {
+                        // fire first renew succeed event.
+                        firstJobRequestRenewed.TrySetResult(0);
+                    }
+
+                    if (encounteringError > 0)
+                    {
+                        encounteringError = 0;
+                        HostContext.WritePerfCounter("JobRenewRecovered");
+                    }
+
+                    // renew again after 60 sec delay
+                    await HostContext.Delay(TimeSpan.FromSeconds(60), token);
+                }
+                catch (TaskOrchestrationJobNotFoundException)
+                {
+                    // no need for retry. the job is not valid anymore.
+                    Trace.Info($"TaskAgentJobNotFoundException received when renew job {jobId}, job is no longer valid, stop renew job request.");
+                    return;
+                }
+                catch (OperationCanceledException) when (token.IsCancellationRequested)
+                {
+                    // OperationCanceledException may caused by http timeout or _lockRenewalTokenSource.Cance();
+                    // Stop renew only on cancellation token fired.
+                    Trace.Info($"job renew has been cancelled, stop renew job {jobId}.");
+                    return;
+                }
+                catch (Exception ex)
+                {
+                    Trace.Error($"Catch exception during renew runner job {jobId}.");
+                    Trace.Error(ex);
+                    encounteringError++;
+
+                    // retry
+                    TimeSpan remainingTime = TimeSpan.Zero;
+                    if (!firstJobRequestRenewed.Task.IsCompleted)
+                    {
+                        // retry 5 times every 10 sec for the first renew
+                        if (firstRenewRetryLimit-- > 0)
+                        {
+                            remainingTime = TimeSpan.FromSeconds(10);
+                        }
+                    }
+                    else
+                    {
+                        // retry till reach lockeduntil + 5 mins extra buffer.
+                        remainingTime = request.LockedUntil.Value + TimeSpan.FromMinutes(5) - DateTime.UtcNow;
+                    }
+
+                    if (remainingTime > TimeSpan.Zero)
+                    {
+                        TimeSpan delayTime;
+                        if (!firstJobRequestRenewed.Task.IsCompleted)
+                        {
+                            Trace.Info($"Retrying lock renewal for job {jobId}. The first job renew request has failed.");
+                            delayTime = BackoffTimerHelper.GetRandomBackoff(TimeSpan.FromSeconds(1), TimeSpan.FromSeconds(10));
+                        }
+                        else
+                        {
+                            Trace.Info($"Retrying lock renewal for job {jobId}. Job is valid until {request.LockedUntil.Value}.");
+                            if (encounteringError > 5)
+                            {
+                                delayTime = BackoffTimerHelper.GetRandomBackoff(TimeSpan.FromSeconds(15), TimeSpan.FromSeconds(30));
+                            }
+                            else
+                            {
+                                delayTime = BackoffTimerHelper.GetRandomBackoff(TimeSpan.FromSeconds(5), TimeSpan.FromSeconds(15));
+                            }
+                        }
+
+                        try
+                        {
+                            // back-off before next retry.
+                            await HostContext.Delay(delayTime, token);
+                        }
+                        catch (OperationCanceledException) when (token.IsCancellationRequested)
+                        {
+                            Trace.Info($"job renew has been cancelled, stop renew job {jobId}.");
+                        }
+                    }
+                    else
+                    {
+                        Trace.Info($"Lock renewal has run out of retry, stop renew lock for job {jobId}.");
+                        HostContext.WritePerfCounter("JobRenewReachLimit");
+                        return;
+                    }
+                }
+            }
+        }
+
+        private async Task RenewJobRequestAsync(IRunnerServer runnerServer, int poolId, long requestId, Guid lockToken, string orchestrationId, TaskCompletionSource<int> firstJobRequestRenewed, CancellationToken token)
+        {
             TaskAgentJobRequest request = null;
             int firstRenewRetryLimit = 5;
             int encounteringError = 0;
@@ -840,12 +969,10 @@ namespace GitHub.Runner.Listener
                 var systemConnection = message.Resources.Endpoints.SingleOrDefault(x => string.Equals(x.Name, WellKnownServiceEndpointNames.SystemVssConnection));
                 ArgUtil.NotNull(systemConnection, nameof(systemConnection));
 
-                var jobServer = HostContext.GetService<IJobServer>();
-                VssCredentials jobServerCredential = VssUtil.GetVssCredential(systemConnection);
-                VssConnection jobConnection = VssUtil.CreateConnection(systemConnection.Url, jobServerCredential);
-
-                await jobServer.ConnectAsync(jobConnection);
+                var server = await InitializeJobServerAsync(systemConnection);
 
+                if (server is IJobServer jobServer)
+                {
                     var timeline = await jobServer.GetTimelineAsync(message.Plan.ScopeIdentifier, message.Plan.PlanType, message.Plan.PlanId, message.Timeline.Id, CancellationToken.None);
 
                     var updatedRecords = new List<TimelineRecord>();
@@ -926,6 +1053,11 @@ namespace GitHub.Runner.Listener
                         await jobServer.UpdateTimelineRecordsAsync(message.Plan.ScopeIdentifier, message.Plan.PlanType, message.Plan.PlanId, message.Timeline.Id, updatedRecords, CancellationToken.None);
                     }
                 }
+                else
+                {
+                    Trace.Info("Job server does not support log upload yet.");
+                }
+            }
             catch (Exception ex)
             {
                 // Ignore any error during log upload since it's best effort
@@ -943,6 +1075,12 @@ namespace GitHub.Runner.Listener
                 return;
             }
 
+            if (this._isRunServiceJob)
+            {
+                Trace.Verbose($"Skip FinishAgentRequest call from Listener because MessageType is {message.MessageType}");
+                return;
+            }
+
             var runnerServer = HostContext.GetService<IRunnerServer>();
             int completeJobRequestRetryLimit = 5;
             List<Exception> exceptions = new();
@@ -979,7 +1117,9 @@ namespace GitHub.Runner.Listener
         }
 
         // log an error issue to job level timeline record
-        private async Task LogWorkerProcessUnhandledException(IJobServer jobServer, Pipelines.AgentJobRequestMessage message, string errorMessage)
+        private async Task LogWorkerProcessUnhandledException(IRunnerService server, Pipelines.AgentJobRequestMessage message, string errorMessage)
+        {
+            if (server is IJobServer jobServer)
             {
                 try
                 {
@@ -1025,9 +1165,17 @@ namespace GitHub.Runner.Listener
                     Trace.Error(ex);
                 }
             }
+            else
+            {
+                Trace.Info("Job server does not support handling unhandled exception yet, error message: {0}", errorMessage);
+                return;
+            }
+        }
 
         // raise job completed event to fail the job.
-        private async Task ForceFailJob(IJobServer jobServer, Pipelines.AgentJobRequestMessage message)
+        private async Task ForceFailJob(IRunnerService server, Pipelines.AgentJobRequestMessage message)
+        {
+            if (server is IJobServer jobServer)
             {
                 try
                 {
@@ -1040,6 +1188,47 @@ namespace GitHub.Runner.Listener
                     Trace.Error(ex);
                 }
             }
+            else if (server is IRunServer runServer)
+            {
+                try
+                {
+                    await runServer.CompleteJobAsync(message.Plan.PlanId, message.JobId, TaskResult.Failed, outputs: null, stepResults: null, CancellationToken.None);
+                }
+                catch (Exception ex)
+                {
+                    Trace.Error("Fail to raise job completion back to service.");
+                    Trace.Error(ex);
+                }
+            }
+            else
+            {
+                throw new NotSupportedException($"Server type {server.GetType().FullName} is not supported.");
+            }
+        }
+
+        private async Task<IRunnerService> InitializeJobServerAsync(ServiceEndpoint systemConnection)
+        {
+            if (this._isRunServiceJob)
+            {
+                return await GetRunServerAsync(systemConnection);
+            }
+            else
+            {
+                var jobServer = HostContext.GetService<IJobServer>();
+                VssCredentials jobServerCredential = VssUtil.GetVssCredential(systemConnection);
+                VssConnection jobConnection = VssUtil.CreateConnection(systemConnection.Url, jobServerCredential);
+                await jobServer.ConnectAsync(jobConnection);
+                return jobServer;
+            }
+        }
+
+        private async Task<IRunServer> GetRunServerAsync(ServiceEndpoint systemConnection)
+        {
+            var runServer = HostContext.GetService<IRunServer>();
+            VssCredentials jobServerCredential = VssUtil.GetVssCredential(systemConnection);
+            await runServer.ConnectAsync(systemConnection.Url, jobServerCredential);
+            return runServer;
+        }
 
         private class WorkerDispatcher : IDisposable
         {

src/Runner.Listener/Runner.cs

@@ -9,6 +9,7 @@ using System.Threading;
 using System.Threading.Tasks;
 using GitHub.DistributedTask.WebApi;
 using GitHub.Runner.Common;
+using GitHub.Runner.Common.Util;
 using GitHub.Runner.Listener.Check;
 using GitHub.Runner.Listener.Configuration;
 using GitHub.Runner.Sdk;
@@ -502,7 +503,7 @@ namespace GitHub.Runner.Listener
                                 }
                             }
                             // Broker flow
-                            else if (string.Equals(message.MessageType, JobRequestMessageTypes.RunnerJobRequest, StringComparison.OrdinalIgnoreCase))
+                            else if (MessageUtil.IsRunServiceJob(message.MessageType))
                             {
                                 if (autoUpdateInProgress || runOnceJobReceived)
                                 {

src/Runner.Worker/JobRunner.cs

@@ -6,6 +6,7 @@ using System.Net.Http;
 using System.Text;
 using System.Threading;
 using System.Threading.Tasks;
+using GitHub.DistributedTask.Pipelines;
 using GitHub.DistributedTask.WebApi;
 using GitHub.Runner.Common;
 using GitHub.Runner.Common.Util;
@@ -19,7 +20,7 @@ namespace GitHub.Runner.Worker
     [ServiceLocator(Default = typeof(JobRunner))]
     public interface IJobRunner : IRunnerService
     {
-        Task<TaskResult> RunAsync(Pipelines.AgentJobRequestMessage message, CancellationToken jobRequestCancellationToken);
+        Task<TaskResult> RunAsync(AgentJobRequestMessage message, CancellationToken jobRequestCancellationToken);
     }
 
     public sealed class JobRunner : RunnerService, IJobRunner
@@ -28,7 +29,7 @@ namespace GitHub.Runner.Worker
         private RunnerSettings _runnerSettings;
         private ITempDirectoryManager _tempDirectoryManager;
 
-        public async Task<TaskResult> RunAsync(Pipelines.AgentJobRequestMessage message, CancellationToken jobRequestCancellationToken)
+        public async Task<TaskResult> RunAsync(AgentJobRequestMessage message, CancellationToken jobRequestCancellationToken)
         {
             // Validate parameters.
             Trace.Entering();
@@ -42,7 +43,7 @@ namespace GitHub.Runner.Worker
             IRunnerService server = null;
 
             ServiceEndpoint systemConnection = message.Resources.Endpoints.Single(x => string.Equals(x.Name, WellKnownServiceEndpointNames.SystemVssConnection, StringComparison.OrdinalIgnoreCase));
-            if (string.Equals(message.MessageType, JobRequestMessageTypes.RunnerJobRequest, StringComparison.OrdinalIgnoreCase))
+            if (MessageUtil.IsRunServiceJob(message.MessageType))
             {
                 var runServer = HostContext.GetService<IRunServer>();
                 VssCredentials jobServerCredential = VssUtil.GetVssCredential(systemConnection);

src/Sdk/RSWebApi/Contracts/RenewJobRequest.cs

@@ -0,0 +1,15 @@
+using System;
+using System.Runtime.Serialization;
+
+namespace GitHub.Actions.RunService.WebApi
+{
+    [DataContract]
+    public class RenewJobRequest
+    {
+        [DataMember(Name = "planId", EmitDefaultValue = false)]
+        public Guid PlanID { get; set; }
+
+        [DataMember(Name = "jobId", EmitDefaultValue = false)]
+        public Guid JobID { get; set; }
+    }
+}

src/Sdk/RSWebApi/Contracts/RenewJobResponse.cs

@@ -0,0 +1,16 @@
+using System;
+using System.Runtime.Serialization;
+
+namespace Sdk.RSWebApi.Contracts
+{
+    [DataContract]
+    public class RenewJobResponse
+    {
+        [DataMember]
+        public DateTime LockedUntil
+        {
+            get;
+            internal set;
+        }
+    }
+}

src/Sdk/RSWebApi/RunServiceHttpClient.cs

@@ -8,6 +8,7 @@ using GitHub.DistributedTask.WebApi;
 using GitHub.Services.Common;
 using GitHub.Services.OAuth;
 using GitHub.Services.WebApi;
+using Sdk.RSWebApi.Contracts;
 using Sdk.WebApi.WebApi;
 
 namespace GitHub.Actions.RunService.WebApi
@@ -103,5 +104,28 @@ namespace GitHub.Actions.RunService.WebApi
                     content: requestContent,
                     cancellationToken: cancellationToken);
         }
+
+        public Task<RenewJobResponse> RenewJobAsync(
+            Uri requestUri,
+            Guid planId,
+            Guid jobId,
+            CancellationToken cancellationToken = default)
+        {
+            HttpMethod httpMethod = new HttpMethod("POST");
+            var payload = new RenewJobRequest()
+            {
+                PlanID = planId,
+                JobID = jobId
+            };
+
+            requestUri = new Uri(requestUri, "renewjob");
+
+            var requestContent = new ObjectContent<RenewJobRequest>(payload, new VssJsonMediaTypeFormatter(true));
+            return SendAsync<RenewJobResponse>(
+                httpMethod,
+                requestUri,
+                content: requestContent,
+                cancellationToken: cancellationToken);
+        }
     }
 }

src/Test/L0/Listener/JobDispatcherL0.cs

@@ -1,12 +1,17 @@
-using System;
+using System;
 using System.Collections.Generic;
 using System.Reflection;
 using System.Threading;
 using System.Threading.Tasks;
+using GitHub.DistributedTask.ObjectTemplating.Tokens;
+using GitHub.DistributedTask.Pipelines;
+using GitHub.DistributedTask.Pipelines.ContextData;
 using GitHub.DistributedTask.WebApi;
 using GitHub.Runner.Listener;
+using GitHub.Runner.Listener.Configuration;
 using GitHub.Services.WebApi;
 using Moq;
+using Sdk.RSWebApi.Contracts;
 using Xunit;
 
 using Pipelines = GitHub.DistributedTask.Pipelines;
@@ -18,6 +23,8 @@ namespace GitHub.Runner.Common.Tests.Listener
         private Mock<IProcessChannel> _processChannel;
         private Mock<IProcessInvoker> _processInvoker;
         private Mock<IRunnerServer> _runnerServer;
+
+        private Mock<IRunServer> _runServer;
         private Mock<IConfigurationStore> _configurationStore;
 
         public JobDispatcherL0()
@@ -25,6 +32,7 @@ namespace GitHub.Runner.Common.Tests.Listener
             _processChannel = new Mock<IProcessChannel>();
             _processInvoker = new Mock<IProcessInvoker>();
             _runnerServer = new Mock<IRunnerServer>();
+            _runServer = new Mock<IRunServer>();
             _configurationStore = new Mock<IConfigurationStore>();
         }
 
@@ -139,7 +147,7 @@ namespace GitHub.Runner.Common.Tests.Listener
                 var jobDispatcher = new JobDispatcher();
                 jobDispatcher.Initialize(hc);
 
-                await jobDispatcher.RenewJobRequestAsync(poolId, requestId, Guid.Empty, Guid.NewGuid().ToString(), firstJobRequestRenewed, cancellationTokenSource.Token);
+                await jobDispatcher.RenewJobRequestAsync(It.IsAny<AgentJobRequestMessage>(), It.IsAny<ServiceEndpoint>(), poolId, requestId, Guid.Empty, Guid.NewGuid().ToString(), firstJobRequestRenewed, cancellationTokenSource.Token);
 
                 Assert.True(firstJobRequestRenewed.Task.IsCompletedSuccessfully);
                 _runnerServer.Verify(x => x.RenewAgentRequestAsync(It.IsAny<int>(), It.IsAny<long>(), It.IsAny<Guid>(), It.IsAny<string>(), It.IsAny<CancellationToken>()), Times.Exactly(5));
@@ -197,7 +205,7 @@ namespace GitHub.Runner.Common.Tests.Listener
                 var jobDispatcher = new JobDispatcher();
                 jobDispatcher.Initialize(hc);
 
-                await jobDispatcher.RenewJobRequestAsync(poolId, requestId, Guid.Empty, Guid.NewGuid().ToString(), firstJobRequestRenewed, cancellationTokenSource.Token);
+                await jobDispatcher.RenewJobRequestAsync(It.IsAny<AgentJobRequestMessage>(), It.IsAny<ServiceEndpoint>(), poolId, requestId, Guid.Empty, Guid.NewGuid().ToString(), firstJobRequestRenewed, cancellationTokenSource.Token);
 
                 Assert.True(firstJobRequestRenewed.Task.IsCompletedSuccessfully, "First renew should succeed.");
                 Assert.False(cancellationTokenSource.IsCancellationRequested);
@@ -205,6 +213,75 @@ namespace GitHub.Runner.Common.Tests.Listener
             }
         }
 
+        [Fact]
+        [Trait("Level", "L0")]
+        [Trait("Category", "Runner")]
+        public async void DispatcherRenewJobOnRunServiceStopOnJobNotFoundExceptions()
+        {
+            //Arrange
+            using (var hc = new TestHostContext(this))
+            {
+                int poolId = 1;
+                Int64 requestId = 1000;
+                int count = 0;
+
+                var trace = hc.GetTrace(nameof(DispatcherRenewJobOnRunServiceStopOnJobNotFoundExceptions));
+                TaskCompletionSource<int> firstJobRequestRenewed = new();
+                CancellationTokenSource cancellationTokenSource = new();
+
+                TaskAgentJobRequest request = new();
+                PropertyInfo lockUntilProperty = request.GetType().GetProperty("LockedUntil", BindingFlags.Instance | BindingFlags.NonPublic | BindingFlags.Public);
+                Assert.NotNull(lockUntilProperty);
+                lockUntilProperty.SetValue(request, DateTime.UtcNow.AddMinutes(5));
+
+                hc.SetSingleton<IRunServer>(_runServer.Object);
+                hc.SetSingleton<IConfigurationStore>(_configurationStore.Object);
+                _configurationStore.Setup(x => x.GetSettings()).Returns(new RunnerSettings() { PoolId = 1 });
+                _ = _runServer.Setup(x => x.RenewJobAsync(It.IsAny<Guid>(), It.IsAny<Guid>(), It.IsAny<CancellationToken>()))
+                            .Returns(() =>
+                            {
+                                count++;
+                                if (!firstJobRequestRenewed.Task.IsCompletedSuccessfully)
+                                {
+                                    trace.Info("First renew happens.");
+                                }
+
+                                if (count < 5)
+                                {
+                                    var response = new RenewJobResponse()
+                                    {
+                                        LockedUntil = request.LockedUntil.Value
+                                    };
+                                    return Task.FromResult<RenewJobResponse>(response);
+                                }
+                                else if (count == 5)
+                                {
+                                    cancellationTokenSource.CancelAfter(10000);
+                                    throw new TaskOrchestrationJobNotFoundException("");
+                                }
+                                else
+                                {
+                                    throw new InvalidOperationException("Should not reach here.");
+                                }
+                            });
+
+
+                var jobDispatcher = new JobDispatcher();
+                jobDispatcher.Initialize(hc);
+                EnableRunServiceJobForJobDispatcher(jobDispatcher);
+
+                // Set the value of the _isRunServiceJob field to true
+                var isRunServiceJobField = typeof(JobDispatcher).GetField("_isRunServiceJob", BindingFlags.NonPublic | BindingFlags.Instance);
+                isRunServiceJobField.SetValue(jobDispatcher, true);
+
+                await jobDispatcher.RenewJobRequestAsync(GetAgentJobRequestMessage(), GetServiceEndpoint(), poolId, requestId, Guid.Empty, Guid.NewGuid().ToString(), firstJobRequestRenewed, cancellationTokenSource.Token);
+
+                Assert.True(firstJobRequestRenewed.Task.IsCompletedSuccessfully, "First renew should succeed.");
+                Assert.False(cancellationTokenSource.IsCancellationRequested);
+                _runServer.Verify(x => x.RenewJobAsync(It.IsAny<Guid>(), It.IsAny<Guid>(), It.IsAny<CancellationToken>()), Times.Exactly(5));
+            }
+        }
+
         [Fact]
         [Trait("Level", "L0")]
         [Trait("Category", "Runner")]
@@ -256,7 +333,7 @@ namespace GitHub.Runner.Common.Tests.Listener
                 var jobDispatcher = new JobDispatcher();
                 jobDispatcher.Initialize(hc);
 
-                await jobDispatcher.RenewJobRequestAsync(poolId, requestId, Guid.Empty, Guid.NewGuid().ToString(), firstJobRequestRenewed, cancellationTokenSource.Token);
+                await jobDispatcher.RenewJobRequestAsync(It.IsAny<AgentJobRequestMessage>(), It.IsAny<ServiceEndpoint>(), poolId, requestId, Guid.Empty, Guid.NewGuid().ToString(), firstJobRequestRenewed, cancellationTokenSource.Token);
 
                 Assert.True(firstJobRequestRenewed.Task.IsCompletedSuccessfully, "First renew should succeed.");
                 Assert.False(cancellationTokenSource.IsCancellationRequested);
@@ -312,8 +389,9 @@ namespace GitHub.Runner.Common.Tests.Listener
                 var jobDispatcher = new JobDispatcher();
                 jobDispatcher.Initialize(hc);
 
+
                 // Act
-                await jobDispatcher.RenewJobRequestAsync(0, 0, Guid.Empty, Guid.NewGuid().ToString(), firstJobRequestRenewed, cancellationTokenSource.Token);
+                await jobDispatcher.RenewJobRequestAsync(It.IsAny<AgentJobRequestMessage>(), It.IsAny<ServiceEndpoint>(), 0, 0, Guid.Empty, Guid.NewGuid().ToString(), firstJobRequestRenewed, cancellationTokenSource.Token);
 
                 // Assert
                 _configurationStore.Verify(x => x.SaveSettings(It.Is<RunnerSettings>(settings => settings.AgentName == newName)), Times.Once);
@@ -368,7 +446,7 @@ namespace GitHub.Runner.Common.Tests.Listener
                 jobDispatcher.Initialize(hc);
 
                 // Act
-                await jobDispatcher.RenewJobRequestAsync(0, 0, Guid.Empty, Guid.NewGuid().ToString(), firstJobRequestRenewed, cancellationTokenSource.Token);
+                await jobDispatcher.RenewJobRequestAsync(It.IsAny<AgentJobRequestMessage>(), It.IsAny<ServiceEndpoint>(), 0, 0, Guid.Empty, Guid.NewGuid().ToString(), firstJobRequestRenewed, cancellationTokenSource.Token);
 
                 // Assert
                 _configurationStore.Verify(x => x.SaveSettings(It.IsAny<RunnerSettings>()), Times.Never);
@@ -421,7 +499,7 @@ namespace GitHub.Runner.Common.Tests.Listener
                 jobDispatcher.Initialize(hc);
 
                 // Act
-                await jobDispatcher.RenewJobRequestAsync(0, 0, Guid.Empty, Guid.NewGuid().ToString(), firstJobRequestRenewed, cancellationTokenSource.Token);
+                await jobDispatcher.RenewJobRequestAsync(It.IsAny<AgentJobRequestMessage>(), It.IsAny<ServiceEndpoint>(), 0, 0, Guid.Empty, Guid.NewGuid().ToString(), firstJobRequestRenewed, cancellationTokenSource.Token);
 
                 // Assert
                 _configurationStore.Verify(x => x.SaveSettings(It.IsAny<RunnerSettings>()), Times.Never);
@@ -479,7 +557,7 @@ namespace GitHub.Runner.Common.Tests.Listener
                 var jobDispatcher = new JobDispatcher();
                 jobDispatcher.Initialize(hc);
 
-                await jobDispatcher.RenewJobRequestAsync(poolId, requestId, Guid.Empty, Guid.NewGuid().ToString(), firstJobRequestRenewed, cancellationTokenSource.Token);
+                await jobDispatcher.RenewJobRequestAsync(It.IsAny<AgentJobRequestMessage>(), It.IsAny<ServiceEndpoint>(), poolId, requestId, Guid.Empty, Guid.NewGuid().ToString(), firstJobRequestRenewed, cancellationTokenSource.Token);
 
                 Assert.True(firstJobRequestRenewed.Task.IsCompletedSuccessfully, "First renew should succeed.");
                 Assert.True(cancellationTokenSource.IsCancellationRequested);
@@ -536,7 +614,7 @@ namespace GitHub.Runner.Common.Tests.Listener
                 var jobDispatcher = new JobDispatcher();
                 jobDispatcher.Initialize(hc);
 
-                await jobDispatcher.RenewJobRequestAsync(poolId, requestId, Guid.Empty, Guid.NewGuid().ToString(), firstJobRequestRenewed, cancellationTokenSource.Token);
+                await jobDispatcher.RenewJobRequestAsync(It.IsAny<AgentJobRequestMessage>(), It.IsAny<ServiceEndpoint>(), poolId, requestId, Guid.Empty, Guid.NewGuid().ToString(), firstJobRequestRenewed, cancellationTokenSource.Token);
 
                 Assert.False(firstJobRequestRenewed.Task.IsCompletedSuccessfully, "First renew should failed.");
                 Assert.False(cancellationTokenSource.IsCancellationRequested);
@@ -600,7 +678,7 @@ namespace GitHub.Runner.Common.Tests.Listener
                 var jobDispatcher = new JobDispatcher();
                 jobDispatcher.Initialize(hc);
 
-                await jobDispatcher.RenewJobRequestAsync(poolId, requestId, Guid.Empty, Guid.NewGuid().ToString(), firstJobRequestRenewed, cancellationTokenSource.Token);
+                await jobDispatcher.RenewJobRequestAsync(It.IsAny<AgentJobRequestMessage>(), It.IsAny<ServiceEndpoint>(), poolId, requestId, Guid.Empty, Guid.NewGuid().ToString(), firstJobRequestRenewed, cancellationTokenSource.Token);
 
                 Assert.True(firstJobRequestRenewed.Task.IsCompletedSuccessfully, "First renew should succeed.");
                 Assert.False(cancellationTokenSource.IsCancellationRequested);
@@ -659,5 +737,78 @@ namespace GitHub.Runner.Common.Tests.Listener
                 Assert.True(jobDispatcher.RunOnceJobCompleted.Task.Result, "JobDispatcher should set task complete token to 'TRUE' for one time agent.");
             }
         }
+
+        private static void EnableRunServiceJobForJobDispatcher(JobDispatcher jobDispatcher)
+        {
+            // Set the value of the _isRunServiceJob field to true
+            var isRunServiceJobField = typeof(JobDispatcher).GetField("_isRunServiceJob", BindingFlags.NonPublic | BindingFlags.Instance);
+            isRunServiceJobField.SetValue(jobDispatcher, true);
+        }
+
+        private static ServiceEndpoint GetServiceEndpoint()
+        {
+            var serviceEndpoint = new ServiceEndpoint
+            {
+                Authorization = new EndpointAuthorization
+                {
+                    Scheme = EndpointAuthorizationSchemes.OAuth
+                }
+            };
+            serviceEndpoint.Authorization.Parameters.Add("AccessToken", "token");
+            return serviceEndpoint;
+        }
+
+        private static AgentJobRequestMessage GetAgentJobRequestMessage()
+        {
+            var message = new AgentJobRequestMessage(
+                new TaskOrchestrationPlanReference()
+                {
+                    PlanType = "Build",
+                    PlanId = Guid.NewGuid(),
+                    Version = 1
+                },
+                new TimelineReference()
+                {
+                    Id = Guid.NewGuid()
+                },
+                Guid.NewGuid(),
+                "jobDisplayName",
+                "jobName",
+                null,
+                null,
+                new List<TemplateToken>(),
+                new Dictionary<string, VariableValue>()
+                {
+                    {
+                        "variables",
+                        new VariableValue()
+                        {
+                            IsSecret = false,
+                            Value = "variables"
+                        }
+                    }
+                },
+                new List<MaskHint>()
+                {
+                    new MaskHint()
+                    {
+                        Type = MaskType.Variable,
+                        Value = "maskHints"
+                    }
+                },
+                new JobResources(),
+                new DictionaryContextData(),
+                new WorkspaceOptions(),
+                new List<JobStep>(),
+                new List<string>()
+                {
+                    "fileTable"
+                },
+                null,
+                new List<TemplateToken>(),
+                new ActionsEnvironmentReference("env")
+            );
+            return message;
+        }
     }
 }