Add remote persistent worker support (#787)

Summary:
Closes #776.

Implements support for persistent workers in remote builds using the Bazel remote execution protocol and the approach documented in the Bazel remote persistent workers proposal:
https://github.com/bazelbuild/proposals/blob/main/designs/2021-03-06-remote-persistent-workers.md

The Bazel remote persistent worker protocol includes an automatic fallback in cases where the remote execution system does not yet support persistent workers. To that end actions take the shape
```
WORKER WORKER_ARGS... REQUEST_ARGS_FILE
```
The remote execution system separates worker arguments on the command-line from request arguments in the response file and adds the `--persistent_worker` flag.

Pull Request resolved: #787

Test Plan: Test added in D68157749.

Reviewed By: IanChilds

Differential Revision: D65213349

Pulled By: KapJI

fbshipit-source-id: 79d219a06c08fad7d16976c6c9e5a6815dd073df

Author: aherrmann

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,8 @@ 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 +356,35 @@ 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::ActionMismatchedOutputs,
+                        "Remote persistent worker command should not produce outputs."
+                    ));
+                }
+                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,8 @@ 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 +83,9 @@ 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 +112,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,
         },
     }
 }