feat(tvix): Switch to jemalloc r/8222

Switch tvix to using jemalloc as the default global allocator on
supported (eg, non-msvc) platforms. This gives a pretty reasonable
performance boost basically for free:

int                     time:   [24.361 µs 24.386 µs 24.418 µs]
                        change: [-19.355% -18.859% -18.527%] (p = 0.00 < 0.05)
                        Performance has improved.

merge small attrs       time:   [37.201 µs 37.328 µs 37.442 µs]
                        change: [-24.609% -24.266% -23.982%] (p = 0.00 < 0.05)
                        Performance has improved.

merge large attrs with small attrs
                        time:   [20.030 ms 20.135 ms 20.251 ms]
                        change: [-4.2811% -3.2549% -2.3807%] (p = 0.00 < 0.05)
                        Performance has improved.

hello outpath           time:   [967.91 ms 974.07 ms 983.33 ms]
                        change: [-5.5008% -4.4552% -3.4099%] (p = 0.00 < 0.05)
                        Performance has improved.

Change-Id: I6c6e6e3295ec2fca01ea28dc37bcb201cd811767
Reviewed-on: https://cl.tvl.fyi/c/depot/+/10851
Tested-by: BuildkiteCI
Reviewed-by: sterni <sternenseemann@systemli.org>
Autosubmit: aspen <root@gws.fyi>
Reviewed-by: flokli <flokli@flokli.de>

8 files changed, 132 insertions, 0 deletions

diff --git a/tvix/Cargo.lock b/tvix/Cargo.lock
index 4f66da69389a..584290084bde 100644
--- a/tvix/Cargo.lock
+++ b/tvix/Cargo.lock
@@ -3661,6 +3661,26 @@ dependencies = [
 ]
 
 [[package]]
