fix jemalloc linking in CI

.envrc

@@ -1,6 +1,6 @@
 #!/usr/bin/env bash
 
-use flake
+use flake ".#${DIRENV_DEVSHELL:-default}"
 
 PATH_add bin
 

.github/workflows/ci.yml

@@ -35,6 +35,11 @@ env:
     # Custom nix binary cache if fork is being used
     ATTIC_ENDPOINT: ${{ vars.ATTIC_ENDPOINT }}
     ATTIC_PUBLIC_KEY: ${{ vars.ATTIC_PUBLIC_KEY }}
+    # Use the all-features devshell instead of default, to ensure that features
+    # match between nix and cargo
+    DIRENV_DEVSHELL: all-features
+    # Get error output from nix that we can actually use
+    NIX_CONFIG: show-trace = true
 
 permissions:
     packages: write
@@ -92,7 +97,7 @@ jobs:
                   echo 'source $HOME/.nix-profile/share/nix-direnv/direnvrc' > "$HOME/.direnvrc"
                   nix profile install --impure --inputs-from . nixpkgs#direnv nixpkgs#nix-direnv
                   direnv allow
-                  nix develop --command true
+                  nix develop .#all-features --command true
 
             - name: Run CI tests
               run: |
@@ -202,7 +207,7 @@ jobs:
                   echo 'source $HOME/.nix-profile/share/nix-direnv/direnvrc' > "$HOME/.direnvrc"
                   nix profile install --impure --inputs-from . nixpkgs#direnv nixpkgs#nix-direnv
                   direnv allow
-                  nix develop --command true
+                  nix develop .#all-features --command true
 
             - name: Build static ${{ matrix.target }}
               run: |

bin/complement

@@ -21,7 +21,7 @@ toplevel="$(git rev-parse --show-toplevel)"
 
 pushd "$toplevel" > /dev/null
 
-bin/nix-build-and-cache just .#complement
+bin/nix-build-and-cache just .#static-complement
 
 docker load < result
 popd > /dev/null

bin/nix-build-and-cache

@@ -70,7 +70,7 @@ ci() {
         --inputs-from "$toplevel"
 
         # Keep sorted
-        "$toplevel#devShells.x86_64-linux.default"
+        "$toplevel#devShells.x86_64-linux.all-features"
         attic#default
         nixpkgs#direnv
         nixpkgs#jq

flake.nix

@@ -16,6 +16,7 @@
     inputs.flake-utils.lib.eachDefaultSystem (system:
     let
       pkgsHost = inputs.nixpkgs.legacyPackages.${system};
+      pkgsHostStatic = pkgsHost.pkgsStatic;
 
       # The Rust toolchain to use
       toolchain = inputs.fenix.packages.${system}.fromToolchainFile {
@@ -25,7 +26,8 @@
         sha256 = "sha256-+syqAd2kX8KVa8/U2gz3blIQTTsYYt3U63xBWaGOSc8";
       };
 
-      scope = pkgs: pkgs.lib.makeScope pkgs.newScope (self: {
+      mkScope = pkgs: pkgs.lib.makeScope pkgs.newScope (self: {
+        inherit pkgs;
         book = self.callPackage ./nix/pkgs/book {};
         complement = self.callPackage ./nix/pkgs/complement {};
         craneLib = ((inputs.crane.mkLib pkgs).overrideToolchain toolchain);
@@ -41,7 +43,59 @@
         });
       });
 
-      scopeHost = (scope pkgsHost);
+      scopeHost = mkScope pkgsHost;
+      scopeHostStatic = mkScope pkgsHostStatic;
+
+      mkDevShell = scope: scope.pkgs.mkShell {
+        env = scope.main.env // {
+          # Rust Analyzer needs to be able to find the path to default crate
+          # sources, and it can read this environment variable to do so. The
+          # `rust-src` component is required in order for this to work.
+          RUST_SRC_PATH = "${toolchain}/lib/rustlib/src/rust/library";
+
+          # Convenient way to access a pinned version of Complement's source
+          # code.
+          COMPLEMENT_SRC = inputs.complement.outPath;
+
+          # Needed for Complement
+          CGO_CFLAGS = "-I${scope.pkgs.olm}/include";
+          CGO_LDFLAGS = "-L${scope.pkgs.olm}/lib";
+        };
+
+        # Development tools
+        packages = [
+          # Always use nightly rustfmt because most of its options are unstable
+          #
+          # This needs to come before `toolchain` in this list, otherwise
+          # `$PATH` will have stable rustfmt instead.
+          inputs.fenix.packages.${system}.latest.rustfmt
+
+          toolchain
+        ]
+        ++ (with pkgsHost.pkgs; [
+          engage
+          cargo-audit
+
+          # Needed for producing Debian packages
+          cargo-deb
+
+          # Needed for Complement
+          go
+
+          # Needed for our script for Complement
+          jq
+
+          # Needed for finding broken markdown links
+          lychee
+
+          # Useful for editing the book locally
+          mdbook
+        ])
+        ++ scope.main.propagatedBuildInputs
+        ++ scope.main.nativeBuildInputs;
+
+        meta.broken = scope.main.meta.broken;
+      };
     in
     {
       packages = {
@@ -64,6 +118,7 @@
         book = scopeHost.book;
 
         complement = scopeHost.complement;
+        static-complement = scopeHostStatic.complement;
       }
       //
       builtins.listToAttrs
@@ -79,7 +134,7 @@
                       config = crossSystem;
                     };
                   }).pkgsStatic;
