Support distributed ThinLTO for Linux cdylibs

cargo/private/cargo_build_script.bzl

@@ -531,7 +531,10 @@ def _cargo_build_script_impl(ctx):
     # Pull in env vars which may be required for the cc_toolchain to work (e.g. on OSX, the SDK version).
     # We hope that the linker env is sufficient for the whole cc_toolchain.
     if use_cc_toolchain:
-        cc_toolchain, feature_configuration = find_cc_toolchain(ctx)
+        # Build-script native objects are opaque to Bazel: no per-object LTO
+        # backend actions can be registered for them. Keep them as machine code
+        # when the Rust link uses distributed ThinLTO.
+        cc_toolchain, feature_configuration = find_cc_toolchain(ctx, ["thin_lto"])
     else:
         cc_toolchain, feature_configuration = None, None
     linker, _, link_args, linker_env = get_linker_and_args(ctx, "bin", toolchain, cc_toolchain, feature_configuration, None)

extensions/pyo3/private/BUILD.bazel

@@ -14,6 +14,7 @@ bzl_library(
     srcs = glob(["*.bzl"]),
     visibility = ["//extensions/pyo3:__pkg__"],
     deps = [
+        "@rules_cc//cc/common",
         "@rules_python//python:defs_bzl",
         "@rules_rust//rust:bzl_lib",
     ],

extensions/pyo3/private/pyo3_toolchain.bzl

@@ -1,6 +1,7 @@
 """PyO3 Toolchains"""
 
 load("@bazel_skylib//rules:common_settings.bzl", "BuildSettingInfo")
+load("@rules_cc//cc/common:cc_info.bzl", "CcInfo")
 load("@rules_rust//rust:defs.bzl", "rust_common")
 
 PYO3_TOOLCHAIN = "//extensions/pyo3:toolchain_type"
@@ -209,6 +210,9 @@ def _current_rust_pyo3_toolchain_impl(ctx):
     if rust_common.crate_group_info in target:
         providers.append(target[rust_common.crate_group_info])
 
+    if CcInfo in target:
+        providers.append(target[CcInfo])
+
     return providers
 
 current_rust_pyo3_toolchain = rule(
@@ -243,6 +247,9 @@ def _current_rust_pyo3_introspection_toolchain_impl(ctx):
     if rust_common.crate_group_info in target:
         providers.append(target[rust_common.crate_group_info])
 
+    if CcInfo in target:
+        providers.append(target[CcInfo])
+
     return providers
 
 current_rust_pyo3_introspection_toolchain = rule(

rust/private/rustc.bzl

@@ -966,7 +966,10 @@ def _will_emit_object_file(emit):
 def _supports_distributed_thin_lto(ctx, toolchain, crate_info):
     """Whether `crate_info` can participate in distributed ThinLTO."""
     return (
-        crate_info.type in ("bin", "lib", "rlib") and
+        (
+            crate_info.type in ("bin", "lib", "rlib") or
+            (crate_info.type == "cdylib" and toolchain.target_os == "linux")
+        ) and
         toolchain.target_arch not in ("wasm32", "wasm64") and
         not toolchain._bootstrapping and
         not is_no_std(ctx, toolchain, crate_info.is_test)
@@ -975,6 +978,11 @@ def _supports_distributed_thin_lto(ctx, toolchain, crate_info):
 def _remove_codegen_units(flag):
     return None if flag.startswith("-Ccodegen-units") else flag
 
+def _cdylib_native_link_args(file):
+    if file.basename.split("-")[1] == "whole":
+        return ["-Wl,--whole-archive", file.path, "-Wl,--no-whole-archive"]
+    return [file.path]
+
 def _should_add_oso_prefix(toolchain):
     """Whether to add -oso_prefix to strip absolute paths from N_OSO entries.
 
@@ -1443,6 +1451,12 @@ def construct_arguments(
 
     if linker_plugin_lto:
         rustc_flags.add("-Clinker-plugin-lto")
+        if crate_info.type in ("lib", "rlib"):
+            # Preserve bundled native archives as separate linker inputs when
+            # the Rust object is replaced by a distributed LTO backend output.
+            rustc_flags.add("-Zpacked-bundled-libs")
+            if inject_allow_features_guardrail:
+                rustc_flags.add("-Zallow-features=")
     else:
         rustc_flags.add_all(construct_lto_arguments(ctx, toolchain, crate_info))
     _add_codegen_units_flags(toolchain, emit, rustc_flags)
@@ -2011,8 +2025,9 @@ def rustc_compile(
             ).format(ctx.label),
         )
     use_cc_common_link = experimental_use_cc_common_link or (
-        distributed_thin_lto and crate_info.type == "bin"
+        distributed_thin_lto and crate_info.type in ("bin", "cdylib")
     )
+    packed_bundled_libs = distributed_thin_lto and crate_info.type in ("lib", "rlib")
 
     scan_msvc_archive_object = (
         rust_toolchain.target_abi == "msvc" and
@@ -2070,8 +2085,24 @@ def rustc_compile(
     # Metadata is emitted by a separate -Zno-codegen action. The full action
     # emits the linked crate unless cc_common.link does, and emits output_o
     # when output_o is declared.
+    cdylib_export_file = None
+    cdylib_symbols_file = None
+    cdylib_native_dir = None
+    cdylib_native_params = None
+    if distributed_thin_lto and crate_info.type == "cdylib":
+        # Keep rustc's ELF export policy and symbol roots while delegating code
+        # generation and the actual shared link to the C++ toolchain.
+        cdylib_export_file = ctx.actions.declare_file(crate_info.output.basename + ".exports", sibling = crate_info.output)
+        cdylib_symbols_file = ctx.actions.declare_file(crate_info.output.basename + ".symbols.o", sibling = crate_info.output)
+        cdylib_native_dir = ctx.actions.declare_directory(crate_info.output.basename + ".native", sibling = crate_info.output)
+        cdylib_native_params = ctx.actions.declare_file(crate_info.output.basename + ".native.params", sibling = crate_info.output)
+        native_args = ctx.actions.args()
+        native_args.add_all([cdylib_native_dir], map_each = _cdylib_native_link_args)
+        native_args.set_param_file_format("shell")
+        ctx.actions.write(output = cdylib_native_params, content = native_args)
+
     emit = []
-    if not use_cc_common_link:
+    if not use_cc_common_link or cdylib_export_file:
         emit.append("link")
     if output_o:
         emit.append(("obj", output_o))
@@ -2184,6 +2215,11 @@ def rustc_compile(
     )
 
     args_metadata = None
+    if cdylib_export_file:
+        args.process_wrapper_flags.add("--rustc-cdylib-export-file", cdylib_export_file)
+        args.process_wrapper_flags.add("--rustc-cdylib-symbols-file", cdylib_symbols_file)
+        args.process_wrapper_flags.add_all([cdylib_native_dir], before_each = "--rustc-cdylib-native-dir", expand_directories = False)
+
     if build_metadata:
         metadata_emit = ["link"]
         args_metadata, _ = construct_arguments(
@@ -2221,8 +2257,9 @@ def rustc_compile(
     # this is the final list of env vars
     env.update(env_from_args)
 
-    if build_metadata and inject_allow_features_guardrail:
-        # RUSTC_BOOTSTRAP=1 is required for -Zno-codegen on stable/beta rustc, and
+    if (build_metadata or packed_bundled_libs) and inject_allow_features_guardrail:
+        # RUSTC_BOOTSTRAP=1 is required for -Zno-codegen and
+        # -Zpacked-bundled-libs on stable/beta rustc, and
         # must be set on both the metadata and full actions for SVH compatibility
         # (since RUSTC_BOOTSTRAP affects the crate hash). Skipped on nightly
         # toolchains (where -Zno-codegen works without bootstrap) and when the
@@ -2247,6 +2284,8 @@ def rustc_compile(
 
     # The action might generate extra output that we don't want to include in the `DefaultInfo` files.
     action_outputs = list(outputs)
+    if cdylib_export_file:
+        action_outputs.extend([cdylib_export_file, cdylib_symbols_file, cdylib_native_dir])
     if output_o and output_o not in action_outputs:
         action_outputs.append(output_o)
     if rustc_output:
@@ -2360,8 +2399,9 @@ def rustc_compile(
         # Wrap the main `.o` file into a compilation output suitable for
         # cc_common.link. The main `.o` file is useful in both PIC and non-PIC
         # modes.
-        cco_args["objects"] = depset([output_o])
-        cco_args["pic_objects"] = depset([output_o])
+        link_objects = [output_o] + ([cdylib_symbols_file] if cdylib_symbols_file else [])
+        cco_args["objects"] = depset(link_objects)
+        cco_args["pic_objects"] = depset(link_objects)
         if distributed_thin_lto:
             cco_args["lto_compilation_context"] = cc_common.create_lto_compilation_context(
                 objects = {
@@ -2457,6 +2497,8 @@ def rustc_compile(
             output_type = "executable" if crate_info.type == "bin" else "dynamic_library",
             additional_outputs = additional_linker_outputs,
             variables_extension = variables_extension,
+            additional_inputs = [cdylib_export_file, cdylib_native_params, cdylib_native_dir] if cdylib_export_file else [],
+            user_link_flags = ["-Wl,--version-script=" + cdylib_export_file.path, "@" + cdylib_native_params.path] if cdylib_export_file else [],
         )
 
         if rust_toolchain.target_os == "linux" and cc_helper.should_create_per_object_debug_info(feature_configuration, ctx.fragments.cpp):

test/unit/lto/lto_test_suite.bzl

@@ -6,7 +6,7 @@ load("@rules_cc//cc:cc_binary.bzl", "cc_binary")
 load("@rules_cc//cc:cc_library.bzl", "cc_library")
 load("@rules_cc//cc:cc_test.bzl", "cc_test")
 load("@rules_cc//cc/common:cc_info.bzl", "CcInfo")
-load("//rust:defs.bzl", "rust_binary", "rust_library", "rust_library_group", "rust_proc_macro")
+load("//rust:defs.bzl", "rust_binary", "rust_library", "rust_library_group", "rust_proc_macro", "rust_shared_library")
 load(
     "//test/unit:common.bzl",
     "assert_action_mnemonic",
@@ -149,6 +149,8 @@ def _distributed_thin_lto_library(ctx):
     assert_argv_contains(env, action, "--emit=link")
     assert_argv_contains_prefix(env, action, "--emit=obj=")
     assert_argv_contains(env, action, "-Clinker-plugin-lto")
+    assert_argv_contains(env, action, "-Zpacked-bundled-libs")
+    assert_argv_contains(env, action, "-Zallow-features=")
     assert_argv_contains_prefix_not(env, action, "-Clto")
     assert_argv_contains_prefix_not(env, action, "-Cembed-bitcode")
 
@@ -210,6 +212,30 @@ _distributed_thin_lto_binary_test = analysistest.make(
     config_settings = _DISTRIBUTED_THIN_LTO_CONFIG_SETTINGS,
 )
 
+def _distributed_thin_lto_cdylib(ctx):
+    env = analysistest.begin(ctx)
+    target = analysistest.target_under_test(env)
+    actions = _assert_distributed_thin_lto_link(env, target)
+    rustc_action = actions["Rustc"]
+    assert_argv_contains(env, rustc_action, "--crate-type=cdylib")
+    assert_argv_contains(env, rustc_action, "-Clinker-plugin-lto")
+    assert_argv_contains_prefix(env, rustc_action, "--emit=obj=")
+    assert_argv_contains(env, rustc_action, "--emit=link")
+    assert_argv_contains(env, rustc_action, "--rustc-cdylib-export-file")
+    assert_argv_contains_prefix_not(env, rustc_action, "-Clto")
+    assert_argv_contains(env, actions["CppLink"], "-shared")
+    for mnemonic in ["CppLTOIndexing", "CppLink"]:
+        assert_argv_contains_prefix(env, actions[mnemonic], "-Wl,--version-script=")
+        asserts.true(env, any([file.extension == "exports" for file in actions[mnemonic].inputs.to_list()]))
+    asserts.true(env, any([file.basename == "libdistributed_cdylib.so" for file in actions["CppLink"].outputs.to_list()]))
+    asserts.equals(env, ["libdistributed_cdylib.so"], [file.basename for file in target[DefaultInfo].files.to_list()])
+    return analysistest.end(env)
+
+_distributed_thin_lto_cdylib_test = analysistest.make(
+    _distributed_thin_lto_cdylib,
+    config_settings = _DISTRIBUTED_THIN_LTO_CONFIG_SETTINGS,
+)
+
 def _distributed_thin_lto_global_allocator(ctx):
     env = analysistest.begin(ctx)
     target = analysistest.target_under_test(env)
@@ -443,6 +469,45 @@ def lto_test_suite(name):
         ],
     )
 
+    write_file(
+        name = "crate_cdylib",
+        out = "cdylib.rs",
+        content = [
+            "extern \"C\" { fn native_add(left: usize, right: usize) -> usize; }",
+            "#[no_mangle]",
+            "pub extern \"C\" fn cdylib_add(left: usize, right: usize) -> usize {",
+            "    unsafe { native_add(distributed_lib::add(left, right), 1) }",
+            "}",
+            "",
+        ],
+    )
+
+    write_file(
+        name = "cdylib_native_src",
+        out = "cdylib_native.cc",
+        content = [
+            "#include <cstdint>",
+            "extern \"C\" uintptr_t native_add(uintptr_t left, uintptr_t right) { return left + right; }",
+            "",
+        ],
+    )
+
+    write_file(
+        name = "cdylib_test_src",
+        out = "cdylib_test.cc",
+        content = [
+            "#include <cstdint>",
+            "#include <dlfcn.h>",
+            "extern \"C\" uintptr_t cdylib_add(uintptr_t, uintptr_t);",
+            "extern \"C\" uintptr_t distributed_add(uintptr_t, uintptr_t);",
+            "int main() {",
+            "    return cdylib_add(2, 2) == 5 && distributed_add(2, 2) == 4",
+            "        && dlsym(RTLD_DEFAULT, \"native_add\") == nullptr ? 0 : 1;",
+            "}",
+            "",
+        ],
+    )
+
     rust_library(
         name = "lib",
         srcs = [":lib.rs"],
@@ -496,6 +561,29 @@ def lto_test_suite(name):
         tags = ["manual"],
     )
 
+    cc_library(
+        name = "cdylib_native",
+        srcs = [":cdylib_native.cc"],
+        tags = ["manual"],
+    )
+
+    rust_shared_library(
+        name = "distributed_cdylib",
+        srcs = [":cdylib.rs"],
+        edition = "2021",
+        deps = [":cdylib_native", ":distributed_lib_group"],
+        tags = ["manual"],
+    )
+
+    cc_test(
+        name = "distributed_cdylib_runtime_test",
+        srcs = [":cdylib_test.cc"],
+        linkopts = ["-ldl"],
+        deps = [":distributed_cdylib"],
+        tags = ["manual"],
+        target_compatible_with = ["@platforms//os:linux"],
+    )
+
     rust_binary(
         name = "distributed_global_allocator_bin",
         srcs = [":global_allocator_bin"],
@@ -589,6 +677,11 @@ def lto_test_suite(name):
         target_under_test = ":distributed_bin",
     )
 
+    _distributed_thin_lto_cdylib_test(
+        name = "distributed_thin_lto_cdylib_test",
+        target_under_test = ":distributed_cdylib",
+    )
+
     _distributed_thin_lto_global_allocator_test(
         name = "distributed_thin_lto_global_allocator_test",
         target_under_test = ":distributed_global_allocator_bin",
@@ -655,6 +748,7 @@ def lto_test_suite(name):
             ":lto_proc_macro_test",
             ":distributed_thin_lto_library_test",
             ":distributed_thin_lto_binary_test",
+            ":distributed_thin_lto_cdylib_test",
             ":distributed_thin_lto_global_allocator_test",
             ":distributed_thin_lto_cc_binary_test",
             ":distributed_thin_lto_shared_backends_test",