tcmalloc rust bindings

.bazelrc

@@ -23,6 +23,9 @@ build --incompatible_disallow_empty_glob
 # for details.
 common --incompatible_config_setting_private_default_visibility
 
+# cc static lib for rust
+build --experimental_cc_static_library
+
 # Link Google Test against Abseil and RE2.
 build --define='absl=1'
 

BUILD

@@ -11,3 +11,41 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+
+load("@bazel_skylib//rules:native_binary.bzl", "native_binary")
+load("@rules_cc//cc:defs.bzl", "cc_import")
+load("@rules_rust_bindgen//:defs.bzl", "rust_bindgen_toolchain")
+
+native_binary(
+    name = "clang",
+    src = "@llvm_toolchain_llvm//:bin/clang",
+    visibility = ["//tcmalloc_rs:__subpackages__"],
+)
+
+cc_import(
+    name = "libclang",
+    shared_library = "@llvm_toolchain_llvm//:libclang",
+    visibility = ["//tcmalloc_rs:__subpackages__"],
+)
+
+cc_import(
+    name = "libstdcxx",
+    static_library = "@llvm_toolchain_llvm//:lib/x86_64-unknown-linux-gnu/libc++.a",
+    visibility = ["//tcmalloc_rs:__subpackages__"],
+)
+
+rust_bindgen_toolchain(
+    name = "rust_bindgen_toolchain",
+    bindgen = "@rules_rust_bindgen//3rdparty:bindgen",
+    clang = ":clang",
+    libclang = ":libclang",
+    libstdcxx = ":libstdcxx",
+    visibility = ["//tcmalloc_rs:__subpackages__"],
+)
+
+toolchain(
+    name = "default_bindgen_toolchain",
+    toolchain = ":rust_bindgen_toolchain",
+    toolchain_type = "@rules_rust_bindgen//:toolchain_type",
+    visibility = ["//tcmalloc_rs:__subpackages__"],
+)

MODULE.bazel

