Add remote persistent worker support

.github/actions/build_example_persistent_worker/action.yml

@@ -0,0 +1,16 @@
+name: build_example_persistent_worker
+inputs:
+  buildbuddyApiKey:
+    description: "The API key for BuildBuddy remote cache and execution."
+    required: true
+runs:
+  using: composite
+  steps:
+  - name: Build examples/persistent_worker directory
+    env:
+      BUILDBUDDY_API_KEY: ${{ inputs.buildbuddyApiKey }}
+    run: |-
+      cd examples/persistent_worker
+      export PATH="$RUNNER_TEMP/artifacts:$PATH"
+      ./test.sh
+    shell: bash

.github/workflows/build-and-test.yml

@@ -2,6 +2,7 @@ name: Build and test
 on:
   push:
   pull_request:
+  workflow_dispatch: # allows manual triggering
 jobs:
   linux-build-and-test:
     runs-on: 4-core-ubuntu
@@ -69,6 +70,9 @@ jobs:
         $RUNNER_TEMP/artifacts/buck2 test //... -v 2
     - uses: ./.github/actions/build_example_conan
     - uses: ./.github/actions/build_example_no_prelude
+    - uses: ./.github/actions/build_example_persistent_worker
+      with:
+        buildbuddyApiKey: ${{ secrets.BUILDBUDDY_API_KEY }}
     - uses: ./.github/actions/setup_reindeer
     - uses: ./.github/actions/build_bootstrap
   windows-build-examples:

app/buck2_action_impl/src/actions/impls/run.rs

@@ -40,6 +40,7 @@ use buck2_core::execution_types::executor_config::RemoteExecutorDependency;
 use buck2_core::fs::artifact_path_resolver::ArtifactFs;
 use buck2_core::fs::buck_out_path::BuildArtifactPath;
 use buck2_core::fs::paths::forward_rel_path::ForwardRelativePathBuf;
+use buck2_error::buck2_error;
 use buck2_error::BuckErrorContext;
 use buck2_events::dispatch::span_async_simple;
 use buck2_execute::artifact::fs::ExecutorFs;
@@ -263,6 +264,7 @@ struct UnpackedWorkerValues<'v> {
     id: WorkerId,
     concurrency: Option<usize>,
     streaming: bool,
+    supports_bazel_remote_persistent_worker_protocol: bool,
 }
 
 struct UnpackedRunActionValues<'v> {
