From 513ad2b4cdc093a9c55d5d9ab71d2c4245f5b57c Mon Sep 17 00:00:00 2001 From: ishu599 Date: Thu, 3 Sep 2026 20:19:39 +0000 Subject: [PATCH] Fix run service job renewal retry --- src/Runner.Listener/JobDispatcher.cs | 11 +-- src/Test/L0/Listener/JobDispatcherL0.cs | 96 +++++++++++++++++++++++++ 2 files changed, 103 insertions(+), 4 deletions(-) diff --git a/src/Runner.Listener/JobDispatcher.cs b/src/Runner.Listener/JobDispatcher.cs index dcc896633aa..a6b6dc4c6b8 100644 --- a/src/Runner.Listener/JobDispatcher.cs +++ b/src/Runner.Listener/JobDispatcher.cs @@ -747,7 +747,7 @@ internal async Task RenewJobRequestAsync(Pipelines.AgentJobRequestMessage messag private async Task RenewJobRequestAsync(IRunServer runServer, Guid planId, Guid jobId, TaskCompletionSource firstJobRequestRenewed, CancellationToken token) { - TaskAgentJobRequest request = null; + DateTime? lastSuccessfulLockedUntil = null; int firstRenewRetryLimit = 5; int encounteringError = 0; @@ -759,6 +759,7 @@ private async Task RenewJobRequestAsync(IRunServer runServer, Guid planId, Guid { var renewResponse = await runServer.RenewJobAsync(planId, jobId, token); Trace.Info($"Successfully renew job {jobId}, job is valid till {renewResponse.LockedUntil}"); + lastSuccessfulLockedUntil = renewResponse.LockedUntil; if (!firstJobRequestRenewed.Task.IsCompleted) { @@ -807,7 +808,7 @@ private async Task RenewJobRequestAsync(IRunServer runServer, Guid planId, Guid else { // retry till reach lockeduntil + 5 mins extra buffer. - remainingTime = request.LockedUntil.Value + TimeSpan.FromMinutes(5) - DateTime.UtcNow; + remainingTime = lastSuccessfulLockedUntil.Value + TimeSpan.FromMinutes(5) - DateTime.UtcNow; } if (remainingTime > TimeSpan.Zero) @@ -820,7 +821,7 @@ private async Task RenewJobRequestAsync(IRunServer runServer, Guid planId, Guid } else { - Trace.Info($"Retrying lock renewal for job {jobId}. Job is valid until {request.LockedUntil.Value}."); + Trace.Info($"Retrying lock renewal for job {jobId}. Job is valid until {lastSuccessfulLockedUntil.Value}."); if (encounteringError > 5) { delayTime = BackoffTimerHelper.GetRandomBackoff(TimeSpan.FromSeconds(15), TimeSpan.FromSeconds(30)); @@ -854,6 +855,7 @@ private async Task RenewJobRequestAsync(IRunServer runServer, Guid planId, Guid private async Task RenewJobRequestAsync(IRunnerServer runnerServer, int poolId, long requestId, Guid lockToken, string orchestrationId, TaskCompletionSource firstJobRequestRenewed, CancellationToken token) { TaskAgentJobRequest request = null; + DateTime? lastSuccessfulLockedUntil = null; int firstRenewRetryLimit = 5; int encounteringError = 0; @@ -865,6 +867,7 @@ private async Task RenewJobRequestAsync(IRunnerServer runnerServer, int poolId, { request = await runnerServer.RenewAgentRequestAsync(poolId, requestId, lockToken, orchestrationId, token); Trace.Info($"Successfully renew job request {requestId}, job is valid till {request.LockedUntil.Value}"); + lastSuccessfulLockedUntil = request.LockedUntil; if (!firstJobRequestRenewed.Task.IsCompleted) { @@ -936,7 +939,7 @@ private async Task RenewJobRequestAsync(IRunnerServer runnerServer, int poolId, } else { - Trace.Info($"Retrying lock renewal for jobrequest {requestId}. Job is valid until {request.LockedUntil.Value}."); + Trace.Info($"Retrying lock renewal for jobrequest {requestId}. Job is valid until {lastSuccessfulLockedUntil.Value}."); if (encounteringError > 5) { delayTime = BackoffTimerHelper.GetRandomBackoff(TimeSpan.FromSeconds(15), TimeSpan.FromSeconds(30)); diff --git a/src/Test/L0/Listener/JobDispatcherL0.cs b/src/Test/L0/Listener/JobDispatcherL0.cs index a160ffba2c0..72a1f56b464 100644 --- a/src/Test/L0/Listener/JobDispatcherL0.cs +++ b/src/Test/L0/Listener/JobDispatcherL0.cs @@ -1,6 +1,7 @@ using System; using System.Collections.Generic; using System.Reflection; +using System.Net.Http; using System.Threading; using System.Threading.Tasks; using GitHub.DistributedTask.ObjectTemplating.Tokens; @@ -344,6 +345,101 @@ public async void DispatcherRenewJobRequestStopOnJobTokenExpiredExceptions() } } + [Fact] + [Trait("Level", "L0")] + [Trait("Category", "Runner")] + public async Task DispatcherRenewJobOnRunServiceRetryAfterGenericException() + { + // Arrange + using (var hc = new TestHostContext(this)) + { + int poolId = 1; + Int64 requestId = 1000; + int count = 0; + + TaskCompletionSource firstJobRequestRenewed = new(); + CancellationTokenSource cancellationTokenSource = new(); + + hc.SetSingleton(_runServer.Object); + hc.SetSingleton(_configurationStore.Object); + _configurationStore.Setup(x => x.GetSettings()) + .Returns(new RunnerSettings() { PoolId = 1 }); + + _ = _runServer.Setup(x => x.RenewJobAsync( + It.IsAny(), + It.IsAny(), + It.IsAny())) + .Returns(() => + { + count++; + + if (count == 1) + { + // First renewal succeeds. + return Task.FromResult(new RenewJobResponse + { + LockedUntil = DateTime.UtcNow.AddMinutes(5) + }); + } + else if (count == 2) + { + // Second renewal fails with a generic exception. + throw new HttpRequestException("Simulated renewal failure."); + } + else if (count == 3) + { + // Retry succeeds. + cancellationTokenSource.Cancel(); + return Task.FromResult(new RenewJobResponse + { + LockedUntil = DateTime.UtcNow.AddMinutes(5) + }); + } + + 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); + + // Act + await jobDispatcher.RenewJobRequestAsync( + GetAgentJobRequestMessage(), + GetServiceEndpoint(), + poolId, + requestId, + Guid.Empty, + Guid.NewGuid().ToString(), + firstJobRequestRenewed, + cancellationTokenSource.Token); + + // Assert + Assert.True( + firstJobRequestRenewed.Task.IsCompletedSuccessfully, + "First renew should succeed."); + + Assert.Equal(3, count); + Assert.True( + cancellationTokenSource.IsCancellationRequested, + "Cancellation should only happen after the retry succeeds."); + + _runServer.Verify( + x => x.RenewJobAsync( + It.IsAny(), + It.IsAny(), + It.IsAny()), + Times.Exactly(3)); + } + } + [Fact] [Trait("Level", "L0")] [Trait("Category", "Runner")]