Skip to content

Commit 2281642

Browse files
EdSchoutenEdSchouten
committed
Decompose the remote worker client
Just like the remote execution client, make it so that the base implementation is only capable of sending byte slices across. We move the Protobuf message handling into a decorator for the Executor. This does require us to make a minor change to the API of Executor, as it should now be possible for the Executor to reject actions as well. We solve this by adding an additional 'error' return value.
1 parent 31b00e1 commit 2281642

7 files changed

Lines changed: 228 additions & 127 deletions

File tree

‎cmd/bonanza_builder/main.go‎

Lines changed: 18 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -151,7 +151,7 @@ func main() {
151151
}
152152
client, err := remoteworker.NewClient(
153153
remoteWorkerClient,
154-
executor,
154+
remoteworker.NewProtoExecutor(executor),
155155
clock.SystemClock,
156156
random.CryptoThreadSafeGenerator,
157157
platformPrivateKeys,
@@ -282,14 +282,14 @@ func (e *builderExecutor) CheckReadiness(ctx context.Context) error {
282282
return nil
283283
}
284284

285-
func (e *builderExecutor) Execute(ctx context.Context, action *model_build_pb.Action, executionTimeout time.Duration, executionEvents chan<- proto.Message) (proto.Message, time.Duration, remoteworker_pb.CurrentState_Completed_Result) {
285+
func (e *builderExecutor) Execute(ctx context.Context, action *model_build_pb.Action, executionTimeout time.Duration, executionEvents chan<- *emptypb.Empty) (*model_build_pb.Result, time.Duration, remoteworker_pb.CurrentState_Completed_Result, error) {
286286
namespace, err := object.NewNamespace(action.Namespace)
287287
if err != nil {
288288
return &model_build_pb.Result{
289289
Failure: &model_build_pb.Result_Failure{
290290
Status: status.Convert(util.StatusWrap(err, "Invalid namespace")).Proto(),
291291
},
292-
}, 0, remoteworker_pb.CurrentState_Completed_FAILED
292+
}, 0, remoteworker_pb.CurrentState_Completed_FAILED, nil
293293
}
294294
instanceName := namespace.InstanceName
295295
buildSpecificationReference, err := model_core.NewDecodableLocalReferenceFromWeakProto(namespace.ReferenceFormat, action.BuildSpecificationReference)
@@ -298,7 +298,7 @@ func (e *builderExecutor) Execute(ctx context.Context, action *model_build_pb.Ac
298298
Failure: &model_build_pb.Result_Failure{
299299
Status: status.Convert(util.StatusWrap(err, "Invalid build specification reference")).Proto(),
300300
},
301-
}, 0, remoteworker_pb.CurrentState_Completed_FAILED
301+
}, 0, remoteworker_pb.CurrentState_Completed_FAILED, nil
302302
}
303303
buildSpecificationEncoder, err := encoding.NewBinaryEncoderFromProto(
304304
action.BuildSpecificationEncoders,
@@ -309,7 +309,7 @@ func (e *builderExecutor) Execute(ctx context.Context, action *model_build_pb.Ac
309309
Failure: &model_build_pb.Result_Failure{
310310
Status: status.Convert(util.StatusWrap(err, "Invalid build specification encoder")).Proto(),
311311
},
312-
}, 0, remoteworker_pb.CurrentState_Completed_FAILED
312+
}, 0, remoteworker_pb.CurrentState_Completed_FAILED, nil
313313
}
314314

315315
// Perform the build.
@@ -415,7 +415,7 @@ func (e *builderExecutor) Execute(ctx context.Context, action *model_build_pb.Ac
415415
Failure: &model_build_pb.Result_Failure{
416416
Status: status.Convert(err).Proto(),
417417
},
418-
}, 0, remoteworker_pb.CurrentState_Completed_FAILED
418+
}, 0, remoteworker_pb.CurrentState_Completed_FAILED, nil
419419
}
420420
patcher := key.Patcher
421421