@@ -320,6 +322,7 @@ impl RunAction {
             id: WorkerId(worker.id),
             concurrency: worker.concurrency(),
             streaming: worker.streaming(),
+            supports_bazel_remote_persistent_worker_protocol: worker.supports_bazel_remote_persistent_worker_protocol(),
         });
 
         Ok(UnpackedRunActionValues {
@@ -352,11 +355,32 @@ impl RunAction {
                 .exe
                 .add_to_command_line(&mut worker_rendered, &mut cli_ctx)?;
             worker.exe.visit_artifacts(artifact_visitor)?;
+            let worker_key = if worker.supports_bazel_remote_persistent_worker_protocol {
+                let mut worker_visitor = SimpleCommandLineArtifactVisitor::new();
+                worker.exe.visit_artifacts(&mut worker_visitor)?;
+                if !worker_visitor.outputs.is_empty() {
+                    // TODO[AH] create appropriate error enum value.
+                    return Err(buck2_error!(
+                        buck2_error::ErrorTag::Input,
+                        "remote persistent worker command should not produce an output"
+                    ));
+                }
+                let worker_inputs: Vec<&ArtifactGroupValues> = worker_visitor
+                    .inputs()
+                    .map(|group| action_execution_ctx.artifact_values(group))
+                    .collect();
+                let (_, worker_digest) =
+                    metadata_content(fs.fs(), &worker_inputs, action_execution_ctx.digest_config())?;
+                Some(worker_digest)
+            } else {
+                None
+            };
             Some(WorkerSpec {
                 exe: worker_rendered,
                 id: worker.id,
                 concurrency: worker.concurrency,
                 streaming: worker.streaming,
+                remote_key: worker_key,
             })
         } else {
             None

app/buck2_build_api/src/actions/execute/action_executor.rs

@@ -761,6 +761,7 @@ mod tests {
                 CommandGenerationOptions {
                     path_separator: PathSeparatorKind::Unix,
                     output_paths_behavior: Default::default(),
+                    use_bazel_protocol_remote_persistent_workers: false,
                 },
                 Default::default(),
             ),

app/buck2_build_api/src/interpreter/rule_defs/command_executor_config.rs

@@ -86,6 +86,7 @@ pub fn register_command_executor_config(builder: &mut GlobalsBuilder) {
     /// * `allow_hybrid_fallbacks_on_failure`: Whether to allow fallbacks when the result is failure (i.e. the command failed on the primary, but the infra worked)
     /// * `use_windows_path_separators`: Whether to use Windows path separators in command line arguments
     /// * `use_persistent workers`: Whether to use persistent workers for local execution if they are available
+    /// * `use_bazel_protocol_remote_persistent_workers`: Whether to use persistent workers for remote execution via the Bazel remote persistent worker protocol if they are available
     /// * `allow_cache_uploads`: Whether to upload local actions to the RE cache
     /// * `max_cache_upload_mebibytes`: Maximum size to upload in cache uploads
     /// * `experimental_low_pass_filter`: Whether to use the experimental low pass filter
@@ -111,6 +112,7 @@ pub fn register_command_executor_config(builder: &mut GlobalsBuilder) {
         #[starlark(default = false, require = named)] allow_hybrid_fallbacks_on_failure: bool,
         #[starlark(default = false, require = named)] use_windows_path_separators: bool,
         #[starlark(default = false, require = named)] use_persistent_workers: bool,
+        #[starlark(default = false, require = named)] use_bazel_protocol_remote_persistent_workers: bool,
         #[starlark(default = false, require = named)] allow_cache_uploads: bool,
         #[starlark(default = NoneOr::None, require = named)] max_cache_upload_mebibytes: NoneOr<
             i32,
@@ -323,6 +325,7 @@ pub fn register_command_executor_config(builder: &mut GlobalsBuilder) {
                         PathSeparatorKind::Unix
                     },
                     output_paths_behavior,
+                    use_bazel_protocol_remote_persistent_workers,
                 },
             }
         };

app/buck2_build_api/src/interpreter/rule_defs/provider/builtin/worker_info.rs

@@ -50,6 +50,8 @@ pub struct WorkerInfoGen<V: ValueLifetimeless> {
     pub concurrency: ValueOfUncheckedGeneric<V, NoneOr<usize>>,
     // Whether to always run actions using this worker via the streaming API
     pub streaming: ValueOfUncheckedGeneric<V, bool>,
+    // Bazel remote persistent worker protocol capable worker.
+    pub supports_bazel_remote_persistent_worker_protocol: ValueOfUncheckedGeneric<V, bool>,
 
     pub id: u64,
 }
@@ -68,6 +70,7 @@ fn worker_info_creator(globals: &mut GlobalsBuilder) {
             ValueOf<'v, usize>,
         >,
         #[starlark(require = named, default = NoneType)] streaming: Value<'v>,
+        #[starlark(require = named, default = false)] supports_bazel_remote_persistent_worker_protocol: bool,
         eval: &mut Evaluator<'v, '_, '_>,
     ) -> starlark::Result<WorkerInfo<'v>> {
         let heap = eval.heap();
@@ -79,6 +82,7 @@ fn worker_info_creator(globals: &mut GlobalsBuilder) {
             id,
             concurrency: heap.alloc_typed_unchecked(concurrency).cast(),
             streaming: ValueOfUnchecked::new(streaming),
+            supports_bazel_remote_persistent_worker_protocol: heap.alloc_typed_unchecked(supports_bazel_remote_persistent_worker_protocol).cast(),
         })
     }
 }
@@ -105,6 +109,13 @@ impl<'v, V: ValueLike<'v>> WorkerInfoGen<V> {
             .into_option()
             .unwrap_or(false)
     }
+
+    pub fn supports_bazel_remote_persistent_worker_protocol(&self) -> bool {
+        self.supports_bazel_remote_persistent_worker_protocol
+            .to_value()
+            .unpack()
+            .expect("validated at construction")
+    }
 }
 
 fn validate_worker_info<'v, V>(info: &WorkerInfoGen<V>) -> buck2_error::Result<()>

