Merge pull request #221851 from Ma27/postgresql-jit-support

postgresql: implement opt-in JIT support

nixos/doc/manual/release-notes/rl-2305.section.md

@@ -273,6 +273,16 @@ 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;
+      enableJIT = true;
+    };
+  }
+  ```
+
 - `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:

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"`.

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"
@@ -42,6 +51,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 +446,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}";
 

nixos/tests/all-tests.nix

@@ -565,6 +565,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 {};

nixos/tests/postgresql-jit.nix

@@ -0,0 +1,48 @@
+{ system ? builtins.currentSystem
+, config ? {}
+, pkgs ? import ../.. { inherit system config; }
+}:
+
+with import ../lib/testing-python.nix { inherit system pkgs; };
+
+let
+  inherit (pkgs) lib;
+  packages = 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;
+        enableJIT = true;
+        package = pkgs.${packageName};
+        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

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)))

nixos/tests/postgresql.nix

@@ -137,7 +137,7 @@ let
       maintainers = [ zagy ];
     };
 
-    machine = {...}:
+    nodes.machine = {...}:
       {
         services.postgresql = {
           enable = true;

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,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 ''
+        # 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 +129,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 ""
+
+        ${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"
+        ''}
       '';
 
-    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 +185,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 +193,32 @@ let
 
     doInstallCheck = false; # needs a running daemon?
 
-    disallowedReferences = [ stdenv.cc ];
+    disallowedReferences = [ stdenv'.cc ];
 
-    passthru = {
-      inherit readline psqlSchema;
+    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; };
+        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 +229,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,8 +288,7 @@ 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
@@ -250,4 +333,15 @@ in self: {
       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

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;
   };
 })

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

pkgs/top-level/all-packages.nix

@@ -25473,13 +25473,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 { };