postgresql: always build with JIT enabled

This changes the build to always enable JIT - but to only enable it at
run-time, when required. This keeps the runtime closure small without
JIT, but allows enabling it without a rebuild. We can do this, because
JIT is actually built as a shared module, which is loaded at run-time.
We put it into a -jit output and only link it into the environment when
requested.

Under the hood, this uses withPackages and adds the "JIT package" -
thus, to be able to use withPackages on top of that, we also need to be
able to apply withPackages repeatedly.

This cuts down the number of NixOS tests in half, because we don't need
to run it for every version with and without JIT anymore. There really
is no point in running everything with llvmjit.so in place, when the
queries are not making use of it anyway.

Also, we only need to build each extension once and not twice, further
reducing the number of rebuilds required for PRs touching postgresql.

nixos/tests/postgresql/anonymizer.nix

@@ -20,7 +20,6 @@ let
           services.postgresql = {
             inherit package;
             enable = true;
-            enableJIT = lib.hasInfix "-jit-" package.name;
             extensions = ps: [ ps.anonymizer ];
             settings.shared_preload_libraries = [ "anon" ];
           };

nixos/tests/postgresql/citus.nix

@@ -35,7 +35,6 @@ let
           services.postgresql = {
             inherit package;
             enable = true;
-            enableJIT = lib.hasInfix "-jit-" package.name;
             extensions =
               ps: with ps; [
                 citus

nixos/tests/postgresql/pgjwt.nix

@@ -24,7 +24,6 @@ let
           services.postgresql = {
             inherit package;
             enable = true;
-            enableJIT = lib.hasInfix "-jit-" package.name;
             extensions =
               ps: with ps; [
                 pgjwt

nixos/tests/postgresql/pgvecto-rs.nix

@@ -38,7 +38,6 @@ let
           services.postgresql = {
             inherit package;
             enable = true;
-            enableJIT = lib.hasInfix "-jit-" package.name;
             extensions =
               ps: with ps; [
                 pgvecto-rs

nixos/tests/postgresql/postgresql-jit.nix

@@ -51,5 +51,4 @@ let
 in
 genTests {
   inherit makeTestFor;
-  filter = n: _: lib.hasSuffix "_jit" n;
 }

nixos/tests/postgresql/postgresql-tls-client-cert.nix

@@ -51,7 +51,6 @@ let
           services.postgresql = {
             inherit package;
             enable = true;
-            enableJIT = lib.hasInfix "-jit-" package.name;
             enableTCPIP = true;
             ensureUsers = [
               {

nixos/tests/postgresql/postgresql-wal-receiver.nix

@@ -32,7 +32,6 @@ let
           services.postgresql = {
             inherit package;
             enable = true;
-            enableJIT = lib.hasInfix "-jit-" package.name;
             settings = {
               max_replication_slots = 10;
               max_wal_senders = 10;

nixos/tests/postgresql/postgresql.nix

@@ -54,7 +54,6 @@ let
           services.postgresql = {
             inherit package;
             enable = true;
-            enableJIT = lib.hasInfix "-jit-" package.name;
             # TODO(@Ma27) split this off into its own VM test and move a few other
             # extension tests to use postgresqlTestExtension.
             extensions = ps: with ps; [ plv8 ];

nixos/tests/postgresql/timescaledb.nix

@@ -54,7 +54,6 @@ let
           services.postgresql = {
             inherit package;
             enable = true;
-            enableJIT = lib.hasInfix "-jit-" package.name;
             extensions =
               ps: with ps; [
                 timescaledb

nixos/tests/postgresql/tsja.nix

@@ -21,7 +21,6 @@ let
           services.postgresql = {
             inherit package;
             enable = true;
-            enableJIT = lib.hasInfix "-jit-" package.name;
             extensions =
               ps: with ps; [
                 tsja

nixos/tests/postgresql/wal2json.nix

@@ -17,7 +17,6 @@ let
         services.postgresql = {
           inherit package;
           enable = true;
-          enableJIT = lib.hasInfix "-jit-" package.name;
           extensions = with package.pkgs; [ wal2json ];
           settings = {
             wal_level = "logical";

pkgs/by-name/pg/pg_checksums/package.nix

@@ -1,13 +1,13 @@
 {
   lib,
-  stdenv,
+  clangStdenv,
   fetchFromGitHub,
   libxslt,
   docbook_xsl,
   postgresql,
 }:
 
-stdenv.mkDerivation rec {
+clangStdenv.mkDerivation rec {
   pname = "pg_checksums";
   version = "1.2";
 

pkgs/by-name/pg/pgcopydb/package.nix

@@ -1,6 +1,6 @@
 {
   lib,
-  stdenv,
+  clangStdenv,
   fetchFromGitHub,
   libkrb5,
   openssl,
@@ -13,7 +13,7 @@
   zlib,
 }:
 
-stdenv.mkDerivation (finalAttrs: {
+clangStdenv.mkDerivation (finalAttrs: {
   pname = "pgcopydb";
   version = "0.15";
 
@@ -38,7 +38,7 @@ stdenv.mkDerivation (finalAttrs: {
       sqlite
       zlib
     ]
-    ++ lib.optionals stdenv.hostPlatform.isLinux [
+    ++ lib.optionals clangStdenv.hostPlatform.isLinux [
       pam
     ];
 

pkgs/servers/sql/postgresql/default.nix

@@ -22,12 +22,10 @@ let
       version: path:
       let
         attrName = if jitSupport then "${version}_jit" else version;
+        postgresql = import path { inherit self; };
+        attrValue = if jitSupport then postgresql.withJIT else postgresql.withoutJIT;
       in
-      self.lib.nameValuePair attrName (
+      self.lib.nameValuePair attrName attrValue
-        import path {
-          inherit jitSupport self;
-        }
-      )
     ) versions;
 
   libpq = self.callPackage ./libpq.nix { };
@@ -35,7 +33,8 @@ let
 in
 {
   # variations without and with JIT
-  postgresqlVersions = mkAttributes false // mkAttributes true;
+  postgresqlVersions = mkAttributes false;
+  postgresqlJitVersions = mkAttributes true;
 
   inherit libpq;
 }

pkgs/servers/sql/postgresql/generic.nix

@@ -66,7 +66,11 @@ let
       icu,
 
       # JIT
-      jitSupport, # not default on purpose, this is set via "_jit or not" attributes
+      jitSupport ?
+        stdenv.hostPlatform.canExecute stdenv.buildPlatform
+        # Building with JIT in pkgsStatic fails like this:
+        #   fatal error: 'stdio.h' file not found
+        && !stdenv.hostPlatform.isStatic,
       llvmPackages,
       nukeReferences,
       overrideCC,
@@ -135,10 +139,8 @@ let
 
       dlSuffix = if olderThan "16" then ".so" else stdenv.hostPlatform.extensions.sharedLibrary;
 
-      pname = "postgresql";
-
       stdenv' =
-        if jitSupport && !stdenv.cc.isClang then
+        if !stdenv.cc.isClang then
           overrideCC llvmPackages.stdenv (
             llvmPackages.stdenv.cc.override {
               # LLVM bintools are not used by default, but are needed to make -flto work below.
@@ -150,7 +152,7 @@ let
     in
     stdenv'.mkDerivation (finalAttrs: {
       inherit version;
-      pname = pname + lib.optionalString jitSupport "-jit";
+      pname = "postgresql";
 
       src = fetchFromGitHub {
         owner = "postgres";
@@ -162,8 +164,6 @@ let
 
       __structuredAttrs = true;
 
-      hardeningEnable = lib.optionals (!stdenv'.cc.isClang) [ "pie" ];
-
       outputs =
         [
           "out"
@@ -172,51 +172,67 @@ let
           "lib"
           "man"
         ]
+        ++ lib.optionals jitSupport [ "jit" ]
         ++ lib.optionals perlSupport [ "plperl" ]
         ++ lib.optionals pythonSupport [ "plpython3" ]
         ++ lib.optionals tclSupport [ "pltcl" ];
-      outputChecks = {
-        out = {
-          disallowedReferences = [
-            "dev"
-            "doc"
-            "man"
-          ];
-          disallowedRequisites = [
-            stdenv'.cc
-            llvmPackages.llvm.out
-          ] ++ (map lib.getDev (builtins.filter (drv: drv ? "dev") finalAttrs.buildInputs));
-        };
 
-        lib = {
+      outputChecks =
-          disallowedReferences = [
+        {
-            "out"
+          out = {
-            "dev"
+            disallowedReferences = [
-            "doc"
+              "dev"
-            "man"
+              "doc"
-          ];
+              "man"
-          disallowedRequisites = [
+            ] ++ lib.optionals jitSupport [ "jit" ];
-            stdenv'.cc
+            disallowedRequisites = [
-            llvmPackages.llvm.out
+              stdenv'.cc
-          ] ++ (map lib.getDev (builtins.filter (drv: drv ? "dev") finalAttrs.buildInputs));
+              llvmPackages.llvm.out
-        };
+            ] ++ (map lib.getDev (builtins.filter (drv: drv ? "dev") finalAttrs.buildInputs));
+          };
 
-        doc = {
+          lib = {
-          disallowedReferences = [
+            disallowedReferences = [
-            "out"
+              "out"
-            "dev"
+              "dev"
-            "man"
+              "doc"
-          ];
+              "man"
-        };
+            ] ++ lib.optionals jitSupport [ "jit" ];
+            disallowedRequisites = [
+              stdenv'.cc
+              llvmPackages.llvm.out
+            ] ++ (map lib.getDev (builtins.filter (drv: drv ? "dev") finalAttrs.buildInputs));
+          };
 
-        man = {
+          doc = {
-          disallowedReferences = [
+            disallowedReferences = [
-            "out"
+              "out"
-            "dev"
+              "dev"
-            "doc"
+              "man"
-          ];
+            ] ++ lib.optionals jitSupport [ "jit" ];
+          };
+
+          man = {
+            disallowedReferences = [
+              "out"
+              "dev"
+              "doc"
+            ] ++ lib.optionals jitSupport [ "jit" ];
+          };
+        }
+        // lib.optionalAttrs jitSupport {
+          jit = {
+            disallowedReferences = [
+              "dev"
+              "doc"
+              "man"
+            ];
+            disallowedRequisites = [
+              stdenv'.cc
+              llvmPackages.llvm.out
+            ] ++ (map lib.getDev (builtins.filter (drv: drv ? "dev") finalAttrs.buildInputs));
+          };
         };
-      };
 
       strictDeps = true;
 
@@ -271,9 +287,7 @@ let
           # flags will remove unused sections from all shared libraries and binaries - including
           # those paths. This avoids a lot of circular dependency problems with different outputs,
           # and allows splitting them cleanly.
-          CFLAGS =
+          CFLAGS = "-fdata-sections -ffunction-sections -flto";
-            "-fdata-sections -ffunction-sections"
-            + (if stdenv'.cc.isClang then " -flto" else " -fmerge-constants -Wl,--gc-sections");
         }
         // lib.optionalAttrs perlSupport { PERL = lib.getExe perl; }
         // lib.optionalAttrs pythonSupport { PYTHON = lib.getExe python3; }
@@ -364,11 +378,25 @@ let
 
       installTargets = [ "install-world" ];
 
-      postPatch = ''
+      postPatch =
-        substituteInPlace "src/Makefile.global.in" --subst-var out
+        ''
-        substituteInPlace "src/common/config_info.c" --subst-var dev
+          substituteInPlace "src/Makefile.global.in" --subst-var out
-        cat ${./pg_config.env.mk} >> src/common/Makefile
+          substituteInPlace "src/common/config_info.c" --subst-var dev
-      '';
+          cat ${./pg_config.env.mk} >> src/common/Makefile
+        ''
+        # This check was introduced upstream to prevent calls to "exit" inside libpq.
+        # However, this doesn't work reliably with static linking, see this and following:
+        # https://postgr.es/m/flat/20210703001639.GB2374652%40rfd.leadboat.com#52584ca4bd3cb9dac376f3158c419f97
+        # Thus, disable the check entirely, as it would currently fail with this:
+        # > libpq.so.5.17:                  U atexit
+        # > libpq.so.5.17:                  U pthread_exit
+        # > libpq must not be calling any function which invokes exit
+        # Don't mind the fact that this checks libpq.**so** in pkgsStatic - that's correct, since PostgreSQL
+        # still needs a shared library internally.
+        + lib.optionalString (atLeast "15" && stdenv'.hostPlatform.isStatic) ''
+          substituteInPlace src/interfaces/libpq/Makefile \
+            --replace-fail "echo 'libpq must not be calling any function which invokes exit'; exit 1;" "echo;"
+        '';
 
       postInstall =
         ''
@@ -415,6 +443,9 @@ let
 
           # Stop lib depending on the -dev output of llvm
           remove-references-to -t ${llvmPackages.llvm.dev} "$out/lib/llvmjit${dlSuffix}"
+
+          moveToOutput "lib/bitcode" "$jit"
+          moveToOutput "lib/llvmjit*" "$jit"
         ''
         + lib.optionalString stdenv'.hostPlatform.isDarwin ''
           # The darwin specific Makefile for PGXS contains a reference to the postgres
@@ -466,17 +497,14 @@ let
       passthru =
         let
           this = self.callPackage generic args;
-          jitToggle = this.override {
-            jitSupport = !jitSupport;
-          };
         in
         {
           inherit dlSuffix;
 
           psqlSchema = lib.versions.major version;
 
-          withJIT = if jitSupport then this else jitToggle;
+          withJIT = this.withPackages (_: [ this.jit ]);
-          withoutJIT = if jitSupport then jitToggle else this;
+          withoutJIT = this;
 
           pkgs =
             let
@@ -586,18 +614,38 @@ let
         inherit installedExtensions;
         inherit (postgresql)
           pg_config
+          pkgs
           psqlSchema
           version
           ;
 
         withJIT = postgresqlWithPackages {
-          inherit buildEnv lib makeBinaryWrapper;
+          inherit
-          postgresql = postgresql.withJIT;
+            buildEnv
-        } f;
+            lib
+            makeBinaryWrapper
+            postgresql
+            ;
+        } (_: installedExtensions ++ [ postgresql.jit ]);
         withoutJIT = postgresqlWithPackages {
-          inherit buildEnv lib makeBinaryWrapper;
+          inherit
-          postgresql = postgresql.withoutJIT;
+            buildEnv
-        } f;
+            lib
+            makeBinaryWrapper
+            postgresql
+            ;
+        } (_: lib.remove postgresql.jit installedExtensions);
+
+        withPackages =
+          f':
+          postgresqlWithPackages {
+            inherit
+              buildEnv
+              lib
+              makeBinaryWrapper
+              postgresql
+              ;
+          } (ps: installedExtensions ++ f' ps);
       };
     };
 

pkgs/top-level/all-packages.nix

@@ -12455,6 +12455,7 @@ with pkgs;
 
   inherit (import ../servers/sql/postgresql pkgs)
     postgresqlVersions
+    postgresqlJitVersions
     libpq
     ;
 
@@ -12464,7 +12465,9 @@ with pkgs;
     postgresql_15
     postgresql_16
     postgresql_17
+    ;
 
+  inherit (postgresqlJitVersions)
     postgresql_13_jit
     postgresql_14_jit
     postgresql_15_jit