app/buck2_build_api_tests/src/interpreter/rule_defs/provider/builtin/worker_info.rs

@@ -24,7 +24,7 @@ fn run_display() {
         .run_starlark_bzl_test(
             r#"
 def test():
-    assert_eq('WorkerInfo(exe=cmd_args("x"), concurrency=None, streaming=None)', str(WorkerInfo(exe="x")))
+    assert_eq('WorkerInfo(exe=cmd_args("x"), concurrency=None, streaming=None, supports_bazel_remote_persistent_worker_protocol=False)', str(WorkerInfo(exe="x")))
 "#,
         )
         .unwrap();

app/buck2_core/src/execution_types/executor_config.rs

@@ -283,6 +283,7 @@ impl Default for CacheUploadBehavior {
 pub struct CommandGenerationOptions {
     pub path_separator: PathSeparatorKind,
     pub output_paths_behavior: OutputPathsBehavior,
+    pub use_bazel_protocol_remote_persistent_workers: bool,
 }
 
 #[derive(Debug, Eq, PartialEq, Hash, Allocative, Clone)]
@@ -314,6 +315,7 @@ impl CommandExecutorConfig {
             options: CommandGenerationOptions {
                 path_separator: PathSeparatorKind::system_default(),
                 output_paths_behavior: Default::default(),
+                use_bazel_protocol_remote_persistent_workers: false,
             },
         })
     }

app/buck2_execute/src/execute/command_executor.rs

@@ -20,7 +20,6 @@ use buck2_core::fs::artifact_path_resolver::ArtifactFs;
 use buck2_core::fs::project_rel_path::ProjectRelativePath;
 use buck2_core::fs::project_rel_path::ProjectRelativePathBuf;
 use buck2_directory::directory::fingerprinted_directory::FingerprintedDirectory;
-#[cfg(fbcode_build)]
 use buck2_error::buck2_error;
 use buck2_error::BuckErrorContext;
 use buck2_futures::cancellation::CancellationContext;
@@ -187,15 +186,45 @@ impl CommandExecutor {
                 }
                 CommandExecutionInput::ScratchPath(_) => None,
             });
+            let mut platform = self.0.re_platform.clone();
+            let args = if self.0.options.use_bazel_protocol_remote_persistent_workers
+                && let Some(worker) = request.worker()
+                && let Some(key) = worker.remote_key.as_ref()
+            {
+                platform.properties.push(RE::Property {
+                    name: "persistentWorkerKey".to_owned(),
+                    value: key.to_string(),
+                });
+                // TODO[AH] Ideally, Buck2 could generate an argfile on the fly.
+                for arg in request.args() {
+                    if !(arg.starts_with("@")
+                        || arg.starts_with("-flagfile")
+                        || arg.starts_with("--flagfile"))
+                    {
+                        return Err(buck2_error!(
+                            buck2_error::ErrorTag::Input,
+                            "Remote persistent worker arguments must be passed as `@argfile`, `-flagfile=argfile`, or `--flagfile=argfile`."
+                        ));
+                    }
+                }
+                worker
+                    .exe
+                    .iter()
+                    .chain(request.args().iter())
+                    .cloned()
+                    .collect()
+            } else {
+                request.all_args_vec()
+            };
             let action = re_create_action(
-                request.all_args_vec(),
+                args,
                 request.paths().output_paths(),
                 request.working_directory(),
                 request.env(),
                 input_digest,
                 action_metadata_blobs,
                 request.timeout(),
-                self.0.re_platform.clone(),
+                platform,
                 false,
                 digest_config,
                 self.0.options.output_paths_behavior,

app/buck2_execute/src/execute/request.rs

@@ -278,12 +278,13 @@ impl CommandExecutionPaths {
 #[derive(Copy, Clone, Dupe, Debug, Display, Allocative, Hash, PartialEq, Eq)]
 pub struct WorkerId(pub u64);
 
-#[derive(Clone)]
+#[derive(Clone, Debug)]
 pub struct WorkerSpec {
     pub id: WorkerId,
     pub exe: Vec<String>,
     pub concurrency: Option<usize>,
     pub streaming: bool,
+    pub remote_key: Option<TrackedFileDigest>,
 }
 
 /// The data contains the information about the command to be executed.

app/buck2_execute/src/lib.rs

@@ -14,6 +14,7 @@
 #![feature(try_trait_v2)]
 #![feature(used_with_arg)]
 #![feature(trait_upcasting)]
+#![feature(let_chains)]
 
 pub mod artifact;
 pub mod artifact_utils;

app/buck2_server/src/daemon/common.rs

@@ -500,6 +500,7 @@ pub fn get_default_executor_config(host_platform: HostPlatformOverride) -> Comma
         options: CommandGenerationOptions {
             path_separator: get_default_path_separator(host_platform),
             output_paths_behavior: Default::default(),
+            use_bazel_protocol_remote_persistent_workers: false,
         },
     }
 }