-                scopeCrossStatic = scope pkgsCrossStatic;
+                scopeCrossStatic = mkScope pkgsCrossStatic;
               in
               [
                 # An output for a statically-linked binary
@@ -138,54 +193,15 @@
           )
         );
 
-      devShells.default = pkgsHost.mkShell {
-        env = scopeHost.main.env // {
-          # Rust Analyzer needs to be able to find the path to default crate
-          # sources, and it can read this environment variable to do so. The
-          # `rust-src` component is required in order for this to work.
-          RUST_SRC_PATH = "${toolchain}/lib/rustlib/src/rust/library";
-
-          # Convenient way to access a pinned version of Complement's source
-          # code.
-          COMPLEMENT_SRC = inputs.complement.outPath;
-        };
-
-        # Development tools
-        packages = [
-          # Always use nightly rustfmt because most of its options are unstable
-          #
-          # This needs to come before `toolchain` in this list, otherwise
-          # `$PATH` will have stable rustfmt instead.
-          inputs.fenix.packages.${system}.latest.rustfmt
-
-          toolchain
-        ]
-        ++ (with pkgsHost; [
-          engage
-          cargo-audit
-
-          # Needed for producing Debian packages
-          cargo-deb
-
-          # Needed for Complement
-          go
-          olm
-
-          # Needed for our script for Complement
-          jq
-
-          # Needed for finding broken markdown links
-          lychee
-
-          # Useful for editing the book locally
-          mdbook
-        ])
-        ++ (if !pkgsHost.stdenv.isDarwin then [
-            # Needed for building with io_uring
-            pkgsHost.liburing
-        ] else [])
-        ++
-        scopeHost.main.nativeBuildInputs;
-      };
+      devShells.default = mkDevShell scopeHostStatic;
+      devShells.all-features = mkDevShell
+        (scopeHostStatic.overrideScope (final: prev: {
+          main = prev.main.override { all_features = true; };
+        }));
+      devShells.no-features = mkDevShell
+        (scopeHostStatic.overrideScope (final: prev: {
+          main = prev.main.override { default_features = false; };
+        }));
+      devShells.dynamic = mkDevShell scopeHost;
     });
 }

nix/pkgs/main/default.nix

@@ -6,20 +6,71 @@
 , pkgsBuildHost
 , rocksdb
 , rust
+, rust-jemalloc-sys
 , stdenv
 
 # Options (keep sorted)
 , default_features ? true
+, all_features ? false
 , features ? []
 , profile ? "release"
 }:
 
 let
+# We perform default-feature unification in nix, because some of the dependencies
+# on the nix side depend on feature values.
+workspaceMembers = builtins.map (member: "${inputs.self}/src/${member}")
+  (builtins.attrNames (builtins.readDir "${inputs.self}/src"));
+crateFeatures = path:
+  let manifest = lib.importTOML "${path}/Cargo.toml"; in
+  lib.remove "default" (lib.attrNames manifest.features) ++
+  lib.attrNames
+    (lib.filterAttrs
+      (_: dependency: dependency.optional or false)
+      manifest.dependencies);
+crateDefaultFeatures = path:
+  (lib.importTOML "${path}/Cargo.toml").features.default;
+allDefaultFeatures = lib.unique
+  (lib.flatten (builtins.map crateDefaultFeatures workspaceMembers));
+allFeatures = lib.unique
+  (lib.flatten (builtins.map crateFeatures workspaceMembers));
+features' = lib.unique
+  (features ++
+    lib.optionals default_features allDefaultFeatures ++
+    lib.optionals all_features allFeatures);
+
+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;
-    };
+    rocksdb' = (rocksdb.override {
+      jemalloc = rust-jemalloc-sys';
+      # rocksdb fails to build with prefixed jemalloc, which is required on
+      # darwin due to [1]. In this case, fall back to building rocksdb with
+      # libc malloc. This should not cause conflicts, because all of the
+      # jemalloc symbols are prefixed.
+      #
+      # [1]: https://github.com/tikv/jemallocator/blob/ab0676d77e81268cd09b059260c75b38dbef2d51/jemalloc-sys/src/env.rs#L17
+      enableJemalloc = featureEnabled "jemalloc" && !stdenv.isDarwin;
+    }).overrideAttrs (old: {
+      # TODO: static rocksdb fails to build on darwin
+      # build log at <https://girlboss.ceo/~strawberry/pb/JjGH>
+      meta.broken = stdenv.hostPlatform.isStatic && stdenv.isDarwin;
+    });
   in
   {
     CARGO_PROFILE = profile;
@@ -60,6 +111,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
@@ -82,13 +135,10 @@ craneLib.buildPackage ( commonAttrs // {
     env = buildDepsOnlyEnv;
   });
 
-  cargoExtraArgs = ""
-    + lib.optionalString
-      (!default_features)
-      "--no-default-features "
+  cargoExtraArgs = "--no-default-features "
     + lib.optionalString
-      (features != [])
-      "--features " + (builtins.concatStringsSep "," features);
+      (features' != [])
+      "--features " + (builtins.concatStringsSep "," features');
 
   # This is redundant with CI
   cargoTestCommand = "";
@@ -105,4 +155,11 @@ craneLib.buildPackage ( commonAttrs // {
   };
 
   meta.mainProgram = commonAttrs.pname;
+  # Dynamically-linked jemalloc is broken on linux due to link-order problems,
+  # where the symbols are being resolved to libc malloc/free before jemalloc is
+  # loaded. This problem does not occur on darwin for unknown reasons.
+  meta.broken =
+    stdenv.isLinux &&
+    !stdenv.hostPlatform.isStatic &&
+    (featureEnabled "jemalloc");
 })