Add test for persistent remote workers on BuildBuddy

Summary:
Part of #787

Includes an example setup that works with
- local builds without persistent worker
- local builds with persistent worker (Buck2 protocol)
- remote builds without persistent worker

The demo worker included in the example in this PR distinguishes between Buck2 worker, Bazel remote worker, and one-shot modes depending on whether Buck2's WORKER_SOCKET, Bazel's --persistent_worker flag, or neither is set.

The example includes a README with detailed instructions how to test this feature.
- remote builds with persistent worker (Bazel protocol)

Reviewed By: scottcao

Differential Revision: D68157749

fbshipit-source-id: 51e2e247c75e0ca9736ddc0a5f383e662edee298

.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
@@ -51,7 +52,7 @@ jobs:
     - uses: ./.github/actions/setup_reindeer
     - uses: ./.github/actions/build_bootstrap
   linux-build-examples:
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-22.04
     steps:
     - uses: actions/[email protected]
     - uses: ./.github/actions/setup_linux_env
@@ -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:

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,28 @@
+load("defs.bzl", "demo", "worker")
+
+oncall("build_infra")
+
+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)
+]

examples/persistent_worker/README.md

@@ -0,0 +1,145 @@
+# Persistent Worker Demo
+
+At the time of writing (2024-09-25) Buck2 supports persistent workers for local
+builds through a dedicated Buck2 persistent worker gRPC protocol. However, Buck2
+does not support persistent workers for builds that use remote execution. This
+demo is part of a patch-set that adds support for remote persistent workers to
+Buck2, see [#776].
+
+[#776]: https://github.com/facebook/buck2/issues/776
+
+## Requirements
+
+This demo uses BuildBuddy remote execution to demonstrate remote persistent
+workers. You will need an API token for at least a free open source account. You
+can use [direnv] to set up the environment:
+
+Credentials for [BuildBuddy] stored in `.envrc.private`:
+
+```
+export BUILDBUDDY_API_KEY=...
+```
+
+On CI the API key is not available for pipelines initiated from forks of the
+main Buck2 repository. The corresponding tests will be skipped in that case. A
+Meta engineer can manually initiate a pipeline run with the token set.
+
+[direnv]: https://direnv.net/
+[BuildBuddy]: https://www.buildbuddy.io/
+
+## Local Build
+
+Configure a local build without persistent workers:
+
+```
+$ cd examples/persistent_worker
+$ echo '<file:.buckconfig.no-workers>' > .buckconfig.local
+```
+
+Run a clean build:
+
+```
+$ buck2 clean; buck2 build : -vstderr
+...
+stderr for root//:demo-7 (demo):
+...
+ONE-SHOT START
+...
+```
+
+## Local Persistent Worker
+
+Configure a local build with persistent workers:
+
+```
+$ cd examples/persistent_worker
+$ echo '<file:.buckconfig.local-persistent-workers>' > .buckconfig.local
+```
+
+Run a clean build:
+
+```
+$ buck2 clean; buck2 build : -vstderr
+...
+stderr for root//:demo-7 (demo):
+...
+Buck2 persistent worker ...
+...
+```
+
+## Remote Execution
+
+Configure a remote build without persistent workers:
+
+```
+$ cd examples/persistent_worker
+$ echo '<file:.buckconfig.buildbuddy>' > .buckconfig.local
+```
+
+Run a clean build:
+
+```
+$ buck2 clean; buck2 build : -vstderr
+...
+stderr for root//:demo-7 (demo):
+...
+ONE-SHOT START
+...
+```
+
+## Remote Persistent Worker
+
+Configure a remote build with persistent workers:
+
+```
+$ cd examples/persistent_worker
+$ echo '<file:.buckconfig.buildbuddy-persistent-workers>' > .buckconfig.local
+```
+
+Run a clean build:
+
+```
+$ buck2 clean; buck2 build : -vstderr
+...
+stderr for root//:demo-7 (demo):
+...
+Bazel persistent worker ...
+...
+```
+
+## Protocol
+
+### Starlark
+
+A Buck2 persistent worker is created by a rule that emits the `WorkerInfo`
+provider. Setting `remote = True` on this provider indicates that this worker is
+remote execution capable.
+
+Buck2 actions indicate that they can utilize a persistent worker by setting the
+`exe` parameter to `ctx.actions.run` to `WorkerRunInfo(worker, exe)`, where
+`worker` is a `WorkerInfo` provider, and `exe` defines the fallback executable
+for non persistent-worker execution.
+
+Buck2 actions that want to utilize a remote persistent worker must pass
+command-line arguments in an argument file specified as `@argfile`,
+`-flagfile=argfile`, or `--flagfile=argfile` on the command-line.
+
+### Local Persistent Worker
+
+A locally executed Buck2 persistent worker falls under the
+[Buck2 persistent worker protocol](./proto/buck2/worker.proto): It is started
+and managed by Buck2 and passed a file path in the `WORKER_SOCKET` environment
+variable where it should create a gRPC Unix domain socket to serve worker
+requests over. Multiple requests may be sent in parallel and expected to be
+served at the same time depending on the `concurrency` attribute of the
+`WorkerInfo` provider.
+
+### Remote Persistent Worker
+
+A remotely executed Buck2 persistent worker falls under the
+[Bazel persistent worker protocol](./proto/bazel/worker_protocol.proto): It is
+started and managed by the remote execution system. Work requests are sent as
+length prefixed protobuf objects to the standard input of the worker process.
+Work responses are expected as length prefixed protobuf objects on the standard
+output of the worker process. The worker process may not use standard output for
+anything else.

examples/persistent_worker/defs.bzl

@@ -0,0 +1,63 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under both the MIT license found in the
+# LICENSE-MIT file in the root directory of this source tree and the Apache
+# License, Version 2.0 found in the LICENSE-APACHE file in the root directory
+# of this source tree.
+
+load("@prelude//utils:argfile.bzl", "at_argfile")
+
+def _worker_impl(ctx: AnalysisContext) -> list[Provider]:
+    return [
+        DefaultInfo(),
+        WorkerInfo(
+            exe = ctx.attrs.worker[RunInfo].args,
+            concurrency = None,
+            supports_bazel_remote_persistent_worker_protocol = True,
+        ),
+    ]
+
+worker = rule(
+    impl = _worker_impl,
+    attrs = {
+        "worker": attrs.dep(providers = [RunInfo]),
+    },
+)
+
+def _demo_impl(ctx: AnalysisContext) -> list[Provider]:
+    output = ctx.actions.declare_output(ctx.label.name)
+    argfile = at_argfile(
+        actions = ctx.actions,
+        name = "demo." + ctx.label.name + ".args",
+        args = cmd_args(output.as_output()),
+    )
+    ctx.actions.run(
+        cmd_args(argfile),
+        category = "demo",
+        env = {
+            # modify this value to force an action rerun even if caching is enabled.
+            # `--no-remote-cache` does not have the desired effect, because it also causes
+            # the action to be omitted from `buck2 log what-ran`, which interferes with the
+            # test setup.
+            "CACHE_SILO_KEY": read_root_config("build", "cache_silo_key", "0"),
+        },
+        exe = WorkerRunInfo(
+            worker = ctx.attrs._worker[WorkerInfo],
+            exe = ctx.attrs._one_shot[RunInfo].args,
+        ),
+    )
+    return [DefaultInfo(default_output = output)]
+
+demo = rule(
+    impl = _demo_impl,
+    attrs = {
+        "_one_shot": attrs.exec_dep(
+            default = "//:one_shot",
+            providers = [RunInfo],
+        ),
+        "_worker": attrs.exec_dep(
+            default = "//:worker",
+            providers = [WorkerInfo],
+        ),
+    },
+)

examples/persistent_worker/one_shot.py

@@ -0,0 +1,31 @@
+#!/usr/bin/env python3
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under both the MIT license found in the
+# LICENSE-MIT file in the root directory of this source tree and the Apache
+# License, Version 2.0 found in the LICENSE-APACHE file in the root directory
+# of this source tree.
+
+import argparse
+import os
+import sys
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        fromfile_prefix_chars="@", prog="one_shot", description="One-shot command"
+    )
+    parser.add_argument("outfile", type=argparse.FileType("w"), help="Output file.")
+
+    args = parser.parse_args()
+
+    print("one-shot.py", file=sys.stderr)
+    print("ONE-SHOT START", file=sys.stderr)
+    name = os.path.basename(args.outfile.name)
+    args.outfile.write(name + "\n")
+    args.outfile.close()
+    print("ONE-SHOT END", file=sys.stderr)
+
+
+if __name__ == "__main__":
+    main()