reduce number of jobs

Add DELETE /lfc/prewarm route which handles ongoing prewarm
cancellation, update API spec, add prewarm Cancelled state
Add offload Cancelled state when LFC is not initialized

.github/actionlint.yml

@@ -4,6 +4,7 @@ self-hosted-runner:
     - large
     - large-arm64
     - small
+    - small-debug
     - small-metal
     - small-arm64
     - unit-perf

.github/workflows/_build-and-test-locally.yml

@@ -338,67 +338,3 @@ jobs:
       - name: Merge and upload coverage data
         if: inputs.build-type == 'debug'
         uses: ./.github/actions/save-coverage-data
-
-  regress-tests:
-    # Don't run regression tests on debug arm64 builds
-    if: inputs.build-type != 'debug' || inputs.arch != 'arm64'
-    permissions:
-      id-token: write # aws-actions/configure-aws-credentials
-      contents: read
-      statuses: write
-    needs: [ build-neon ]
-    runs-on: ${{ fromJSON(format('["self-hosted", "{0}"]', inputs.arch == 'arm64' && 'large-arm64' || 'large-metal')) }}
-    container:
-      image: ${{ inputs.build-tools-image }}
-      credentials:
-        username: ${{ github.actor }}
-        password: ${{ secrets.GITHUB_TOKEN }}
-      # for changed limits, see comments on `options:` earlier in this file
-      options: --init --shm-size=512mb --ulimit memlock=67108864:67108864
-    strategy:
-      fail-fast: false
-      matrix: ${{ fromJSON(format('{{"include":{0}}}', inputs.test-cfg)) }}
-    steps:
-      - name: Harden the runner (Audit all outbound calls)
-        uses: step-security/harden-runner@4d991eb9b905ef189e4c376166672c3f2f230481 # v2.11.0
-        with:
-          egress-policy: audit
-
-      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
-        with:
-          submodules: true
-
-      - name: Pytest regression tests
-        continue-on-error: ${{ matrix.lfc_state == 'with-lfc' && inputs.build-type == 'debug' }}
-        uses: ./.github/actions/run-python-test-set
-        timeout-minutes: ${{ (inputs.build-type == 'release' && inputs.sanitizers != 'enabled') && 75 || 180 }}
-        with:
-          build_type: ${{ inputs.build-type }}
-          test_selection: regress
-          needs_postgres_source: true
-          run_with_real_s3: true
-          real_s3_bucket: neon-github-ci-tests
-          real_s3_region: eu-central-1
-          rerun_failed: ${{ inputs.rerun-failed }}
-          pg_version: ${{ matrix.pg_version }}
-          sanitizers: ${{ inputs.sanitizers }}
-          aws-oidc-role-arn: ${{ vars.DEV_AWS_OIDC_ROLE_ARN }}
-          # `--session-timeout` is equal to (timeout-minutes - 10 minutes) * 60 seconds.
-          # Attempt to stop tests gracefully to generate test reports
-          # until they are forcibly stopped by the stricter `timeout-minutes` limit.
-          extra_params: --session-timeout=${{ (inputs.build-type == 'release' && inputs.sanitizers != 'enabled') && 3000 || 10200 }} --count=${{ inputs.test-run-count }}
-                        ${{ inputs.test-selection != '' && format('-k "{0}"', inputs.test-selection) || '' }}
-        env:
-          TEST_RESULT_CONNSTR: ${{ secrets.REGRESS_TEST_RESULT_CONNSTR_NEW }}
-          CHECK_ONDISK_DATA_COMPATIBILITY: nonempty
-          BUILD_TAG: ${{ inputs.build-tag }}
-          PAGESERVER_VIRTUAL_FILE_IO_ENGINE: tokio-epoll-uring
-          USE_LFC: ${{ matrix.lfc_state == 'with-lfc' && 'true' || 'false' }}
-
-      # Temporary disable this step until we figure out why it's so flaky
-      # Ref https://github.com/neondatabase/neon/issues/4540
-      - name: Merge and upload coverage data
-        if: |
-          false &&
-          inputs.build-type == 'debug' && matrix.pg_version == 'v16'
-        uses: ./.github/actions/save-coverage-data

.github/workflows/proxy-benchmark.yml

@@ -13,12 +13,6 @@ on:
     #        │ │ │ │ ┌───────────── day of the week (0 - 6 or SUN-SAT)
     # - cron: '0 5 * * *' # Runs at 5 UTC once a day
   workflow_dispatch: # adds an ability to run this manually
-    inputs:
-      commit_hash:
-        type: string
-        description: 'The long neon repo commit hash for the system under test (proxy) to be tested.'
-        required: false
-        default: ''
 
 defaults:
   run:
@@ -39,14 +33,20 @@ jobs:
       contents: write
       pull-requests: write
     runs-on: [ self-hosted, unit-perf-aws-arm ]
+    timeout-minutes: 60  # 1h timeout
     container:
       image: ghcr.io/neondatabase/build-tools:pinned-bookworm
       credentials:
         username: ${{ github.actor }}
         password: ${{ secrets.GITHUB_TOKEN }}
       options: --init
-    timeout-minutes: 60  # 1h timeout
     steps:
+    - name: Checkout proxy-bench Repo
+      uses: actions/checkout@v4
+      with:
+        repository: neondatabase/proxy-bench
+        path: proxy-bench
+
     - name: Set up the environment which depends on $RUNNER_TEMP on nvme drive
       id: set-env
       shell: bash -euxo pipefail {0}
@@ -54,64 +54,19 @@ jobs:
         PROXY_BENCH_PATH=$(realpath ./proxy-bench)
         {
           echo "PROXY_BENCH_PATH=$PROXY_BENCH_PATH"
-          echo "NEON_DIR=${GITHUB_WORKSPACE}/"
-          echo "NEON_PROXY_PATH=${GITHUB_WORKSPACE}/bin/proxy"
+          echo "NEON_DIR=${RUNNER_TEMP}/neon"
+          echo "NEON_PROXY_PATH=${RUNNER_TEMP}/neon/bin/proxy"
           echo "TEST_OUTPUT=${PROXY_BENCH_PATH}/test_output"
-          echo "DOCKER_COMPOSE_FILE=${PROXY_BENCH_PATH}/docker-compose.yml"
           echo ""
         } >> "$GITHUB_ENV"
 
