Respect #[global_allocator] in cc_common.link builds

.bazelci/presubmit.yml

@@ -497,13 +497,33 @@ tasks:
     platform: ubuntu2004
     working_directory: test/cc_common_link
     build_targets:
+      - "--"
       - "//..."
+      # The with_global_alloc directory is a repository on its own tested in the 'Build via cc_common.link using a global allocator' task.
+      - "-//with_global_alloc/..."
     test_targets:
+      - "--"
       - "//..."
+      # The with_global_alloc directory is a repository on its own tested in the 'Build via cc_common.link using a global allocator' task.
+      - "-//with_global_alloc/..."
     build_flags:
       - "--@rules_rust//rust/settings:experimental_use_cc_common_link=True"
     test_flags:
       - "--@rules_rust//rust/settings:experimental_use_cc_common_link=True"
+  cc_common_link_with_global_alloc_ubuntu2004:
+    name: Build via cc_common.link using s global allocator
+    platform: ubuntu2004
+    working_directory: test/cc_common_link/with_global_alloc
+    build_targets:
+      - "//..."
+    test_targets:
+      - "//..."
+    build_flags:
+      - "--@rules_rust//rust/settings:experimental_use_cc_common_link=True"
+      - "--@rules_rust//rust/settings:experimental_use_global_allocator=True"
+    test_flags:
+      - "--@rules_rust//rust/settings:experimental_use_cc_common_link=True"
+      - "--@rules_rust//rust/settings:experimental_use_global_allocator=True"
   android_examples_ubuntu2004:
     name: Android Examples
     platform: ubuntu2004

rust/private/repository_utils.bzl

