From 2282fa73a1d654c3c7b7429d8de49947c11bd0e1 Mon Sep 17 00:00:00 2001
From: Maximilian Bosch
Date: Sat, 18 Mar 2023 09:54:54 +0100
Subject: postgresql: implement opt-in JIT support

Closes #150801

Note: I decided against resuming directly on #150801 because the
conflict was too big (and resolving it seemed too error-prone to me).
Also the `this`-refactoring could be done in an easier manner, i.e. by
exposing JIT attributes with the correct configuration. More on that
below.

This patch creates variants of the `postgresql*`-packages with JIT[1]
support. Please note that a lot of the work was derived from previous
patches filed by other contributors, namely dasJ, andir and abbradar,
hence the co-authored-by tags below.

Effectively, the following things have changed:

* For JIT variants an LLVM-backed stdenv with clang is now used as
  suggested by dasJ[2]. We need LLVM and CLang[3] anyways to build the
  JIT-part, so no need to mix this up with GCC's stdenv. Also, using the
  `dev`-output of LLVM and clang's stdenv for building (and adding llvm
  libs as build-inputs) seems more cross friendly to me (which will
  become useful when cross-building for JIT-variants will actually be
  supported).

* Plugins inherit the build flags from the Makefiles in
  `$out/lib/pgxs/src` (e.g. `-Werror=unguarded-availability-new`). Since
  some of the flags are clang-specific (and stem from the use of the
  CLang stdenv) and don't work on gcc, the stdenv of `pkgs.postgresql`
  is passed to the plugins. I.e., plugins for non-JIT variants are built
  with a gcc stdenv on Linux and plugins for JIT variants with a clang
  stdenv.

  Since `plv8` hard-codes `gcc` as `$CC` in its Makefile[4], I marked it
  as broken for JIT-variants of postgresql only.

