tcmalloc rust bindings

Author: jayakasadev

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

@@ -28,3 +28,39 @@ bazel_dep(name = "protobuf", version = "27.5", repo_name = "com_google_protobuf"
 bazel_dep(name = "re2", version = "2024-02-01", repo_name = "com_googlesource_code_re2")
 bazel_dep(name = "rules_cc", version = "0.0.9")
 bazel_dep(name = "rules_fuzzing", version = "0.5.2")
+bazel_dep(name = "toolchains_llvm", version = "1.3.0")
+
+bazel_dep(name = "rules_rust", version = "0.59.2")
+bazel_dep(name = "rules_rust_bindgen", version = "0.59.2")
+
+
+# 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",
+    llvm_version = "17.0.6",
+    cxx_standard = {"": "c++19"},
+)
+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,42 @@ 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.
+### Building
+To build the `tcmalloc` rust bindings, you need to turn off the `alwayslink` flag on `cc_library` targets. 
+This is done with the `--//tcmalloc:alwayslink=false` flag.
+
+`//tcmalloc:tcmalloc` with `alwayslink=1` will produce `libtcmalloc.lo` which rustc does not recognize as a valid static library ([related](https://github.com/bazelbuild/rules_rust/issues/325#issuecomment-706548453)).
+Optionally turning off `alwayslink` for `tcmalloc_sys` will allow `bindgen` to successfully build.
+
+#### 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 --//tcmalloc:alwayslink=false`
+
+### 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 --//tcmalloc:alwayslink=false`.
+
 ## What's Next
 
 *   Read our [overview](overview.md), if you haven't already. The overview

tcmalloc/BUILD

@@ -18,14 +18,28 @@
 # https://github.com/google/tcmalloc/tree/master/docs/design.md for a high-level description of
 # how this malloc works.
 
+load("@bazel_skylib//rules:common_settings.bzl", "bool_flag")
 load("//tcmalloc:copts.bzl", "TCMALLOC_DEFAULT_COPTS")
 load("//tcmalloc:variants.bzl", "create_tcmalloc_benchmark", "create_tcmalloc_libraries", "create_tcmalloc_testsuite")
 
 package(default_visibility = ["//visibility:private"])
 
 licenses(["notice"])
 
-exports_files(["LICENSE"])
+exports_files(["LICENSE", "tcmalloc.h"])
+
+bool_flag(
+    name = "alwayslink",
+    build_setting_default = True,
+)
+
+config_setting(
+    name = "unlinked",
+    flag_values = {
+        ":alwayslink": "false",
+    },
+    visibility = ["//visibility:public"],
+)
 
 config_setting(
     name = "llvm",
@@ -108,7 +122,10 @@ cc_library(
         "@com_google_absl//absl/time",
         "@com_google_absl//absl/types:span",
     ],
-    alwayslink = 1,
+    alwayslink = select({
+        ":unlinked": 0,
+        "//conditions:default": 1,
+    }),
 )
 
 cc_library(
@@ -334,7 +351,10 @@ create_tcmalloc_libraries(
         "@com_google_absl//absl/types:optional",
         "@com_google_absl//absl/types:span",
     ],
-    alwayslink = 1,
+    alwayslink = select({
+        ":unlinked": 0,
+        "//conditions:default": 1,
+    }),
 )
 
 # TEMPORARY. WILL BE REMOVED.
@@ -435,7 +455,10 @@ cc_library(
         "@com_google_absl//absl/time",
         "@com_google_absl//absl/types:span",
     ],
-    alwayslink = 1,
+    alwayslink = select({
+        ":unlinked": 0,
+        "//conditions:default": 1,
+    }),
 )
 
 # TCMalloc with 256k pages is usually faster but fragmentation is higher.  See
@@ -458,7 +481,10 @@ cc_library(
         "@com_google_absl//absl/time",
         "@com_google_absl//absl/types:span",
     ],
-    alwayslink = 1,
+    alwayslink = select({
+        ":unlinked": 0,
+        "//conditions:default": 1,
+    }),
 )
 
 cc_library(

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,75 @@
+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"],
+)
+
+rust_bindgen_library(
+    name = "tcmalloc_sys",
+    cc_lib = select({
+        ":xlarge": "//tcmalloc:tcmalloc_256k_pages",
+        ":large": "//tcmalloc:tcmalloc_large_pages",
+        "//conditions:default": "//tcmalloc",
+    }),
+    header = "//tcmalloc:tcmalloc.h",
+    clang_flags = [
+        # "-v",  # enable for debugging
+        "-x",
+        "c++",
+        "-std=c++17",
+    ],
+    bindgen_flags = [
+        "--allowlist-function=TCMallocInternalMalloc",
+        "--allowlist-function=TCMallocInternalFree",
+        "--allowlist-function=TCMallocInternalCalloc",
+        "--allowlist-function=TCMallocInternalRealloc",
+        "--allowlist-function=malloc_usable_size",
+        "--use-core", # don't compile with std, allows baremetal envs
+    ],
+    rustc_flags = [
+        "-Ccodegen-units=1",
+        "-Copt-level=3",
+        "-Cstrip=symbols",
+    ],
+)
+
+rust_library(
+    name = "tcmalloc_rs",
+    deps = [":tcmalloc_sys"],
+    srcs = ["lib.rs"],
+    visibility = ["//visibility:public"],
+)
+
+
+# bazel test //tcmalloc_rs/... --extra_toolchains=@llvm_toolchain//:cc-toolchain-x86_64-linux --//tcmalloc:alwayslink=false
+rust_test(
+    name = "test",
+    size = "small",
+    crate = ":tcmalloc_rs",
+)
+
+# bazel run //tcmalloc_rs:main --extra_toolchains=@llvm_toolchain//:cc-toolchain-x86_64-linux --//tcmalloc:alwayslink=false
+rust_binary(
+    name = "main",
+    srcs = ["main.rs"],
+    deps = [":tcmalloc_rs"],
+)