@@ -434,7 +434,7 @@ func (e *builderExecutor) Execute(ctx context.Context, action *model_build_pb.Ac
434434
Failure: &model_build_pb.Result_Failure{
435435
Status: status.Convert(err).Proto(),
436436
},
437-
}, 0, remoteworker_pb.CurrentState_Completed_FAILED
437+
}, 0, remoteworker_pb.CurrentState_Completed_FAILED, nil
438438
}
439439
patcher.Merge(value.Patcher)
440440
}
@@ -472,7 +472,7 @@ func (e *builderExecutor) Execute(ctx context.Context, action *model_build_pb.Ac
472472
Failure: &model_build_pb.Result_Failure{
473473
Status: status.Convert(err).Proto(),
474474
},
475-
}, 0, remoteworker_pb.CurrentState_Completed_FAILED
475+
}, 0, remoteworker_pb.CurrentState_Completed_FAILED, nil
476476
}
477477
if err := dependencyTreeBuilder.PushChild(
478478
model_core.NewPatchedMessage(
@@ -488,7 +488,7 @@ func (e *builderExecutor) Execute(ctx context.Context, action *model_build_pb.Ac
488488
Failure: &model_build_pb.Result_Failure{
489489
Status: status.Convert(err).Proto(),
490490
},
491-
}, 0, remoteworker_pb.CurrentState_Completed_FAILED
491+
}, 0, remoteworker_pb.CurrentState_Completed_FAILED, nil
492492
}
493493
}
494494
dependencies, err := dependencyTreeBuilder.FinalizeList()
@@ -497,7 +497,7 @@ func (e *builderExecutor) Execute(ctx context.Context, action *model_build_pb.Ac
497497
Failure: &model_build_pb.Result_Failure{
498498
Status: status.Convert(err).Proto(),
499499
},
500-
}, 0, remoteworker_pb.CurrentState_Completed_FAILED
500+
}, 0, remoteworker_pb.CurrentState_Completed_FAILED, nil
501501
}
502502
patcher.Merge(dependencies.Patcher)
503503

@@ -519,7 +519,7 @@ func (e *builderExecutor) Execute(ctx context.Context, action *model_build_pb.Ac
519519
Failure: &model_build_pb.Result_Failure{
520520
Status: status.Convert(err).Proto(),
521521
},
522-
}, 0, remoteworker_pb.CurrentState_Completed_FAILED
522+
}, 0, remoteworker_pb.CurrentState_Completed_FAILED, nil
523523
}
524524
}
525525
evaluations, err := evaluationTreeBuilder.FinalizeList()
@@ -528,7 +528,7 @@ func (e *builderExecutor) Execute(ctx context.Context, action *model_build_pb.Ac
528528
Failure: &model_build_pb.Result_Failure{
529529
Status: status.Convert(err).Proto(),
530530
},
531-
}, 0, remoteworker_pb.CurrentState_Completed_FAILED
531+
}, 0, remoteworker_pb.CurrentState_Completed_FAILED, nil
532532
}
533533
var evaluationsReference *model_core_pb.WeakDecodableReference
534534
if len(evaluations.Message) > 0 {
@@ -542,7 +542,7 @@ func (e *builderExecutor) Execute(ctx context.Context, action *model_build_pb.Ac
542542
Failure: &model_build_pb.Result_Failure{
543543
Status: status.Convert(err).Proto(),
544544
},
545-
}, 0, remoteworker_pb.CurrentState_Completed_FAILED
545+
}, 0, remoteworker_pb.CurrentState_Completed_FAILED, nil
546546
}
547547
createdEvaluationsReference := createdEvaluations.Value.GetLocalReference()
548548
if err := dag.UploadDAG(
@@ -562,7 +562,7 @@ func (e *builderExecutor) Execute(ctx context.Context, action *model_build_pb.Ac
562562
Failure: &model_build_pb.Result_Failure{
563563
Status: status.Convert(errCompute).Proto(),
564564
},
565-
}, 0, remoteworker_pb.CurrentState_Completed_FAILED
565+
}, 0, remoteworker_pb.CurrentState_Completed_FAILED, nil
566566
}
567567
evaluationsReference = model_core.DecodableLocalReferenceToWeakProto(
568568
model_core.CopyDecodable(
@@ -583,7 +583,7 @@ func (e *builderExecutor) Execute(ctx context.Context, action *model_build_pb.Ac
583583
Status: status.Convert(err).Proto(),
584584
},
585585
EvaluationsReference: evaluationsReference,
586-
}, 0, remoteworker_pb.CurrentState_Completed_FAILED
586+
}, 0, remoteworker_pb.CurrentState_Completed_FAILED, nil
587587
}
588588
marshaledKey, err := model_core.MarshalTopLevelMessage(keyAny)
589589
if err != nil {
@@ -592,7 +592,7 @@ func (e *builderExecutor) Execute(ctx context.Context, action *model_build_pb.Ac
592592
Status: status.Convert(err).Proto(),
593593
},
594594
EvaluationsReference: evaluationsReference,
595-
}, 0, remoteworker_pb.CurrentState_Completed_FAILED
595+
}, 0, remoteworker_pb.CurrentState_Completed_FAILED, nil
596596
}
597597
marshaledStackTraceKeys = append(marshaledStackTraceKeys, marshaledKey)
598598
}
@@ -602,12 +602,12 @@ func (e *builderExecutor) Execute(ctx context.Context, action *model_build_pb.Ac
602602
Status: status.Convert(errCompute).Proto(),
603603
},
604604
EvaluationsReference: evaluationsReference,
605-
}, 0, remoteworker_pb.CurrentState_Completed_FAILED
605+
}, 0, remoteworker_pb.CurrentState_Completed_FAILED, nil
606606
}
607607
return &model_build_pb.Result{
608608
Failure: &model_build_pb.Result_Failure{
609609
Status: status.Newf(codes.Internal, "TODO: %s", value).Proto(),
610610
},
611611
EvaluationsReference: evaluationsReference,
612-
}, 0, remoteworker_pb.CurrentState_Completed_FAILED
612+
}, 0, remoteworker_pb.CurrentState_Completed_FAILED, nil
613613
}

