From a7fe434086d3a48d91310b4c471e637f0f1b9793 Mon Sep 17 00:00:00 2001
From: Benjamin Lee <benjamin@computer.surgery>
Date: Thu, 23 May 2024 13:32:43 -0700
Subject: [PATCH] only link to one jemalloc build

Without setting JEMALLOC_OVERRIDE, we end up linking to two different
jemalloc builds. Once dynamically, as a transitive dependency through
rocksdb, and a second time to the static jemalloc that tikv-jemalloc-sys
builds.
---
 flake.nix                 |  4 ++--
 nix/pkgs/main/default.nix | 22 +++++++++++++++++++++-
 2 files changed, 23 insertions(+), 3 deletions(-)

diff --git a/flake.nix b/flake.nix
index 553ca023..ff47a625 100644
--- a/flake.nix
+++ b/flake.nix
@@ -184,8 +184,8 @@
             # Needed for building with io_uring
             pkgsHost.liburing
         ] else [])
-        ++
-        scopeHost.main.nativeBuildInputs;
+        ++ scopeHost.main.propagatedBuildInputs
+        ++ scopeHost.main.nativeBuildInputs;
       };
     });
 }
diff --git a/nix/pkgs/main/default.nix b/nix/pkgs/main/default.nix
index ec2aebc2..e4873e42 100644
--- a/nix/pkgs/main/default.nix
+++ b/nix/pkgs/main/default.nix
@@ -6,6 +6,7 @@
 , pkgsBuildHost
 , rocksdb
 , rust
+, rust-jemalloc-sys
 , stdenv
 
 # Options (keep sorted)
@@ -15,10 +16,27 @@
 }:
 
 let
+featureEnabled = feature : builtins.elem feature features;
+
+# This derivation will set the JEMALLOC_OVERRIDE variable, causing the
+# tikv-jemalloc-sys crate to use the nixpkgs jemalloc instead of building it's
+# own. In order for this to work, we need to set flags on the build that match
+# whatever flags tikv-jemalloc-sys was going to use. These are dependent on
+# which features we enable in tikv-jemalloc-sys.
+rust-jemalloc-sys' = (rust-jemalloc-sys.override {
+  # tikv-jemalloc-sys/unprefixed_malloc_on_supported_platforms feature
+  unprefixed = true;
+}).overrideAttrs (old: {
+  configureFlags = old.configureFlags ++
+    # tikv-jemalloc-sys/profiling feature
+    lib.optional (featureEnabled "jemalloc_prof") "--enable-prof";
+});
+
 buildDepsOnlyEnv =
   let
     rocksdb' = rocksdb.override {
-      enableJemalloc = builtins.elem "jemalloc" features;
+      jemalloc = rust-jemalloc-sys';
+      enableJemalloc = featureEnabled "jemalloc";
     };
   in
   {
@@ -60,6 +78,8 @@ commonAttrs = {
       ];
     };
 
+    buildInputs = lib.optional (featureEnabled "jemalloc") rust-jemalloc-sys';
+
     nativeBuildInputs = [
       # bindgen needs the build platform's libclang. Apparently due to "splicing
       # weirdness", pkgs.rustPlatform.bindgenHook on its own doesn't quite do the