app/buck2_test/src/orchestrator.rs

@@ -1191,6 +1191,7 @@ impl<'b> BuckTestOrchestrator<'b> {
             options: CommandGenerationOptions {
                 path_separator: PathSeparatorKind::system_default(),
                 output_paths_behavior: Default::default(),
+                use_bazel_protocol_remote_persistent_workers: false,
             },
         };
         let CommandExecutorResponse {
@@ -1733,6 +1734,7 @@ impl<'a> Execute2RequestExpander<'a> {
                     id: WorkerId(worker.id),
                     concurrency: worker.concurrency(),
                     streaming: worker.streaming(),
+                    remote_key: None,
                 })
             }
             _ => None,

examples/persistent_worker/.buckconfig

@@ -0,0 +1,17 @@
+[cells]
+root = .
+prelude = prelude
+toolchains = toolchains
+none = none
+
+[cell_aliases]
+config = prelude
+fbcode = none
+fbsource = none
+buck = none
+
+[external_cells]
+prelude = bundled
+
+[parser]
+target_platform_detector_spec = target:root//...->prelude//platforms:default

examples/persistent_worker/.buckconfig.buildbuddy

@@ -0,0 +1,13 @@
+[buck2]
+digest_algorithms = SHA256
+
+[buck2_re_client]
+engine_address       = grpc://remote.buildbuddy.io
+action_cache_address = grpc://remote.buildbuddy.io
+cas_address          = grpc://remote.buildbuddy.io
+tls                  = true
+http_headers         = \
+  x-buildbuddy-api-key:$BUILDBUDDY_API_KEY
+
+[build]
+execution_platforms = root//platforms:buildbuddy

examples/persistent_worker/.buckconfig.buildbuddy-persistent-workers

@@ -0,0 +1,13 @@
+[buck2]
+digest_algorithms = SHA256
+
+[buck2_re_client]
+engine_address       = grpc://remote.buildbuddy.io
+action_cache_address = grpc://remote.buildbuddy.io
+cas_address          = grpc://remote.buildbuddy.io
+tls                  = true
+http_headers         = \
+  x-buildbuddy-api-key:$BUILDBUDDY_API_KEY
+
+[build]
+execution_platforms = root//platforms:buildbuddy-persistent-workers

examples/persistent_worker/.buckconfig.local-persistent-workers

@@ -0,0 +1,2 @@
+[build]
+execution_platforms = root//platforms:local-persistent-workers

examples/persistent_worker/.buckconfig.no-workers

@@ -0,0 +1,2 @@
+[build]
+execution_platforms = root//platforms:local

examples/persistent_worker/.envrc

@@ -0,0 +1,3 @@
+# specify the following:
+# - BUILDBUDDY_API_KEY
+source_env_if_exists .envrc.private

examples/persistent_worker/.gitignore

@@ -0,0 +1,4 @@
+.buckconfig.local
+.direnv
+.envrc.private
+prelude

examples/persistent_worker/BUCK

@@ -0,0 +1,26 @@
+load("defs.bzl", "demo", "worker")
+
+python_binary(
+    name = "one_shot",
+    main = "one_shot.py",
+)
+
+python_binary(
+    name = "worker_py",
+    main = "persistent_worker.py",
+    deps = [
+        "//proto/bazel:worker_protocol_pb2",
+        "//proto/buck2:worker_pb2",
+    ],
+)
+
+worker(
+    name = "worker",
+    visibility = ["PUBLIC"],
+    worker = ":worker_py",
+)
+
+[
+    demo(name = "demo-" + str(i))
+    for i in range(4)
+]