@@ -29,3 +29,38 @@ bazel_dep(name = "fuzztest", version = "20241028.0", dev_dependency = True, repo
 bazel_dep(name = "google_benchmark", version = "1.9.2", dev_dependency = True, repo_name = "com_github_google_benchmark")
 bazel_dep(name = "googletest", version = "1.13.0", dev_dependency = True, repo_name = "com_google_googletest")
 bazel_dep(name = "rules_fuzzing", version = "0.5.2", dev_dependency = True)
+
+bazel_dep(name = "toolchains_llvm", version = "1.4.0")
+bazel_dep(name = "rules_rust", version = "0.61.0")
+bazel_dep(name = "rules_rust_bindgen", version = "0.61.0")
+
+# Configure and register the toolchain.
+# https://github.com/bazel-contrib/toolchains_llvm/blob/master/toolchain/internal/llvm_distributions.bzl
+llvm = use_extension("@toolchains_llvm//toolchain/extensions:llvm.bzl", "llvm", dev_dependency = True)
+
+# LLVM toolchain.
+llvm.toolchain(
+    name = "llvm_toolchain",
+    cxx_standard = {"": "c++17"},
+    llvm_version = "17.0.6",
+)
+use_repo(llvm, "llvm_toolchain", "llvm_toolchain_llvm")
+
+rust = use_extension("@rules_rust//rust:extensions.bzl", "rust", dev_dependency = True)
+rust.toolchain(
+    edition = "2021",
+    versions = [
+        "1.85.0",
+    ],
+)
+use_repo(rust, "rust_toolchains")
+
+register_toolchains(
+    "@rust_toolchains//:all",
+    dev_dependency = True,
+)
+
+register_toolchains(
+    "//:default_bindgen_toolchain",
+    dev_dependency = True,
+)

docs/quickstart.md

@@ -254,6 +254,36 @@ build --cxxopt='-std=c++17'
 
 Congratulations! You've created your first binary using TCMalloc.
 
+## TCMalloc-rs
+This is the rust wrapper for using `TCMalloc` as a global allocator.
+
+#### Pagesize
+By default, `tcmalloc` build with support for `8k` pagesize.
+
+To build `tcmalloc` with `32k` pagesizes, build `tcmalloc_rs` with `--//tcmalloc_rs:pagesize=32k`
+
+To build `tcmalloc` with `512k` pagesizes, build `tcmalloc_rs` with `--//tcmalloc_rs:pagesize=512k`
+### Running Tests
+`bazel run //tcmalloc_rs:test --extra_toolchains=@llvm_toolchain//:cc-toolchain-x86_64-linux`
+
+### Usage
+Add the bazel dependency:
+```Starlark
+//tcmalloc_rs
+```
+To use the allocator, add the following to your code:
+```rust
+// To set TCMalloc as the global allocator add this to your project:
+use tcmalloc_rs;
+
+#[global_allocator]
+static GLOBAL: tcmalloc_rs::TCMalloc = tcmalloc_rs::TCMalloc;
+```
+
+#### Running the example
+Take a look at the code in `tcmalloc_rs/main.rs` and run
+`bazel run //tcmalloc_rs:main --extra_toolchains=@llvm_toolchain//:cc-toolchain-x86_64-linux`.
+
 ## What's Next
 
 *   Read our [overview](overview.md), if you haven't already. The overview

tcmalloc/BUILD

@@ -25,7 +25,11 @@ package(default_visibility = ["//visibility:private"])
 
 licenses(["notice"])
 
-exports_files(["LICENSE"])
+exports_files([
+    "LICENSE",
+    "tcmalloc.h",
+    "malloc_extension.h",
+])
 
 config_setting(
     name = "llvm",

tcmalloc/selsan/BUILD

@@ -38,7 +38,10 @@ cc_library(
         "//tcmalloc/internal:logging",
         "@com_google_absl//absl/base:core_headers",
     ],
-    alwayslink = 1,
+    alwayslink = select({
+        "//tcmalloc:unlinked": 0,
+        "//conditions:default": 1,
+    }),
 )
 
 cc_test(

tcmalloc_rs/BUILD

@@ -0,0 +1,129 @@
+load("@bazel_skylib//rules:common_settings.bzl", "string_flag")
+load("@rules_rust//rust:defs.bzl", "rust_binary", "rust_library", "rust_test")
+load("@rules_rust_bindgen//:defs.bzl", "rust_bindgen_library")
+
+string_flag(
+    name = "pagesize",
+    build_setting_default = "8k",
+)
+
+config_setting(
+    name = "large",
+    flag_values = {
+        ":pagesize": "32k",
+    },
+    visibility = ["//visibility:public"],
+)
+
+config_setting(
+    name = "xlarge",
+    flag_values = {
+        ":pagesize": "256k",
+    },
+    visibility = ["//visibility:public"],
+)
+
+cc_static_library(
+    name = "tcmalloc",
+    visibility = ["//visibility:public"],
+    deps = select({
+        ":xlarge": ["//tcmalloc:tcmalloc_256k_pages"],
+        ":large": ["//tcmalloc:tcmalloc_large_pages"],
+        "//conditions:default": ["//tcmalloc"],
+    }),
+)
+
+rust_bindgen_library(
+    name = "tcmalloc_sys_extensions",
+    bindgen_flags = [
+        "--allowlist-function=.*MallocExtension_.*",
+        "--allowlist-type=.*MallocExtension_Ownership",
+        "--allowlist-type=.*MallocExtension_Property",
+        "--blocklist-function=.*MallocExtension_GetProfileSamplingRate",
+        "--blocklist-function=.*MallocExtension_SetProfileSamplingRate",
+        "--blocklist-function=.*MallocExtension_GetGuardedSamplingRate",
+        "--blocklist-function=.*MallocExtension_SetGuardedSamplingRate",
+        "--blocklist-function=.*MallocExtension_GetSkipSubreleaseInterval",
+        "--blocklist-function=.*MallocExtension_SetSkipSubreleaseInterval",
+        "--opaque-type=std::optional",  # prevents unsafe comparisons which could result in overflow
+        "--opaque-type=std::basic_string",  # prevents unsafe comparisons which could result in overflow
+        "--opaque-type=std::basic_string_view",  # allow bindgen to use __BindgenOpaqueArray to repr std::string
+        "--opaque-type=std::unique_ptr",  # prevents unsafe comparisons which could result in overflow
+        "--use-core",  # don't compile with std, allows baremetal envs
+    ],
+    cc_lib = "//tcmalloc:malloc_extension",
+    clang_flags = [
+        # "-v",  # enable for debugging
+        "-x",
+        "c++",
+        "-std=c++17",
+    ],
+    header = "//tcmalloc:malloc_extension.h",
+    alwayslink = 1,
+)
+
+rust_bindgen_library(
+    name = "tcmalloc_sys",
+    bindgen_flags = [
+        "--allowlist-function=malloc",
+        "--allowlist-function=free",
+        "--allowlist-function=calloc",
+        "--allowlist-function=realloc",
+        "--allowlist-function=malloc_usable_size",
+        # comment out the above and uncomment the two lines below to generate bindings for everything
+        # "--opaque-type=std::.*",
+        # "--blocklist-item=std::value",
+        "--use-core",  # don't compile with std, allows baremetal envs
+    ],
+    cc_lib = ":tcmalloc",
+    clang_flags = [
+        # "-v",  # enable for debugging
+        "-x",
+        "c++",
+        "-std=c++17",
+    ],
+    header = "//tcmalloc:tcmalloc.h",
+)
+
+rust_library(
+    name = "tcmalloc_rs",
+    srcs = ["lib.rs"],
+    visibility = ["//visibility:public"],
+    deps = [":tcmalloc_sys"],
+)
+
+rust_library(
+    name = "tcmalloc_extensions_rs",
+    srcs = ["extension.rs"],
+    visibility = ["//visibility:public"],
+    deps = [":tcmalloc_sys_extensions"],
+)
+
+# bazel test //tcmalloc_rs/... --extra_toolchains=@llvm_toolchain//:cc-toolchain-x86_64-linux
+rust_test(
+    name = "test_malloc",
+    size = "small",
+    crate = ":tcmalloc_rs",
+)
+
+rust_test(
+    name = "test_malloc_extension",
+    size = "small",
+    crate = ":tcmalloc_extensions_rs",
+    deps = [":tcmalloc_rs"],
+)
+
+# bazel run //tcmalloc_rs:main --extra_toolchains=@llvm_toolchain//:cc-toolchain-x86_64-linux
+rust_binary(
+    name = "main",
+    srcs = ["main.rs"],
+    rustc_flags = [
+        "-Ccodegen-units=1",
+        "-Copt-level=3",
+        "-Cstrip=symbols",
+    ],
+    deps = [
+        ":tcmalloc_extensions_rs",
+        ":tcmalloc_rs",
+    ],
+)