test(nix): supautils site profile swap vm test

.github/workflows/ami-release-nix.yml

@@ -11,6 +11,7 @@ on:
       - flake.lock
       - flake.nix
       - nix/packages/build-ami.nix
+      - nix/**
   workflow_dispatch:
 
 permissions:

.github/workflows/site-env-manifest.yml

@@ -0,0 +1,70 @@
+name: Publish site-env manifest
+
+on:
+  pull_request:
+    paths:
+      - .github/workflows/site-env-manifest.yml
+      - flake.lock
+      - flake.nix
+      - nix/**
+  push:
+    branches:
+      - develop
+      - release/*
+    paths:
+      - .github/workflows/site-env-manifest.yml
+      - flake.lock
+      - flake.nix
+      - nix/**
+  workflow_dispatch:
+
+permissions:
+  contents: write
+
+jobs:
+  manifest:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout Repo
+        uses: supabase/postgres/.github/actions/shared-checkout@HEAD
+
+      - name: Install nix
+        uses: ./.github/actions/nix-install-ephemeral
+
+      - name: Resolve git sha
+        id: resolve-git-sha
+        uses: ./.github/actions/resolve-git-sha
+        with:
+          flake_input: ${{ vars.GIT_SHA_FROM_FLAKE_INPUT }}
+
+      - name: Build manifest
+        run: |
+          echo '{}' > site-env-manifest.json
+          for system in aarch64-linux x86_64-linux; do
+            for name in site-env-15 site-env-17 site-env-orioledb-17; do
+              path=$(nix eval --raw ".#legacyPackages.${system}.${name}.outPath")
+              jq --arg k "${system}/${name}" --arg v "$path" '.[$k] = $v' \
+                site-env-manifest.json > tmp.json && mv tmp.json site-env-manifest.json
+            done
+          done
+          cat site-env-manifest.json
+
+      - name: Resolve release tag
+        id: tag
+        run: |
+          if [[ "${{ github.event_name }}" == "pull_request" ]]; then
+            echo "tag=pr-${{ github.event.pull_request.number }}" >> "$GITHUB_OUTPUT"
+            echo "prerelease=true" >> "$GITHUB_OUTPUT"
+          else
+            echo "tag=site-env-manifest-${{ steps.resolve-git-sha.outputs.sha }}" >> "$GITHUB_OUTPUT"
+            echo "prerelease=false" >> "$GITHUB_OUTPUT"
+          fi
+
+      - name: Publish manifest release
+        uses: softprops/action-gh-release@a06a81a03ee405af7f2048a818ed3f03bbf83c7b # v2.5.0
+        with:
+          name: ${{ steps.tag.outputs.tag }}
+          tag_name: ${{ steps.tag.outputs.tag }}
+          target_commitish: ${{ steps.resolve-git-sha.outputs.sha }}
+          prerelease: ${{ steps.tag.outputs.prerelease }}
+          files: site-env-manifest.json

ansible/files/postgresql_config/postgresql.conf.j2

@@ -692,7 +692,7 @@ jit_provider = 'llvmjit'		# JIT library to use
 
 # - Other Defaults -
 
-#dynamic_library_path = '$libdir'
+dynamic_library_path = '/nix/var/nix/profiles/site/pg-extensions:$libdir'
 #gin_fuzzy_search_limit = 0
 
 #------------------------------------------------------------------------------

ansible/tasks/stage2-setup-postgres.yml

@@ -77,6 +77,18 @@
           nix-env --set {{ postgres_env_path.stdout }}
         "
 
+    - name: Resolve site env store path
+      ansible.builtin.shell: |
+        . /nix/var/nix/profiles/default/etc/profile.d/nix-daemon.sh &&
+        nix build --no-link --print-out-paths github:supabase/postgres/{{ git_commit_sha }}#site-env-{{ postgresql_major_version }}
+      register: site_env_path
+
+    - name: Install site env from nix binary cache
+      ansible.builtin.shell: |
+        . /nix/var/nix/profiles/default/etc/profile.d/nix-daemon.sh &&
+        nix-env --profile /nix/var/nix/profiles/site --set {{ site_env_path.stdout }}
+
+
     - name: Install supascan for baseline validation
       ansible.builtin.shell: |
         sudo -u ubuntu bash -c ". /nix/var/nix/profiles/default/etc/profile.d/nix-daemon.sh && nix profile install github:supabase/postgres/{{ git_commit_sha }}#supascan"

nix/checks.nix

@@ -939,6 +939,48 @@
             wal-g-3
             ;
           devShell = self'.devShells.default;
+          site =
+            let
+              system = pkgs.pkgsLinux.stdenv.hostPlatform.system;
+              site-env-17 = self.packages.${system}."site-env-17";
+              psql_17 = self.legacyPackages.${system}."psql_17".bin;
+              pgConf = pkgs.writeText "postgresql-test.conf" ''
+                dynamic_library_path = '/nix/var/nix/profiles/site/pg-extensions:$libdir'
+                session_preload_libraries = 'supautils'
+                listen_addresses = 'localhost'
+                unix_socket_directories = '/tmp'
+              '';
+            in
+            pkgs.testers.runNixOSTest {
+              name = "site";
+              nodes.machine =
+                { ... }:
+                {
+                  environment.systemPackages = [
+                    site-env-17
+                  ];
+                  users.users.postgres = {
+                    isSystemUser = true;
+                    group = "postgres";
+                    shell = pkgs.bash;
+                  };
+                  users.groups.postgres = { };
+                };
+              testScript = ''
+                machine.succeed("nix-store --realise ${site-env-17} && nix-env --profile /nix/var/nix/profiles/site --set ${site-env-17}")
+                machine.succeed("[ \"$(readlink -f /nix/var/nix/profiles/site)\" = \"${site-env-17}\" ]")
+
+                # idempotent
+                machine.succeed("nix-store --realise ${site-env-17} && nix-env --profile /nix/var/nix/profiles/site --set ${site-env-17}")
+
+                # postgres can load supautils via the site profile's dynamic_library_path
+                machine.succeed("install -d -o postgres -g postgres /tmp/pgdata")
+                machine.succeed("su postgres -c '${psql_17}/bin/initdb -D /tmp/pgdata'")
+                machine.succeed("install -o postgres -g postgres ${pgConf} /tmp/pgdata/postgresql.conf")
+                machine.succeed("su postgres -c '${psql_17}/bin/pg_ctl -D /tmp/pgdata -l /tmp/pg.log start'")
+                machine.succeed("su postgres -c '${psql_17}/bin/psql -h localhost -d postgres -c \"select 1\"'")
+              '';
+            };
         }
         // (import ./ext/tests {
           inherit self;

nix/ext/tests/supautils-site-profile.nix

@@ -0,0 +1,143 @@
+# supautils is session-preloaded, so each new backend dlopens it again through
+# dynamic_library_path while open sessions keep the copy they loaded. Adding the
+# site path to an instance takes a reload; later site profile flips take nothing.
+{ self, pkgs }:
+let
+  testLib = import ./lib.nix { inherit self pkgs; };
+  system = pkgs.pkgsLinux.stdenv.hostPlatform.system;
+  pgBin = "${self.packages.${system}."psql_17/bin"}/bin";
+  newSiteEnv = self.packages.${system}.site-env-17;
+  # Same supautils source under another name: a distinct store path to swap from.
+  oldSupautils = self.legacyPackages.${system}.psql_17.exts.supautils.overrideAttrs (_: {
+    name = "supautils-previous";
+  });
+  oldSiteEnv = pkgs.pkgsLinux.buildEnv {
+    name = "site-env-17-previous";
+    paths = [ oldSupautils ];
+    postBuild = ''
+      mkdir -p $out/pg-extensions
+      ln -s ../lib/supautils.so $out/pg-extensions/supautils.so
+    '';
+  };
+in
+pkgs.testers.runNixOSTest {
+  name = "supautils-site-profile";
+  nodes.server =
+    { ... }:
+    {
+      imports = [ (testLib.makeSupabaseTestConfig { majorVersion = "17"; }) ];
+      virtualisation.additionalPaths = [
+        oldSiteEnv
+        newSiteEnv
+      ];
+    };
+  testScript = ''
+    PG_BIN = "${pgBin}"
+    CONNECT = "-h localhost -U supabase_admin"
+    CONF = "/var/lib/postgresql/data/postgresql.conf"
+    sessions = 0
+
+    def sql(query, db="postgres"):
+      return server.succeed(
+        f"psql -U supabase_admin -d {db} -X -t -A -v ON_ERROR_STOP=1 -c \"{query}\""
+      ).strip()
+
+    def mapped_supautils(pid):
+      return server.succeed(
+        f"grep -o '/nix/store/[^ ]*/supautils.so' /proc/{pid}/maps | sort -u"
+      ).strip()
+
+    def new_session_supautils():
+      global sessions
+      sessions += 1
+      name = f"session-{sessions}"
+      server.succeed(
+        f"systemd-run --unit={name} -E PGAPPNAME={name} "
+        f"{PG_BIN}/psql {CONNECT} -d postgres -X -c 'select pg_sleep(120)'"
+      )
+      query = f"select pid from pg_stat_activity where application_name = '{name}'"
+      server.wait_until_succeeds(f"psql -U supabase_admin -d postgres -X -tA -c \"{query}\" | grep -q .")
+      pid = sql(query)
+      mapped = mapped_supautils(pid)
+      sql(f"select pg_terminate_backend({pid})")
+      return mapped
+
+    def start_bench(name, flags):
+      log = f"/tmp/{name}.log"
+      server.succeed(
+        f"systemd-run --unit={name} -p RemainAfterExit=yes "
+        f"-p StandardOutput=file:{log} -p StandardError=file:{log} -E PGAPPNAME={name} "
+        f"{PG_BIN}/pgbench {CONNECT} -n -S {flags} -R 50 -T 60 bench"
+      )
+
+    def finish_bench(name):
+      server.wait_until_succeeds(f"systemctl show {name} -p SubState --value | grep -qx exited", timeout=120)
+      log = server.succeed(f"cat /tmp/{name}.log")
+      status = server.succeed(f"systemctl show {name} -p ExecMainStatus --value").strip()
+      assert status == "0", f"{name} exited with {status}:\n{log}"
+      assert "number of failed transactions: 0 " in log, f"{name} had failures:\n{log}"
+
+    def set_site_profile(path):
+      server.succeed(f"nix-store --realise {path} && nix-env --profile /nix/var/nix/profiles/site --set {path}")
+
+    def resolve(path):
+      return server.succeed(f"readlink -f {path}").strip()
+
+    start_all()
+    server.wait_for_unit("supabase-db-init.service")
+    old_lib = resolve("${oldSiteEnv}/pg-extensions/supautils.so")
+    new_lib = resolve("${newSiteEnv}/pg-extensions/supautils.so")
+    assert old_lib != new_lib, "old and new site envs ship the same supautils"
+
+    with subtest("Without the site path, sessions load supautils from $libdir"):
+      server.succeed(f"cp {CONF} /tmp/postgresql.conf.site")
+      server.succeed(f"sed -i '/^dynamic_library_path/d' {CONF}")
+      sql("select pg_reload_conf()")
+      server.wait_until_succeeds("[ \"$(psql -U supabase_admin -d postgres -X -tA -c 'show dynamic_library_path')\" = '$libdir' ]")
+      set_site_profile("${oldSiteEnv}")
+      base_lib = new_session_supautils()
+      assert base_lib != old_lib, f"$libdir already serves the old site supautils: {base_lib}"
+      postmaster = server.succeed("head -1 /var/lib/postgresql/data/postmaster.pid").strip()
+      started_at = sql("select pg_postmaster_start_time()")
+
+    with subtest("Start persistent and reconnecting load"):
+      server.succeed("createdb -U supabase_admin bench")
+      server.succeed(f"{PG_BIN}/pgbench {CONNECT} -i -q bench")
+      start_bench("bench-persistent", "-c 2")
+      start_bench("bench-reconnect", "-C -c 2")
+      query = "select pid from pg_stat_activity where application_name = 'bench-persistent'"
+      server.wait_until_succeeds(f"[ $(psql -U supabase_admin -d postgres -X -tA -c \"{query}\" | wc -l) -eq 2 ]")
+      persistent = sql(query).splitlines()
+      for pid in persistent:
+        assert mapped_supautils(pid) == base_lib
+
+    with subtest("Deploying the site path takes effect on reload"):
+      server.succeed(f"cp /tmp/postgresql.conf.site {CONF}")
+      assert new_session_supautils() == base_lib, "site path applied before the reload"
+      sql("select pg_reload_conf()")
+      server.wait_until_succeeds(
+        "psql -U supabase_admin -d postgres -X -tA -c 'show dynamic_library_path' | grep -q /nix/var/nix/profiles/site/pg-extensions"
+      )
+      assert new_session_supautils() == old_lib, "a new session did not load supautils from the site profile"
+      for pid in persistent:
+        assert mapped_supautils(pid) == base_lib, f"persistent backend {pid} lost its supautils"
+
+    with subtest("Flip the site profile without a reload"):
+      set_site_profile("${newSiteEnv}")
+      assert new_session_supautils() == new_lib, "a new session did not load the new supautils"
+      sql("select pg_reload_conf()")
+      assert new_session_supautils() == new_lib, "a reload changed what new sessions load"
+      for pid in persistent:
+        assert mapped_supautils(pid) == base_lib, f"persistent backend {pid} lost its supautils"
+
+    with subtest("Roll back the site profile"):
+      server.succeed("nix-env -p /nix/var/nix/profiles/site --rollback")
+      assert new_session_supautils() == old_lib
+
+    with subtest("Postmaster never restarted and the load saw no failures"):
+      assert server.succeed("head -1 /var/lib/postgresql/data/postmaster.pid").strip() == postmaster
+      assert sql("select pg_postmaster_start_time()") == started_at
+      finish_bench("bench-persistent")
+      finish_bench("bench-reconnect")
+  '';
+}

