From a94387dcdc818a47e263d65090f7c08c37a7a899 Mon Sep 17 00:00:00 2001
From: strawberry <strawberry@puppygock.gay>
Date: Wed, 10 Apr 2024 20:31:00 -0400
Subject: [PATCH] maybe add multiple outputs support to nix flake for jemalloc
 and hardened_malloc

Signed-off-by: strawberry <strawberry@puppygock.gay>
---
 Cargo.lock |   1 +
 Cargo.toml |   3 +-
 flake.nix  | 124 ++++++++++++++++++++++++++++++++++++++++++++++++++++-
 3 files changed, 124 insertions(+), 4 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index f5c438b5..a60c52d4 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2482,6 +2482,7 @@ dependencies = [
  "libz-sys",
  "lz4-sys",
  "pkg-config",
+ "tikv-jemalloc-sys",
  "zstd-sys",
 ]
 
diff --git a/Cargo.toml b/Cargo.toml
index 6073c05d..9bc2ccb6 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -326,12 +326,11 @@ default = [
     "gzip_compression",
     "brotli_compression",
     "zstd_compression",
-    "hardened_malloc",
 ]
 backend_sqlite = ["sqlite"]
 backend_rocksdb = ["rocksdb"]
 rocksdb = ["rust-rocksdb", "num_cpus"]
-jemalloc = ["tikv-jemalloc-ctl", "tikv-jemallocator"]
+jemalloc = ["tikv-jemalloc-ctl", "tikv-jemallocator", "rust-rocksdb/jemalloc"]
 sqlite = ["rusqlite", "parking_lot", "thread_local", "num_cpus"]
 systemd = ["sd-notify"]
 sentry_telemetry = ["sentry", "sentry-tracing", "sentry-tower"]
diff --git a/flake.nix b/flake.nix
index 26a8aac4..46f11b6e 100644
--- a/flake.nix
+++ b/flake.nix
@@ -41,7 +41,7 @@
         let
           version = "9.0.0";
         in
-        pkgs.rocksdb.overrideAttrs (old: {
+        (pkgs.rocksdb.overrideAttrs (old: {
           inherit version;
           src = pkgs.fetchFromGitHub {
             owner = "girlbossceo";
@@ -49,7 +49,7 @@
             rev = "449768a833b79c267c584b5ab1d50e73db6faf9d";
             hash = "sha256-MjmGfAlZ5WC2+hFH6nEUprqBjO8xiTQh2HJIqQ5mIg8=";
           };
-        });
+        }));
 
       # Nix-accessible `Cargo.toml`
       cargoToml = builtins.fromTOML (builtins.readFile ./Cargo.toml);
@@ -185,6 +185,50 @@
         meta.mainProgram = cargoToml.package.name;
       };
 
+      packageJemalloc = pkgs: builder pkgs {
+        src = nix-filter {
+          root = ./.;
+          include = [
+            "src"
+            "Cargo.toml"
+            "Cargo.lock"
+          ];
+        };
+
+        rocksdb' = pkgs.rocksdb.override { enableJemalloc = true; };
+
+        # This is redundant with CI
+        doCheck = false;
+
+        env = env pkgs;
+        nativeBuildInputs = nativeBuildInputs pkgs;
+
+        meta.mainProgram = cargoToml.package.name;
+
+        buildFeatures = [ "jemalloc" ];
+      };
+
+      packageHMalloc = pkgs: builder pkgs {
+        src = nix-filter {
+          root = ./.;
+          include = [
+            "src"
+            "Cargo.toml"
+            "Cargo.lock"
+          ];
+        };
+
+        # This is redundant with CI
+        doCheck = false;
+
+        env = env pkgs;
+        nativeBuildInputs = nativeBuildInputs pkgs;
+
+        meta.mainProgram = cargoToml.package.name;
+
+        buildFeatures = [ "hardened_malloc" ];
+      };
+
       mkOciImage = pkgs: package:
         pkgs.dockerTools.buildImage {
           name = package.pname;
@@ -204,11 +248,55 @@
             ];
           };
         };
+
+        mkOciImageJemalloc = pkgs: packageJemalloc:
+          pkgs.dockerTools.buildImage {
+            name = "${package.pname}-jemalloc";
+            tag = "main";
+            copyToRoot = [
+              pkgs.dockerTools.caCertificates
+            ];
+            config = {
+              # Use the `tini` init system so that signals (e.g. ctrl+c/SIGINT)
+              # are handled as expected
+              Entrypoint = [
+                "${pkgs.lib.getExe' pkgs.tini "tini"}"
+                "--"
+              ];
+              Cmd = [
+                "${pkgs.lib.getExe package}"
+              ];
+            };
+          };
+
+          mkOciImageHMalloc = pkgs: packageHMalloc:
+            pkgs.dockerTools.buildImage {
+            name = "${package.pname}-hmalloc";
+              tag = "main";
+              copyToRoot = [
+                pkgs.dockerTools.caCertificates
+              ];
+              config = {
+                # Use the `tini` init system so that signals (e.g. ctrl+c/SIGINT)
+                # are handled as expected
+                Entrypoint = [
+                  "${pkgs.lib.getExe' pkgs.tini "tini"}"
+                  "--"
+                ];
+                Cmd = [
+                  "${pkgs.lib.getExe package}"
+                ];
+              };
+            };
     in
     {
       packages = {
         default = package pkgsHost;
+        jemalloc = packageJemalloc pkgsHost;
+        hmalloc = packageHMalloc pkgsHost;
         oci-image = mkOciImage pkgsHost self.packages.${system}.default;
+        oci-image-jemalloc = mkOciImageJemalloc pkgsHost self.packages.${system}.default;
+        oci-image-hmalloc = mkOciImageHMalloc pkgsHost self.packages.${system}.default;
 
         book =
           let
@@ -261,6 +349,18 @@
                   value = package pkgsCrossStatic;
                 }
 
+                # An output for a statically-linked binary with jemalloc
+                {
+                  name = "${binaryName}-jemalloc";
+                  value = packageJemalloc pkgsCrossStatic;
+                }
+
+                # An output for a statically-linked binary with hardened_malloc
+                {
+                  name = "${binaryName}-hmalloc";
+                  value = packageHMalloc pkgsCrossStatic;
+                }
+
                 # An output for an OCI image based on that binary
                 {
                   name = "oci-image-${crossSystem}";
@@ -268,11 +368,31 @@
                     pkgsCrossStatic
                     self.packages.${system}.${binaryName};
                 }
+
+                # An output for an OCI image based on that binary with jemalloc
+                {
+                  name = "oci-image-${crossSystem}-jemalloc";
+                  value = mkOciImageJemalloc
+                    pkgsCrossStatic
+                    self.packages.${system}.${binaryName};
+                }
+
+                # An output for an OCI image based on that binary with hardened_malloc
+                {
+                  name = "oci-image-${crossSystem}-hmalloc";
+                  value = mkOciImageHMalloc
+                    pkgsCrossStatic
+                    self.packages.${system}.${binaryName};
+                }
               ]
             )
             [
               "x86_64-unknown-linux-musl"
+              "x86_64-unknown-linux-musl-jemalloc"
+              "x86_64-unknown-linux-musl-hmalloc"
               "aarch64-unknown-linux-musl"
+              "aarch64-unknown-linux-musl-jemalloc"
+              "aarch64-unknown-linux-musl-hmalloc"
             ]
           )
         );