perf: add custom memory allocator for improved performance (#2021)

Author: baszalmstra

.github/workflows/rust.yml

@@ -113,7 +113,8 @@ jobs:
       CARGO_PROFILE_OPT_LEVEL: ${{ (github.ref == 'refs/heads/main' || startsWith(github.ref, 'refs/tags/')) && 's' || '0' }}
 
       # Sets the default features which are always used also when `--no-default` is used.
-      REQUIRED_FEATURES: recipe-generation
+      # Enable the `performance` feature for release builds to use optimized memory allocators.
+      REQUIRED_FEATURES: ${{ (github.ref == 'refs/heads/main' || startsWith(github.ref, 'refs/tags/')) && 'recipe-generation,performance' || 'recipe-generation' }}
 
     steps:
       - name: Checkout source code

Cargo.lock

@@ -12,7 +12,7 @@ default-run = "rattler-build"
 rust-version = "1.89.0"
 
 [workspace]
-members = ["rust-tests", "crates/rattler_build_recipe_generator"]
+members = ["rust-tests", "crates/rattler_build_recipe_generator", "crates/rattler_build_allocator"]
 
 [workspace.package]
 authors = ["rattler-build contributors <[email protected]>"]
@@ -52,6 +52,10 @@ zip = { version = "6.0.0", default-features = false, features = ["deflate"] }
 
 [features]
 default = ['rustls-tls', 'recipe-generation', 's3']
+# Enable optimized memory allocators for release builds.
+# This provides significant performance improvements, especially on Linux musl.
+# Disabled by default due to increased compile time.
+performance = ["rattler_build_allocator"]
 native-tls = [
   'reqwest/native-tls',
   'rattler/native-tls',
@@ -219,6 +223,7 @@ dialoguer = "0.12.0"
 rattler_build_recipe_generator = { path = "crates/rattler_build_recipe_generator", optional = true, features = [
   "cli",
 ] }
+rattler_build_allocator = { path = "crates/rattler_build_allocator", optional = true }
 
 [target.'cfg(not(target_os = "windows"))'.dependencies]
 sha2 = { workspace = true, features = ["asm"] }

Cargo.toml

@@ -0,0 +1,19 @@
+[package]
+name = "rattler_build_allocator"
+version = "0.1.0"
+authors.workspace = true
+edition.workspace = true
+license.workspace = true
+repository.workspace = true
+description = "A crate that provides the best memory allocator for different platforms"
+
+[lib]
+doctest = false
+
+# Use mimalloc on Windows
+[target.'cfg(windows)'.dependencies]
+mimalloc = "0.1.43"
+
+# Use jemalloc on supported unix platforms (excluding OpenBSD and FreeBSD)
+[target.'cfg(all(not(windows), not(target_os = "openbsd"), not(target_os = "freebsd"), any(target_arch = "x86_64", target_arch = "aarch64", target_arch = "powerpc64")))'.dependencies]
+tikv-jemallocator = "0.6.0"

crates/rattler_build_allocator/Cargo.toml

@@ -0,0 +1,28 @@
+//! This crate provides the best memory allocator for different platforms.
+//!
+//! On Windows, we use mimalloc because it provides good performance.
+//! On most Unix platforms (Linux, macOS), we use jemalloc for its excellent
+//! performance characteristics with multi-threaded applications.
+//!
+//! This crate is designed to be used as a dependency that, when included,
+//! automatically sets the global allocator. Simply add this crate as a
+//! dependency and the allocator will be configured.
+
+// Use mimalloc on Windows
+#[cfg(windows)]
+#[global_allocator]
+static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
+
+// Use jemalloc on supported unix platforms
+#[cfg(all(
+    not(windows),
+    not(target_os = "openbsd"),
+    not(target_os = "freebsd"),
+    any(
+        target_arch = "x86_64",
+        target_arch = "aarch64",
+        target_arch = "powerpc64"
+    )
+))]
+#[global_allocator]
+static GLOBAL: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;

crates/rattler_build_allocator/src/lib.rs

@@ -1,5 +1,10 @@
 //! This is the main entry point for the `rattler-build` binary.
 
+// Use custom allocators for improved performance when the `performance` feature is enabled.
+// This must be at the crate root to set the global allocator.
+#[cfg(feature = "performance")]
+use rattler_build_allocator as _;
+
 use std::{
     fs::File,
     io::{self, IsTerminal},