Pass BillingOwnerId through Acquire/Complete calls (#3689)

* Pass BillingOwnerId through Acquire/Complete calls

* add param to test
This commit is contained in:
Luke Tomlinson
2025-02-03 15:15:54 -05:00
committed by GitHub
parent 52dc98b10f
commit a61328a7e7
10 changed files with 39 additions and 12 deletions

View File

@@ -18,7 +18,7 @@ namespace GitHub.Runner.Common
{
Task ConnectAsync(Uri serverUrl, VssCredentials credentials);
Task<AgentJobRequestMessage> GetJobMessageAsync(string id, CancellationToken token);
Task<AgentJobRequestMessage> GetJobMessageAsync(string id, string billingOwnerId, CancellationToken token);
Task CompleteJobAsync(
Guid planId,
@@ -29,6 +29,7 @@ namespace GitHub.Runner.Common
IList<Annotation> jobAnnotations,
string environmentUrl,
IList<Telemetry> telemetry,
string billingOwnerId,
CancellationToken token);
Task<RenewJobResponse> RenewJobAsync(Guid planId, Guid jobId, CancellationToken token);
@@ -58,11 +59,11 @@ namespace GitHub.Runner.Common
}
}
public Task<AgentJobRequestMessage> GetJobMessageAsync(string id, CancellationToken cancellationToken)
public Task<AgentJobRequestMessage> GetJobMessageAsync(string id, string billingOwnerId, CancellationToken cancellationToken)
{
CheckConnection();
return RetryRequest<AgentJobRequestMessage>(
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<Annotation> jobAnnotations,
string environmentUrl,
IList<Telemetry> 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<RenewJobResponse> RenewJobAsync(Guid planId, Guid jobId, CancellationToken cancellationToken)

View File

@@ -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)
{

View File

@@ -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 (

View File

@@ -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; }
}
}

View File

@@ -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)

View File

@@ -246,6 +246,13 @@ namespace GitHub.DistributedTask.Pipelines
set;
}
[DataMember(EmitDefaultValue = false)]
public String BillingOwnerId
{
get;
set;
}
/// <summary>
/// Gets the collection of variables associated with the current context.
/// </summary>

View File

@@ -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; }
}
}

View File

@@ -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; }
}
}

View File

@@ -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<Annotation> jobAnnotations,
string environmentUrl,
IList<Telemetry> 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");

View File

@@ -36,20 +36,23 @@ namespace GitHub.Runner.Common.Tests.Listener
_configurationStore = new Mock<IConfigurationStore>();
}
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<string, VariableValue>(), new List<MaskHint>(), new Pipelines.JobResources(), new Pipelines.ContextData.DictionaryContextData(), new Pipelines.WorkspaceOptions(), new List<Pipelines.ActionStep>(), 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<String>(), It.IsAny<String>(), "spawnclient 1 2", null, It.IsAny<CancellationToken>()))