‎cmd/bonanza_worker/main.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -199,7 +199,7 @@ func main() {
199199

200200
client, err := remoteworker.NewClient(
201201
schedulerClient,
202-
executor,
202+
remoteworker.NewProtoExecutor(executor),
203203
clock.SystemClock,
204204
random.CryptoThreadSafeGenerator,
205205
platformPrivateKeys,

‎pkg/model/command/local_executor.go‎

Lines changed: 16 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -128,7 +128,7 @@ func NewLocalExecutor(
128128
environmentVariables map[string]string,
129129
buildDirectoryOwnerUserID uint32,
130130
buildDirectoryOwnerGroupID uint32,
131-
) remoteworker.Executor[*model_command_pb.Action] {
131+
) remoteworker.Executor[*model_command_pb.Action, proto.Message, *model_command_pb.Result] {
132132
return &localExecutor{
133133
objectDownloader: objectDownloader,
134134
parsedObjectPool: parsedObjectPool,
@@ -196,12 +196,12 @@ func captureLog(ctx context.Context, buildDirectory virtual.PrepopulatedDirector
196196
return model_core.PatchedMessage[*model_filesystem_pb.FileContents, dag.ObjectContentsWalker]{}, status.Error(codes.InvalidArgument, "File is of an incorrect type")
197197
}
198198

199-
func (e *localExecutor) Execute(ctx context.Context, action *model_command_pb.Action, executionTimeout time.Duration, executionEvents chan<- proto.Message) (proto.Message, time.Duration, remoteworker_pb.CurrentState_Completed_Result) {
199+
func (e *localExecutor) Execute(ctx context.Context, action *model_command_pb.Action, executionTimeout time.Duration, executionEvents chan<- proto.Message) (*model_command_pb.Result, time.Duration, remoteworker_pb.CurrentState_Completed_Result, error) {
200200
namespace, err := object.NewNamespace(action.Namespace)
201201
if err != nil {
202202
return &model_command_pb.Result{
203203
Status: status.Convert(util.StatusWrap(err, "Invalid namespace")).Proto(),
204-
}, 0, remoteworker_pb.CurrentState_Completed_FAILED
204+
}, 0, remoteworker_pb.CurrentState_Completed_FAILED, nil
205205
}
206206
parsedObjectPoolIngester := model_parser.NewParsedObjectPoolIngester(
207207
e.parsedObjectPool,
@@ -219,7 +219,7 @@ func (e *localExecutor) Execute(ctx context.Context, action *model_command_pb.Ac
219219
if err != nil {
220220
return &model_command_pb.Result{
221221
Status: status.Convert(util.StatusWrap(err, "Invalid command encoders")).Proto(),
222-
}, 0, remoteworker_pb.CurrentState_Completed_FAILED
222+
}, 0, remoteworker_pb.CurrentState_Completed_FAILED, nil
223223
}
224224
commandReader := model_parser.LookupParsedObjectReader[object.LocalReference](
225225
parsedObjectPoolIngester,
@@ -233,13 +233,13 @@ func (e *localExecutor) Execute(ctx context.Context, action *model_command_pb.Ac
233233
if err != nil {
234234
return &model_command_pb.Result{
235235
Status: status.Convert(util.StatusWrap(err, "Invalid command reference")).Proto(),
236-
}, 0, remoteworker_pb.CurrentState_Completed_FAILED
236+
}, 0, remoteworker_pb.CurrentState_Completed_FAILED, nil
237237
}
238238
command, err := commandReader.ReadParsedObject(ctx, commandReference)
239239
if err != nil {
240240
return &model_command_pb.Result{
241241
Status: status.Convert(util.StatusWrap(err, "Failed to read command")).Proto(),
242-
}, 0, remoteworker_pb.CurrentState_Completed_FAILED
242+
}, 0, remoteworker_pb.CurrentState_Completed_FAILED, nil
243243
}
244244

245245
// Convert arguments and environment variables stored in B-trees
@@ -266,14 +266,14 @@ func (e *localExecutor) Execute(ctx context.Context, action *model_command_pb.Ac
266266
if !ok {
267267
return &model_command_pb.Result{
268268
Status: status.New(codes.InvalidArgument, "Invalid leaf element in arguments").Proto(),
269-
}, 0, remoteworker_pb.CurrentState_Completed_FAILED
269+
}, 0, remoteworker_pb.CurrentState_Completed_FAILED, nil
270270
}
271271
arguments = append(arguments, level.Leaf)
272272
}
273273
if errIter != nil {
274274
return &model_command_pb.Result{
275275
Status: status.Convert(util.StatusWrap(err, "Failed to iterate arguments")).Proto(),
276-
}, 0, remoteworker_pb.CurrentState_Completed_FAILED
276+
}, 0, remoteworker_pb.CurrentState_Completed_FAILED, nil
277277
}
278278

279279
environmentVariables := map[string]string{}
@@ -296,14 +296,14 @@ func (e *localExecutor) Execute(ctx context.Context, action *model_command_pb.Ac
296296
if !ok {
297297
return &model_command_pb.Result{
298298
Status: status.New(codes.InvalidArgument, "Invalid leaf entry in environment variables").Proto(),
299-
}, 0, remoteworker_pb.CurrentState_Completed_FAILED
299+
}, 0, remoteworker_pb.CurrentState_Completed_FAILED, nil
300300
}
301301
environmentVariables[level.Leaf.Name] = level.Leaf.Value
302302
}
303303
if errIter != nil {
304304
return &model_command_pb.Result{
305305
Status: status.Convert(util.StatusWrap(err, "Failed to iterate environment variables")).Proto(),
306-
}, 0, remoteworker_pb.CurrentState_Completed_FAILED
306+
}, 0, remoteworker_pb.CurrentState_Completed_FAILED, nil
307307
}
308308

309309
// Error logger to terminate execution and capture I/O error events.
@@ -318,7 +318,7 @@ func (e *localExecutor) Execute(ctx context.Context, action *model_command_pb.Ac
318318
if err != nil {
319319
return &model_command_pb.Result{
320320
Status: status.Convert(util.StatusWrap(err, "Invalid file creation parameters")).Proto(),
321-
}, 0, remoteworker_pb.CurrentState_Completed_FAILED
321+
}, 0, remoteworker_pb.CurrentState_Completed_FAILED, nil
322322
}
323323
directoryCreationParameters, err := model_filesystem.NewDirectoryCreationParametersFromProto(
324324
command.Message.DirectoryCreationParameters,
@@ -327,7 +327,7 @@ func (e *localExecutor) Execute(ctx context.Context, action *model_command_pb.Ac
327327
if err != nil {
328328
return &model_command_pb.Result{
329329
Status: status.Convert(util.StatusWrap(err, "Invalid file creation parameters")).Proto(),
330-
}, 0, remoteworker_pb.CurrentState_Completed_FAILED
330+
}, 0, remoteworker_pb.CurrentState_Completed_FAILED, nil
331331
}
332332
directoryEncoder := directoryCreationParameters.GetEncoder()
333333

@@ -356,7 +356,7 @@ func (e *localExecutor) Execute(ctx context.Context, action *model_command_pb.Ac
356356
if err != nil {
357357
return &model_command_pb.Result{
358358
Status: status.Convert(util.StatusWrap(err, "Invalid input root reference")).Proto(),
359-
}, 0, remoteworker_pb.CurrentState_Completed_FAILED
359+
}, 0, remoteworker_pb.CurrentState_Completed_FAILED, nil
360360
}
361361
if err := buildDirectory.CreateChildren(map[path.Component]virtual.InitialChild{
362362
inputRootDirectoryComponent: virtual.InitialChild{}.FromDirectory(
@@ -408,7 +408,7 @@ func (e *localExecutor) Execute(ctx context.Context, action *model_command_pb.Ac
408408
}, false); err != nil {
409409
return &model_command_pb.Result{
410410
Status: status.Convert(util.StatusWrap(err, "Failed to create initial children of build directory")).Proto(),
411-
}, 0, remoteworker_pb.CurrentState_Completed_FAILED
411+
}, 0, remoteworker_pb.CurrentState_Completed_FAILED, nil
412412
}
413413

414414
// If the command requires a stable input root path, we should
@@ -425,7 +425,7 @@ func (e *localExecutor) Execute(ctx context.Context, action *model_command_pb.Ac
425425
if err := e.topLevelDirectory.AddChild(ctx, buildDirectoryName, virtual.DirectoryChild{}.FromDirectory(buildDirectory)); err != nil {
426426
return &model_command_pb.Result{
427427
Status: status.Convert(util.StatusWrap(err, "Failed to attach build directory")).Proto(),
428-
}, 0, remoteworker_pb.CurrentState_Completed_FAILED
428+
}, 0, remoteworker_pb.CurrentState_Completed_FAILED, nil
429429
}
430430
defer e.topLevelDirectory.RemoveChild(buildDirectoryName)
431431

@@ -595,7 +595,7 @@ func (e *localExecutor) Execute(ctx context.Context, action *model_command_pb.Ac
595595
setError(util.StatusWrap(err, "Failed to marshal outputs"))
596596
}
597597
}
598-
return resultMessage, virtualExecutionDuration, resultCode
598+
return resultMessage, virtualExecutionDuration, resultCode, nil
599599
}
600600

601601
type prepopulatedCapturableDirectoryOptions struct {

‎pkg/remoteworker/BUILD.bazel‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ go_library(
55
srcs = [
66
"client.go",
77
"executor.go",
8+
"proto_executor.go",
89
],
910
importpath = "bonanza.build/pkg/remoteworker",
1011
visibility = ["//visibility:public"],
@@ -25,5 +26,6 @@ go_library(
2526
"@org_golang_google_protobuf//types/known/anypb",
2627
"@org_golang_google_protobuf//types/known/durationpb",
2728
"@org_golang_google_protobuf//types/known/emptypb",
29+
"@org_golang_x_sync//errgroup",
2830
],
2931
)

0 commit comments

Comments
 (0)