From a61328a7e7c41cb37a64e4db60deefc44a51d8e0 Mon Sep 17 00:00:00 2001 From: Luke Tomlinson Date: Mon, 3 Feb 2025 15:15:54 -0500 Subject: [PATCH] Pass BillingOwnerId through Acquire/Complete calls (#3689) * Pass BillingOwnerId through Acquire/Complete calls * add param to test --- src/Runner.Common/RunServer.cs | 10 ++++++---- src/Runner.Listener/JobDispatcher.cs | 2 +- src/Runner.Listener/Runner.cs | 2 +- src/Runner.Listener/RunnerJobRequestRef.cs | 5 +++++ src/Runner.Worker/JobRunner.cs | 2 +- .../DTPipelines/Pipelines/AgentJobRequestMessage.cs | 7 +++++++ src/Sdk/RSWebApi/Contracts/AcquireJobRequest.cs | 3 +++ src/Sdk/RSWebApi/Contracts/CompleteJobRequest.cs | 3 +++ src/Sdk/RSWebApi/RunServiceHttpClient.cs | 6 +++++- src/Test/L0/Listener/JobDispatcherL0.cs | 11 +++++++---- 10 files changed, 39 insertions(+), 12 deletions(-) diff --git a/src/Runner.Common/RunServer.cs b/src/Runner.Common/RunServer.cs index 8a7f4d5b4..a343bf0e3 100644 --- a/src/Runner.Common/RunServer.cs +++ b/src/Runner.Common/RunServer.cs @@ -18,7 +18,7 @@ namespace GitHub.Runner.Common { Task ConnectAsync(Uri serverUrl, VssCredentials credentials); - Task GetJobMessageAsync(string id, CancellationToken token); + Task GetJobMessageAsync(string id, string billingOwnerId, CancellationToken token); Task CompleteJobAsync( Guid planId, @@ -29,6 +29,7 @@ namespace GitHub.Runner.Common IList jobAnnotations, string environmentUrl, IList telemetry, + string billingOwnerId, CancellationToken token); Task RenewJobAsync(Guid planId, Guid jobId, CancellationToken token); @@ -58,11 +59,11 @@ namespace GitHub.Runner.Common } } - public Task GetJobMessageAsync(string id, CancellationToken cancellationToken) + public Task GetJobMessageAsync(string id, string billingOwnerId, CancellationToken cancellationToken) { CheckConnection(); return RetryRequest( - async () => await _runServiceHttpClient.GetJobMessageAsync(requestUri, id, VarUtil.OS, cancellationToken), cancellationToken, + async () => await _runServiceHttpClient.GetJobMessageAsync(requestUri, id, VarUtil.OS, billingOwnerId, cancellationToken), cancellationToken, shouldRetry: ex => ex is not TaskOrchestrationJobNotFoundException && // HTTP status 404 ex is not TaskOrchestrationJobAlreadyAcquiredException && // HTTP status 409 @@ -78,11 +79,12 @@ namespace GitHub.Runner.Common IList jobAnnotations, string environmentUrl, IList telemetry, + string billingOwnerId, CancellationToken cancellationToken) { CheckConnection(); return RetryRequest( - async () => await _runServiceHttpClient.CompleteJobAsync(requestUri, planId, jobId, result, outputs, stepResults, jobAnnotations, environmentUrl, telemetry, cancellationToken), cancellationToken); + async () => await _runServiceHttpClient.CompleteJobAsync(requestUri, planId, jobId, result, outputs, stepResults, jobAnnotations, environmentUrl, telemetry, billingOwnerId, cancellationToken), cancellationToken); } public Task RenewJobAsync(Guid planId, Guid jobId, CancellationToken cancellationToken) diff --git a/src/Runner.Listener/JobDispatcher.cs b/src/Runner.Listener/JobDispatcher.cs index 1b8196091..b908c8b07 100644 --- a/src/Runner.Listener/JobDispatcher.cs +++ b/src/Runner.Listener/JobDispatcher.cs @@ -1206,7 +1206,7 @@ namespace GitHub.Runner.Listener jobAnnotations.Add(annotation.Value); } - await runServer.CompleteJobAsync(message.Plan.PlanId, message.JobId, TaskResult.Failed, outputs: null, stepResults: null, jobAnnotations: jobAnnotations, environmentUrl: null, telemetry: null, CancellationToken.None); + await runServer.CompleteJobAsync(message.Plan.PlanId, message.JobId, TaskResult.Failed, outputs: null, stepResults: null, jobAnnotations: jobAnnotations, environmentUrl: null, telemetry: null, billingOwnerId: message.BillingOwnerId, CancellationToken.None); } catch (Exception ex) { diff --git a/src/Runner.Listener/Runner.cs b/src/Runner.Listener/Runner.cs index 94b86b9f9..cf83dedae 100644 --- a/src/Runner.Listener/Runner.cs +++ b/src/Runner.Listener/Runner.cs @@ -584,7 +584,7 @@ namespace GitHub.Runner.Listener await runServer.ConnectAsync(new Uri(messageRef.RunServiceUrl), creds); try { - jobRequestMessage = await runServer.GetJobMessageAsync(messageRef.RunnerRequestId, messageQueueLoopTokenSource.Token); + jobRequestMessage = await runServer.GetJobMessageAsync(messageRef.RunnerRequestId, messageRef.BillingOwnerId, messageQueueLoopTokenSource.Token); _acquireJobThrottler.Reset(); } catch (Exception ex) when ( diff --git a/src/Runner.Listener/RunnerJobRequestRef.cs b/src/Runner.Listener/RunnerJobRequestRef.cs index df8a4d793..a74662dd2 100644 --- a/src/Runner.Listener/RunnerJobRequestRef.cs +++ b/src/Runner.Listener/RunnerJobRequestRef.cs @@ -7,9 +7,14 @@ namespace GitHub.Runner.Listener { [DataMember(Name = "id")] public string Id { get; set; } + [DataMember(Name = "runner_request_id")] public string RunnerRequestId { get; set; } + [DataMember(Name = "run_service_url")] public string RunServiceUrl { get; set; } + + [DataMember(Name = "billing_owner_id")] + public string BillingOwnerId { get; set; } } } diff --git a/src/Runner.Worker/JobRunner.cs b/src/Runner.Worker/JobRunner.cs index 7ab506a53..7c96ef84d 100644 --- a/src/Runner.Worker/JobRunner.cs +++ b/src/Runner.Worker/JobRunner.cs @@ -318,7 +318,7 @@ namespace GitHub.Runner.Worker { try { - await runServer.CompleteJobAsync(message.Plan.PlanId, message.JobId, result, jobContext.JobOutputs, jobContext.Global.StepsResult, jobContext.Global.JobAnnotations, environmentUrl, telemetry, default); + await runServer.CompleteJobAsync(message.Plan.PlanId, message.JobId, result, jobContext.JobOutputs, jobContext.Global.StepsResult, jobContext.Global.JobAnnotations, environmentUrl, telemetry, billingOwnerId: message.BillingOwnerId, default); return result; } catch (Exception ex) diff --git a/src/Sdk/DTPipelines/Pipelines/AgentJobRequestMessage.cs b/src/Sdk/DTPipelines/Pipelines/AgentJobRequestMessage.cs index 070d86ee2..e6ecbf450 100644 --- a/src/Sdk/DTPipelines/Pipelines/AgentJobRequestMessage.cs +++ b/src/Sdk/DTPipelines/Pipelines/AgentJobRequestMessage.cs @@ -246,6 +246,13 @@ namespace GitHub.DistributedTask.Pipelines set; } + [DataMember(EmitDefaultValue = false)] + public String BillingOwnerId + { + get; + set; + } + /// /// Gets the collection of variables associated with the current context. /// diff --git a/src/Sdk/RSWebApi/Contracts/AcquireJobRequest.cs b/src/Sdk/RSWebApi/Contracts/AcquireJobRequest.cs index c075d318e..020094b91 100644 --- a/src/Sdk/RSWebApi/Contracts/AcquireJobRequest.cs +++ b/src/Sdk/RSWebApi/Contracts/AcquireJobRequest.cs @@ -10,5 +10,8 @@ namespace GitHub.Actions.RunService.WebApi [DataMember(Name = "runnerOS", EmitDefaultValue = false)] public string RunnerOS { get; set; } + + [DataMember(Name = "billingOwnerId", EmitDefaultValue = false)] + public string BillingOwnerId { get; set; } } } diff --git a/src/Sdk/RSWebApi/Contracts/CompleteJobRequest.cs b/src/Sdk/RSWebApi/Contracts/CompleteJobRequest.cs index ff02ab658..a9ba71a57 100644 --- a/src/Sdk/RSWebApi/Contracts/CompleteJobRequest.cs +++ b/src/Sdk/RSWebApi/Contracts/CompleteJobRequest.cs @@ -32,5 +32,8 @@ namespace GitHub.Actions.RunService.WebApi [DataMember(Name = "environmentUrl", EmitDefaultValue = false)] public string EnvironmentUrl { get; set; } + + [DataMember(Name = "billingOwnerId", EmitDefaultValue = false)] + public string BillingOwnerId { get; set; } } } diff --git a/src/Sdk/RSWebApi/RunServiceHttpClient.cs b/src/Sdk/RSWebApi/RunServiceHttpClient.cs index f4d0c539f..66bd08a98 100644 --- a/src/Sdk/RSWebApi/RunServiceHttpClient.cs +++ b/src/Sdk/RSWebApi/RunServiceHttpClient.cs @@ -71,13 +71,15 @@ namespace GitHub.Actions.RunService.WebApi Uri requestUri, string messageId, string runnerOS, + string billingOwnerId, CancellationToken cancellationToken = default) { HttpMethod httpMethod = new HttpMethod("POST"); var payload = new AcquireJobRequest { JobMessageId = messageId, - RunnerOS = runnerOS + RunnerOS = runnerOS, + BillingOwnerId = billingOwnerId, }; requestUri = new Uri(requestUri, "acquirejob"); @@ -128,6 +130,7 @@ namespace GitHub.Actions.RunService.WebApi IList jobAnnotations, string environmentUrl, IList telemetry, + string billingOwnerId, CancellationToken cancellationToken = default) { HttpMethod httpMethod = new HttpMethod("POST"); @@ -141,6 +144,7 @@ namespace GitHub.Actions.RunService.WebApi Annotations = jobAnnotations, EnvironmentUrl = environmentUrl, Telemetry = telemetry, + BillingOwnerId = billingOwnerId, }; requestUri = new Uri(requestUri, "completejob"); diff --git a/src/Test/L0/Listener/JobDispatcherL0.cs b/src/Test/L0/Listener/JobDispatcherL0.cs index cc50c1804..3b26233a4 100644 --- a/src/Test/L0/Listener/JobDispatcherL0.cs +++ b/src/Test/L0/Listener/JobDispatcherL0.cs @@ -36,20 +36,23 @@ namespace GitHub.Runner.Common.Tests.Listener _configurationStore = new Mock(); } - private Pipelines.AgentJobRequestMessage CreateJobRequestMessage() + private Pipelines.AgentJobRequestMessage CreateJobRequestMessage(string billingOwnerId = null) { TaskOrchestrationPlanReference plan = new(); TimelineReference timeline = null; Guid jobId = Guid.NewGuid(); var result = new Pipelines.AgentJobRequestMessage(plan, timeline, jobId, "someJob", "someJob", null, null, null, new Dictionary(), new List(), new Pipelines.JobResources(), new Pipelines.ContextData.DictionaryContextData(), new Pipelines.WorkspaceOptions(), new List(), null, null, null, null, null); result.ContextData["github"] = new Pipelines.ContextData.DictionaryContextData(); + result.BillingOwnerId = billingOwnerId; return result; } - [Fact] + [Theory] [Trait("Level", "L0")] [Trait("Category", "Runner")] - public async void DispatchesJobRequest() + [InlineData(null)] + [InlineData("billingOwnerId")] + public async void DispatchesJobRequest(string billingOwnerId) { //Arrange using (var hc = new TestHostContext(this)) @@ -65,7 +68,7 @@ namespace GitHub.Runner.Common.Tests.Listener jobDispatcher.Initialize(hc); var ts = new CancellationTokenSource(); - Pipelines.AgentJobRequestMessage message = CreateJobRequestMessage(); + Pipelines.AgentJobRequestMessage message = CreateJobRequestMessage(billingOwnerId); string strMessage = JsonUtility.ToString(message); _processInvoker.Setup(x => x.ExecuteAsync(It.IsAny(), It.IsAny(), "spawnclient 1 2", null, It.IsAny()))