-    - name: Determine commit hash
-      id: commit_hash
-      shell: bash -euxo pipefail {0}
-      env:
-        INPUT_COMMIT_HASH: ${{ github.event.inputs.commit_hash }}
-      run: |
-        if [[ -z "${INPUT_COMMIT_HASH}" ]]; then
-          COMMIT_HASH=$(curl -s https://api.github.com/repos/neondatabase/neon/commits/main | jq -r '.sha')
-          echo "COMMIT_HASH=$COMMIT_HASH" >> $GITHUB_ENV
-          echo "commit_hash=$COMMIT_HASH" >> "$GITHUB_OUTPUT"
-          echo "COMMIT_HASH_TYPE=latest" >> $GITHUB_ENV
-        else
-          COMMIT_HASH="${INPUT_COMMIT_HASH}"
-          echo "COMMIT_HASH=$COMMIT_HASH" >> $GITHUB_ENV
-          echo "commit_hash=$COMMIT_HASH" >> "$GITHUB_OUTPUT"
-          echo "COMMIT_HASH_TYPE=manual" >> $GITHUB_ENV
-        fi
-    - name: Checkout the neon repository at given commit hash
-      uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
-      with:
-        ref: ${{ steps.commit_hash.outputs.commit_hash }}
-
-    - name: Print GITHUB_WORKSPACE
-      run: echo "$GITHUB_WORKSPACE"
-
-    - name: List parent dir of workspace
-      run: ls -la /__w/neon
-
-    - name: List all env vars
-      run: env
-
-    - name: Checkout proxy-bench Repo
-      uses: actions/checkout@v4
-      with:
-        repository: neondatabase/proxy-bench
-        path: proxy-bench
-
     - name: Cache poetry deps
       uses: actions/cache@v4
       with:
         path: ~/.cache/pypoetry/virtualenvs
         key: v2-${{ runner.os }}-${{ runner.arch }}-python-deps-bookworm-${{ hashFiles('poetry.lock') }}
 
-    - name: DEBUG List files for debugging
-      working-directory: ${{ env.NEON_DIR }}
-      run: |
-        pwd
-        ls -la
-
     - name: Install Python deps
-      working-directory: ${{ env.NEON_DIR }}
       shell: bash -euxo pipefail {0}
       run: ./scripts/pysync
 
@@ -122,14 +77,14 @@ jobs:
 
     - name: Run proxy-bench
       working-directory: ${{ env.PROXY_BENCH_PATH }}
-      run: ./run.sh --bare-metal
+      run: ./run.sh --with-grafana --bare-metal
 
     - name: Ingest Bench Results
       if: always()
       working-directory: ${{ env.NEON_DIR }}
       run: |
         mkdir -p $TEST_OUTPUT
-        python3 ./scripts/proxy_bench_results_ingest.py --out $TEST_OUTPUT
+        python $NEON_DIR/scripts/proxy_bench_results_ingest.py --out $TEST_OUTPUT
 
     - name: Push Metrics to Proxy perf database
       shell: bash -euxo pipefail {0}
@@ -154,4 +109,4 @@ jobs:
         fi
         if [[ -d "${PROXY_BENCH_PATH}/test_output" ]]; then
           rm -rf ${PROXY_BENCH_PATH}/test_output
-        fi
+        fi

compute/etc/sql_exporter/checkpoints_req.17.sql

@@ -1 +1 @@
-SELECT num_requested AS checkpoints_req FROM pg_stat_checkpointer;
+SELECT num_requested AS checkpoints_req FROM pg_catalog.pg_stat_checkpointer;

compute/etc/sql_exporter/checkpoints_req.sql

@@ -1 +1 @@
-SELECT checkpoints_req FROM pg_stat_bgwriter;
+SELECT checkpoints_req FROM pg_catalog.pg_stat_bgwriter;

compute/etc/sql_exporter/checkpoints_timed.sql

@@ -1 +1 @@
-SELECT checkpoints_timed FROM pg_stat_bgwriter;
+SELECT checkpoints_timed FROM pg_catalog.pg_stat_bgwriter;

compute/etc/sql_exporter/compute_backpressure_throttling_seconds_total.sql

@@ -1 +1 @@
-SELECT (neon.backpressure_throttling_time()::float8 / 1000000) AS throttled;
+SELECT (neon.backpressure_throttling_time()::pg_catalog.float8 / 1000000) AS throttled;

compute/etc/sql_exporter/compute_current_lsn.sql

@@ -1,4 +1,4 @@
 SELECT CASE
-  WHEN pg_catalog.pg_is_in_recovery() THEN (pg_last_wal_replay_lsn() - '0/0')::FLOAT8
-  ELSE (pg_current_wal_lsn() - '0/0')::FLOAT8
+  WHEN pg_catalog.pg_is_in_recovery() THEN (pg_catalog.pg_last_wal_replay_lsn() - '0/0')::pg_catalog.FLOAT8
+  ELSE (pg_catalog.pg_current_wal_lsn() - '0/0')::pg_catalog.FLOAT8
 END AS lsn;

compute/etc/sql_exporter/compute_logical_snapshot_files.sql

@@ -1,7 +1,7 @@
 SELECT
-  (SELECT setting FROM pg_settings WHERE name = 'neon.timeline_id') AS timeline_id,
+  (SELECT setting FROM pg_catalog.pg_settings WHERE name = 'neon.timeline_id') AS timeline_id,
   -- Postgres creates temporary snapshot files of the form %X-%X.snap.%d.tmp.
   -- These temporary snapshot files are renamed to the actual snapshot files
   -- after they are completely built. We only WAL-log the completely built
   -- snapshot files
-  (SELECT COUNT(*) FROM pg_ls_dir('pg_logical/snapshots') AS name WHERE name LIKE '%.snap') AS num_logical_snapshot_files;
+  (SELECT COUNT(*) FROM pg_catalog.pg_ls_dir('pg_logical/snapshots') AS name WHERE name LIKE '%.snap') AS num_logical_snapshot_files;

compute/etc/sql_exporter/compute_logical_snapshots_bytes.15.sql

@@ -1,7 +1,7 @@
 SELECT
-  (SELECT current_setting('neon.timeline_id')) AS timeline_id,
+  (SELECT pg_catalog.current_setting('neon.timeline_id')) AS timeline_id,
   -- Postgres creates temporary snapshot files of the form %X-%X.snap.%d.tmp.
   -- These temporary snapshot files are renamed to the actual snapshot files
   -- after they are completely built. We only WAL-log the completely built
   -- snapshot files
-  (SELECT COALESCE(sum(size), 0) FROM pg_ls_logicalsnapdir() WHERE name LIKE '%.snap') AS logical_snapshots_bytes;
+  (SELECT COALESCE(pg_catalog.sum(size), 0) FROM pg_catalog.pg_ls_logicalsnapdir() WHERE name LIKE '%.snap') AS logical_snapshots_bytes;

compute/etc/sql_exporter/compute_logical_snapshots_bytes.sql

@@ -1,9 +1,9 @@
 SELECT
-  (SELECT setting FROM pg_settings WHERE name = 'neon.timeline_id') AS timeline_id,
+  (SELECT setting FROM pg_catalog.pg_settings WHERE name = 'neon.timeline_id') AS timeline_id,
   -- Postgres creates temporary snapshot files of the form %X-%X.snap.%d.tmp.
   -- These temporary snapshot files are renamed to the actual snapshot files
   -- after they are completely built. We only WAL-log the completely built
   -- snapshot files
-  (SELECT COALESCE(sum((pg_stat_file('pg_logical/snapshots/' || name, missing_ok => true)).size), 0)
-    FROM (SELECT * FROM pg_ls_dir('pg_logical/snapshots') WHERE pg_ls_dir LIKE '%.snap') AS name
+  (SELECT COALESCE(pg_catalog.sum((pg_catalog.pg_stat_file('pg_logical/snapshots/' || name, missing_ok => true)).size), 0)
+   FROM (SELECT * FROM pg_catalog.pg_ls_dir('pg_logical/snapshots') WHERE pg_ls_dir LIKE '%.snap') AS name
   ) AS logical_snapshots_bytes;

compute/etc/sql_exporter/compute_max_connections.sql

@@ -1 +1 @@
-SELECT current_setting('max_connections') as max_connections;
+SELECT pg_catalog.current_setting('max_connections') AS max_connections;

compute/etc/sql_exporter/compute_pg_oldest_frozen_xid_age.sql

@@ -1,4 +1,4 @@
 SELECT datname database_name,
-  age(datfrozenxid) frozen_xid_age
-FROM pg_database
+   pg_catalog.age(datfrozenxid) frozen_xid_age
+FROM pg_catalog.pg_database
 ORDER BY frozen_xid_age DESC LIMIT 10;

compute/etc/sql_exporter/compute_pg_oldest_mxid_age.sql

@@ -1,4 +1,4 @@
 SELECT datname database_name,
-  mxid_age(datminmxid) min_mxid_age
-FROM pg_database
+  pg_catalog.mxid_age(datminmxid) min_mxid_age
+FROM pg_catalog.pg_database
 ORDER BY min_mxid_age DESC LIMIT 10;

compute/etc/sql_exporter/compute_receive_lsn.sql

@@ -1,4 +1,4 @@
 SELECT CASE
-  WHEN pg_catalog.pg_is_in_recovery() THEN (pg_last_wal_receive_lsn() - '0/0')::FLOAT8
+  WHEN pg_catalog.pg_is_in_recovery() THEN (pg_catalog.pg_last_wal_receive_lsn() - '0/0')::pg_catalog.FLOAT8
   ELSE 0
 END AS lsn;

compute/etc/sql_exporter/compute_subscriptions_count.sql

@@ -1 +1 @@
-SELECT subenabled::text AS enabled, count(*) AS subscriptions_count FROM pg_subscription GROUP BY subenabled;
+SELECT subenabled::pg_catalog.text AS enabled, pg_catalog.count(*) AS subscriptions_count FROM pg_catalog.pg_subscription GROUP BY subenabled;

compute/etc/sql_exporter/connection_counts.sql

@@ -1 +1 @@
-SELECT datname, state, count(*) AS count FROM pg_stat_activity WHERE state <> '' GROUP BY datname, state;
+SELECT datname, state, pg_catalog.count(*) AS count FROM pg_catalog.pg_stat_activity WHERE state <> '' GROUP BY datname, state;

compute/etc/sql_exporter/db_total_size.sql

@@ -1,5 +1,5 @@
-SELECT sum(pg_database_size(datname)) AS total
-FROM pg_database
+SELECT pg_catalog.sum(pg_catalog.pg_database_size(datname)) AS total
+FROM pg_catalog.pg_database
 -- Ignore invalid databases, as we will likely have problems with
 -- getting their size from the Pageserver.
 WHERE datconnlimit != -2;

compute/etc/sql_exporter/lfc_approximate_working_set_size_windows.autoscaling.sql

@@ -3,6 +3,6 @@
 -- minutes.
 
 SELECT
-  x::text as duration_seconds,
+  x::pg_catalog.text AS duration_seconds,
   neon.approximate_working_set_size_seconds(x) AS size
 FROM (SELECT generate_series * 60 AS x FROM generate_series(1, 60)) AS t (x);

compute/etc/sql_exporter/lfc_approximate_working_set_size_windows.sql

@@ -3,6 +3,6 @@
 
 SELECT
   x AS duration,
-  neon.approximate_working_set_size_seconds(extract('epoch' FROM x::interval)::int) AS size FROM (
+  neon.approximate_working_set_size_seconds(extract('epoch' FROM x::pg_catalog.interval)::pg_catalog.int4) AS size FROM (
     VALUES ('5m'), ('15m'), ('1h')
   ) AS t (x);

compute/etc/sql_exporter/lfc_cache_size_limit.sql

@@ -1 +1 @@
-SELECT pg_size_bytes(current_setting('neon.file_cache_size_limit')) AS lfc_cache_size_limit;
+SELECT pg_catalog.pg_size_bytes(pg_catalog.current_setting('neon.file_cache_size_limit')) AS lfc_cache_size_limit;

compute/etc/sql_exporter/logical_slot_restart_lsn.sql

@@ -1,3 +1,3 @@
-SELECT slot_name, (restart_lsn - '0/0')::FLOAT8 as restart_lsn
-FROM pg_replication_slots
+SELECT slot_name, (restart_lsn - '0/0')::pg_catalog.FLOAT8 AS restart_lsn
+FROM pg_catalog.pg_replication_slots
 WHERE slot_type = 'logical';

compute/etc/sql_exporter/max_cluster_size.sql

@@ -1 +1 @@
-SELECT setting::int AS max_cluster_size FROM pg_settings WHERE name = 'neon.max_cluster_size';
+SELECT setting::pg_catalog.int4 AS max_cluster_size FROM pg_catalog.pg_settings WHERE name = 'neon.max_cluster_size';

compute/etc/sql_exporter/pg_stats_userdb.sql

@@ -1,13 +1,13 @@
 -- We export stats for 10 non-system databases. Without this limit it is too
 -- easy to abuse the system by creating lots of databases.
 
-SELECT pg_database_size(datname) AS db_size,
+SELECT pg_catalog.pg_database_size(datname) AS db_size,
   deadlocks,
   tup_inserted AS inserted,
   tup_updated AS updated,
   tup_deleted AS deleted,
   datname
-FROM pg_stat_database
+FROM pg_catalog.pg_stat_database
 WHERE datname IN (
   SELECT datname FROM pg_database
   -- Ignore invalid databases, as we will likely have problems with

compute/etc/sql_exporter/replication_delay_bytes.sql

@@ -3,4 +3,4 @@
 -- replay LSN may have advanced past the receive LSN we are using for the
 -- calculation.
 
-SELECT GREATEST(0, pg_wal_lsn_diff(pg_last_wal_receive_lsn(), pg_last_wal_replay_lsn())) AS replication_delay_bytes;
+SELECT GREATEST(0, pg_catalog.pg_wal_lsn_diff(pg_catalog.pg_last_wal_receive_lsn(), pg_catalog.pg_last_wal_replay_lsn())) AS replication_delay_bytes;

compute/etc/sql_exporter/replication_delay_seconds.sql

@@ -1,5 +1,5 @@
 SELECT
   CASE
-    WHEN pg_last_wal_receive_lsn() = pg_last_wal_replay_lsn() THEN 0
-    ELSE GREATEST(0, EXTRACT (EPOCH FROM now() - pg_last_xact_replay_timestamp()))
+    WHEN pg_catalog.pg_last_wal_receive_lsn() = pg_catalog.pg_last_wal_replay_lsn() THEN 0
+    ELSE GREATEST(0, EXTRACT (EPOCH FROM pg_catalog.now() - pg_catalog.pg_last_xact_replay_timestamp()))
   END AS replication_delay_seconds;

compute/etc/sql_exporter/retained_wal.sql

@@ -1,10 +1,10 @@
 SELECT
   slot_name,
-  pg_wal_lsn_diff(
+  pg_catalog.pg_wal_lsn_diff(
     CASE
-      WHEN pg_is_in_recovery() THEN pg_last_wal_replay_lsn()
-      ELSE pg_current_wal_lsn()
+      WHEN pg_catalog.pg_is_in_recovery() THEN pg_catalog.pg_last_wal_replay_lsn()
+      ELSE pg_catalog.pg_current_wal_lsn()
     END,
-    restart_lsn)::FLOAT8 AS retained_wal
-FROM pg_replication_slots
+    restart_lsn)::pg_catalog.FLOAT8 AS retained_wal
+FROM pg_catalog.pg_replication_slots
 WHERE active = false;

compute/etc/sql_exporter/wal_is_lost.sql

@@ -4,4 +4,4 @@ SELECT
     WHEN wal_status = 'lost' THEN 1
     ELSE 0
   END AS wal_is_lost
-FROM pg_replication_slots;
+FROM pg_catalog.pg_replication_slots;

compute_tools/src/bin/compute_ctl.rs

@@ -279,7 +279,7 @@ fn main() -> Result<()> {
         config,
     )?;
 
-    let exit_code = compute_node.run()?;
+    let exit_code = compute_node.run().context("running compute node")?;
 
     scenario.teardown();
 

compute_tools/src/checker.rs

@@ -24,9 +24,9 @@ pub async fn check_writability(compute: &ComputeNode) -> Result<()> {
     });
 
     let query = "
-    INSERT INTO health_check VALUES (1, now())
+    INSERT INTO public.health_check VALUES (1, pg_catalog.now())
         ON CONFLICT (id) DO UPDATE
-         SET updated_at = now();";
+         SET updated_at = pg_catalog.now();";
 
     match client.simple_query(query).await {
         Result::Ok(result) => {

compute_tools/src/compute.rs

@@ -32,6 +32,7 @@ use std::sync::{Arc, Condvar, Mutex, RwLock};
 use std::time::{Duration, Instant};
 use std::{env, fs};
 use tokio::{spawn, sync::watch, task::JoinHandle, time};
+use tokio_util::sync::CancellationToken;
 use tracing::{Instrument, debug, error, info, instrument, warn};
 use url::Url;
 use utils::id::{TenantId, TimelineId};
@@ -192,6 +193,7 @@ pub struct ComputeState {
     pub startup_span: Option<tracing::span::Span>,
 
     pub lfc_prewarm_state: LfcPrewarmState,
+    pub lfc_prewarm_token: CancellationToken,
     pub lfc_offload_state: LfcOffloadState,
 
     /// WAL flush LSN that is set after terminating Postgres and syncing safekeepers if
@@ -217,6 +219,7 @@ impl ComputeState {
             lfc_offload_state: LfcOffloadState::default(),
             terminate_flush_lsn: None,
             promote_state: None,
+            lfc_prewarm_token: CancellationToken::new(),
         }
     }
 
@@ -583,7 +586,7 @@ impl ComputeNode {
         // that can affect `compute_ctl` and prevent it from properly configuring the database schema.
         // Unset them via connection string options before connecting to the database.
         // N.B. keep it in sync with `ZENITH_OPTIONS` in `get_maintenance_client()`.
-        const EXTRA_OPTIONS: &str = "-c role=cloud_admin -c default_transaction_read_only=off -c search_path=public -c statement_timeout=0 -c pgaudit.log=none";
+        const EXTRA_OPTIONS: &str = "-c role=cloud_admin -c default_transaction_read_only=off -c search_path='' -c statement_timeout=0 -c pgaudit.log=none";
         let options = match conn_conf.get_options() {
             // Allow the control plane to override any options set by the
             // compute
@@ -1884,7 +1887,7 @@ impl ComputeNode {
 
                     // It doesn't matter what were the options before, here we just want
                     // to connect and create a new superuser role.
-                    const ZENITH_OPTIONS: &str = "-c role=zenith_admin -c default_transaction_read_only=off -c search_path=public -c statement_timeout=0";
+                    const ZENITH_OPTIONS: &str = "-c role=zenith_admin -c default_transaction_read_only=off -c search_path='' -c statement_timeout=0";
                     zenith_admin_conf.options(ZENITH_OPTIONS);
 
                     let mut client =
@@ -2339,13 +2342,13 @@ impl ComputeNode {
         let result = client
             .simple_query(
                 "SELECT
-    row_to_json(pg_stat_statements)
+    pg_catalog.row_to_json(pss)
 FROM
-    pg_stat_statements
+    public.pg_stat_statements pss
 WHERE
-    userid != 'cloud_admin'::regrole::oid
+    pss.userid != 'cloud_admin'::pg_catalog.regrole::pg_catalog.oid
 ORDER BY
-    (mean_exec_time + mean_plan_time) DESC
+    (pss.mean_exec_time + pss.mean_plan_time) DESC
 LIMIT 100",
             )
             .await;
@@ -2473,11 +2476,11 @@ LIMIT 100",
 
         // check the role grants first - to gracefully handle read-replicas.
         let select = "SELECT privilege_type
-            FROM pg_namespace
-                JOIN LATERAL (SELECT * FROM aclexplode(nspacl) AS x) acl ON true
-                JOIN pg_user users ON acl.grantee = users.usesysid
-            WHERE users.usename = $1
-                AND nspname = $2";
+            FROM pg_catalog.pg_namespace
+                JOIN LATERAL (SELECT * FROM aclexplode(nspacl) AS x) AS acl ON true
+                JOIN pg_catalog.pg_user users ON acl.grantee = users.usesysid
+            WHERE users.usename OPERATOR(pg_catalog.=) $1::pg_catalog.name
+                AND nspname OPERATOR(pg_catalog.=) $2::pg_catalog.name";
         let rows = db_client
             .query(select, &[role_name, schema_name])
             .await
@@ -2546,8 +2549,9 @@ LIMIT 100",
                 .await
                 .with_context(|| format!("Failed to execute query: {query}"))?;
         } else {
-            let query =
-                format!("CREATE EXTENSION IF NOT EXISTS {ext_name} WITH VERSION {quoted_version}");
+            let query = format!(
+                "CREATE EXTENSION IF NOT EXISTS {ext_name} WITH SCHEMA public VERSION {quoted_version}"
+            );
             db_client
                 .simple_query(&query)
                 .await

compute_tools/src/compute_prewarm.rs

@@ -7,7 +7,8 @@ use http::StatusCode;
 use reqwest::Client;
 use std::mem::replace;
 use std::sync::Arc;
-use tokio::{io::AsyncReadExt, spawn};
+use tokio::{io::AsyncReadExt, select, spawn};
+use tokio_util::sync::CancellationToken;
 use tracing::{error, info};
 
 #[derive(serde::Serialize, Default)]
@@ -92,34 +93,35 @@ impl ComputeNode {
     /// If there is a prewarm request ongoing, return `false`, `true` otherwise.
     /// Has a failpoint "compute-prewarm"
     pub fn prewarm_lfc(self: &Arc<Self>, from_endpoint: Option<String>) -> bool {
+        let token: CancellationToken;
         {
-            let state = &mut self.state.lock().unwrap().lfc_prewarm_state;
-            if let LfcPrewarmState::Prewarming = replace(state, LfcPrewarmState::Prewarming) {
+            let state = &mut self.state.lock().unwrap();
+            token = state.lfc_prewarm_token.clone();
+            if let LfcPrewarmState::Prewarming =
+                replace(&mut state.lfc_prewarm_state, LfcPrewarmState::Prewarming)
+            {
                 return false;
             }
         }
         crate::metrics::LFC_PREWARMS.inc();
 
-        let cloned = self.clone();
+        let this = self.clone();
         spawn(async move {
-            let state = match cloned.prewarm_impl(from_endpoint).await {
-                Ok(true) => LfcPrewarmState::Completed,
-                Ok(false) => {
-                    info!(
-                        "skipping LFC prewarm because LFC state is not found in endpoint storage"
-                    );
-                    LfcPrewarmState::Skipped
-                }
+            let prewarm_state = match this.prewarm_impl(from_endpoint, token).await {
+                Ok(state) => state,
                 Err(err) => {
                     crate::metrics::LFC_PREWARM_ERRORS.inc();
                     error!(%err, "could not prewarm LFC");
-                    LfcPrewarmState::Failed {
-                        error: format!("{err:#}"),
-                    }
+                    let error = format!("{err:#}");
+                    LfcPrewarmState::Failed { error }
                 }
             };
 
-            cloned.state.lock().unwrap().lfc_prewarm_state = state;
+            let state = &mut this.state.lock().unwrap();
+            if let LfcPrewarmState::Cancelled = prewarm_state {
+                state.lfc_prewarm_token = CancellationToken::new();
+            }
+            state.lfc_prewarm_state = prewarm_state;
         });
         true
     }
@@ -132,47 +134,70 @@ impl ComputeNode {
 
     /// Request LFC state from endpoint storage and load corresponding pages into Postgres.
     /// Returns a result with `false` if the LFC state is not found in endpoint storage.
-    async fn prewarm_impl(&self, from_endpoint: Option<String>) -> Result<bool> {
-        let EndpointStoragePair { url, token } = self.endpoint_storage_pair(from_endpoint)?;
+    async fn prewarm_impl(
+        &self,
+        from_endpoint: Option<String>,
+        token: CancellationToken,
+    ) -> Result<LfcPrewarmState> {
+        let EndpointStoragePair {
+            url,
+            token: storage_token,
+        } = self.endpoint_storage_pair(from_endpoint)?;
 
         #[cfg(feature = "testing")]
-        fail::fail_point!("compute-prewarm", |_| {
-            bail!("prewarm configured to fail because of a failpoint")
-        });
+        fail::fail_point!("compute-prewarm", |_| bail!("compute-prewarm failpoint"));
 
         info!(%url, "requesting LFC state from endpoint storage");
-        let request = Client::new().get(&url).bearer_auth(token);
-        let res = request.send().await.context("querying endpoint storage")?;
-        match res.status() {
+        let request = Client::new().get(&url).bearer_auth(storage_token);
+        let response = select! {
+            _ = token.cancelled() => return Ok(LfcPrewarmState::Cancelled),
+            response = request.send() => response
+        }
+        .context("querying endpoint storage")?;
+
+        match response.status() {
             StatusCode::OK => (),
-            StatusCode::NOT_FOUND => {
-                return Ok(false);
-            }
+            StatusCode::NOT_FOUND => return Ok(LfcPrewarmState::Skipped),
             status => bail!("{status} querying endpoint storage"),
         }
 
         let mut uncompressed = Vec::new();
-        let lfc_state = res
-            .bytes()
-            .await
-            .context("getting request body from endpoint storage")?;
-        ZstdDecoder::new(lfc_state.iter().as_slice())
-            .read_to_end(&mut uncompressed)
-            .await
-            .context("decoding LFC state")?;
+        let lfc_state = select! {
+            _ = token.cancelled() => return Ok(LfcPrewarmState::Cancelled),
+            lfc_state = response.bytes() => lfc_state
+        }
+        .context("getting request body from endpoint storage")?;
+
+        let mut decoder = ZstdDecoder::new(lfc_state.iter().as_slice());
+        select! {
+            _ = token.cancelled() => return Ok(LfcPrewarmState::Cancelled),
+            read = decoder.read_to_end(&mut uncompressed) => read
+        }
+        .context("decoding LFC state")?;
+
         let uncompressed_len = uncompressed.len();
+        info!(%url, "downloaded LFC state, uncompressed size {uncompressed_len}");
 
-        info!(%url, "downloaded LFC state, uncompressed size {uncompressed_len}, loading into Postgres");
-
-        ComputeNode::get_maintenance_client(&self.tokio_conn_conf)
+        // Client connection and prewarm info querying are fast and therefore don't need
+        // cancellation
+        let client = ComputeNode::get_maintenance_client(&self.tokio_conn_conf)
             .await
-            .context("connecting to postgres")?
-            .query_one("select neon.prewarm_local_cache($1)", &[&uncompressed])
-            .await
-            .context("loading LFC state into postgres")
-            .map(|_| ())?;
+            .context("connecting to postgres")?;
+        let pg_token = client.cancel_token();
 
-        Ok(true)
+        let params: Vec<&(dyn postgres_types::ToSql + Sync)> = vec![&uncompressed];
+        select! {
+            res = client.query_one("select neon.prewarm_local_cache($1)", &params) => res,
+            _ = token.cancelled() => {
+                pg_token.cancel_query(postgres::NoTls).await
+                    .context("cancelling neon.prewarm_local_cache()")?;
+                return Ok(LfcPrewarmState::Cancelled)
+            }
+        }
+        .context("loading LFC state into postgres")
+        .map(|_| ())?;
+
+        Ok(LfcPrewarmState::Completed)
     }
 
     /// If offload request is ongoing, return false, true otherwise
@@ -200,20 +225,20 @@ impl ComputeNode {
 
     async fn offload_lfc_with_state_update(&self) {
         crate::metrics::LFC_OFFLOADS.inc();
-
-        let Err(err) = self.offload_lfc_impl().await else {
-            self.state.lock().unwrap().lfc_offload_state = LfcOffloadState::Completed;
-            return;
+        let state = match self.offload_lfc_impl().await {
+            Ok(state) => state,
+            Err(err) => {
+                crate::metrics::LFC_OFFLOAD_ERRORS.inc();
+                error!(%err, "could not offload LFC");
+                let error = format!("{err:#}");
+                LfcOffloadState::Failed { error }
+            }
         };
 
-        crate::metrics::LFC_OFFLOAD_ERRORS.inc();
-        error!(%err, "could not offload LFC state to endpoint storage");
-        self.state.lock().unwrap().lfc_offload_state = LfcOffloadState::Failed {
-            error: format!("{err:#}"),
-        };
+        self.state.lock().unwrap().lfc_offload_state = state;
     }
 
-    async fn offload_lfc_impl(&self) -> Result<()> {
+    async fn offload_lfc_impl(&self) -> Result<LfcOffloadState> {
         let EndpointStoragePair { url, token } = self.endpoint_storage_pair(None)?;
         info!(%url, "requesting LFC state from Postgres");
 
@@ -228,7 +253,7 @@ impl ComputeNode {
             .context("deserializing LFC state")?;
         let Some(state) = state else {
             info!(%url, "empty LFC state, not exporting");
-            return Ok(());
+            return Ok(LfcOffloadState::Skipped);
         };
 
         let mut compressed = Vec::new();
@@ -242,7 +267,7 @@ impl ComputeNode {
 
         let request = Client::new().put(url).bearer_auth(token).body(compressed);
         match request.send().await {
-            Ok(res) if res.status() == StatusCode::OK => Ok(()),
+            Ok(res) if res.status() == StatusCode::OK => Ok(LfcOffloadState::Completed),
             Ok(res) => bail!(
                 "Request to endpoint storage failed with status: {}",
                 res.status()
@@ -250,4 +275,8 @@ impl ComputeNode {
             Err(err) => Err(err).context("writing to endpoint storage"),
         }
     }
+
+    pub fn cancel_prewarm(self: &Arc<Self>) {
+        self.state.lock().unwrap().lfc_prewarm_token.cancel();
+    }
 }

compute_tools/src/compute_promote.rs

@@ -78,7 +78,7 @@ impl ComputeNode {
         const RETRIES: i32 = 20;
         for i in 0..=RETRIES {
             let row = client
-                .query_one("SELECT pg_last_wal_replay_lsn()", &[])
+                .query_one("SELECT pg_catalog.pg_last_wal_replay_lsn()", &[])
                 .await
                 .context("getting last replay lsn")?;
             let lsn: u64 = row.get::<usize, postgres_types::PgLsn>(0).into();
@@ -103,7 +103,7 @@ impl ComputeNode {
             .await
             .context("setting safekeepers")?;
         client
-            .query("SELECT pg_reload_conf()", &[])
+            .query("SELECT pg_catalog.pg_reload_conf()", &[])
             .await
             .context("reloading postgres config")?;
 
@@ -113,7 +113,7 @@ impl ComputeNode {
         });
 
         let row = client
-            .query_one("SELECT * FROM pg_promote()", &[])
+            .query_one("SELECT * FROM pg_catalog.pg_promote()", &[])
             .await
             .context("pg_promote")?;
         if !row.get::<usize, bool>(0) {

compute_tools/src/http/openapi_spec.yaml

@@ -139,6 +139,15 @@ paths:
             application/json:
               schema:
                 $ref: "#/components/schemas/LfcPrewarmState"
+    delete:
+      tags:
+        - Prewarm
+      summary: Cancel ongoing LFC prewarm
+      description: ""
+      operationId: cancelLfcPrewarm
+      responses:
+        202:
+          description: Prewarm cancelled
 
   /lfc/offload:
     post:
@@ -636,7 +645,7 @@ components:
       properties:
         status:
           description: LFC offload status
-          enum: [not_offloaded, offloading, completed, failed]
+          enum: [not_offloaded, offloading, completed, skipped, failed]
           type: string
         error:
           description: LFC offload error, if any

compute_tools/src/http/routes/lfc.rs

@@ -46,3 +46,8 @@ pub(in crate::http) async fn offload(compute: Compute) -> Response {
         )
     }
 }
+
+pub(in crate::http) async fn cancel_prewarm(compute: Compute) -> StatusCode {
+    compute.cancel_prewarm();
+    StatusCode::ACCEPTED
+}

compute_tools/src/http/server.rs

@@ -99,7 +99,12 @@ impl From<&Server> for Router<Arc<ComputeNode>> {
                     );
 
                 let authenticated_router = Router::<Arc<ComputeNode>>::new()
-                    .route("/lfc/prewarm", get(lfc::prewarm_state).post(lfc::prewarm))
+                    .route(
+                        "/lfc/prewarm",
+                        get(lfc::prewarm_state)
+                            .post(lfc::prewarm)
+                            .delete(lfc::cancel_prewarm),
+                    )
                     .route("/lfc/offload", get(lfc::offload_state).post(lfc::offload))
                     .route("/promote", post(promote::promote))
                     .route("/check_writability", post(check_writability::is_writable))

compute_tools/src/installed_extensions.rs

@@ -19,7 +19,7 @@ async fn list_dbs(client: &mut Client) -> Result<Vec<String>, PostgresError> {
         .query(
             "SELECT datname FROM pg_catalog.pg_database
                 WHERE datallowconn
-                AND datconnlimit <> - 2
+                AND datconnlimit OPERATOR(pg_catalog.<>) (OPERATOR(pg_catalog.-) 2::pg_catalog.int4)
                 LIMIT 500",
             &[],
         )
@@ -67,7 +67,7 @@ pub async fn get_installed_extensions(
 
         let extensions: Vec<(String, String, i32)> = client
             .query(
-                "SELECT extname, extversion, extowner::integer FROM pg_catalog.pg_extension",
+                "SELECT extname, extversion, extowner::pg_catalog.int4 FROM pg_catalog.pg_extension",
                 &[],
             )
             .await?

compute_tools/src/migration.rs

@@ -76,7 +76,7 @@ impl<'m> MigrationRunner<'m> {
         self.client
             .simple_query("CREATE SCHEMA IF NOT EXISTS neon_migration")
             .await?;
-        self.client.simple_query("CREATE TABLE IF NOT EXISTS neon_migration.migration_id (key INT NOT NULL PRIMARY KEY, id bigint NOT NULL DEFAULT 0)").await?;
+        self.client.simple_query("CREATE TABLE IF NOT EXISTS neon_migration.migration_id (key pg_catalog.int4 NOT NULL PRIMARY KEY, id pg_catalog.int8 NOT NULL DEFAULT 0)").await?;
         self.client
             .simple_query(
                 "INSERT INTO neon_migration.migration_id VALUES (0, 0) ON CONFLICT DO NOTHING",

compute_tools/src/migrations/0002-alter_roles.sql

@@ -15,17 +15,17 @@ DO $$
 DECLARE
     role_name text;
 BEGIN
-    FOR role_name IN SELECT rolname FROM pg_roles WHERE pg_has_role(rolname, '{privileged_role_name}', 'member')
+    FOR role_name IN SELECT rolname FROM pg_catalog.pg_roles WHERE pg_catalog.pg_has_role(rolname, '{privileged_role_name}', 'member')
     LOOP
-        RAISE NOTICE 'EXECUTING ALTER ROLE % INHERIT', quote_ident(role_name);
-        EXECUTE 'ALTER ROLE ' || quote_ident(role_name) || ' INHERIT';
+        RAISE NOTICE 'EXECUTING ALTER ROLE % INHERIT', pg_catalog.quote_ident(role_name);
+        EXECUTE pg_catalog.format('ALTER ROLE %I INHERIT;', role_name);
     END LOOP;
 
-    FOR role_name IN SELECT rolname FROM pg_roles
+    FOR role_name IN SELECT rolname FROM pg_catalog.pg_roles
         WHERE
-            NOT pg_has_role(rolname, '{privileged_role_name}', 'member') AND NOT starts_with(rolname, 'pg_')
+            NOT pg_catalog.pg_has_role(rolname, '{privileged_role_name}', 'member') AND NOT pg_catalog.starts_with(rolname, 'pg_')
     LOOP
-        RAISE NOTICE 'EXECUTING ALTER ROLE % NOBYPASSRLS', quote_ident(role_name);
-        EXECUTE 'ALTER ROLE ' || quote_ident(role_name) || ' NOBYPASSRLS';
+        RAISE NOTICE 'EXECUTING ALTER ROLE % NOBYPASSRLS', pg_catalog.quote_ident(role_name);
+        EXECUTE pg_catalog.format('ALTER ROLE %I NOBYPASSRLS;', role_name);
     END LOOP;
 END $$;

compute_tools/src/migrations/0003-grant_pg_create_subscription_to_privileged_role.sql

@@ -1,6 +1,6 @@
 DO $$
 BEGIN
-    IF (SELECT setting::numeric >= 160000 FROM pg_settings WHERE name = 'server_version_num') THEN
+    IF (SELECT setting::pg_catalog.numeric >= 160000 FROM pg_catalog.pg_settings WHERE name = 'server_version_num') THEN
         EXECUTE 'GRANT pg_create_subscription TO {privileged_role_name}';
     END IF;
 END $$;

compute_tools/src/migrations/0009-revoke_replication_for_previously_allowed_roles.sql

@@ -5,9 +5,9 @@ DO $$
 DECLARE
     role_name TEXT;
 BEGIN
-    FOR role_name IN SELECT rolname FROM pg_roles WHERE rolreplication IS TRUE
+    FOR role_name IN SELECT rolname FROM pg_catalog.pg_roles WHERE rolreplication IS TRUE
     LOOP
-        RAISE NOTICE 'EXECUTING ALTER ROLE % NOREPLICATION', quote_ident(role_name);
-        EXECUTE 'ALTER ROLE ' || quote_ident(role_name) || ' NOREPLICATION';
+        RAISE NOTICE 'EXECUTING ALTER ROLE % NOREPLICATION', pg_catalog.quote_ident(role_name);
+        EXECUTE pg_catalog.format('ALTER ROLE %I NOREPLICATION;', role_name);
     END LOOP;
 END $$;

compute_tools/src/migrations/0010-grant_snapshot_synchronization_funcs_to_privileged_role.sql

@@ -1,6 +1,6 @@
 DO $$
 BEGIN
-    IF (SELECT setting::numeric >= 160000 FROM pg_settings WHERE name = 'server_version_num') THEN
+    IF (SELECT setting::pg_catalog.numeric >= 160000 FROM pg_catalog.pg_settings WHERE name OPERATOR(pg_catalog.=) 'server_version_num'::pg_catalog.text) THEN
        EXECUTE 'GRANT EXECUTE ON FUNCTION pg_export_snapshot TO {privileged_role_name}';
        EXECUTE 'GRANT EXECUTE ON FUNCTION pg_log_standby_snapshot TO {privileged_role_name}';
     END IF;

compute_tools/src/migrations/tests/0001-add_bypass_rls_to_privileged_role.sql

@@ -2,7 +2,7 @@ DO $$
 DECLARE
     bypassrls boolean;
 BEGIN
-    SELECT rolbypassrls INTO bypassrls FROM pg_roles WHERE rolname = 'neon_superuser';
+    SELECT rolbypassrls INTO bypassrls FROM pg_catalog.pg_roles WHERE rolname = 'neon_superuser';
     IF NOT bypassrls THEN
         RAISE EXCEPTION 'neon_superuser cannot bypass RLS';
     END IF;

compute_tools/src/migrations/tests/0002-alter_roles.sql

@@ -4,8 +4,8 @@ DECLARE
 BEGIN
     FOR role IN
         SELECT rolname AS name, rolinherit AS inherit
-        FROM pg_roles
-        WHERE pg_has_role(rolname, 'neon_superuser', 'member')
+        FROM pg_catalog.pg_roles
+        WHERE pg_catalog.pg_has_role(rolname, 'neon_superuser', 'member')
     LOOP
         IF NOT role.inherit THEN
             RAISE EXCEPTION '% cannot inherit', quote_ident(role.name);
@@ -14,12 +14,12 @@ BEGIN
 
     FOR role IN
         SELECT rolname AS name, rolbypassrls AS bypassrls
-        FROM pg_roles
-        WHERE NOT pg_has_role(rolname, 'neon_superuser', 'member')
-            AND NOT starts_with(rolname, 'pg_')
+        FROM pg_catalog.pg_roles
+        WHERE NOT pg_catalog.pg_has_role(rolname, 'neon_superuser', 'member')
+            AND NOT pg_catalog.starts_with(rolname, 'pg_')
     LOOP
         IF role.bypassrls THEN
-            RAISE EXCEPTION  '% can bypass RLS', quote_ident(role.name);
+            RAISE EXCEPTION  '% can bypass RLS', pg_catalog.quote_ident(role.name);
         END IF;
     END LOOP;
 END $$;

compute_tools/src/migrations/tests/0003-grant_pg_create_subscription_to_privileged_role.sql

@@ -1,10 +1,10 @@
 DO $$
 BEGIN
-    IF (SELECT current_setting('server_version_num')::numeric < 160000) THEN
+    IF (SELECT pg_catalog.current_setting('server_version_num')::pg_catalog.numeric < 160000) THEN
         RETURN;
     END IF;
 
-    IF NOT (SELECT pg_has_role('neon_superuser', 'pg_create_subscription', 'member')) THEN
+    IF NOT (SELECT pg_catalog.pg_has_role('neon_superuser', 'pg_create_subscription', 'member')) THEN
         RAISE EXCEPTION 'neon_superuser cannot execute pg_create_subscription';
     END IF;
 END $$;

compute_tools/src/migrations/tests/0004-grant_pg_monitor_to_privileged_role.sql

@@ -2,12 +2,12 @@ DO $$
 DECLARE
     monitor record;
 BEGIN
-    SELECT pg_has_role('neon_superuser', 'pg_monitor', 'member') AS member,
+    SELECT pg_catalog.pg_has_role('neon_superuser', 'pg_monitor', 'member') AS member,
             admin_option AS admin
         INTO monitor
-        FROM pg_auth_members
-        WHERE roleid = 'pg_monitor'::regrole
-            AND member = 'neon_superuser'::regrole;
+        FROM pg_catalog.pg_auth_members
+        WHERE roleid = 'pg_monitor'::pg_catalog.regrole
+            AND member = 'neon_superuser'::pg_catalog.regrole;
 
     IF monitor IS NULL THEN
         RAISE EXCEPTION 'no entry in pg_auth_members for neon_superuser and pg_monitor';

compute_tools/src/migrations/tests/0010-grant_snapshot_synchronization_funcs_to_privileged_role.sql

@@ -2,11 +2,11 @@ DO $$
 DECLARE
     can_execute boolean;
 BEGIN
-    SELECT bool_and(has_function_privilege('neon_superuser', oid, 'execute'))
+    SELECT pg_catalog.bool_and(pg_catalog.has_function_privilege('neon_superuser', oid, 'execute'))
        INTO can_execute
-       FROM pg_proc
+       FROM pg_catalog.pg_proc
        WHERE proname IN ('pg_export_snapshot', 'pg_log_standby_snapshot')
-           AND pronamespace = 'pg_catalog'::regnamespace;
+           AND pronamespace = 'pg_catalog'::pg_catalog.regnamespace;
     IF NOT can_execute THEN
         RAISE EXCEPTION 'neon_superuser cannot execute both pg_export_snapshot and pg_log_standby_snapshot';
     END IF;

compute_tools/src/migrations/tests/0011-grant_pg_show_replication_origin_status_to_privileged_role.sql

@@ -2,9 +2,9 @@ DO $$
 DECLARE
     can_execute boolean;
 BEGIN
-    SELECT has_function_privilege('neon_superuser', oid, 'execute')
+    SELECT pg_catalog.has_function_privilege('neon_superuser', oid, 'execute')
        INTO can_execute
-       FROM pg_proc
+       FROM pg_catalog.pg_proc
        WHERE proname = 'pg_show_replication_origin_status'
            AND pronamespace = 'pg_catalog'::regnamespace;
     IF NOT can_execute THEN

compute_tools/src/migrations/tests/0012-grant_pg_signal_backend_to_privileged_role.sql

@@ -2,10 +2,10 @@ DO $$
 DECLARE
     signal_backend record;
 BEGIN
-    SELECT pg_has_role('neon_superuser', 'pg_signal_backend', 'member') AS member,
+    SELECT pg_catalog.pg_has_role('neon_superuser', 'pg_signal_backend', 'member') AS member,
             admin_option AS admin
         INTO signal_backend
-        FROM pg_auth_members
+        FROM pg_catalog.pg_auth_members
         WHERE roleid = 'pg_signal_backend'::regrole
             AND member = 'neon_superuser'::regrole;
 

compute_tools/src/monitor.rs

@@ -407,9 +407,9 @@ fn get_database_stats(cli: &mut Client) -> anyhow::Result<(f64, i64)> {
     // like `postgres_exporter` use it to query Postgres statistics.
     // Use explicit 8 bytes type casts to match Rust types.
     let stats = cli.query_one(
-        "SELECT coalesce(sum(active_time), 0.0)::float8 AS total_active_time,
-            coalesce(sum(sessions), 0)::bigint AS total_sessions
-        FROM pg_stat_database
+        "SELECT pg_catalog.coalesce(pg_catalog.sum(active_time), 0.0)::pg_catalog.float8 AS total_active_time,
+            pg_catalog.coalesce(pg_catalog.sum(sessions), 0)::pg_catalog.bigint AS total_sessions
+        FROM pg_catalog.pg_stat_database
         WHERE datname NOT IN (
                 'postgres',
                 'template0',
@@ -445,11 +445,11 @@ fn get_backends_state_change(cli: &mut Client) -> anyhow::Result<Option<DateTime
     let mut last_active: Option<DateTime<Utc>> = None;
     // Get all running client backends except ourself, use RFC3339 DateTime format.
     let backends = cli.query(
-        "SELECT state, to_char(state_change, 'YYYY-MM-DD\"T\"HH24:MI:SS.US\"Z\"') AS state_change
+        "SELECT state, pg_catalog.to_char(state_change, 'YYYY-MM-DD\"T\"HH24:MI:SS.US\"Z\"'::pg_catalog.text) AS state_change
                 FROM pg_stat_activity
-                    WHERE backend_type = 'client backend'
-                    AND pid != pg_backend_pid()
-                    AND usename != 'cloud_admin';", // XXX: find a better way to filter other monitors?
+                    WHERE backend_type OPERATOR(pg_catalog.=) 'client backend'::pg_catalog.text
+                    AND pid OPERATOR(pg_catalog.!=) pg_catalog.pg_backend_pid()
+                    AND usename OPERATOR(pg_catalog.!=) 'cloud_admin'::pg_catalog.name;", // XXX: find a better way to filter other monitors?
         &[],
     );
 

compute_tools/src/pg_helpers.rs

@@ -299,9 +299,9 @@ pub async fn get_existing_dbs_async(
         .query_raw::<str, &String, &[String; 0]>(
             "SELECT
                 datname AS name,
-                (SELECT rolname FROM pg_roles WHERE oid = datdba) AS owner,
+                (SELECT rolname FROM pg_catalog.pg_roles WHERE oid OPERATOR(pg_catalog.=) datdba) AS owner,
                 NOT datallowconn AS restrict_conn,
-                datconnlimit = - 2 AS invalid
+                datconnlimit OPERATOR(pg_catalog.=) (OPERATOR(pg_catalog.-) 2) AS invalid
             FROM
                 pg_catalog.pg_database;",
             &[],

compute_tools/src/spec_apply.rs

@@ -82,7 +82,7 @@ impl ComputeNode {
                 info!("Checking if drop subscription operation was already performed for timeline_id: {}", timeline_id);
 
                 drop_subscriptions_done = match
-                    client.query("select 1 from neon.drop_subscriptions_done where timeline_id = $1", &[&timeline_id.to_string()]).await {
+                    client.query("select 1 from neon.drop_subscriptions_done where timeline_id OPERATOR(pg_catalog.=) $1", &[&timeline_id.to_string()]).await {
                     Ok(result) => !result.is_empty(),
                     Err(e) =>
                     {
@@ -1142,7 +1142,9 @@ async fn get_operations<'a>(
             if let Some(libs) = spec.cluster.settings.find("shared_preload_libraries") {
                 if libs.contains("pg_stat_statements") {
                     return Ok(Box::new(once(Operation {
-                        query: String::from("CREATE EXTENSION IF NOT EXISTS pg_stat_statements"),
+                        query: String::from(
+                            "CREATE EXTENSION IF NOT EXISTS pg_stat_statements WITH SCHEMA public",
+                        ),
                         comment: Some(String::from("create system extensions")),
                     })));
                 }
@@ -1150,11 +1152,13 @@ async fn get_operations<'a>(
             Ok(Box::new(empty()))
         }
         ApplySpecPhase::CreatePgauditExtension => Ok(Box::new(once(Operation {
-            query: String::from("CREATE EXTENSION IF NOT EXISTS pgaudit"),
+            query: String::from("CREATE EXTENSION IF NOT EXISTS pgaudit WITH SCHEMA public"),
             comment: Some(String::from("create pgaudit extensions")),
         }))),
         ApplySpecPhase::CreatePgauditlogtofileExtension => Ok(Box::new(once(Operation {
-            query: String::from("CREATE EXTENSION IF NOT EXISTS pgauditlogtofile"),
+            query: String::from(
+                "CREATE EXTENSION IF NOT EXISTS pgauditlogtofile WITH SCHEMA public",
+            ),
             comment: Some(String::from("create pgauditlogtofile extensions")),
         }))),
         // Disable pgaudit logging for postgres database.
@@ -1178,7 +1182,7 @@ async fn get_operations<'a>(
                 },
                 Operation {
                     query: String::from(
-                        "UPDATE pg_extension SET extrelocatable = true WHERE extname = 'neon'",
+                        "UPDATE pg_catalog.pg_extension SET extrelocatable = true WHERE extname OPERATOR(pg_catalog.=) 'neon'::pg_catalog.name AND extrelocatable OPERATOR(pg_catalog.=) false",
                     ),
                     comment: Some(String::from("compat/fix: make neon relocatable")),
                 },

compute_tools/src/sql/add_availabilitycheck_tables.sql

@@ -3,16 +3,17 @@ BEGIN
     IF NOT EXISTS(
         SELECT 1
         FROM pg_catalog.pg_tables
-        WHERE tablename = 'health_check'
+        WHERE tablename::pg_catalog.name OPERATOR(pg_catalog.=) 'health_check'::pg_catalog.name
+        AND schemaname::pg_catalog.name OPERATOR(pg_catalog.=) 'public'::pg_catalog.name
     )
     THEN
-    CREATE TABLE health_check (
-        id serial primary key,
-        updated_at timestamptz default now()
+    CREATE TABLE public.health_check (
+        id pg_catalog.int4 primary key generated by default as identity,
+        updated_at pg_catalog.timestamptz default pg_catalog.now()
     );
-    INSERT INTO health_check VALUES (1, now())
+    INSERT INTO public.health_check VALUES (1, pg_catalog.now())
         ON CONFLICT (id) DO UPDATE
-         SET updated_at = now();
+         SET updated_at = pg_catalog.now();
     END IF;
 END
 $$

compute_tools/src/sql/anon_ext_fn_reassign.sql

@@ -2,10 +2,11 @@ DO $$
 DECLARE
     query varchar;
 BEGIN
-    FOR query IN SELECT 'ALTER FUNCTION '||nsp.nspname||'.'||p.proname||'('||pg_get_function_identity_arguments(p.oid)||') OWNER TO {db_owner};'
-    FROM pg_proc p
-        JOIN pg_namespace nsp ON p.pronamespace = nsp.oid
-    WHERE nsp.nspname = 'anon' LOOP
+    FOR query IN
+    SELECT pg_catalog.format('ALTER FUNCTION %I(%s) OWNER TO {db_owner};', p.oid::regproc, pg_catalog.pg_get_function_identity_arguments(p.oid))
+    FROM pg_catalog.pg_proc p
+        WHERE p.pronamespace OPERATOR(pg_catalog.=) 'anon'::regnamespace::oid
+    LOOP
         EXECUTE query;
     END LOOP;
 END

compute_tools/src/sql/create_privileged_role.sql

@@ -1,6 +1,6 @@
 DO $$
     BEGIN
-        IF NOT EXISTS (SELECT FROM pg_catalog.pg_roles WHERE rolname = '{privileged_role_name}')
+        IF NOT EXISTS (SELECT FROM pg_catalog.pg_roles WHERE rolname OPERATOR(pg_catalog.=) '{privileged_role_name}'::pg_catalog.name)
         THEN
             CREATE ROLE {privileged_role_name} {privileges} IN ROLE pg_read_all_data, pg_write_all_data;
         END IF;

compute_tools/src/sql/default_grants.sql

@@ -4,14 +4,14 @@ $$
         IF EXISTS(
             SELECT nspname
             FROM pg_catalog.pg_namespace
-            WHERE nspname = 'public'
+            WHERE nspname OPERATOR(pg_catalog.=) 'public'
         ) AND
-           current_setting('server_version_num')::int / 10000 >= 15
+           pg_catalog.current_setting('server_version_num')::int OPERATOR(pg_catalog./) 10000 OPERATOR(pg_catalog.>=) 15
         THEN
             IF EXISTS(
                 SELECT rolname
                 FROM pg_catalog.pg_roles
-                WHERE rolname = 'web_access'
+                WHERE rolname OPERATOR(pg_catalog.=) 'web_access'
             )
             THEN
                 GRANT CREATE ON SCHEMA public TO web_access;
@@ -20,7 +20,7 @@ $$
         IF EXISTS(
             SELECT nspname
             FROM pg_catalog.pg_namespace
-            WHERE nspname = 'public'
+            WHERE nspname OPERATOR(pg_catalog.=) 'public'
         )
         THEN
             ALTER DEFAULT PRIVILEGES IN SCHEMA public GRANT ALL ON TABLES TO neon_superuser WITH GRANT OPTION;

compute_tools/src/sql/drop_subscriptions.sql

@@ -2,11 +2,17 @@ DO ${outer_tag}$
 DECLARE
     subname TEXT;
 BEGIN
-    LOCK TABLE pg_subscription IN ACCESS EXCLUSIVE MODE;
-    FOR subname IN SELECT pg_subscription.subname FROM pg_subscription WHERE subdbid = (SELECT oid FROM pg_database WHERE datname = {datname_str}) LOOP
-        EXECUTE format('ALTER SUBSCRIPTION %I DISABLE;', subname);
-        EXECUTE format('ALTER SUBSCRIPTION %I SET (slot_name = NONE);', subname);
-        EXECUTE format('DROP SUBSCRIPTION %I;', subname);
+    LOCK TABLE pg_catalog.pg_subscription IN ACCESS EXCLUSIVE MODE;
+    FOR subname IN
+        SELECT pg_subscription.subname
+        FROM pg_catalog.pg_subscription
+        WHERE subdbid OPERATOR(pg_catalog.=) (
+            SELECT oid FROM pg_database WHERE datname OPERATOR(pg_catalog.=) {datname_str}::pg_catalog.name
+        )
+    LOOP
+        EXECUTE pg_catalog.format('ALTER SUBSCRIPTION %I DISABLE;', subname);
+        EXECUTE pg_catalog.format('ALTER SUBSCRIPTION %I SET (slot_name = NONE);', subname);
+        EXECUTE pg_catalog.format('DROP SUBSCRIPTION %I;', subname);
     END LOOP;
 END;
 ${outer_tag}$;

compute_tools/src/sql/finalize_drop_subscriptions.sql

@@ -3,19 +3,19 @@ BEGIN
     IF NOT EXISTS(
         SELECT 1
         FROM pg_catalog.pg_tables
-        WHERE tablename = 'drop_subscriptions_done'
-        AND schemaname = 'neon'
+        WHERE tablename OPERATOR(pg_catalog.=) 'drop_subscriptions_done'::pg_catalog.name
+        AND schemaname OPERATOR(pg_catalog.=) 'neon'::pg_catalog.name
     )
     THEN
         CREATE TABLE neon.drop_subscriptions_done
-        (id serial primary key, timeline_id text);
+        (id pg_catalog.int4 primary key generated by default as identity, timeline_id pg_catalog.text);
     END IF;
 
     -- preserve the timeline_id of the last drop_subscriptions run
     -- to ensure that the cleanup of a timeline is executed only once.
     -- use upsert to avoid the table bloat in case of cascade branching (branch of a branch)
-    INSERT INTO neon.drop_subscriptions_done VALUES (1, current_setting('neon.timeline_id'))
+    INSERT INTO neon.drop_subscriptions_done VALUES (1, pg_catalog.current_setting('neon.timeline_id'))
     ON CONFLICT (id) DO UPDATE
-    SET timeline_id = current_setting('neon.timeline_id');
+    SET timeline_id = pg_catalog.current_setting('neon.timeline_id')::pg_catalog.text;
 END
 $$

compute_tools/src/sql/pre_drop_role_revoke_privileges.sql

@@ -15,15 +15,15 @@ BEGIN
         WHERE schema_name IN ('public')
     LOOP
         FOR grantor IN EXECUTE
-            format(
-                'SELECT DISTINCT rtg.grantor FROM information_schema.role_table_grants AS rtg WHERE grantee = %s',
+            pg_catalog.format(
+                'SELECT DISTINCT rtg.grantor FROM information_schema.role_table_grants AS rtg WHERE grantee OPERATOR(pg_catalog.=) %s',
                 -- N.B. this has to be properly dollar-escaped with `pg_quote_dollar()`
                 quote_literal({role_name})
             )
         LOOP
-            EXECUTE format('SET LOCAL ROLE %I', grantor);
+            EXECUTE pg_catalog.format('SET LOCAL ROLE %I', grantor);
 
-            revoke_query := format(
+            revoke_query := pg_catalog.format(
                 'REVOKE ALL PRIVILEGES ON ALL TABLES IN SCHEMA %I FROM %I GRANTED BY %I',
                 schema,
                 -- N.B. this has to be properly dollar-escaped with `pg_quote_dollar()`

compute_tools/src/sql/set_public_schema_owner.sql

@@ -5,17 +5,17 @@ DO ${outer_tag}$
         IF EXISTS(
             SELECT nspname
             FROM pg_catalog.pg_namespace
-            WHERE nspname = 'public'
+            WHERE nspname OPERATOR(pg_catalog.=) 'public'::pg_catalog.name
         )
         THEN
             SELECT nspowner::regrole::text
             FROM pg_catalog.pg_namespace
-            WHERE nspname = 'public'
+            WHERE nspname OPERATOR(pg_catalog.=) 'public'::pg_catalog.text
             INTO schema_owner;
 
-            IF schema_owner = 'cloud_admin' OR schema_owner = 'zenith_admin'
+            IF schema_owner OPERATOR(pg_catalog.=) 'cloud_admin'::pg_catalog.text OR schema_owner OPERATOR(pg_catalog.=) 'zenith_admin'::pg_catalog.text
             THEN
-                EXECUTE format('ALTER SCHEMA public OWNER TO %I', {db_owner});
+                EXECUTE pg_catalog.format('ALTER SCHEMA public OWNER TO %I', {db_owner});
             END IF;
         END IF;
     END

compute_tools/src/sql/unset_template_for_drop_dbs.sql

@@ -3,10 +3,10 @@ DO ${outer_tag}$
         IF EXISTS(
             SELECT 1
             FROM pg_catalog.pg_database
-            WHERE datname = {datname}
+            WHERE datname OPERATOR(pg_catalog.=) {datname}::pg_catalog.name
         )
         THEN
-            EXECUTE format('ALTER DATABASE %I is_template false', {datname});
+            EXECUTE pg_catalog.format('ALTER DATABASE %I is_template false', {datname});
         END IF;
     END
 ${outer_tag}$;

libs/compute_api/src/responses.rs

@@ -68,11 +68,15 @@ pub enum LfcPrewarmState {
     /// We tried to fetch the corresponding LFC state from the endpoint storage,
     /// but received `Not Found 404`. This should normally happen only during the
     /// first endpoint start after creation with `autoprewarm: true`.
+    /// This may also happen if LFC is turned off or not initialized
     ///
     /// During the orchestrated prewarm via API, when a caller explicitly
     /// provides the LFC state key to prewarm from, it's the caller responsibility
     /// to handle this status as an error state in this case.
     Skipped,
+    /// LFC prewarm was cancelled. Some pages in LFC cache may be prewarmed if query
+    /// has started working before cancellation
+    Cancelled,
 }
 
 impl Display for LfcPrewarmState {
@@ -83,6 +87,7 @@ impl Display for LfcPrewarmState {
             LfcPrewarmState::Completed => f.write_str("Completed"),
             LfcPrewarmState::Skipped => f.write_str("Skipped"),
             LfcPrewarmState::Failed { error } => write!(f, "Error({error})"),
+            LfcPrewarmState::Cancelled => f.write_str("Cancelled"),
         }
     }
 }
@@ -97,6 +102,7 @@ pub enum LfcOffloadState {
     Failed {
         error: String,
     },
+    Skipped,
 }
 
 #[derive(Serialize, Debug, Clone, PartialEq)]

libs/walproposer/src/api_bindings.rs

@@ -341,6 +341,34 @@ extern "C-unwind" fn log_internal(
     }
 }
 
+/* BEGIN_HADRON */
+extern "C" fn reset_safekeeper_statuses_for_metrics(wp: *mut WalProposer, num_safekeepers: u32) {
+    unsafe {
+        let callback_data = (*(*wp).config).callback_data;
+        let api = callback_data as *mut Box<dyn ApiImpl>;
+        if api.is_null() {
+            return;
+        }
+        (*api).reset_safekeeper_statuses_for_metrics(&mut (*wp), num_safekeepers);
+    }
+}
+
+extern "C" fn update_safekeeper_status_for_metrics(
+    wp: *mut WalProposer,
+    sk_index: u32,
+    status: u8,
+) {
+    unsafe {
+        let callback_data = (*(*wp).config).callback_data;
+        let api = callback_data as *mut Box<dyn ApiImpl>;
+        if api.is_null() {
+            return;
+        }
+        (*api).update_safekeeper_status_for_metrics(&mut (*wp), sk_index, status);
+    }
+}
+/* END_HADRON */
+
 #[derive(Debug, PartialEq)]
 pub enum Level {
     Debug5,
@@ -414,6 +442,10 @@ pub(crate) fn create_api() -> walproposer_api {
         finish_sync_safekeepers: Some(finish_sync_safekeepers),
         process_safekeeper_feedback: Some(process_safekeeper_feedback),
         log_internal: Some(log_internal),
+        /* BEGIN_HADRON */
+        reset_safekeeper_statuses_for_metrics: Some(reset_safekeeper_statuses_for_metrics),
+        update_safekeeper_status_for_metrics: Some(update_safekeeper_status_for_metrics),
+        /* END_HADRON */
     }
 }
 
@@ -451,6 +483,8 @@ pub fn empty_shmem() -> crate::bindings::WalproposerShmemState {
         replica_promote: false,
         min_ps_feedback: empty_feedback,
         wal_rate_limiter: empty_wal_rate_limiter,
+        num_safekeepers: 0,
+        safekeeper_status: [0; 32],
     }
 }
 

libs/walproposer/src/walproposer.rs

@@ -159,6 +159,21 @@ pub trait ApiImpl {
     fn after_election(&self, _wp: &mut WalProposer) {
         todo!()
     }
+
+    /* BEGIN_HADRON */
+    fn reset_safekeeper_statuses_for_metrics(&self, _wp: &mut WalProposer, _num_safekeepers: u32) {
+        // Do nothing for testing purposes.
+    }
+
+    fn update_safekeeper_status_for_metrics(
+        &self,
+        _wp: &mut WalProposer,
+        _sk_index: u32,
+        _status: u8,
+    ) {
+        // Do nothing for testing purposes.
+    }
+    /* END_HADRON */
 }
 
 #[derive(Debug)]

pgxn/neon/neon_perf_counters.c

@@ -391,6 +391,12 @@ neon_get_perf_counters(PG_FUNCTION_ARGS)
 	neon_per_backend_counters totals = {0};
 	metric_t   *metrics;
 
+	/* BEGIN_HADRON */
+	WalproposerShmemState *wp_shmem;
+	uint32 num_safekeepers;
+	uint32 num_active_safekeepers;
+	/* END_HADRON */
+
 	/* We put all the tuples into a tuplestore in one go. */
 	InitMaterializedSRF(fcinfo, 0);
 
@@ -437,11 +443,32 @@ neon_get_perf_counters(PG_FUNCTION_ARGS)
 		// Not ideal but piggyback our databricks counters into the neon perf counters view
 		// so that we don't need to introduce neon--1.x+1.sql to add a new view.
 		{
+		// Keeping this code in its own block to work around the C90 "don't mix declarations and code" rule when we define
+		// the `databricks_metrics` array in the next block. Yes, we are seriously dealing with C90 rules in 2025.
+
+		// Read safekeeper status from wal proposer shared memory first.
+		// Note that we are taking a mutex when reading from walproposer shared memory so that the total safekeeper count is
+		// consistent with the active wal acceptors count. Assuming that we don't query this view too often the mutex should
+		// not be a huge deal.
+		wp_shmem = GetWalpropShmemState();
+		SpinLockAcquire(&wp_shmem->mutex);
+		num_safekeepers = wp_shmem->num_safekeepers;
+		num_active_safekeepers = 0;
+		for (int i = 0; i < num_safekeepers; i++) {
+			if (wp_shmem->safekeeper_status[i] == 1) {
+				num_active_safekeepers++;
+			}
+		}
+		SpinLockRelease(&wp_shmem->mutex);
+	}
+	{
 			metric_t databricks_metrics[] = {
 				{"sql_index_corruption_count", false, 0, (double) pg_atomic_read_u32(&databricks_metrics_shared->index_corruption_count)},
 				{"sql_data_corruption_count", false, 0, (double) pg_atomic_read_u32(&databricks_metrics_shared->data_corruption_count)},
 				{"sql_internal_error_count", false, 0, (double) pg_atomic_read_u32(&databricks_metrics_shared->internal_error_count)},
 				{"ps_corruption_detected", false, 0, (double) pg_atomic_read_u32(&databricks_metrics_shared->ps_corruption_detected)},
+				{"num_active_safekeepers", false, 0.0, (double) num_active_safekeepers},
+				{"num_configured_safekeepers", false, 0.0, (double) num_safekeepers},
 				{NULL, false, 0, 0},
 			};
 			for (int i = 0; databricks_metrics[i].name != NULL; i++)

pgxn/neon/walproposer.c

@@ -154,7 +154,9 @@ WalProposerCreate(WalProposerConfig *config, walproposer_api api)
 		wp->safekeeper[wp->n_safekeepers].state = SS_OFFLINE;
 		wp->safekeeper[wp->n_safekeepers].active_state = SS_ACTIVE_SEND;
 		wp->safekeeper[wp->n_safekeepers].wp = wp;
-
+		/* BEGIN_HADRON */
+		wp->safekeeper[wp->n_safekeepers].index = wp->n_safekeepers;
+		/* END_HADRON */
 		{
 			Safekeeper *sk = &wp->safekeeper[wp->n_safekeepers];
 			int			written = 0;
@@ -183,6 +185,10 @@ WalProposerCreate(WalProposerConfig *config, walproposer_api api)
 	if (wp->safekeepers_generation > INVALID_GENERATION && wp->config->proto_version < 3)
 		wp_log(FATAL, "enabling generations requires protocol version 3");
 	wp_log(LOG, "using safekeeper protocol version %d", wp->config->proto_version);
+	
+	/* BEGIN_HADRON */
+	wp->api.reset_safekeeper_statuses_for_metrics(wp, wp->n_safekeepers);
+	/* END_HADRON */
 
 	/* Fill the greeting package */
 	wp->greetRequest.pam.tag = 'g';
@@ -355,6 +361,10 @@ ShutdownConnection(Safekeeper *sk)
 	sk->state = SS_OFFLINE;
 	sk->streamingAt = InvalidXLogRecPtr;
 
+	/* BEGIN_HADRON */
+	sk->wp->api.update_safekeeper_status_for_metrics(sk->wp, sk->index, 0);
+	/* END_HADRON */
+
 	MembershipConfigurationFree(&sk->greetResponse.mconf);
 	if (sk->voteResponse.termHistory.entries)
 		pfree(sk->voteResponse.termHistory.entries);
@@ -1530,6 +1540,10 @@ StartStreaming(Safekeeper *sk)
 	sk->active_state = SS_ACTIVE_SEND;
 	sk->streamingAt = sk->startStreamingAt;
 
+	/* BEGIN_HADRON */
+	sk->wp->api.update_safekeeper_status_for_metrics(sk->wp, sk->index, 1);
+	/* END_HADRON */
+
 	/*
 	 * Donors can only be in SS_ACTIVE state, so we potentially update the
 	 * donor when we switch one to SS_ACTIVE.

pgxn/neon/walproposer.h

@@ -432,6 +432,10 @@ typedef struct WalproposerShmemState
 	/* BEGIN_HADRON */
 	/* The WAL rate limiter */
 	WalRateLimiter wal_rate_limiter;
+	/* Number of safekeepers in the config */
+	uint32 num_safekeepers;
+	/* Per-safekeeper status flags: 0=inactive, 1=active */
+	uint8 safekeeper_status[MAX_SAFEKEEPERS];
 	/* END_HADRON */
 } WalproposerShmemState;
 
@@ -483,6 +487,11 @@ typedef struct Safekeeper
 	char const *host;
 	char const *port;
 
+	/* BEGIN_HADRON */
+	/* index of this safekeeper in the WalProposer array */
+	uint32 index;
+	/* END_HADRON */
+
 	/*
 	 * connection string for connecting/reconnecting.
 	 *
@@ -731,6 +740,23 @@ typedef struct walproposer_api
 	 * handled by elog().
 	 */
 	void		(*log_internal) (WalProposer *wp, int level, const char *line);
+
+	/*
+	 * BEGIN_HADRON
+	 * APIs manipulating shared memory state used for Safekeeper quorum health metrics.
+	 */
+
+	/*
+	 * Reset the safekeeper statuses in shared memory for metric purposes.
+	 */
+	void		(*reset_safekeeper_statuses_for_metrics) (WalProposer *wp, uint32 num_safekeepers);
+
+	/*
+	 * Update the safekeeper status in shared memory for metric purposes.
+	 */
+	void		(*update_safekeeper_status_for_metrics) (WalProposer *wp, uint32 sk_index, uint8 status);
+
+	/* END_HADRON */
 } walproposer_api;
 
 /*

pgxn/neon/walproposer_pg.c

@@ -2261,6 +2261,27 @@ GetNeonCurrentClusterSize(void)
 }
 uint64		GetNeonCurrentClusterSize(void);
 
+/* BEGIN_HADRON */
+static void
+walprop_pg_reset_safekeeper_statuses_for_metrics(WalProposer *wp, uint32 num_safekeepers)
+{
+	WalproposerShmemState* shmem = wp->api.get_shmem_state(wp);
+	SpinLockAcquire(&shmem->mutex);
+	shmem->num_safekeepers = num_safekeepers;
+	memset(shmem->safekeeper_status, 0, sizeof(shmem->safekeeper_status));
+	SpinLockRelease(&shmem->mutex);
+}
+
+static void
+walprop_pg_update_safekeeper_status_for_metrics(WalProposer *wp, uint32 sk_index, uint8 status)
+{
+	WalproposerShmemState* shmem = wp->api.get_shmem_state(wp);
+	Assert(sk_index < MAX_SAFEKEEPERS);
+	SpinLockAcquire(&shmem->mutex);
+	shmem->safekeeper_status[sk_index] = status;
+	SpinLockRelease(&shmem->mutex);
+}
+/* END_HADRON */
 
 static const walproposer_api walprop_pg = {
 	.get_shmem_state = walprop_pg_get_shmem_state,
@@ -2294,4 +2315,6 @@ static const walproposer_api walprop_pg = {
 	.finish_sync_safekeepers = walprop_pg_finish_sync_safekeepers,
 	.process_safekeeper_feedback = walprop_pg_process_safekeeper_feedback,
 	.log_internal = walprop_pg_log_internal,
+	.reset_safekeeper_statuses_for_metrics = walprop_pg_reset_safekeeper_statuses_for_metrics,
+	.update_safekeeper_status_for_metrics = walprop_pg_update_safekeeper_status_for_metrics,
 };

storage_controller/src/persistence.rs

@@ -471,11 +471,17 @@ impl Persistence {
         &self,
         input_node_id: NodeId,
         input_https_port: Option<u16>,
+        input_grpc_addr: Option<String>,
+        input_grpc_port: Option<u16>,
     ) -> DatabaseResult<()> {
         use crate::schema::nodes::dsl::*;
         self.update_node(
             input_node_id,
-            listen_https_port.eq(input_https_port.map(|x| x as i32)),
+            (
+                listen_https_port.eq(input_https_port.map(|x| x as i32)),
+                listen_grpc_addr.eq(input_grpc_addr),
+                listen_grpc_port.eq(input_grpc_port.map(|x| x as i32)),
+            ),
         )
         .await
     }

storage_controller/src/service.rs

@@ -7813,7 +7813,7 @@ impl Service {
             register_req.listen_https_port,
             register_req.listen_pg_addr,
             register_req.listen_pg_port,
-            register_req.listen_grpc_addr,
+            register_req.listen_grpc_addr.clone(),
             register_req.listen_grpc_port,
             register_req.availability_zone_id.clone(),
             self.config.use_https_pageserver_api,
@@ -7848,6 +7848,8 @@ impl Service {
                     .update_node_on_registration(
                         register_req.node_id,
                         register_req.listen_https_port,
+                        register_req.listen_grpc_addr,
+                        register_req.listen_grpc_port,
                     )
                     .await?
             }

test_runner/fixtures/endpoint/http.py

@@ -78,20 +78,26 @@ class EndpointHttpClient(requests.Session):
         json: dict[str, str] = res.json()
         return json
 
-    def prewarm_lfc(self, from_endpoint_id: str | None = None):
+    def prewarm_lfc(self, from_endpoint_id: str | None = None) -> dict[str, str]:
         """
         Prewarm LFC cache from given endpoint and wait till it finishes or errors
         """
         params = {"from_endpoint": from_endpoint_id} if from_endpoint_id else dict()
         self.post(self.prewarm_url, params=params).raise_for_status()
-        self.prewarm_lfc_wait()
+        return self.prewarm_lfc_wait()
 
-    def prewarm_lfc_wait(self):
+    def cancel_prewarm_lfc(self):
+        """
+        Cancel LFC prewarm if any is ongoing
+        """
+        self.delete(self.prewarm_url).raise_for_status()
+
+    def prewarm_lfc_wait(self) -> dict[str, str]:
         """
         Wait till LFC prewarm returns with error or success.
         If prewarm was not requested before calling this function, it will error
         """
-        statuses = "failed", "completed", "skipped"
+        statuses = "failed", "completed", "skipped", "cancelled"
 
         def prewarmed():
             json = self.prewarm_lfc_status()
@@ -101,6 +107,7 @@ class EndpointHttpClient(requests.Session):
         wait_until(prewarmed, timeout=60)
         res = self.prewarm_lfc_status()
         assert res["status"] != "failed", res
+        return res
 
     def offload_lfc_status(self) -> dict[str, str]:
         res = self.get(self.offload_url)
@@ -108,29 +115,31 @@ class EndpointHttpClient(requests.Session):
         json: dict[str, str] = res.json()
         return json
 
-    def offload_lfc(self):
+    def offload_lfc(self) -> dict[str, str]:
         """
         Offload LFC cache to endpoint storage and wait till offload finishes or errors
         """
         self.post(self.offload_url).raise_for_status()
-        self.offload_lfc_wait()
+        return self.offload_lfc_wait()
 
-    def offload_lfc_wait(self):
+    def offload_lfc_wait(self) -> dict[str, str]:
         """
         Wait till LFC offload returns with error or success.
         If offload was not requested before calling this function, it will error
         """
+        statuses = "failed", "completed", "skipped"
 
         def offloaded():
             json = self.offload_lfc_status()
             status, err = json["status"], json.get("error")
-            assert status in ["failed", "completed"], f"{status}, {err=}"
+            assert status in statuses, f"{status}, {err=}"
 
         wait_until(offloaded, timeout=60)
         res = self.offload_lfc_status()
         assert res["status"] != "failed", res
+        return res
 
-    def promote(self, promote_spec: dict[str, Any], disconnect: bool = False):
+    def promote(self, promote_spec: dict[str, Any], disconnect: bool = False) -> dict[str, str]:
         url = f"http://localhost:{self.external_port}/promote"
         if disconnect:
             try:  # send first request to start promote and disconnect

test_runner/fixtures/neon_fixtures.py

@@ -262,7 +262,6 @@ class PgProtocol:
         # pooler does not support statement_timeout
         # Check if the hostname contains the string 'pooler'
         hostname = result.get("host", "")
-        log.info(f"Hostname: {hostname}")
         options = result.get("options", "")
         if "statement_timeout" not in options and "pooler" not in hostname:
             options = f"-cstatement_timeout=120s {options}"

test_runner/regress/test_compaction.py

@@ -863,7 +863,6 @@ def test_pageserver_compaction_circuit_breaker(neon_env_builder: NeonEnvBuilder)
     assert not env.pageserver.log_contains(".*Circuit breaker failure ended.*")
 
 
-@pytest.mark.skip(reason="Lakebase mode")
 def test_ps_corruption_detection_feedback(neon_env_builder: NeonEnvBuilder):
     """
     Test that when the pageserver detects corruption during image layer creation,
@@ -890,7 +889,9 @@ def test_ps_corruption_detection_feedback(neon_env_builder: NeonEnvBuilder):
     timeline_id = env.initial_timeline
 
     pageserver_http = env.pageserver.http_client()
-    workload = Workload(env, tenant_id, timeline_id)
+    workload = Workload(
+        env, tenant_id, timeline_id, endpoint_opts={"config_lines": ["neon.lakebase_mode=true"]}
+    )
     workload.init()
 
     # Enable the failpoint that will cause image layer creation to fail due to a (simulated) detected

test_runner/regress/test_lfc_prewarm.py

@@ -1,6 +1,6 @@
 import random
-import threading
 from enum import StrEnum
+from threading import Thread
 from time import sleep
 from typing import Any
 
@@ -47,19 +47,23 @@ def offload_lfc(method: PrewarmMethod, client: EndpointHttpClient, cur: Cursor)
         # With autoprewarm, we need to be sure LFC was offloaded after all writes
         # finish, so we sleep. Otherwise we'll have less prewarmed pages than we want
         sleep(AUTOOFFLOAD_INTERVAL_SECS)
-        client.offload_lfc_wait()
-        return
+        offload_res = client.offload_lfc_wait()
+        log.info(offload_res)
+        return offload_res
 
     if method == PrewarmMethod.COMPUTE_CTL:
         status = client.prewarm_lfc_status()
         assert status["status"] == "not_prewarmed"
         assert "error" not in status
-        client.offload_lfc()
+        offload_res = client.offload_lfc()
+        log.info(offload_res)
         assert client.prewarm_lfc_status()["status"] == "not_prewarmed"
+
         parsed = prom_parse(client)
         desired = {OFFLOAD_LABEL: 1, PREWARM_LABEL: 0, OFFLOAD_ERR_LABEL: 0, PREWARM_ERR_LABEL: 0}
         assert parsed == desired, f"{parsed=} != {desired=}"
-        return
+
+        return offload_res
 
     raise AssertionError(f"{method} not in PrewarmMethod")
 
@@ -68,21 +72,30 @@ def prewarm_endpoint(
     method: PrewarmMethod, client: EndpointHttpClient, cur: Cursor, lfc_state: str | None
 ):
     if method == PrewarmMethod.AUTOPREWARM:
-        client.prewarm_lfc_wait()
+        prewarm_res = client.prewarm_lfc_wait()
+        log.info(prewarm_res)
     elif method == PrewarmMethod.COMPUTE_CTL:
-        client.prewarm_lfc()
+        prewarm_res = client.prewarm_lfc()
+        log.info(prewarm_res)
+        return prewarm_res
     elif method == PrewarmMethod.POSTGRES:
         cur.execute("select neon.prewarm_local_cache(%s)", (lfc_state,))
 
 
-def check_prewarmed(
+def check_prewarmed_contains(
     method: PrewarmMethod, client: EndpointHttpClient, desired_status: dict[str, str | int]
 ):
     if method == PrewarmMethod.AUTOPREWARM:
-        assert client.prewarm_lfc_status() == desired_status
+        prewarm_status = client.prewarm_lfc_status()
+        for k in desired_status:
+            assert desired_status[k] == prewarm_status[k]
+
         assert prom_parse(client)[PREWARM_LABEL] == 1
     elif method == PrewarmMethod.COMPUTE_CTL:
-        assert client.prewarm_lfc_status() == desired_status
+        prewarm_status = client.prewarm_lfc_status()
+        for k in desired_status:
+            assert desired_status[k] == prewarm_status[k]
+
         desired = {OFFLOAD_LABEL: 0, PREWARM_LABEL: 1, PREWARM_ERR_LABEL: 0, OFFLOAD_ERR_LABEL: 0}
         assert prom_parse(client) == desired
 
@@ -149,9 +162,6 @@ def test_lfc_prewarm(neon_simple_env: NeonEnv, method: PrewarmMethod):
     log.info(f"Used LFC size: {lfc_used_pages}")
     pg_cur.execute("select * from neon.get_prewarm_info()")
     total, prewarmed, skipped, _ = pg_cur.fetchall()[0]
-    log.info(f"Prewarm info: {total=} {prewarmed=} {skipped=}")
-    progress = (prewarmed + skipped) * 100 // total
-    log.info(f"Prewarm progress: {progress}%")
     assert lfc_used_pages > 10000
     assert total > 0
     assert prewarmed > 0
@@ -161,7 +171,54 @@ def test_lfc_prewarm(neon_simple_env: NeonEnv, method: PrewarmMethod):
     assert lfc_cur.fetchall()[0][0] == n_records * (n_records + 1) / 2
 
     desired = {"status": "completed", "total": total, "prewarmed": prewarmed, "skipped": skipped}
-    check_prewarmed(method, client, desired)
+    check_prewarmed_contains(method, client, desired)
+
+
+@pytest.mark.skipif(not USE_LFC, reason="LFC is disabled, skipping")
+def test_lfc_prewarm_cancel(neon_simple_env: NeonEnv):
+    """
+    Test we can cancel LFC prewarm and prewarm successfully after
+    """
+    env = neon_simple_env
+    n_records = 1000000
+    cfg = [
+        "autovacuum = off",
+        "shared_buffers=1MB",
+        "neon.max_file_cache_size=1GB",
+        "neon.file_cache_size_limit=1GB",
+        "neon.file_cache_prewarm_limit=1000",
+    ]
+    endpoint = env.endpoints.create_start(branch_name="main", config_lines=cfg)
+
+    pg_conn = endpoint.connect()
+    pg_cur = pg_conn.cursor()
+    pg_cur.execute("create schema neon; create extension neon with schema neon")
+    pg_cur.execute("create database lfc")
+
+    lfc_conn = endpoint.connect(dbname="lfc")
+    lfc_cur = lfc_conn.cursor()
+    log.info(f"Inserting {n_records} rows")
+    lfc_cur.execute("create table t(pk integer primary key, payload text default repeat('?', 128))")
+    lfc_cur.execute(f"insert into t (pk) values (generate_series(1,{n_records}))")
+    log.info(f"Inserted {n_records} rows")
+
+    client = endpoint.http_client()
+    method = PrewarmMethod.COMPUTE_CTL
+    offload_lfc(method, client, pg_cur)
+
+    endpoint.stop()
+    endpoint.start()
+
+    thread = Thread(target=lambda: prewarm_endpoint(method, client, pg_cur, None))
+    thread.start()
+    # wait 2 seconds to ensure we cancel prewarm SQL query
+    sleep(2)
+    client.cancel_prewarm_lfc()
+    thread.join()
+    assert client.prewarm_lfc_status()["status"] == "cancelled"
+
+    prewarm_endpoint(method, client, pg_cur, None)
+    assert client.prewarm_lfc_status()["status"] == "completed"
 
 
 @pytest.mark.skipif(not USE_LFC, reason="LFC is disabled, skipping")
@@ -178,9 +235,8 @@ def test_lfc_prewarm_empty(neon_simple_env: NeonEnv):
     cur = conn.cursor()
     cur.execute("create schema neon; create extension neon with schema neon")
     method = PrewarmMethod.COMPUTE_CTL
-    offload_lfc(method, client, cur)
-    prewarm_endpoint(method, client, cur, None)
-    assert client.prewarm_lfc_status()["status"] == "skipped"
+    assert offload_lfc(method, client, cur)["status"] == "skipped"
+    assert prewarm_endpoint(method, client, cur, None)["status"] == "skipped"
 
 
 # autoprewarm isn't needed as we prewarm manually
@@ -251,11 +307,11 @@ def test_lfc_prewarm_under_workload(neon_simple_env: NeonEnv, method: PrewarmMet
 
     workload_threads = []
     for _ in range(n_threads):
-        t = threading.Thread(target=workload)
+        t = Thread(target=workload)
         workload_threads.append(t)
         t.start()
 
-    prewarm_thread = threading.Thread(target=prewarm)
+    prewarm_thread = Thread(target=prewarm)
     prewarm_thread.start()
 
     def prewarmed():

test_runner/regress/test_wal_acceptor.py

@@ -2757,18 +2757,37 @@ def test_timeline_disk_usage_limit(neon_env_builder: NeonEnvBuilder):
     remote_storage_kind = s3_storage()
     neon_env_builder.enable_safekeeper_remote_storage(remote_storage_kind)
 
-    # Set a very small disk usage limit (1KB)
-    neon_env_builder.safekeeper_extra_opts = ["--max-timeline-disk-usage-bytes=1024"]
-
     env = neon_env_builder.init_start()
 
     # Create a timeline and endpoint
     env.create_branch("test_timeline_disk_usage_limit")
-    endpoint = env.endpoints.create_start("test_timeline_disk_usage_limit")
+    endpoint = env.endpoints.create_start(
+        "test_timeline_disk_usage_limit",
+        config_lines=[
+            "neon.lakebase_mode=true",
+        ],
+    )
+
+    # Install the neon extension in the test database. We need it to query perf counter metrics.
+    with closing(endpoint.connect()) as conn:
+        with conn.cursor() as cur:
+            cur.execute("CREATE EXTENSION IF NOT EXISTS neon")
+            # Sanity-check safekeeper connection status in neon_perf_counters in the happy case.
+            cur.execute(
+                "SELECT value FROM neon_perf_counters WHERE metric = 'num_active_safekeepers'"
+            )
+            assert cur.fetchone() == (1,), "Expected 1 active safekeeper"
+            cur.execute(
+                "SELECT value FROM neon_perf_counters WHERE metric = 'num_configured_safekeepers'"
+            )
+            assert cur.fetchone() == (1,), "Expected 1 configured safekeeper"
 
     # Get the safekeeper
     sk = env.safekeepers[0]
 
+    # Restart the safekeeper with a very small disk usage limit (1KB)
+    sk.stop().start(["--max-timeline-disk-usage-bytes=1024"])
+
     # Inject a failpoint to stop WAL backup
     with sk.http_client() as http_cli:
         http_cli.configure_failpoints([("backup-lsn-range-pausable", "pause")])
@@ -2794,6 +2813,18 @@ def test_timeline_disk_usage_limit(neon_env_builder: NeonEnvBuilder):
     wait_until(error_logged)
     log.info("Found expected error message in compute log, resuming.")
 
+    with closing(endpoint.connect()) as conn:
+        with conn.cursor() as cur:
+            # Confirm that neon_perf_counters also indicates that there are no active safekeepers
+            cur.execute(
+                "SELECT value FROM neon_perf_counters WHERE metric = 'num_active_safekeepers'"
+            )
+            assert cur.fetchone() == (0,), "Expected 0 active safekeepers"
+            cur.execute(
+                "SELECT value FROM neon_perf_counters WHERE metric = 'num_configured_safekeepers'"
+            )
+            assert cur.fetchone() == (1,), "Expected 1 configured safekeeper"
+
     # Sanity check that the hanging insert is indeed still hanging. Otherwise means the circuit breaker we
     # implemented didn't work as expected.
     time.sleep(2)

vendor/revisions.json

@@ -1,18 +1,18 @@
 {
   "v17": [
     "17.5",
-    "fa1788475e3146cc9c7c6a1b74f48fd296898fcd"
+    "1e01fcea2a6b38180021aa83e0051d95286d9096"
   ],
   "v16": [
     "16.9",
-    "9b9cb4b3e33347aea8f61e606bb6569979516de5"
+    "a42351fcd41ea01edede1daed65f651e838988fc"
   ],
   "v15": [
     "15.13",
-    "aaaeff2550d5deba58847f112af9b98fa3a58b00"
+    "2aaab3bb4a13557aae05bb2ae0ef0a132d0c4f85"
   ],
   "v14": [
     "14.18",
-    "c9f9fdd0113b52c0bd535afdb09d3a543aeee25f"
+    "2155cb165d05f617eb2c8ad7e43367189b627703"
   ]
 }