* Added a test-matrix to confirm that JIT works fine on each
  `pkgs.postgresql_*_jit` (thanks Andi for the original test in
  #124804!).

* For each postgresql version, a new attribute
  `postgresql_<version>_jit` (and a corresponding
  `postgresqlPackages<version>JitPackages`) are now exposed for better
  discoverability and prebuilt artifacts in the binary cache.

* In #150801 the `this`-argument was replaced by an internal recursion.
  I decided against this approach because it'd blow up the diff even
  more which makes the readability way harder and also harder to revert
  this if necessary.

  Instead, it is made sure that `this` always points to the correct
  variant of `postgresql` and re-using that in an additional
  `.override {}`-expression is trivial because the JIT-variant is
  exposed in `all-packages.nix`.

* I think the changes are sufficiently big to actually add myself as
  maintainer here.

* Added `libxcrypt` to `buildInputs` for versions <v13. While
  building things with an LLVM stdenv, these versions complained that
  the extern `crypt()` symbol can't be found. Not sure what this is
  exactly about, but since we want to switch to libxcrypt for `crypt()`
  usage anyways[5] I decided to add it. For >=13 it's not relevant
  anymore anyways[6].

* JIT support doesn't work with cross-compilation. It is attempted to
  build LLVM-bytecode (`%.bc` is the corresponding `make(1)`-rule) for
  each sub-directory in `backend/` for the JIT apparently, but with a
  $(CLANG) that can produce binaries for the build, not the host-platform.

  I managed to get a cross-build with JIT support working with
  `depsBuildBuild = [ llvmPackages.clang ] ++ buildInputs`, but
  considering that the resulting LLVM IR isn't platform-independent this
  doesn't give you much. In fact, I tried to test the result in a VM-test,
  but as soon as JIT was used to optimize a query, postgres would
  coredump with `Illegal instruction`.

A common concern of the original approach - with llvm as build input -
was the massive increase of closure size. With the new approach of using
the LLVM stdenv directly and patching out references to the clang drv in
`$out` the effective closure size changes are:

    $ nix path-info -Sh $(nix-build -A postgresql_14)
    /nix/store/kssxxqycwa3c7kmwmykwxqvspxxa6r1w-postgresql-14.7	306.4M
    $ nix path-info -Sh $(nix-build -A postgresql_14_jit)
    /nix/store/xc7qmgqrn4h5yr4vmdwy56gs4bmja9ym-postgresql-14.7	689.2M

Most of the increase in closure-size stems from the `lib`-output of
LLVM

    $ nix path-info -Sh /nix/store/5r97sbs5j6mw7qnbg8nhnq1gad9973ap-llvm-11.1.0-lib
    /nix/store/5r97sbs5j6mw7qnbg8nhnq1gad9973ap-llvm-11.1.0-lib	349.8M

which is why this shouldn't be enabled by default.

While this is quite much because of LLVM, it's still a massive
improvement over the simple approach of adding llvm/clang as
build-inputs and building with `--with-llvm`:

    $ nix path-info -Sh $(nix-build -E '
	with import ./. {};
	postgresql.overrideAttrs ({ configureFlags ? [], buildInputs ? [], ... }: {
	  configureFlags = configureFlags ++ [ "--with-llvm" ];
	  buildInputs = buildInputs ++ [ llvm clang ];
	})' -j0)
    /nix/store/i3bd2r21c6c3428xb4gavjnplfqxn27p-postgresql-14.7	  1.6G

Co-authored-by: Andreas Rammhold <andreas@rammhold.de>
Co-authored-by: Janne Heß <janne@hess.ooo>
Co-authored-by: Nikolay Amiantov <ab@fmap.me>

[1] https://www.postgresql.org/docs/current/jit-reason.html
[2] https://github.com/NixOS/nixpkgs/pull/124804#issuecomment-864616931
    & https://github.com/NixOS/nixpkgs/pull/150801#issuecomment-1467868321
[3] This fails with the following error otherwise:
    ```
    configure: error: clang not found, but required when compiling --with-llvm, specify with CLANG=
    ```
[4] https://github.com/plv8/plv8/blob/v3.1.5/Makefile#L14
[5] https://github.com/NixOS/nixpkgs/pull/181764
[6] https://github.com/postgres/postgres/commit/c45643d618e35ec2fe91438df15abd4f3c0d85ca
---
 nixos/doc/manual/release-notes/rl-2305.section.md |  11 ++
 nixos/tests/all-tests.nix                         |   1 +
 nixos/tests/postgresql-jit.nix                    |  50 +++++
 nixos/tests/postgresql-wal-receiver.nix           |   2 +-
 pkgs/servers/sql/postgresql/default.nix           | 211 +++++++++++++++-------
 pkgs/servers/sql/postgresql/ext/plv8/default.nix  |   1 +
 pkgs/servers/sql/postgresql/ext/postgis.nix       |   2 +-
 pkgs/top-level/all-packages.nix                   |  13 ++
 8 files changed, 227 insertions(+), 64 deletions(-)
 create mode 100644 nixos/tests/postgresql-jit.nix

diff --git a/nixos/doc/manual/release-notes/rl-2305.section.md b/nixos/doc/manual/release-notes/rl-2305.section.md
index deeda73d2fb5..2312cea87500 100644
--- a/nixos/doc/manual/release-notes/rl-2305.section.md
+++ b/nixos/doc/manual/release-notes/rl-2305.section.md
@@ -249,6 +249,17 @@ In addition to numerous new and upgraded packages, this release has the followin
 
 - `services.chronyd` is now started with additional systemd sandbox/hardening options for better security.
 
+- PostgreSQL has opt-in support for [JIT compilation](https://www.postgresql.org/docs/current/jit-reason.html). It can be enabled like this:
+  ```nix
+  {
+    services.postgresql = {
+      enable = true;
+      package = pkgs.postgresql_jit;
+      settings."jit" = "on";
+    };
+  }
+  ```
+
 - `services.dhcpcd` service now don't solicit or accept IPv6 Router Advertisements on interfaces that use static IPv6 addresses.
 
 - The module `services.headscale` was refactored to be compliant with [RFC 0042](https://github.com/NixOS/rfcs/blob/master/rfcs/0042-config-option.md). To be precise, this means that the following things have changed:
diff --git a/nixos/tests/all-tests.nix b/nixos/tests/all-tests.nix
index 2c34a3996d0b..f1141667ce0f 100644
--- a/nixos/tests/all-tests.nix
+++ b/nixos/tests/all-tests.nix
@@ -561,6 +561,7 @@ in {
   postfixadmin = handleTest ./postfixadmin.nix {};
   postgis = handleTest ./postgis.nix {};
   postgresql = handleTest ./postgresql.nix {};
+  postgresql-jit = handleTest ./postgresql-jit.nix {};
   postgresql-wal-receiver = handleTest ./postgresql-wal-receiver.nix {};
   powerdns = handleTest ./powerdns.nix {};
   powerdns-admin = handleTest ./powerdns-admin.nix {};
diff --git a/nixos/tests/postgresql-jit.nix b/nixos/tests/postgresql-jit.nix
new file mode 100644
index 000000000000..21682aad96da
--- /dev/null
+++ b/nixos/tests/postgresql-jit.nix
@@ -0,0 +1,50 @@
+{ system ? builtins.currentSystem
+, config ? {}
+, pkgs ? import ../.. { inherit system config; }
+}:
+
+with import ../lib/testing-python.nix { inherit system pkgs; };
+
+let
+  inherit (pkgs) lib;
+  packages = lib.filter
+    (lib.hasSuffix "_jit")
+    (builtins.attrNames (import ../../pkgs/servers/sql/postgresql pkgs));
+
+  mkJitTest = packageName: makeTest {
+    name = "${packageName}";
+    meta.maintainers = with lib.maintainers; [ ma27 ];
+    nodes.machine = { pkgs, lib, ... }: {
+      services.postgresql = {
+        enable = true;
+        package = pkgs.${packageName};
+        settings.jit = "on";
+        initialScript = pkgs.writeText "init.sql" ''
+          create table demo (id int);
+          insert into demo (id) select generate_series(1, 5);
+        '';
+      };
+    };
+    testScript = ''
+      machine.start()
+      machine.wait_for_unit("postgresql.service")
+
+      with subtest("JIT is enabled"):
+          machine.succeed("sudo -u postgres psql <<<'show jit;' | grep 'on'")
+
+      with subtest("Test JIT works fine"):
+          output = machine.succeed(
+              "cat ${pkgs.writeText "test.sql" ''
+                set jit_above_cost = 1;
+                EXPLAIN ANALYZE SELECT CONCAT('jit result = ', SUM(id)) FROM demo;
+                SELECT CONCAT('jit result = ', SUM(id)) from demo;
+              ''} | sudo -u postgres psql"
+          )
+          assert "JIT:" in output
+          assert "jit result = 15" in output
+
+      machine.shutdown()
+    '';
+  };
+in
+lib.genAttrs packages mkJitTest
diff --git a/nixos/tests/postgresql-wal-receiver.nix b/nixos/tests/postgresql-wal-receiver.nix
index ae2708546f5d..b0bd7711dbcd 100644
--- a/nixos/tests/postgresql-wal-receiver.nix
+++ b/nixos/tests/postgresql-wal-receiver.nix
@@ -116,4 +116,4 @@ let
     };
 
 # Maps the generic function over all attributes of PostgreSQL packages
-in builtins.listToAttrs (map makePostgresqlWalReceiverTest (builtins.attrNames (import ../../pkgs/servers/sql/postgresql { })))
+in builtins.listToAttrs (map makePostgresqlWalReceiverTest (builtins.attrNames (import ../../pkgs/servers/sql/postgresql pkgs)))
diff --git a/pkgs/servers/sql/postgresql/default.nix b/pkgs/servers/sql/postgresql/default.nix
index 38f3dfb96f2a..1177024b53db 100644
--- a/pkgs/servers/sql/postgresql/default.nix
+++ b/pkgs/servers/sql/postgresql/default.nix
@@ -14,17 +14,28 @@ let
       , this, self, newScope, buildEnv
 
       # source specification
-      , version, hash, psqlSchema,
+      , version, hash, psqlSchema
 
       # for tests
-      nixosTests, thisAttr
+      , nixosTests, thisAttr
+
+      # JIT
+      , jitSupport ? false
+      , nukeReferences, patchelf, llvmPackages
+      , makeRustPlatform, buildPgxExtension, rustPlatform
+
+      # detection of crypt fails when using llvm stdenv, so we add it manually
+      # for <13 (where it got removed: https://github.com/postgres/postgres/commit/c45643d618e35ec2fe91438df15abd4f3c0d85ca)
+      , libxcrypt
     }:
   let
     atLeast = lib.versionAtLeast version;
+    olderThan = lib.versionOlder version;
     lz4Enabled = atLeast "14";
     zstdEnabled = atLeast "15";
 
-  in stdenv.mkDerivation rec {
+    stdenv' = if jitSupport then llvmPackages.stdenv else stdenv;
+  in stdenv'.mkDerivation rec {
     pname = "postgresql";
     inherit version;
 
@@ -33,7 +44,7 @@ let
       inherit hash;
     };
 
-    hardeningEnable = lib.optionals (!stdenv.cc.isClang) [ "pie" ];
+    hardeningEnable = lib.optionals (!stdenv'.cc.isClang) [ "pie" ];
 
     outputs = [ "out" "lib" "doc" "man" ];
     setOutputFlags = false; # $out retains configureFlags :-/
@@ -45,18 +56,21 @@ let
       libxml2
       icu
     ]
+      ++ lib.optionals (olderThan "13") [ libxcrypt ]
+      ++ lib.optionals jitSupport [ llvmPackages.llvm ]
       ++ lib.optionals lz4Enabled [ lz4 ]
       ++ lib.optionals zstdEnabled [ zstd ]
       ++ lib.optionals enableSystemd [ systemd ]
       ++ lib.optionals gssSupport [ libkrb5 ]
-      ++ lib.optionals (!stdenv.isDarwin) [ libossp_uuid ];
+      ++ lib.optionals (!stdenv'.isDarwin) [ libossp_uuid ];
 
     nativeBuildInputs = [
       makeWrapper
       pkg-config
-    ];
+    ]
+      ++ lib.optionals jitSupport [ llvmPackages.llvm.dev nukeReferences patchelf ];
 
-    enableParallelBuilding = !stdenv.isDarwin;
+    enableParallelBuilding = !stdenv'.isDarwin;
 
     separateDebugInfo = true;
 
@@ -65,7 +79,7 @@ let
     env.NIX_CFLAGS_COMPILE = "-I${libxml2.dev}/include/libxml2";
 
     # Otherwise it retains a reference to compiler and fails; see #44767.  TODO: better.
-    preConfigure = "CC=${stdenv.cc.targetPrefix}cc";
+    preConfigure = "CC=${stdenv'.cc.targetPrefix}cc";
 
     configureFlags = [
       "--with-openssl"
@@ -76,11 +90,12 @@ let
       "--with-system-tzdata=${tzdata}/share/zoneinfo"
       "--enable-debug"
       (lib.optionalString enableSystemd "--with-systemd")
-      (if stdenv.isDarwin then "--with-uuid=e2fs" else "--with-ossp-uuid")
+      (if stdenv'.isDarwin then "--with-uuid=e2fs" else "--with-ossp-uuid")
     ] ++ lib.optionals lz4Enabled [ "--with-lz4" ]
       ++ lib.optionals zstdEnabled [ "--with-zstd" ]
       ++ lib.optionals gssSupport [ "--with-gssapi" ]
-      ++ lib.optionals stdenv.hostPlatform.isRiscV [ "--disable-spinlocks" ];
+      ++ lib.optionals stdenv'.hostPlatform.isRiscV [ "--disable-spinlocks" ]
+      ++ lib.optionals jitSupport [ "--with-llvm" ];
 
     patches = [
       ./patches/disable-resolve_symlinks.patch
@@ -88,7 +103,7 @@ let
       ./patches/hardcode-pgxs-path.patch
       ./patches/specify_pkglibdir_at_runtime.patch
       ./patches/findstring.patch
-    ] ++ lib.optionals stdenv.isLinux [
+    ] ++ lib.optionals stdenv'.isLinux [
       (if atLeast "13" then ./patches/socketdir-in-run-13.patch else ./patches/socketdir-in-run.patch)
     ];
 
@@ -99,6 +114,13 @@ let
     postPatch = ''
       # Hardcode the path to pgxs so pg_config returns the path in $out
       substituteInPlace "src/common/config_info.c" --replace HARDCODED_PGXS_PATH "$out/lib"
+
+      ${lib.optionalString jitSupport ''
+        # Force lookup of jit stuff in $out instead of $lib
+        substituteInPlace src/backend/jit/jit.c --replace pkglib_path \"$out/lib\"
+        substituteInPlace src/backend/jit/llvm/llvmjit.c --replace pkglib_path \"$out/lib\"
+        substituteInPlace src/backend/jit/llvm/llvmjit_inline.cpp --replace pkglib_path \"$out/lib\"
+      ''}
     '';
 
     postInstall =
@@ -109,27 +131,54 @@ let
         moveToOutput "lib/libecpg*" "$out"
 
         # Prevent a retained dependency on gcc-wrapper.
-        substituteInPlace "$out/lib/pgxs/src/Makefile.global" --replace ${stdenv.cc}/bin/ld ld
+        substituteInPlace "$out/lib/pgxs/src/Makefile.global" --replace ${stdenv'.cc}/bin/ld ld
 
         if [ -z "''${dontDisableStatic:-}" ]; then
           # Remove static libraries in case dynamic are available.
           for i in $out/lib/*.a $lib/lib/*.a; do
             name="$(basename "$i")"
-            ext="${stdenv.hostPlatform.extensions.sharedLibrary}"
+            ext="${stdenv'.hostPlatform.extensions.sharedLibrary}"
             if [ -e "$lib/lib/''${name%.a}$ext" ] || [ -e "''${i%.a}$ext" ]; then
               rm "$i"
             fi
           done
         fi
+
+        ${lib.optionalString jitSupport ''
+          # Move the bitcode and libllvmjit.so library out of $lib; otherwise, every client that
+          # depends on libpq.so will also have libLLVM.so in its closure too, bloating it
+          moveToOutput "lib/bitcode" "$out"
+          moveToOutput "lib/llvmjit*" "$out"
+
+          # In the case of JIT support, prevent a retained dependency on clang-wrapper
+          substituteInPlace "$out/lib/pgxs/src/Makefile.global" --replace ${self.llvmPackages.stdenv.cc}/bin/clang clang
+          nuke-refs $out/lib/llvmjit_types.bc $(find $out/lib/bitcode -type f)
+
+          # Stop out depending on the default output of llvm
+          substituteInPlace $out/lib/pgxs/src/Makefile.global \
+            --replace ${self.llvmPackages.llvm.out}/bin "" \
+            --replace '$(LLVM_BINPATH)/' ""
+
+          # Stop out depending on the -dev output of llvm
+          substituteInPlace $out/lib/pgxs/src/Makefile.global \
+            --replace ${self.llvmPackages.llvm.dev}/bin/llvm-config llvm-config \
+            --replace -I${self.llvmPackages.llvm.dev}/include ""
+
+          # Stop lib depending on the -dev output of llvm
+          rpath=$(patchelf --print-rpath $out/lib/llvmjit.so)
+          nuke-refs -e $out $out/lib/llvmjit.so
+          # Restore the correct rpath
+          patchelf $out/lib/llvmjit.so --set-rpath "$rpath"
+        ''}
       '';
 
-    postFixup = lib.optionalString (!stdenv.isDarwin && stdenv.hostPlatform.libc == "glibc")
+    postFixup = lib.optionalString (!stdenv'.isDarwin && stdenv'.hostPlatform.libc == "glibc")
       ''
         # initdb needs access to "locale" command from glibc.
         wrapProgram $out/bin/initdb --prefix PATH ":" ${glibc.bin}/bin
       '';
 
-    doCheck = !stdenv.isDarwin;
+    doCheck = !stdenv'.isDarwin;
     # autodetection doesn't seem to able to find this, but it's there.
     checkTarget = "check";
 
@@ -138,7 +187,7 @@ let
       #     ! ERROR:  could not load library "/build/postgresql-11.5/tmp_install/nix/store/...-postgresql-11.5-lib/lib/libpqwalreceiver.so": Error loading shared library libpq.so.5: No such file or directory (needed by /build/postgresql-11.5/tmp_install/nix/store/...-postgresql-11.5-lib/lib/libpqwalreceiver.so)
       # See also here:
       #     https://git.alpinelinux.org/aports/tree/main/postgresql/disable-broken-tests.patch?id=6d7d32c12e073a57a9e5946e55f4c1fbb68bd442
-      if stdenv.hostPlatform.isMusl then ''
+      if stdenv'.hostPlatform.isMusl then ''
         substituteInPlace src/test/regress/parallel_schedule \
           --replace "subscription" "" \
           --replace "object_address" ""
@@ -146,13 +195,23 @@ let
 
     doInstallCheck = false; # needs a running daemon?
 
-    disallowedReferences = [ stdenv.cc ];
+    disallowedReferences = [ stdenv'.cc ];
 
     passthru = {
-      inherit readline psqlSchema;
+      inherit readline psqlSchema jitSupport;
 
       pkgs = let
-        scope = { postgresql = this; };
+        scope = {
+          postgresql = this;
+          stdenv = stdenv';
+          buildPgxExtension = buildPgxExtension.override {
+            stdenv = stdenv';
+            rustPlatform = makeRustPlatform {
+              stdenv = stdenv';
+              inherit (rustPlatform.rust) rustc cargo;
+            };
+          };
+        };
         newSelf = self // scope;
         newSuper = { callPackage = newScope (scope // this.pkgs); };
       in import ./packages.nix newSelf newSuper;
@@ -163,15 +222,33 @@ let
                      }
                      this.pkgs;
 
-      tests.postgresql = nixosTests.postgresql-wal-receiver.${thisAttr};
+      tests = {
+        postgresql = nixosTests.postgresql-wal-receiver.${thisAttr};
+      } // lib.optionalAttrs jitSupport {
+        postgresql-jit = nixosTests.postgresql-jit.${thisAttr};
+      };
+    } // lib.optionalAttrs jitSupport {
+      inherit (llvmPackages) llvm;
     };
 
     meta = with lib; {
       homepage    = "https://www.postgresql.org";
       description = "A powerful, open source object-relational database system";
       license     = licenses.postgresql;
-      maintainers = with maintainers; [ thoughtpolice danbst globin marsam ivan ];
+      maintainers = with maintainers; [ thoughtpolice danbst globin marsam ivan ma27 ];
       platforms   = platforms.unix;
+
+      # JIT support doesn't work with cross-compilation. It is attempted to build LLVM-bytecode
+      # (`%.bc` is the corresponding `make(1)`-rule) for each sub-directory in `backend/` for
+      # the JIT apparently, but with a $(CLANG) that can produce binaries for the build, not the
+      # host-platform.
+      #
+      # I managed to get a cross-build with JIT support working with
+      # `depsBuildBuild = [ llvmPackages.clang ] ++ buildInputs`, but considering that the
+      # resulting LLVM IR isn't platform-independent this doesn't give you much.
+      # In fact, I tried to test the result in a VM-test, but as soon as JIT was used to optimize
+      # a query, postgres would coredump with `Illegal instruction`.
+      broken = jitSupport && (stdenv.hostPlatform != stdenv.buildPlatform);
     };
   };
 
@@ -204,50 +281,60 @@ let
     passthru.psqlSchema = postgresql.psqlSchema;
   };
 
-in self: {
+  mkPackages = self: {
+    postgresql_11 = self.callPackage generic {
+      version = "11.19";
+      psqlSchema = "11.1"; # should be 11, but changing it is invasive
+      hash = "sha256-ExCeK3HxE5QFwnIB2jczphrOcu4cIo2cnwMg4GruFMI=";
+      this = self.postgresql_11;
+      thisAttr = "postgresql_11";
+      inherit self;
+    };
 
-  postgresql_11 = self.callPackage generic {
-    version = "11.19";
-    psqlSchema = "11.1"; # should be 11, but changing it is invasive
-    hash = "sha256-ExCeK3HxE5QFwnIB2jczphrOcu4cIo2cnwMg4GruFMI=";
-    this = self.postgresql_11;
-    thisAttr = "postgresql_11";
-    inherit self;
-  };
+    postgresql_12 = self.callPackage generic {
+      version = "12.14";
+      psqlSchema = "12";
+      hash = "sha256-eFYQI304LIQtNW40cTjljAb/6uJA5swLUqxevMMNBD4=";
+      this = self.postgresql_12;
+      thisAttr = "postgresql_12";
+      inherit self;
+    };
 
-  postgresql_12 = self.callPackage generic {
-    version = "12.14";
-    psqlSchema = "12";
-    hash = "sha256-eFYQI304LIQtNW40cTjljAb/6uJA5swLUqxevMMNBD4=";
-    this = self.postgresql_12;
-    thisAttr = "postgresql_12";
-    inherit self;
-  };
+    postgresql_13 = self.callPackage generic {
+      version = "13.10";
+      psqlSchema = "13";
+      hash = "sha256-W7z1pW2FxE86iwWPtGhi/0nLyRg00H4pXQLm3jwhbfI=";
+      this = self.postgresql_13;
+      thisAttr = "postgresql_13";
+      inherit self;
+    };
 
-  postgresql_13 = self.callPackage generic {
-    version = "13.10";
-    psqlSchema = "13";
-    hash = "sha256-W7z1pW2FxE86iwWPtGhi/0nLyRg00H4pXQLm3jwhbfI=";
-    this = self.postgresql_13;
-    thisAttr = "postgresql_13";
-    inherit self;
-  };
+    postgresql_14 = self.callPackage generic {
+      version = "14.7";
+      psqlSchema = "14";
+      hash = "sha256-zvYPAJj6gQHBVG9CVORbcir1QxM3lFs3ryBwB2MNszE=";
+      this = self.postgresql_14;
+      thisAttr = "postgresql_14";
+      inherit self;
+    };
 
-  postgresql_14 = self.callPackage generic {
-    version = "14.7";
-    psqlSchema = "14";
-    hash = "sha256-zvYPAJj6gQHBVG9CVORbcir1QxM3lFs3ryBwB2MNszE=";
-    this = self.postgresql_14;
-    thisAttr = "postgresql_14";
-    inherit self;
+    postgresql_15 = self.callPackage generic {
+      version = "15.2";
+      psqlSchema = "15";
+      hash = "sha256-maIXH8PWtbX1a3V6ejy4XVCaOOQnOAXe8jlB7SuEaMc=";
+      this = self.postgresql_15;
+      thisAttr = "postgresql_15";
+      inherit self;
+    };
   };
 
-  postgresql_15 = self.callPackage generic {
-    version = "15.2";
-    psqlSchema = "15";
-    hash = "sha256-maIXH8PWtbX1a3V6ejy4XVCaOOQnOAXe8jlB7SuEaMc=";
-    this = self.postgresql_15;
-    thisAttr = "postgresql_15";
-    inherit self;
-  };
-}
+in self:
+  let packages = mkPackages self; in
+  packages
+  // self.lib.mapAttrs'
+    (attrName: postgres: self.lib.nameValuePair "${attrName}_jit" (postgres.override rec {
+      jitSupport = true;
+      thisAttr = "${attrName}_jit";
+      this = self.${thisAttr};
+    }))
+    packages
diff --git a/pkgs/servers/sql/postgresql/ext/plv8/default.nix b/pkgs/servers/sql/postgresql/ext/plv8/default.nix
index 931a5deb40f1..be743401f2c9 100644
--- a/pkgs/servers/sql/postgresql/ext/plv8/default.nix
+++ b/pkgs/servers/sql/postgresql/ext/plv8/default.nix
@@ -138,5 +138,6 @@ stdenv.mkDerivation (finalAttrs: {
     maintainers = with maintainers; [ marsam ];
     platforms = [ "x86_64-linux" ];
     license = licenses.postgresql;
+    broken = postgresql.jitSupport;
   };
 })
diff --git a/pkgs/servers/sql/postgresql/ext/postgis.nix b/pkgs/servers/sql/postgresql/ext/postgis.nix
index 154cdf0c5a2f..18ac6bfc2a0f 100644
--- a/pkgs/servers/sql/postgresql/ext/postgis.nix
+++ b/pkgs/servers/sql/postgresql/ext/postgis.nix
@@ -26,7 +26,7 @@ stdenv.mkDerivation rec {
 
   buildInputs = [ libxml2 postgresql geos proj gdal json_c protobufc ]
                 ++ lib.optional stdenv.isDarwin libiconv;
-  nativeBuildInputs = [ perl pkg-config ];
+  nativeBuildInputs = [ perl pkg-config ] ++ lib.optional postgresql.jitSupport postgresql.llvm;
   dontDisableStatic = true;
 
   # postgis config directory assumes /include /lib from the same root for json-c library
diff --git a/pkgs/top-level/all-packages.nix b/pkgs/top-level/all-packages.nix
index 1b72f6026898..4b3c8da2a5e5 100644
--- a/pkgs/top-level/all-packages.nix
+++ b/pkgs/top-level/all-packages.nix
@@ -25444,13 +25444,26 @@ with pkgs;
     postgresql_13
     postgresql_14
     postgresql_15
+
+    postgresql_11_jit
+    postgresql_12_jit
+    postgresql_13_jit
+    postgresql_14_jit
+    postgresql_15_jit
   ;
   postgresql = postgresql_14.override { this = postgresql; };
+  postgresql_jit = postgresql_14_jit.override { this = postgresql_jit; };
   postgresqlPackages = recurseIntoAttrs postgresql.pkgs;
+  postgresqlJitPackages = recurseIntoAttrs postgresql_jit.pkgs;
   postgresql11Packages = recurseIntoAttrs postgresql_11.pkgs;
   postgresql12Packages = recurseIntoAttrs postgresql_12.pkgs;
   postgresql13Packages = recurseIntoAttrs postgresql_13.pkgs;
   postgresql15Packages = recurseIntoAttrs postgresql_15.pkgs;
+  postgresql11JitPackages = recurseIntoAttrs postgresql_11_jit.pkgs;
+  postgresql12JitPackages = recurseIntoAttrs postgresql_12_jit.pkgs;
+  postgresql13JitPackages = recurseIntoAttrs postgresql_13_jit.pkgs;
+  postgresql14JitPackages = recurseIntoAttrs postgresql_14_jit.pkgs;
+  postgresql15JitPackages = recurseIntoAttrs postgresql_15_jit.pkgs;
   postgresql14Packages = postgresqlPackages;
 
   postgresql_jdbc = callPackage ../development/java-modules/postgresql_jdbc { };
-- 
cgit 1.4.1


From 43dbeae02d74eafc91ad269ebe9363ff5aed1842 Mon Sep 17 00:00:00 2001
From: Maximilian Bosch
Date: Mon, 20 Mar 2023 11:15:12 +0100
Subject: postgresql: pass through JIT-enabled variant of non-JIT postgres and
 vice versa

This is useful if your postgresql version is dependant on
`system.stateVersion` and not pinned down manually. Then it's not
necessary to find out which version exactly is in use and define
`package` manually, but just stay with what NixOS provides as default:

    $ nix-instantiate -A postgresql
    /nix/store/82fzmb77mz2b787dgj7mn4a8i4f6l6sn-postgresql-14.7.drv
    $ nix-instantiate -A postgresql_jit
    /nix/store/qsjkb72fcrrfpsszrwbsi9q9wgp39m50-postgresql-14.7.drv
    $ nix-instantiate -A postgresql.withJIT
    /nix/store/qsjkb72fcrrfpsszrwbsi9q9wgp39m50-postgresql-14.7.drv
    $ nix-instantiate -A postgresql.withJIT.withoutJIT
    /nix/store/82fzmb77mz2b787dgj7mn4a8i4f6l6sn-postgresql-14.7.drv

I.e. you can use postgresql with JIT (for complex queries only[1]) like
this:

    services.postgresql = {
      enable = true;
      enableJIT = true;
    };

Performing a new override instead of re-using the `_jit`-variants for
that has the nice property that overlays for the original package apply
to the JIT-enabled variant, i.e.

    with import ./. {
      overlays = [
        (self: super: {
          postgresql = super.postgresql.overrideAttrs (_: { fnord = "snens"; });
        })
      ];
    };
    postgresql.withJIT.fnord

still gives the string `snens` whereas `postgresql_jit` doesn't have the
attribute `fnord` in its derivation.

[1] https://www.postgresql.org/docs/current/runtime-config-query.html#GUC-JIT-ABOVE-COST
---
 nixos/doc/manual/release-notes/rl-2305.section.md |  3 +-
 nixos/modules/services/databases/postgresql.md    | 37 +++++++++++++++++++++++
 nixos/modules/services/databases/postgresql.nix   | 14 ++++++---
 pkgs/servers/sql/postgresql/default.nix           | 11 ++++++-
 4 files changed, 57 insertions(+), 8 deletions(-)

diff --git a/nixos/doc/manual/release-notes/rl-2305.section.md b/nixos/doc/manual/release-notes/rl-2305.section.md
index 2312cea87500..67b70ebbd117 100644
--- a/nixos/doc/manual/release-notes/rl-2305.section.md
+++ b/nixos/doc/manual/release-notes/rl-2305.section.md
@@ -254,8 +254,7 @@ In addition to numerous new and upgraded packages, this release has the followin
   {
     services.postgresql = {
       enable = true;
-      package = pkgs.postgresql_jit;
-      settings."jit" = "on";
+      enableJIT = true;
     };
   }
   ```
diff --git a/nixos/modules/services/databases/postgresql.md b/nixos/modules/services/databases/postgresql.md
index 1805bafe3be3..4d66ee38be42 100644
--- a/nixos/modules/services/databases/postgresql.md
+++ b/nixos/modules/services/databases/postgresql.md
@@ -171,3 +171,40 @@ self: super: {
   };
 }
 ```
+
+## JIT (Just-In-Time compilation) {#module-services-postgres-jit}
+
+[JIT](https://www.postgresql.org/docs/current/jit-reason.html)-support in the PostgreSQL package
+is disabled by default because of the ~300MiB closure-size increase from the LLVM dependency. It
+can be optionally enabled in PostgreSQL with the following config option:
+
+```nix
+{
+  services.postgresql.enableJIT = true;
+}
+```
+
+This makes sure that the [`jit`](https://www.postgresql.org/docs/current/runtime-config-query.html#GUC-JIT)-setting
+is set to `on` and a PostgreSQL package with JIT enabled is used. Further tweaking of the JIT compiler, e.g. setting a different
+query cost threshold via [`jit_above_cost`](https://www.postgresql.org/docs/current/runtime-config-query.html#GUC-JIT-ABOVE-COST)
+can be done manually via [`services.postgresql.settings`](#opt-services.postgresql.settings).
+
+The attribute-names of JIT-enabled PostgreSQL packages are suffixed with `_jit`, i.e. for each `pkgs.postgresql`
+(and `pkgs.postgresql_<major>`) in `nixpkgs` there's also a `pkgs.postgresql_jit` (and `pkgs.postgresql_<major>_jit`).
+Alternatively, a JIT-enabled variant can be derived from a given `postgresql` package via `postgresql.withJIT`.
+This is also useful if it's not clear which attribute from `nixpkgs` was originally used (e.g. when working with
+[`config.services.postgresql.package`](#opt-services.postgresql.package) or if the package was modified via an
+overlay) since all modifications are propagated to `withJIT`. I.e.
+
+```nix
+with import <nixpkgs> {
+  overlays = [
+    (self: super: {
+      postgresql = super.postgresql.overrideAttrs (_: { pname = "foobar"; });
+    })
+  ];
+};
+postgresql.withJIT.pname
+```
+
+evaluates to `"foobar"`.
diff --git a/nixos/modules/services/databases/postgresql.nix b/nixos/modules/services/databases/postgresql.nix
index 7bbe1ad22595..16f76248b3a0 100644
--- a/nixos/modules/services/databases/postgresql.nix
+++ b/nixos/modules/services/databases/postgresql.nix
@@ -42,6 +42,8 @@ in
 
       enable = mkEnableOption (lib.mdDoc "PostgreSQL Server");
 
+      enableJIT = mkEnableOption (lib.mdDoc "JIT support");
+
       package = mkOption {
         type = types.package;
         example = literalExpression "pkgs.postgresql_11";
@@ -435,19 +437,21 @@ in
         log_line_prefix = cfg.logLinePrefix;
         listen_addresses = if cfg.enableTCPIP then "*" else "localhost";
         port = cfg.port;
+        jit = mkDefault (if cfg.enableJIT then "on" else "off");
       };
 
     services.postgresql.package = let
         mkThrow = ver: throw "postgresql_${ver} was removed, please upgrade your postgresql version.";
+        base = if versionAtLeast config.system.stateVersion "22.05" then pkgs.postgresql_14
+            else if versionAtLeast config.system.stateVersion "21.11" then pkgs.postgresql_13
+            else if versionAtLeast config.system.stateVersion "20.03" then pkgs.postgresql_11
+            else if versionAtLeast config.system.stateVersion "17.09" then mkThrow "9_6"
+            else mkThrow "9_5";
     in
       # Note: when changing the default, make it conditional on
       # ‘system.stateVersion’ to maintain compatibility with existing
       # systems!
-      mkDefault (if versionAtLeast config.system.stateVersion "22.05" then pkgs.postgresql_14
-            else if versionAtLeast config.system.stateVersion "21.11" then pkgs.postgresql_13
-            else if versionAtLeast config.system.stateVersion "20.03" then pkgs.postgresql_11
-            else if versionAtLeast config.system.stateVersion "17.09" then mkThrow "9_6"
-            else mkThrow "9_5");
+      mkDefault (if cfg.enableJIT then base.withJIT else base);
 
     services.postgresql.dataDir = mkDefault "/var/lib/postgresql/${cfg.package.psqlSchema}";
 
diff --git a/pkgs/servers/sql/postgresql/default.nix b/pkgs/servers/sql/postgresql/default.nix
index 1177024b53db..82a590e898f9 100644
--- a/pkgs/servers/sql/postgresql/default.nix
+++ b/pkgs/servers/sql/postgresql/default.nix
@@ -197,9 +197,18 @@ let
 
     disallowedReferences = [ stdenv'.cc ];
 
-    passthru = {
+    passthru = let
+      jitToggle = this.override {
+        jitSupport = !jitSupport;
+        this = jitToggle;
+      };
+    in
+    {
       inherit readline psqlSchema jitSupport;
 
+      withJIT = if jitSupport then this else jitToggle;
+      withoutJIT = if jitSupport then jitToggle else this;
+
       pkgs = let
         scope = {
           postgresql = this;
-- 
cgit 1.4.1


From 608cb37533f25a049d9363f7053890fdab6076eb Mon Sep 17 00:00:00 2001
From: Maximilian Bosch
Date: Mon, 20 Mar 2023 13:25:13 +0100
Subject: nixos/tests/postgresql: fix deprecation warning

---
 nixos/tests/postgresql.nix | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/nixos/tests/postgresql.nix b/nixos/tests/postgresql.nix
index 7e0a82c38828..b44849e0a14e 100644
--- a/nixos/tests/postgresql.nix
+++ b/nixos/tests/postgresql.nix
@@ -137,7 +137,7 @@ let
       maintainers = [ zagy ];
     };
 
-    machine = {...}:
+    nodes.machine = {...}:
       {
         services.postgresql = {
           enable = true;
-- 
cgit 1.4.1


From a5a715bb249ee4e531d7e743f5a4e4234dd29346 Mon Sep 17 00:00:00 2001
From: Maximilian Bosch
Date: Fri, 24 Mar 2023 20:41:20 +0100
Subject: postgresql_jit: fix darwin build

So, patchelf works on ELF files and these don't exist on Darwin. I'm not
aware of any other way to eliminate all references of the executable's
body while preserving library paths in the header. This still works, the
only problem is that darwin has a dependency on llvm.dev, so it's
runtime closure is larger.
---
 pkgs/servers/sql/postgresql/default.nix | 12 +++++++-----
 1 file changed, 7 insertions(+), 5 deletions(-)

diff --git a/pkgs/servers/sql/postgresql/default.nix b/pkgs/servers/sql/postgresql/default.nix
index 82a590e898f9..0b50b0ebfc94 100644
--- a/pkgs/servers/sql/postgresql/default.nix
+++ b/pkgs/servers/sql/postgresql/default.nix
@@ -164,11 +164,13 @@ let
             --replace ${self.llvmPackages.llvm.dev}/bin/llvm-config llvm-config \
             --replace -I${self.llvmPackages.llvm.dev}/include ""
 
-          # Stop lib depending on the -dev output of llvm
-          rpath=$(patchelf --print-rpath $out/lib/llvmjit.so)
-          nuke-refs -e $out $out/lib/llvmjit.so
-          # Restore the correct rpath
-          patchelf $out/lib/llvmjit.so --set-rpath "$rpath"
+          ${lib.optionalString (!stdenv'.isDarwin) ''
+            # Stop lib depending on the -dev output of llvm
+            rpath=$(patchelf --print-rpath $out/lib/llvmjit.so)
+            nuke-refs -e $out $out/lib/llvmjit.so
+            # Restore the correct rpath
+            patchelf $out/lib/llvmjit.so --set-rpath "$rpath"
+          ''}
         ''}
       '';
 
-- 
cgit 1.4.1


From e2fb65175228a992f196f3b1700a53e18602e7f6 Mon Sep 17 00:00:00 2001
From: Maximilian Bosch
Date: Sat, 25 Mar 2023 11:34:24 +0100
Subject: nixos/postgresql: fix enableJIT

Make sure that JIT is actually available when using

    services.postgresql = {
      enable = true;
      enableJIT = true;
      package = pkgs.postgresql_15;
    };

The current behavior is counter-intuitive because the docs state that
`enableJIT = true;` is sufficient even though it wasn't in that case
because the declared package doesn't have the LLVM dependency.

Fixed by using `package.withJIT` if `enableJIT = true;` and
`package.jitSupport` is `false`.

Also updated the postgresql-jit test to test for that case.
---
 nixos/modules/services/databases/postgresql.nix | 13 +++++++++++--
 nixos/tests/postgresql-jit.nix                  |  6 ++----
 2 files changed, 13 insertions(+), 6 deletions(-)

diff --git a/nixos/modules/services/databases/postgresql.nix b/nixos/modules/services/databases/postgresql.nix
index 16f76248b3a0..3d55995aba05 100644
--- a/nixos/modules/services/databases/postgresql.nix
+++ b/nixos/modules/services/databases/postgresql.nix
@@ -7,9 +7,18 @@ let
   cfg = config.services.postgresql;
 
   postgresql =
+    let
+      # ensure that
+      #   services.postgresql = {
+      #     enableJIT = true;
+      #     package = pkgs.postgresql_<major>;
+      #   };
+      # works.
+      base = if cfg.enableJIT && !cfg.package.jitSupport then cfg.package.withJIT else cfg.package;
+    in
     if cfg.extraPlugins == []
-      then cfg.package
-      else cfg.package.withPackages (_: cfg.extraPlugins);
+      then base
+      else base.withPackages (_: cfg.extraPlugins);
 
   toStr = value:
     if true == value then "yes"
diff --git a/nixos/tests/postgresql-jit.nix b/nixos/tests/postgresql-jit.nix
index 21682aad96da..baf26b8da2b3 100644
--- a/nixos/tests/postgresql-jit.nix
+++ b/nixos/tests/postgresql-jit.nix
@@ -7,9 +7,7 @@ with import ../lib/testing-python.nix { inherit system pkgs; };
 
 let
   inherit (pkgs) lib;
-  packages = lib.filter
-    (lib.hasSuffix "_jit")
-    (builtins.attrNames (import ../../pkgs/servers/sql/postgresql pkgs));
+  packages = builtins.attrNames (import ../../pkgs/servers/sql/postgresql pkgs);
 
   mkJitTest = packageName: makeTest {
     name = "${packageName}";
@@ -17,8 +15,8 @@ let
     nodes.machine = { pkgs, lib, ... }: {
       services.postgresql = {
         enable = true;
+        enableJIT = true;
         package = pkgs.${packageName};
-        settings.jit = "on";
         initialScript = pkgs.writeText "init.sql" ''
           create table demo (id int);
           insert into demo (id) select generate_series(1, 5);
-- 
cgit 1.4.1


From 9f2b4357c144f4fd9f95694b7624e159f8163c74 Mon Sep 17 00:00:00 2001
From: Maximilian Bosch
Date: Sat, 25 Mar 2023 11:39:11 +0100
Subject: postgresql: save rebuilds of existing packages

...by using `+ lib.optionalString ...` rather than a substitution. That
way the phases don't have additional trailing white-spaces in the
non-JIT case which cause rebuilds.
---
 pkgs/servers/sql/postgresql/default.nix | 58 +++++++++++++++------------------
 1 file changed, 27 insertions(+), 31 deletions(-)

diff --git a/pkgs/servers/sql/postgresql/default.nix b/pkgs/servers/sql/postgresql/default.nix
index 0b50b0ebfc94..48b94841548e 100644
--- a/pkgs/servers/sql/postgresql/default.nix
+++ b/pkgs/servers/sql/postgresql/default.nix
@@ -114,13 +114,11 @@ let
     postPatch = ''
       # Hardcode the path to pgxs so pg_config returns the path in $out
       substituteInPlace "src/common/config_info.c" --replace HARDCODED_PGXS_PATH "$out/lib"
-
-      ${lib.optionalString jitSupport ''
+    '' + lib.optionalString jitSupport ''
         # Force lookup of jit stuff in $out instead of $lib
         substituteInPlace src/backend/jit/jit.c --replace pkglib_path \"$out/lib\"
         substituteInPlace src/backend/jit/llvm/llvmjit.c --replace pkglib_path \"$out/lib\"
         substituteInPlace src/backend/jit/llvm/llvmjit_inline.cpp --replace pkglib_path \"$out/lib\"
-      ''}
     '';
 
     postInstall =
@@ -143,34 +141,32 @@ let
             fi
           done
         fi
-
-        ${lib.optionalString jitSupport ''
-          # Move the bitcode and libllvmjit.so library out of $lib; otherwise, every client that
-          # depends on libpq.so will also have libLLVM.so in its closure too, bloating it
-          moveToOutput "lib/bitcode" "$out"
-          moveToOutput "lib/llvmjit*" "$out"
-
-          # In the case of JIT support, prevent a retained dependency on clang-wrapper
-          substituteInPlace "$out/lib/pgxs/src/Makefile.global" --replace ${self.llvmPackages.stdenv.cc}/bin/clang clang
-          nuke-refs $out/lib/llvmjit_types.bc $(find $out/lib/bitcode -type f)
-
-          # Stop out depending on the default output of llvm
-          substituteInPlace $out/lib/pgxs/src/Makefile.global \
-            --replace ${self.llvmPackages.llvm.out}/bin "" \
-            --replace '$(LLVM_BINPATH)/' ""
-
-          # Stop out depending on the -dev output of llvm
-          substituteInPlace $out/lib/pgxs/src/Makefile.global \
-            --replace ${self.llvmPackages.llvm.dev}/bin/llvm-config llvm-config \
-            --replace -I${self.llvmPackages.llvm.dev}/include ""
-
-          ${lib.optionalString (!stdenv'.isDarwin) ''
-            # Stop lib depending on the -dev output of llvm
-            rpath=$(patchelf --print-rpath $out/lib/llvmjit.so)
-            nuke-refs -e $out $out/lib/llvmjit.so
-            # Restore the correct rpath
-            patchelf $out/lib/llvmjit.so --set-rpath "$rpath"
-          ''}
+      '' + lib.optionalString jitSupport ''
+        # Move the bitcode and libllvmjit.so library out of $lib; otherwise, every client that
+        # depends on libpq.so will also have libLLVM.so in its closure too, bloating it
+        moveToOutput "lib/bitcode" "$out"
+        moveToOutput "lib/llvmjit*" "$out"
+
+        # In the case of JIT support, prevent a retained dependency on clang-wrapper
+        substituteInPlace "$out/lib/pgxs/src/Makefile.global" --replace ${self.llvmPackages.stdenv.cc}/bin/clang clang
+        nuke-refs $out/lib/llvmjit_types.bc $(find $out/lib/bitcode -type f)
+
+        # Stop out depending on the default output of llvm
+        substituteInPlace $out/lib/pgxs/src/Makefile.global \
+          --replace ${self.llvmPackages.llvm.out}/bin "" \
+          --replace '$(LLVM_BINPATH)/' ""
+
+        # Stop out depending on the -dev output of llvm
+        substituteInPlace $out/lib/pgxs/src/Makefile.global \
+          --replace ${self.llvmPackages.llvm.dev}/bin/llvm-config llvm-config \
+          --replace -I${self.llvmPackages.llvm.dev}/include ""
+
+        ${lib.optionalString (!stdenv'.isDarwin) ''
+          # Stop lib depending on the -dev output of llvm
+          rpath=$(patchelf --print-rpath $out/lib/llvmjit.so)
+          nuke-refs -e $out $out/lib/llvmjit.so
+          # Restore the correct rpath
+          patchelf $out/lib/llvmjit.so --set-rpath "$rpath"
         ''}
       '';
 
-- 
cgit 1.4.1