nix/packages/extension-catalog.nix

@@ -110,8 +110,8 @@
             makeWrapper ${self'.packages.site-extensions-resolve}/bin/site-extensions-resolve \
               "$out/bin/site-extensions-resolve" \
               --set PG_EXTENSIONS_CATALOG "$out/share/pg-extensions-catalog.json"
-            makeWrapper ${self'.packages.site-extensions-update}/bin/site-extensions-update \
-              "$out/bin/site-extensions-update" \
+            makeWrapper ${self'.packages.update-site-extensions}/bin/update-site-extensions \
+              "$out/bin/update-site-extensions" \
               --set PG_EXTENSIONS_CATALOG "$out/share/pg-extensions-catalog.json"
           ''
         )
@@ -146,19 +146,16 @@
           '';
         };
 
-        # Takes manifest json as argument.
-        # Downloads paths and installs them as an env into the profile, replacing all existing ones.
-        site-extensions-update = pkgs.writeShellApplication {
-          name = "site-extensions-update";
-          runtimeInputs = [
-            self'.packages.site-extensions-resolve
-            pkgs.nix
-          ];
+        # Downloads the manifest's paths and installs them as the profile's env, replacing all existing ones.
+        # nix comes from the environment.
+        update-site-extensions = pkgs.writeShellApplication {
+          name = "update-site-extensions";
+          runtimeInputs = [ self'.packages.site-extensions-resolve ];
           text = ''
             manifest="''${1:?Usage: $0 path-to/pg-extensions.json}"
-            profile="''${NIX_PROFILE:-/nix/var/nix/profiles/site-extensions}"
+            profile="/nix/var/nix/profiles/site-extensions"
             readarray -t paths < <(site-extensions-resolve "$manifest")
-            nix-store -r --option stalled-download-timeout 120 "''${paths[@]}" >/dev/null
+            nix-store --realise --option stalled-download-timeout 120 "''${paths[@]}" >/dev/null
             nix-env --profile "$profile" --install "''${paths[@]}" --remove-all
           '';
         };

nix/packages/site-env.nix

@@ -9,11 +9,26 @@
       ...
     }:
     let
+      activate = pkgs.writeShellApplication {
+        name = "activate";
+        text = ''
+          echo "Activating site profile."
+        '';
+      };
+
       makeSiteEnv =
         version: extraPaths:
         pkgs.buildEnv {
           name = "site-env-${version}";
-          paths = [ self'.legacyPackages."psql_${version}".exts.supautils ] ++ extraPaths;
+          paths = [
+            self'.legacyPackages."psql_${version}".exts.supautils
+            activate
+          ]
+          ++ extraPaths;
+          postBuild = ''
+            mkdir -p $out/pg-extensions
+            ln -s ../lib/supautils.so $out/pg-extensions/supautils.so
+          '';
         };
 
       siteEnvs = {