@@ -242,6 +242,7 @@ rust_toolchain(
     llvm_profdata = {llvm_profdata_label},
     rustc_lib = "//:rustc_lib",
     allocator_library = {allocator_library},
+    global_allocator_library = {global_allocator_library},
     binary_ext = "{binary_ext}",
     staticlib_ext = "{staticlib_ext}",
     dylib_ext = "{dylib_ext}",
@@ -259,6 +260,7 @@ def BUILD_for_rust_toolchain(
         exec_triple,
         target_triple,
         allocator_library,
+        global_allocator_library,
         default_edition,
         include_rustfmt,
         include_llvm_tools,
@@ -270,6 +272,7 @@ def BUILD_for_rust_toolchain(
         exec_triple (triple): The rust-style target that this compiler runs on
         target_triple (triple): The rust-style target triple of the tool
         allocator_library (str, optional): Target that provides allocator functions when rust_library targets are embedded in a cc_binary.
+        global_allocator_library (str, optional): Target that provides allocator functions global allocator is used with cc_common_link.
         default_edition (str): Default Rust edition.
         include_rustfmt (bool): Whether rustfmt is present in the toolchain.
         include_llvm_tools (bool): Whether llvm-tools are present in the toolchain.
@@ -295,13 +298,17 @@ def BUILD_for_rust_toolchain(
     allocator_library_label = "None"
     if allocator_library:
         allocator_library_label = "\"{allocator_library}\"".format(allocator_library = allocator_library)
+    global_allocator_library_label = "None"
+    if global_allocator_library:
+        global_allocator_library_label = "\"{global_allocator_library}\"".format(global_allocator_library = global_allocator_library)
 
     return _build_file_for_rust_toolchain_template.format(
         toolchain_name = name,
         binary_ext = system_to_binary_ext(target_triple.system),
         staticlib_ext = system_to_staticlib_ext(target_triple.system),
         dylib_ext = system_to_dylib_ext(target_triple.system),
         allocator_library = allocator_library_label,
+        global_allocator_library = global_allocator_library_label,
         stdlib_linkflags = stdlib_linkflags,
         system = target_triple.system,
         default_edition = default_edition,

rust/private/rustc.bzl

@@ -1271,7 +1271,7 @@ def rustc_compile_action(
         # Collect the linking contexts of the standard library and dependencies.
         linking_contexts = [
             malloc_library[CcInfo].linking_context,
-            toolchain.libstd_and_allocator_ccinfo.linking_context,
+            _get_std_and_alloc_info(ctx, toolchain).linking_context,
             toolchain.stdlib_linkflags.linking_context,
         ]
 
@@ -1380,6 +1380,14 @@ def rustc_compile_action(
 
     return providers
 
+def _get_std_and_alloc_info(ctx, toolchain):
+    if is_exec_configuration(ctx):
+        return toolchain.libstd_and_allocator_ccinfo
+    if toolchain._experimental_use_global_allocator:
+        return toolchain.libstd_and_global_allocator_ccinfo
+    else:
+        return toolchain.libstd_and_allocator_ccinfo
+
 def _is_dylib(dep):
     return not bool(dep.static_library or dep.pic_static_library)
 
@@ -1492,9 +1500,11 @@ def establish_cc_info(ctx, attr, crate_info, toolchain, cc_toolchain, feature_co
             else:
                 cc_infos.append(dep.cc_info)
 
-    if crate_info.type in ("rlib", "lib") and toolchain.libstd_and_allocator_ccinfo:
-        # TODO: if we already have an rlib in our deps, we could skip this
-        cc_infos.append(toolchain.libstd_and_allocator_ccinfo)
+    if crate_info.type in ("rlib", "lib"):
+        libstd_and_allocator_cc_info = _get_std_and_alloc_info(ctx, toolchain)
+        if libstd_and_allocator_cc_info:
+            # TODO: if we already have an rlib in our deps, we could skip this
+            cc_infos.append(libstd_and_allocator_cc_info)
 
     return [cc_common.merge_cc_infos(cc_infos = cc_infos)]
 

rust/repositories.bzl

@@ -99,6 +99,7 @@ def rust_register_toolchains(
         dev_components = False,
         edition = None,
         allocator_library = None,
+        global_allocator_library = None,
         iso_date = None,
         register_toolchains = True,
         rustfmt_version = DEFAULT_NIGHTLY_VERSION,
@@ -130,6 +131,7 @@ def rust_register_toolchains(
         dev_components (bool, optional): Whether to download the rustc-dev components (defaults to False). Requires version to be "nightly".
         edition (str, optional): The rust edition to be used by default (2015, 2018, or 2021). If absent, every target is required to specify its `edition` attribute.
         allocator_library (str, optional): Target that provides allocator functions when rust_library targets are embedded in a cc_binary.
+        global_allocator_library (str, optional): Target that provides allocator functions when global allocator is used with cc_common.link.
         iso_date (str, optional):  **Deprecated**: Use `versions` instead.
         register_toolchains (bool): If true, repositories will be generated to produce and register `rust_toolchain` targets.
         rustfmt_version (str, optional): The version of rustfmt.
@@ -219,6 +221,7 @@ def rust_register_toolchains(
             exec_triple = exec_triple,
             extra_target_triples = extra_target_triples,
             allocator_library = allocator_library,
+            global_allocator_library = global_allocator_library,
             iso_date = iso_date,
             register_toolchain = register_toolchains,
             rustfmt_version = rustfmt_version,
@@ -342,6 +345,7 @@ def _rust_toolchain_tools_repository_impl(ctx):
         name = "rust_toolchain",
         exec_triple = exec_triple,
         allocator_library = ctx.attr.allocator_library,
+        global_allocator_library = ctx.attr.global_allocator_library,
         target_triple = target_triple,
         stdlib_linkflags = stdlib_linkflags,
         default_edition = ctx.attr.edition,
@@ -389,6 +393,9 @@ rust_toolchain_tools_repository = repository_rule(
             doc = "The Rust-style target that this compiler runs on",
             mandatory = True,
         ),
+        "global_allocator_library": attr.string(
+            doc = "Target that provides allocator functions when a global allocator is used with cc_common.link.",
+        ),
         "iso_date": attr.string(
             doc = "The date of the tool (or None, if the version is a specific version).",
         ),
@@ -470,6 +477,7 @@ def rust_toolchain_repository(
         target_settings = [],
         channel = None,
         allocator_library = None,
+        global_allocator_library = None,
         iso_date = None,
         rustfmt_version = None,
         edition = None,
@@ -490,6 +498,7 @@ def rust_toolchain_repository(
         target_compatible_with (list, optional): A list of constraints for the target platform for this toolchain.
         target_settings (list, optional): A list of config_settings that must be satisfied by the target configuration in order for this toolchain to be selected during toolchain resolution.
         allocator_library (str, optional): Target that provides allocator functions when rust_library targets are embedded in a cc_binary.
+        global_allocator_library (str, optional): Target that provides allocator functions when a global allocator is used with cc_common.link.
         iso_date (str, optional): The date of the tool.
         rustfmt_version (str, optional):  The version of rustfmt to be associated with the
             toolchain.
@@ -523,6 +532,7 @@ def rust_toolchain_repository(
         name = tools_repo_name,
         exec_triple = exec_triple,
         allocator_library = allocator_library,
+        global_allocator_library = global_allocator_library,
         target_triple = target_triple,
         iso_date = iso_date,
         version = version,
@@ -850,6 +860,7 @@ def rust_repository_set(
         version = None,
         versions = [],
         allocator_library = None,
+        global_allocator_library = None,
         extra_target_triples = [],
         iso_date = None,
         rustfmt_version = None,
@@ -871,6 +882,7 @@ def rust_repository_set(
             per channel. E.g. `["1.65.0", "nightly/2022-11-02", "beta/2020-12-30"]`.
         allocator_library (str, optional): Target that provides allocator functions when rust_library targets are
             embedded in a cc_binary.
+        global_allocator_library (str, optional): Target that provides allocator functions a global allocator is used with cc_common.link.
         extra_target_triples (list, optional): Additional rust-style targets that this set of
             toolchains should support.
         iso_date (str, optional): The date of the tool.
@@ -911,6 +923,7 @@ def rust_repository_set(
         all_toolchain_names.append(rust_toolchain_repository(
             name = toolchain.name,
             allocator_library = allocator_library,
+            global_allocator_library = global_allocator_library,
             auth = auth,
             channel = toolchain.channel.name,
             dev_components = dev_components,

rust/settings/BUILD.bazel

@@ -44,6 +44,14 @@ bool_flag(
     build_setting_default = False,
 )
 
+# A flag to indicate that a global allocator is in use when using --@rules_rust//rust/settings:experimental_use_cc_common_link
+# Users need to specify this flag because rustc generates different set of symbols at link time when a global allocator is in use.
+# When the linking is not done by rustc, the `rust_toolchain` itself provides the appropriate set of symbols.
+bool_flag(
+    name = "experimental_use_global_allocator",
+    build_setting_default = False,
+)
+
 bzl_library(
     name = "bzl_lib",
     srcs = glob(["**/*.bzl"]),

rust/toolchain.bzl

@@ -441,8 +441,18 @@ def _rust_toolchain_impl(ctx):
     pipelined_compilation = ctx.attr._pipelined_compilation[BuildSettingInfo].value
 
     experimental_use_cc_common_link = ctx.attr.experimental_use_cc_common_link[BuildSettingInfo].value
-    if experimental_use_cc_common_link and not ctx.attr.allocator_library:
-        fail("rust_toolchain.experimental_use_cc_common_link requires rust_toolchain.allocator_library to be set")
+    experimental_use_global_allocator = ctx.attr._experimental_use_global_allocator[BuildSettingInfo].value
+    if experimental_use_cc_common_link:
+        if experimental_use_global_allocator and not ctx.attr.global_allocator_library:
+            fail("rust_toolchain.experimental_use_cc_common_link with --@rules_rust//rust/settings:experimental_use_global_allocator " +
+                 "requires rust_toolchain.global_allocator_library to be set")
+        if not ctx.attr.allocator_library:
+            fail("rust_toolchain.experimental_use_cc_common_link requires rust_toolchain.allocator_library to be set")
+    if experimental_use_global_allocator and not experimental_use_cc_common_link:
+        fail(
+            "Using @rules_rust//rust/settings:experimental_use_global_allocator requires" +
+            "--@rules_rust//rust/settings:experimental_use_cc_common_link to be set",
+        )
 
     rust_std = ctx.attr.rust_std
 
@@ -564,6 +574,7 @@ def _rust_toolchain_impl(ctx):
         env = ctx.attr.env,
         exec_triple = exec_triple,
         libstd_and_allocator_ccinfo = _make_libstd_and_allocator_ccinfo(ctx, rust_std, ctx.attr.allocator_library),
+        libstd_and_global_allocator_ccinfo = _make_libstd_and_allocator_ccinfo(ctx, rust_std, ctx.attr.global_allocator_library),
         llvm_cov = ctx.file.llvm_cov,
         llvm_profdata = ctx.file.llvm_profdata,
         make_variables = make_variable_info,
@@ -591,6 +602,7 @@ def _rust_toolchain_impl(ctx):
         _third_party_dir = third_party_dir,
         _pipelined_compilation = pipelined_compilation,
         _experimental_use_cc_common_link = experimental_use_cc_common_link,
+        _experimental_use_global_allocator = experimental_use_global_allocator,
     )
     return [
         toolchain,
@@ -656,6 +668,9 @@ rust_toolchain = rule(
         "extra_rustc_flags": attr.string_list(
             doc = "Extra flags to pass to rustc in non-exec configuration",
         ),
+        "global_allocator_library": attr.label(
+            doc = "Target that provides allocator functions for when a global allocator is present.",
+        ),
         "llvm_cov": attr.label(
             doc = "The location of the `llvm-cov` binary. Can be a direct source or a filegroup containing one item. If None, rust code is not instrumented for coverage.",
             allow_single_file = True,
@@ -733,6 +748,13 @@ rust_toolchain = rule(
         "_cc_toolchain": attr.label(
             default = Label("@bazel_tools//tools/cpp:current_cc_toolchain"),
         ),
+        "_experimental_use_global_allocator": attr.label(
+            default = Label("//rust/settings:experimental_use_global_allocator"),
+            doc = (
+                "Label to a boolean build setting that informs the target build whether a global allocator is being used." +
+                "This flag is only relevant when used together with --@rules_rust//rust/settings:experimental_use_global_allocator."
+            ),
+        ),
         "_pipelined_compilation": attr.label(
             default = Label("//rust/settings:pipelined_compilation"),
         ),

test/cc_common_link/BUILD.bazel

@@ -13,6 +13,12 @@ cc_library(
     visibility = ["//visibility:public"],
 )
 
+cc_library(
+    name = "global_allocator_library",
+    srcs = ["global_allocator_library.cc"],
+    visibility = ["//visibility:public"],
+)
+
 cc_library(
     name = "cclinkstampdep",
     linkstamp = "cclinkstampdep.cc",

test/cc_common_link/global_allocator_library.cc

@@ -0,0 +1,52 @@
+#include <stdint.h>
+
+// This file has some exciting magic to get Rust code linking in a cc_binary.
+// The Rust compiler generates some similar symbol aliases when it links, so we
+// have to do it manually. We mark all our symbols as weak so that linking this
+// via Rust tooling to produce a binary with a Rust main works.
+//
+// It is intended to be used in rust_toolchain.allocator_library.
+//
+// https://github.com/rust-lang/rust/blob/master/library/alloc/src/alloc.rs
+// and https://github.com/rust-lang/rust/blob/master/library/std/src/alloc.rs
+// are the best source of docs I've found on these functions and variables.
+// https://doc.rust-lang.org/std/alloc/index.html talks about how this is
+// intended to be used.
+//
+// Also note
+// https://rust-lang.github.io/unsafe-code-guidelines/layout/scalars.html for
+// the sizes of the various integer types.
+//
+// This file strongly assumes that the default allocator is used. It will
+// not work with any other allocated switched in via `#[global_allocator]`.
+
+// New feature as of https://github.com/rust-lang/rust/pull/88098.
+__attribute__((weak)) uint8_t __rust_alloc_error_handler_should_panic = 0;
+
+extern "C" uint8_t *__rg_alloc(uintptr_t size, uintptr_t align);
+extern "C" __attribute__((weak))
+uint8_t *__rust_alloc(uintptr_t size, uintptr_t align) {
+  return __rg_alloc(size, align);
+}
+extern "C" void __rg_dealloc(uint8_t *ptr, uintptr_t size, uintptr_t align);
+extern "C" __attribute__((weak))
+void __rust_dealloc(uint8_t *ptr, uintptr_t size, uintptr_t align) {
+  __rg_dealloc(ptr, size, align);
+}
+extern "C" uint8_t *__rg_realloc(uint8_t *ptr, uintptr_t old_size, uintptr_t align,
+                       uintptr_t new_size);
+extern "C" __attribute__((weak))
+uint8_t *__rust_realloc(uint8_t *ptr, uintptr_t old_size, uintptr_t align,
+                        uintptr_t new_size) {
+  return __rg_realloc(ptr, old_size, align, new_size);
+}
+extern "C" uint8_t *__rg_alloc_zeroed(uintptr_t size, uintptr_t align);
+extern "C" __attribute__((weak))
+uint8_t *__rust_alloc_zeroed(uintptr_t size, uintptr_t align) {
+  return __rg_alloc_zeroed(size, align);
+}
+extern "C" void __rg_oom(uintptr_t size, uintptr_t align);
+extern "C" __attribute__((weak))
+void __rust_alloc_error_handler(uintptr_t size, uintptr_t align) {
+  __rg_oom(size, align);
+}

test/cc_common_link/with_global_alloc/BUILD.bazel

@@ -0,0 +1,17 @@
+load(
+    "@rules_rust//rust:defs.bzl",
+    "rust_binary",
+    "rust_test",
+)
+
+rust_binary(
+    name = "main",
+    srcs = ["main.rs"],
+    edition = "2021",
+)
+
+rust_test(
+    name = "global_alloc_test",
+    crate = ":main",
+    edition = "2021",
+)

test/cc_common_link/with_global_alloc/WORKSPACE.bazel

@@ -0,0 +1,25 @@
+workspace(name = "test_cc_common_link_with_global_alloc")
+
+local_repository(
+    name = "rules_rust",
+    path = "../../../",
+)
+
+local_repository(
+    name = "test_cc_common_link",
+    path = "../",
+)
+
+load("@rules_rust//rust:repositories.bzl", "rules_rust_dependencies", "rust_register_toolchains")
+
+rules_rust_dependencies()
+
+rust_register_toolchains(
+    allocator_library = "@test_cc_common_link//:allocator_library",
+    edition = "2018",
+    global_allocator_library = "@test_cc_common_link//:global_allocator_library",
+)
+
+load("@bazel_skylib//:workspace.bzl", "bazel_skylib_workspace")
+
+bazel_skylib_workspace()