Implement Local Remote Execution for Rust

This change implements ~40 Rust toolchains that reuse Nix's executables and make them available via Bazel. The new setup is fully reproducible, supports musl and glibc and works from all nix-supported hosts to all sensible crosscompilation targets seamlessly through remote and local execution. For instance, you can now run this command on a MacOS host to build a statically linked executable for x86_64-linux on a remote arm worker: ```bash bazel build \ nativelink \ --extra_execution_platforms=@local-remote-execution//lre-rs:aarch64-unknown-linux-gnu \ --platforms=@local-remote-execution//lre-rs:x86_64-unknown-linux-musl ```
TraceMachina · Dec 3, 2024 · 91ea106 · 91ea106
1 parent ae71bc8
commit 91ea106
Show file tree

Hide file tree

Showing 16 changed files with 1,368 additions and 38 deletions.
diff --git a/.bazelignore b/.bazelignore
@@ -13,3 +13,4 @@ bazel-testlogs
 bazel-nativelink
 local-remote-execution/generated-cc
 local-remote-execution/generated-java
+local-remote-execution/lre-rs
diff --git a/MODULE.bazel b/MODULE.bazel
@@ -6,7 +6,7 @@ module(
 
 bazel_dep(name = "rules_cc", version = "0.0.17")
 bazel_dep(name = "platforms", version = "0.0.10")
-bazel_dep(name = "rules_python", version = "0.36.0")
+bazel_dep(name = "rules_python", version = "0.40.0")
 
 python = use_extension("@rules_python//python/extensions:python.bzl", "python")
 python.toolchain(
@@ -20,6 +20,12 @@ use_repo(python, python = "python_versions")
 
 bazel_dep(name = "rules_rust", version = "0.54.1")
 
+# TODO(aaronmondal): REMOVE
+local_path_override(
+    module_name = "rules_rust",
+    path = "../rules_rust",
+)
+
 rust = use_extension("@rules_rust//rust:extensions.bzl", "rust")
 rust.toolchain(
     edition = "2021",
@@ -40,19 +46,19 @@ rust.toolchain(
         "2024-11-23/rustfmt-nightly-x86_64-unknown-linux-gnu.tar.xz": "ac2a8fff891473220d36c9e62afc2dc0d5aebaefbb8489584bd2446a37d6fd7b",
     },
     versions = [
-        "1.82.0",
-        "nightly/2024-11-23",
+        # "1.82.0",
+        # "nightly/2024-11-23",
     ],
 )
 
-rust_host_tools = use_extension(
-    "@rules_rust//rust:extension.bzl",
-    "rust_host_tools",
-)
-rust_host_tools.host_tools(
-    edition = "2021",
-    version = "1.82.0",
-)
+# rust_host_tools = use_extension(
+#     "@rules_rust//rust:extension.bzl",
+#     "rust_host_tools",
+# )
+# rust_host_tools.host_tools(
+#     edition = "2021",
+#     version = "1.82.0",
+# )
 
 use_repo(rust, "rust_toolchains")
 
@@ -64,23 +70,25 @@ crate.from_cargo(
     supported_platform_triples = [
         "aarch64-apple-darwin",
         "aarch64-unknown-linux-gnu",
+        "aarch64-unknown-linux-musl",
         "arm-unknown-linux-gnueabi",
         "armv7-unknown-linux-gnueabi",
         "x86_64-apple-darwin",
         "x86_64-pc-windows-msvc",
         "x86_64-unknown-linux-gnu",
+        "x86_64-unknown-linux-musl",
     ],
 )
 use_repo(crate, "crates")
 
-rust_analyzer = use_extension(
-    "@rules_rust//tools/rust_analyzer:extension.bzl",
-    "rust_analyzer_dependencies",
-)
-rust_analyzer.rust_analyzer_dependencies()
+# rust_analyzer = use_extension(
+#     "@rules_rust//tools/rust_analyzer:extension.bzl",
+#     "rust_analyzer_dependencies",
+# )
+# rust_analyzer.rust_analyzer_dependencies()
 
-bazel_dep(name = "protobuf", version = "27.5", repo_name = "com_google_protobuf")
-bazel_dep(name = "toolchains_protoc", version = "0.3.3")
+bazel_dep(name = "protobuf", version = "29.0", repo_name = "com_google_protobuf")
+bazel_dep(name = "toolchains_protoc", version = "0.3.4")
 
 protoc = use_extension("@toolchains_protoc//protoc:extensions.bzl", "protoc")
 protoc.toolchain(

diff --git a/flake.nix b/flake.nix
@@ -52,20 +52,53 @@
         system,
         ...
       }: let
+        # C and C++ toolchain configuration. The basis of all toolchains.
+        llvmPackages = pkgs.llvmPackages_19;
+
+        # Rust toolchain configuration.
         stable-rust-version = "1.82.0";
         nightly-rust-version = "2024-11-23";
 
-        llvmPackages = pkgs.llvmPackages_19;
+        # This map translates execution platforms to sensible targets that can
+        # be built on such a platform. For instance, an x86_64-linux execution
+        # platform can target aarch64-linux and musl targets, but not darwin.
+        #
+        # On the Bazel side each key maps to a dedicated toolchain capture the
+        # cross-compile capabilities on the different exec platforms.
+        nixSystemToRustTargets = {
+          "aarch64-darwin" = [
+            "aarch64-apple-darwin"
+            "aarch64-unknown-linux-gnu"
+            "aarch64-unknown-linux-musl"
+            "x86_64-apple-darwin"
+            "x86_64-unknown-linux-gnu"
+            "x86_64-unknown-linux-musl"
+          ];
+          "aarch64-linux" = [
+            "aarch64-unknown-linux-gnu"
+            "aarch64-unknown-linux-musl"
+            "x86_64-unknown-linux-gnu"
+            "x86_64-unknown-linux-musl"
+          ];
+          "x86_64-darwin" = [
+            "aarch64-apple-darwin"
+            "aarch64-unknown-linux-gnu"
+            "aarch64-unknown-linux-musl"
+            "x86_64-apple-darwin"
+            "x86_64-unknown-linux-gnu"
+            "x86_64-unknown-linux-musl"
+          ];
+          "x86_64-linux" = [
+            "aarch64-unknown-linux-gnu"
+            "aarch64-unknown-linux-musl"
+            "x86_64-unknown-linux-gnu"
+            "x86_64-unknown-linux-musl"
+          ];
+        };
 
-        nixSystemToRustTriple = nixSystem:
-          {
-            "x86_64-linux" = "x86_64-unknown-linux-musl";
-            "aarch64-linux" = "aarch64-unknown-linux-musl";
-            "x86_64-darwin" = "x86_64-apple-darwin";
-            "aarch64-darwin" = "aarch64-apple-darwin";
-          }
-          .${nixSystem}
-          or (throw "Unsupported Nix system: ${nixSystem}");
+        nixExecToRustTargets = nixSystem:
+          nixSystemToRustTargets.${nixSystem}
+          or (throw "Unsupported Nix exec platform: ${nixSystem}");
 
         # Calling `pkgs.pkgsCross` changes the host and target platform to the
         # cross-target but leaves the build platform the same as pkgs.
@@ -80,24 +113,33 @@
         #
         # For optimal cache reuse of different crosscompilation toolchains we
         # take our rust toolchain from the host's `pkgs` and remap the rust
-        # target to the target platform of the `pkgsCross` target. This lets us
-        # reuse the same executables (for instance rustc) to build artifacts for
-        # different target platforms.
+        # target to sensible `pkgsCross` targets. This lets us reuse the same
+        # executables (for instance rustc) to build artifacts for different
+        # target platforms.
         stableRustFor = p:
           p.rust-bin.stable.${stable-rust-version}.default.override {
-            targets = [
-              "${nixSystemToRustTriple p.stdenv.targetPlatform.system}"
-            ];
+            targets = nixExecToRustTargets p.stdenv.targetPlatform.system;
           };
 
         nightlyRustFor = p:
           p.rust-bin.nightly.${nightly-rust-version}.default.override {
             extensions = ["llvm-tools"];
-            targets = [
-              "${nixSystemToRustTriple p.stdenv.targetPlatform.system}"
-            ];
+            targets = nixExecToRustTargets p.stdenv.targetPlatform.system;
           };
 
+        # This convenience method maps a nix system to the Rust target triple
+        # that we'd want to target by default. This is only used for Nix builds
+        # and doesn't have a Bazel equivalent.
+        nixSystemToRustTriple = nixSystem:
+          {
+            "x86_64-linux" = "x86_64-unknown-linux-musl";
+            "aarch64-linux" = "aarch64-unknown-linux-musl";
+            "x86_64-darwin" = "x86_64-apple-darwin";
+            "aarch64-darwin" = "aarch64-apple-darwin";
+          }
+          .${nixSystem}
+          or (throw "Unsupported Nix host platform: ${nixSystem}");
+
         craneLibFor = p: (crane.mkLib p).overrideToolchain stableRustFor;
         nightlyCraneLibFor = p: (crane.mkLib p).overrideToolchain nightlyRustFor;
 
@@ -248,6 +290,9 @@
           inherit (pkgs.lre) stdenv;
         };
         createWorker = pkgs.callPackage ./tools/create-worker.nix {inherit buildImage self;};
+        gen-lre-rs = pkgs.callPackage ./tools/lre-rs.nix {
+          inherit nixSystemToRustTargets;
+        };
         buck2-toolchain = let
           buck2-nightly-rust-version = "2024-04-28";
           buck2-nightly-rust = pkgs.rust-bin.nightly.${buck2-nightly-rust-version};
@@ -288,6 +333,14 @@
         lre-cc = pkgs.callPackage ./local-remote-execution/lre-cc.nix {
           inherit buildImage;
         };
+        lre-rs = pkgs.callPackage ./local-remote-execution/lre-rs.nix {
+          inherit buildImage lre-cc;
+          rust-toolchains = [
+            (stableRustFor pkgs)
+            (nightlyRustFor pkgs)
+          ];
+        };
+
         toolchain-drake = buildImage {
           name = "toolchain-drake";
           # imageDigest and sha256 are generated by toolchain-drake.sh for non-reproducible builds.
@@ -365,6 +418,7 @@
             inherit
               local-image-test
               lre-cc
+              lre-rs
               native-cli
               nativelink
               nativelinkCoverageForHost
@@ -376,13 +430,18 @@
               nativelink-x86_64-linux
               publish-ghcr
               ;
+
+            stable-rust = stableRustFor pkgs;
+            nightly-rust = nightlyRustFor pkgs;
+
             default = nativelink;
 
             rbe-autogen-lre-cc = rbe-autogen lre-cc;
             nativelink-worker-lre-cc = createWorker lre-cc;
             lre-java = pkgs.callPackage ./local-remote-execution/lre-java.nix {inherit buildImage;};
             rbe-autogen-lre-java = rbe-autogen lre-java;
             nativelink-worker-lre-java = createWorker lre-java;
+            nativelink-worker-lre-rs = createWorker lre-rs;
             nativelink-worker-siso-chromium = createWorker siso-chromium;
             nativelink-worker-toolchain-drake = createWorker toolchain-drake;
             nativelink-worker-toolchain-buck2 = createWorker toolchain-buck2;
@@ -423,7 +482,7 @@
           Env =
             if pkgs.stdenv.isDarwin
             then [] # Doesn't support Darwin yet.
-            else lre-cc.meta.Env;
+            else (lre-cc.meta.Env ++ lre-rs.meta.Env);
           prefix = "linux";
         };
         nixos.settings = {
@@ -451,6 +510,7 @@
               # Rust
               (stableRustFor pkgs)
               bazel
+              gen-lre-rs
 
               ## Infrastructure
               pkgs.awscli2

diff --git a/local-remote-execution/MODULE.bazel b/local-remote-execution/MODULE.bazel
@@ -14,3 +14,18 @@ bazel_dep(name = "rules_cc", version = "0.0.17")
 
 # Use the starlark implementation of Java rules instead of the builtin ones.
 bazel_dep(name = "rules_java", version = "8.5.1")
+bazel_dep(name = "rules_rust", version = "0.54.1")
+bazel_dep(name = "bazel_skylib", version = "1.7.1")
+
+lre_rs = use_extension("//lre-rs:extension.bzl", "lre_rs")
+use_repo(
+    lre_rs,
+    "lre-rs-nightly-aarch64-darwin",
+    "lre-rs-nightly-aarch64-linux",
+    "lre-rs-nightly-x86_64-darwin",
+    "lre-rs-nightly-x86_64-linux",
+    "lre-rs-stable-aarch64-darwin",
+    "lre-rs-stable-aarch64-linux",
+    "lre-rs-stable-x86_64-darwin",
+    "lre-rs-stable-x86_64-linux",
+)
diff --git a/local-remote-execution/lre-rs.nix b/local-remote-execution/lre-rs.nix
@@ -0,0 +1,27 @@
+{
+  lib,
+  buildImage,
+  lre-cc,
+  rust-toolchains,
+}: let
+  # This environment is shared between toolchain autogen images and the final
+  # toolchain image.
+  Env =
+    lre-cc.meta.Env
+    ++ [
+      # Add all tooling here so that the generated toolchains use `/nix/store/*`
+      # paths instead of `/bin` or `/usr/bin`. This way we're guaranteed to use
+      # binary identical toolchains during local and remote execution.
+      "RUST=${lib.concatStringsSep ":" rust-toolchains}"
+    ];
+in
+  buildImage {
+    name = "lre-rs";
+    maxLayers = 100;
+    config = {inherit Env;};
+    # Attached for passthrough to rbe-configs-gen.
+    meta = {inherit Env;};
+
+    # Don't set a tag here so that the image is tagged by its derivation hash.
+    # tag = null;
+  }