+name = "tikv-jemalloc-sys"
+version = "0.5.4+5.3.0-patched"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9402443cb8fd499b6f327e40565234ff34dbda27460c5b47db0db77443dd85d1"
+dependencies = [
+ "cc",
+ "libc",
+]
+
+[[package]]
+name = "tikv-jemallocator"
+version = "0.5.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "965fe0c26be5c56c94e38ba547249074803efd52adfb66de62107d95aab3eaca"
+dependencies = [
+ "libc",
+ "tikv-jemalloc-sys",
+]
+
+[[package]]
 name = "time"
 version = "0.3.34"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -4181,6 +4201,7 @@ dependencies = [
  "rnix",
  "rustyline",
  "thiserror",
+ "tikv-jemallocator",
  "tokio",
  "tracing",
  "tracing-subscriber",
@@ -4225,6 +4246,7 @@ dependencies = [
  "tabwriter",
  "tempfile",
  "test-strategy",
+ "tikv-jemallocator",
  "toml",
  "tvix-eval-builtin-macros",
 ]
@@ -4265,6 +4287,7 @@ dependencies = [
  "sha2",
  "tempfile",
  "thiserror",
+ "tikv-jemallocator",
  "tokio",
  "tokio-tar",
  "tokio-util",
diff --git a/tvix/Cargo.nix b/tvix/Cargo.nix
index 694356d1eab2..1f92aed2830d 100644
--- a/tvix/Cargo.nix
+++ b/tvix/Cargo.nix
@@ -11190,6 +11190,72 @@ rec {
         ];
         features = { };
       };
+      "tikv-jemalloc-sys" = rec {
+        crateName = "tikv-jemalloc-sys";
+        version = "0.5.4+5.3.0-patched";
+        edition = "2018";
+        links = "jemalloc";
+        sha256 = "1lc5vm1p9dqdvd3mn3264zddnd7z6i95ch3y69prnjgxp0y480ll";
+        authors = [
+          "Alex Crichton <alex@alexcrichton.com>"
+          "Gonzalo Brito Gadeschi <gonzalobg88@gmail.com>"
+          "The TiKV Project Developers"
+        ];
+        dependencies = [
+          {
+            name = "libc";
+            packageId = "libc";
+            usesDefaultFeatures = false;
+          }
+        ];
+        buildDependencies = [
+          {
+            name = "cc";
+            packageId = "cc";
+          }
+        ];
+        features = {
+          "background_threads" = [ "background_threads_runtime_support" ];
+          "default" = [ "background_threads_runtime_support" ];
+        };
+        resolvedDefaultFeatures = [ "background_threads_runtime_support" ];
+      };
+      "tikv-jemallocator" = rec {
+        crateName = "tikv-jemallocator";
+        version = "0.5.4";
+        edition = "2018";
+        sha256 = "1jpanfm9az8hcbg6dyxdabykx03lj0j4g9cbwfa6rig5dg1f0pwn";
+        authors = [
+          "Alex Crichton <alex@alexcrichton.com>"
+          "Gonzalo Brito Gadeschi <gonzalobg88@gmail.com>"
+          "Simon Sapin <simon.sapin@exyr.org>"
+          "Steven Fackler <sfackler@gmail.com>"
+          "The TiKV Project Developers"
+        ];
+        dependencies = [
+          {
+            name = "libc";
+            packageId = "libc";
+            usesDefaultFeatures = false;
+          }
+          {
+            name = "tikv-jemalloc-sys";
+            packageId = "tikv-jemalloc-sys";
+            usesDefaultFeatures = false;
+          }
+        ];
+        features = {
+          "background_threads" = [ "tikv-jemalloc-sys/background_threads" ];
+          "background_threads_runtime_support" = [ "tikv-jemalloc-sys/background_threads_runtime_support" ];
+          "debug" = [ "tikv-jemalloc-sys/debug" ];
+          "default" = [ "background_threads_runtime_support" ];
+          "disable_initial_exec_tls" = [ "tikv-jemalloc-sys/disable_initial_exec_tls" ];
+          "profiling" = [ "tikv-jemalloc-sys/profiling" ];
+          "stats" = [ "tikv-jemalloc-sys/stats" ];
+          "unprefixed_malloc_on_supported_platforms" = [ "tikv-jemalloc-sys/unprefixed_malloc_on_supported_platforms" ];
+        };
+        resolvedDefaultFeatures = [ "background_threads_runtime_support" "default" ];
+      };
       "time" = rec {
         crateName = "time";
         version = "0.3.34";
@@ -13221,6 +13287,11 @@ rec {
             packageId = "thiserror";
           }
           {
+            name = "tikv-jemallocator";
+            packageId = "tikv-jemallocator";
+            target = { target, features }: (!("msvc" == target."env" or null));
+          }
+          {
             name = "tokio";
             packageId = "tokio";
           }
@@ -13414,6 +13485,11 @@ rec {
             name = "tempfile";
             packageId = "tempfile";
           }
+          {
+            name = "tikv-jemallocator";
+            packageId = "tikv-jemallocator";
+            target = { target, features }: (!("msvc" == target."env" or null));
+          }
         ];
         features = {
           "arbitrary" = [ "proptest" "test-strategy" "imbl/proptest" ];
@@ -13535,6 +13611,11 @@ rec {
             packageId = "thiserror";
           }
           {
+            name = "tikv-jemallocator";
+            packageId = "tikv-jemallocator";
+            target = { target, features }: (!("msvc" == target."env" or null));
+          }
+          {
             name = "tokio";
             packageId = "tokio";
           }
diff --git a/tvix/cli/Cargo.toml b/tvix/cli/Cargo.toml
index be85597fa2ae..d3c05172e456 100644
--- a/tvix/cli/Cargo.toml
+++ b/tvix/cli/Cargo.toml
@@ -26,3 +26,6 @@ tracing-subscriber = "0.3.16"
 
 [dependencies.wu-manber]
 git = "https://github.com/tvlfyi/wu-manber.git"
+
+[target.'cfg(not(target_env = "msvc"))'.dependencies]
+tikv-jemallocator = "0.5"
diff --git a/tvix/cli/src/main.rs b/tvix/cli/src/main.rs
index 292a223cbb1a..66ec5a99dba6 100644
--- a/tvix/cli/src/main.rs
+++ b/tvix/cli/src/main.rs
@@ -18,6 +18,13 @@ use tvix_glue::tvix_io::TvixIO;
 use tvix_glue::tvix_store_io::TvixStoreIO;
 use tvix_glue::{builtins::add_derivation_builtins, configure_nix_path};
 
+#[cfg(not(target_env = "msvc"))]
+use tikv_jemallocator::Jemalloc;
+
+#[cfg(not(target_env = "msvc"))]
+#[global_allocator]
+static GLOBAL: Jemalloc = Jemalloc;
+
 #[derive(Parser, Clone)]
 struct Args {
     #[arg(long)]
diff --git a/tvix/eval/Cargo.toml b/tvix/eval/Cargo.toml
index cc590791cc92..4cf8ea146cf8 100644
--- a/tvix/eval/Cargo.toml
+++ b/tvix/eval/Cargo.toml
@@ -42,6 +42,9 @@ pretty_assertions = "1.2.1"
 rstest = "0.19.0"
 tempfile = "3.3.0"
 
+[target.'cfg(not(target_env = "msvc"))'.dev-dependencies]
+tikv-jemallocator = "0.5"
+
 [features]
 default = ["impure", "arbitrary", "nix_tests"]
 
diff --git a/tvix/eval/benches/eval.rs b/tvix/eval/benches/eval.rs
index 57d4eb71b59f..1333f5018cb3 100644
--- a/tvix/eval/benches/eval.rs
+++ b/tvix/eval/benches/eval.rs
@@ -1,5 +1,11 @@
 use criterion::{black_box, criterion_group, criterion_main, Criterion};
 use itertools::Itertools;
+#[cfg(not(target_env = "msvc"))]
+use tikv_jemallocator::Jemalloc;
+
+#[cfg(not(target_env = "msvc"))]
+#[global_allocator]
+static GLOBAL: Jemalloc = Jemalloc;
 
 fn interpret(code: &str) {
     tvix_eval::Evaluation::new_pure().evaluate(code, None);
diff --git a/tvix/glue/Cargo.toml b/tvix/glue/Cargo.toml
index 0afdefeaaa0e..dcc46d599096 100644
--- a/tvix/glue/Cargo.toml
+++ b/tvix/glue/Cargo.toml
@@ -33,6 +33,9 @@ walkdir = "2.4.0"
 [dependencies.wu-manber]
 git = "https://github.com/tvlfyi/wu-manber.git"
 
+[target.'cfg(not(target_env = "msvc"))'.dependencies]
+tikv-jemallocator = "0.5"
+
 [dev-dependencies]
 criterion = { version = "0.5", features = ["html_reports"] }
 hex-literal = "0.4.1"
diff --git a/tvix/glue/benches/eval.rs b/tvix/glue/benches/eval.rs
index 202278c1aa01..9e0154cad787 100644
--- a/tvix/glue/benches/eval.rs
+++ b/tvix/glue/benches/eval.rs
@@ -1,6 +1,8 @@
 use criterion::{black_box, criterion_group, criterion_main, Criterion};
 use lazy_static::lazy_static;
 use std::{env, rc::Rc, sync::Arc, time::Duration};
+#[cfg(not(target_env = "msvc"))]
+use tikv_jemallocator::Jemalloc;
 use tvix_build::buildservice::DummyBuildService;
 use tvix_eval::{builtins::impure_builtins, EvalIO};
 use tvix_glue::{
@@ -11,6 +13,10 @@ use tvix_glue::{
 };
 use tvix_store::utils::construct_services;
 
+#[cfg(not(target_env = "msvc"))]
+#[global_allocator]
+static GLOBAL: Jemalloc = Jemalloc;
+
 lazy_static! {
     static ref TOKIO_RUNTIME: tokio::runtime::Runtime = tokio::runtime::Runtime::new().unwrap();
 }