add the worker as a toolchain

Author: nikhilm

rust/private/rust.bzl

@@ -370,13 +370,6 @@ _rust_common_attrs = {
         allow_single_file = True,
         cfg = "exec",
     ),
-    "_process_wrapper_worker": attr.label(
-        default = "@rustc_worker//file",
-        # Actually we may want the persistent worker to be built by the rust rules themselves, but right now we are just going to use a binary.
-        executable = True,
-        allow_single_file = True,
-        cfg = "exec",
-    ),
 }
 
 _rust_library_attrs = {
@@ -412,6 +405,7 @@ rust_library = rule(
     toolchains = [
         "@io_bazel_rules_rust//rust:toolchain",
         "@bazel_tools//tools/cpp:toolchain_type",
+        "@io_bazel_rules_rust//worker:toolchain_type",
     ],
     doc = """
 Builds a Rust library crate.
@@ -502,6 +496,7 @@ rust_binary = rule(
     toolchains = [
         "@io_bazel_rules_rust//rust:toolchain",
         "@bazel_tools//tools/cpp:toolchain_type",
+        "@io_bazel_rules_rust//worker:toolchain_type",
     ],
     doc = """
 Builds a Rust binary crate.
@@ -600,6 +595,7 @@ rust_test = rule(
     toolchains = [
         "@io_bazel_rules_rust//rust:toolchain",
         "@bazel_tools//tools/cpp:toolchain_type",
+        "@io_bazel_rules_rust//worker:toolchain_type",
     ],
     doc = """
 Builds a Rust test crate.
@@ -748,6 +744,7 @@ rust_test_binary = rule(
     toolchains = [
         "@io_bazel_rules_rust//rust:toolchain",
         "@bazel_tools//tools/cpp:toolchain_type",
+        "@io_bazel_rules_rust//worker:toolchain_type",
     ],
     doc = """
 Builds a Rust test binary, without marking this rule as a Bazel test.
@@ -771,6 +768,7 @@ rust_benchmark = rule(
     toolchains = [
         "@io_bazel_rules_rust//rust:toolchain",
         "@bazel_tools//tools/cpp:toolchain_type",
+        "@io_bazel_rules_rust//worker:toolchain_type",
     ],
     doc = """
 Builds a Rust benchmark test.

rust/private/rustc.bzl

@@ -306,7 +306,8 @@ def construct_arguments(
         out_dir,
         build_env_file,
         build_flags_files,
-        maker_path):
+        maker_path,
+        use_worker=False):
     output_dir = getattr(crate_info.output, "dirname") if hasattr(crate_info.output, "dirname") else None
 
     linker_script = getattr(file, "linker_script") if hasattr(file, "linker_script") else None
@@ -315,8 +316,10 @@ def construct_arguments(
 
     # Wrapper args first
     args = ctx.actions.args()
-    args.set_param_file_format("multiline")
-    args.use_param_file("@%s", use_always=True)
+    if use_worker:
+        # Write the args to a param file that will be used by Bazel to send messages to the worker.
+        args.set_param_file_format("multiline")
+        args.use_param_file("@%s", use_always=True)
 
     if build_env_file != None:
         args.add("--env-file", build_env_file)
@@ -443,6 +446,8 @@ def rustc_compile_action(
                      - DepInfo: The transitive dependencies of this crate.
                      - DefaultInfo: The output file for this crate, and its runfiles.
     """
+    worker_binary = ctx.toolchains["@io_bazel_rules_rust//worker:toolchain_type"].worker_binary
+
     dep_info, build_info = collect_deps(
         ctx.label,
         crate_info.deps,
@@ -478,22 +483,36 @@ def rustc_compile_action(
         build_env_file,
         build_flags_files,
         maker_path = None,
+        use_worker = worker_binary != None,
     )
 
     if hasattr(ctx.attr, "version") and ctx.attr.version != "0.0.0":
         formatted_version = " v{}".format(ctx.attr.version)
     else:
         formatted_version = ""
 
+    if worker_binary != None:
+        executable = worker_binary
+        tools = [ctx.executable._process_wrapper]
+        arguments = [ctx.executable._process_wrapper.path, toolchain.rustc.path, ctx.var["COMPILATION_MODE"], args]
+        execution_requirements = { "supports-workers": "1" }
+    else:
+        # Not all execution platforms support a worker.
+        executable = ctx.executable._process_wrapper
+        tools = []
+        arguments = [args]
+        execution_requirements = {}
+    # TODO: Select if worker is to be enabled.
+
     ctx.actions.run(
-        executable = ctx.executable._process_wrapper_worker,
+        executable = executable,
         inputs = compile_inputs,
         outputs = [crate_info.output],
-        tools = [ctx.executable._process_wrapper],
+        tools = tools,
         env = env,
-        arguments = [ctx.executable._process_wrapper.path, toolchain.rustc.path, ctx.var["COMPILATION_MODE"], args],
+        arguments = arguments,
         mnemonic = "Rustc",
-        execution_requirements = { "supports-workers": "1" },
+        execution_requirements = execution_requirements,
         progress_message = "Compiling Rust {} {}{} ({} files)".format(
             crate_info.type,
             ctx.label.name,

rust/repositories.bzl

@@ -2,6 +2,7 @@ load(":known_shas.bzl", "FILE_KEY_TO_SHA")
 load("//rust/platform:triple_mappings.bzl", "system_to_binary_ext", "system_to_dylib_ext", "system_to_staticlib_ext", "system_to_stdlib_linkflags", "triple_to_constraint_set", "triple_to_system")
 load("@bazel_tools//tools/build_defs/repo:utils.bzl", "maybe")
 load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
+load("//worker:repositories.bzl", "rust_worker_repositories")
 
 DEFAULT_TOOLCHAIN_NAME_PREFIX = "toolchain_for"
 
@@ -93,6 +94,8 @@ def rust_repositories(
         sha256s = sha256s,
     )
 
+    rust_worker_repositories()
+
 def _check_version_valid(version, iso_date, param_prefix = ""):
     """Verifies that the provided rust version and iso_date make sense."""
 

worker/BUILD

@@ -0,0 +1,26 @@
+load("//worker:toolchain.bzl", "worker_toolchain")
+
+toolchain_type(name = "toolchain_type")
+
+worker_toolchain(
+    name = "worker_linux_x86_64",
+    worker_binary = "@rust_worker_linux_x86_64//file",
+)
+
+worker_toolchain(name = "worker_dummy")
+
+toolchain(
+    name = "linux_x86_64",
+    exec_compatible_with = [
+        "@platforms//os:linux",
+        "@platforms//cpu:x86_64",
+    ],
+    toolchain = ":worker_linux_x86_64",
+    toolchain_type = ":toolchain_type",
+)
+
+toolchain(
+    name = "dummy",
+    toolchain = ":worker_dummy",
+    toolchain_type = ":toolchain_type",
+)

worker/repositories.bzl

@@ -0,0 +1,14 @@
+load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_file")
+
+def rust_worker_repositories():
+    http_file(
+        name = "rust_worker_linux_x86_64",
+        executable = True,
+        sha256 = "e3383d062752bcbeeb8a97d36bba0dea244db77b45a083bee718978a0deb5d37",
+        urls = ["https://github.com/nikhilm/rustc-worker/releases/download/v0.2.0/rustc-worker-linux_x86_64"],
+    )
+
+    native.register_toolchains(
+        "@io_bazel_rules_rust//worker:linux_x86_64",
+        "@io_bazel_rules_rust//worker:dummy",
+    )

worker/toolchain.bzl

@@ -0,0 +1,12 @@
+def _worker_toolchain_impl(ctx):
+    toolchain_info = platform_common.ToolchainInfo(
+        worker_binary = ctx.executable.worker_binary,
+    )
+    return [toolchain_info]
+
+worker_toolchain = rule(
+    implementation = _worker_toolchain_impl,
+    attrs = {
+        "worker_binary": attr.label(allow_single_file = True, executable = True, cfg = "exec"),
+    },
+)