Filter obsolete layer files from an older generation from heatmap

## Problem
rest broker needs to respond with the correct cors headers for the api
to be usable from other domains

## Summary of changes
added a code path in rest broker to handle the OPTIONS requests

---------

Co-authored-by: Ruslan Talpa <ruslan.talpa@databricks.com>

.config/hakari.toml

@@ -21,13 +21,14 @@ platforms = [
     # "x86_64-apple-darwin",
     # "x86_64-pc-windows-msvc",
 ]
-
 [final-excludes]
 workspace-members = [
     # vm_monitor benefits from the same Cargo.lock as the rest of our artifacts, but
     # it is built primarly in separate repo neondatabase/autoscaling and thus is excluded
     # from depending on workspace-hack because most of the dependencies are not used.
     "vm_monitor",
+    # subzero-core is a stub crate that should be excluded from workspace-hack
+    "subzero-core",
     # All of these exist in libs and are not usually built independently.
     # Putting workspace hack there adds a bottleneck for cargo builds.
     "compute_api",

.dockerignore

@@ -27,4 +27,4 @@
 !storage_controller/
 !vendor/postgres-*/
 !workspace_hack/
-!build_tools/patches
+!build-tools/patches

.github/actionlint.yml

@@ -31,6 +31,7 @@ config-variables:
   - NEON_PROD_AWS_ACCOUNT_ID
   - PGREGRESS_PG16_PROJECT_ID
   - PGREGRESS_PG17_PROJECT_ID
+  - PREWARM_PROJECT_ID
   - REMOTE_STORAGE_AZURE_CONTAINER
   - REMOTE_STORAGE_AZURE_REGION
   - SLACK_CICD_CHANNEL_ID

.github/actions/prepare-for-subzero/action.yml

@@ -0,0 +1,28 @@
+name: 'Prepare current job for subzero'
+description: >
+  Set git token to access `neondatabase/subzero` from cargo build,
+  and set `CARGO_NET_GIT_FETCH_WITH_CLI=true` env variable to use git CLI
+
+inputs:
+  token:
+    description: 'GitHub token with access to neondatabase/subzero'
+    required: true
+
+runs:
+  using: "composite"
+
+  steps:
+    - name: Set git token for neondatabase/subzero
+      uses: pyTooling/Actions/with-post-step@2307b526df64d55e95884e072e49aac2a00a9afa # v5.1.0
+      env:
+        SUBZERO_ACCESS_TOKEN: ${{ inputs.token }}
+      with:
+        main: |
+          git config --global url."https://x-access-token:${SUBZERO_ACCESS_TOKEN}@github.com/neondatabase/subzero".insteadOf "https://github.com/neondatabase/subzero"
+          cargo add -p proxy subzero-core --git https://github.com/neondatabase/subzero --rev 396264617e78e8be428682f87469bb25429af88a
+        post: |
+          git config --global --unset url."https://x-access-token:${SUBZERO_ACCESS_TOKEN}@github.com/neondatabase/subzero".insteadOf "https://github.com/neondatabase/subzero"
+
+    - name: Set `CARGO_NET_GIT_FETCH_WITH_CLI=true` env variable
+      shell: bash -euxo pipefail {0}
+      run: echo "CARGO_NET_GIT_FETCH_WITH_CLI=true" >> ${GITHUB_ENV}

.github/actions/run-python-test-set/action.yml

@@ -176,7 +176,13 @@ runs:
         fi
 
         if [[ $BUILD_TYPE == "debug" && $RUNNER_ARCH == 'X64' ]]; then
-          cov_prefix=(scripts/coverage "--profraw-prefix=$GITHUB_JOB" --dir=/tmp/coverage run)
+          # We don't use code coverage for regression tests (the step is disabled),
+          # so there's no need to collect it.
+          # Ref https://github.com/neondatabase/neon/issues/4540
+          # cov_prefix=(scripts/coverage "--profraw-prefix=$GITHUB_JOB" --dir=/tmp/coverage run)
+          cov_prefix=()
+          # Explicitly set LLVM_PROFILE_FILE to /dev/null to avoid writing *.profraw files
+          export LLVM_PROFILE_FILE=/dev/null
         else
           cov_prefix=()
         fi

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

@@ -86,6 +86,10 @@ jobs:
         with:
           submodules: true
 
+      - uses: ./.github/actions/prepare-for-subzero
+        with:
+          token: ${{ secrets.CI_ACCESS_TOKEN }}
+
       - name: Set pg 14 revision for caching
         id: pg_v14_rev
         run: echo pg_rev=$(git rev-parse HEAD:vendor/postgres-v14) >> $GITHUB_OUTPUT
@@ -116,7 +120,7 @@ jobs:
           ARCH: ${{ inputs.arch }}
           SANITIZERS: ${{ inputs.sanitizers }}
         run: |
-          CARGO_FLAGS="--locked --features testing"
+          CARGO_FLAGS="--locked --features testing,rest_broker"
           if [[ $BUILD_TYPE == "debug" && $ARCH == 'x64' ]]; then
             cov_prefix="scripts/coverage --profraw-prefix=$GITHUB_JOB --dir=/tmp/coverage run"
             CARGO_PROFILE=""
@@ -150,7 +154,7 @@ jobs:
           secretKey: ${{ secrets.HETZNER_CACHE_SECRET_KEY }}
           use-fallback: false
           path: pg_install/v14
-          key: v1-${{ runner.os }}-${{ runner.arch }}-${{ inputs.build-type }}-pg-${{ steps.pg_v14_rev.outputs.pg_rev }}-bookworm-${{ hashFiles('Makefile', 'build-tools.Dockerfile') }}
+          key: v1-${{ runner.os }}-${{ runner.arch }}-${{ inputs.build-type }}-pg-${{ steps.pg_v14_rev.outputs.pg_rev }}-bookworm-${{ hashFiles('Makefile', 'build-tools/Dockerfile') }}
 
       - name: Cache postgres v15 build
         id: cache_pg_15
@@ -162,7 +166,7 @@ jobs:
           secretKey: ${{ secrets.HETZNER_CACHE_SECRET_KEY }}
           use-fallback: false
           path: pg_install/v15
-          key: v1-${{ runner.os }}-${{ runner.arch }}-${{ inputs.build-type }}-pg-${{ steps.pg_v15_rev.outputs.pg_rev }}-bookworm-${{ hashFiles('Makefile', 'build-tools.Dockerfile') }}
+          key: v1-${{ runner.os }}-${{ runner.arch }}-${{ inputs.build-type }}-pg-${{ steps.pg_v15_rev.outputs.pg_rev }}-bookworm-${{ hashFiles('Makefile', 'build-tools/Dockerfile') }}
 
       - name: Cache postgres v16 build
         id: cache_pg_16
@@ -174,7 +178,7 @@ jobs:
           secretKey: ${{ secrets.HETZNER_CACHE_SECRET_KEY }}
           use-fallback: false
           path: pg_install/v16
-          key: v1-${{ runner.os }}-${{ runner.arch }}-${{ inputs.build-type }}-pg-${{ steps.pg_v16_rev.outputs.pg_rev }}-bookworm-${{ hashFiles('Makefile', 'build-tools.Dockerfile') }}
+          key: v1-${{ runner.os }}-${{ runner.arch }}-${{ inputs.build-type }}-pg-${{ steps.pg_v16_rev.outputs.pg_rev }}-bookworm-${{ hashFiles('Makefile', 'build-tools/Dockerfile') }}
 
       - name: Cache postgres v17 build
         id: cache_pg_17
@@ -186,7 +190,7 @@ jobs:
           secretKey: ${{ secrets.HETZNER_CACHE_SECRET_KEY }}
           use-fallback: false
           path: pg_install/v17
-          key: v1-${{ runner.os }}-${{ runner.arch }}-${{ inputs.build-type }}-pg-${{ steps.pg_v17_rev.outputs.pg_rev }}-bookworm-${{ hashFiles('Makefile', 'build-tools.Dockerfile') }}
+          key: v1-${{ runner.os }}-${{ runner.arch }}-${{ inputs.build-type }}-pg-${{ steps.pg_v17_rev.outputs.pg_rev }}-bookworm-${{ hashFiles('Makefile', 'build-tools/Dockerfile') }}
 
       - name: Build all
         # Note: the Makefile picks up BUILD_TYPE and CARGO_PROFILE from the env variables

.github/workflows/_check-codestyle-rust.yml

@@ -46,6 +46,10 @@ jobs:
         uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
         with:
           submodules: true
+      
+      - uses: ./.github/actions/prepare-for-subzero
+        with:
+          token: ${{ secrets.CI_ACCESS_TOKEN }}
 
       - name: Cache cargo deps
         uses: tespkg/actions-cache@b7bf5fcc2f98a52ac6080eb0fd282c2f752074b1  # v1.8.0

.github/workflows/benchbase_tpcc.yml

@@ -0,0 +1,384 @@
+name: TPC-C like benchmark using benchbase
+
+on:
+  schedule:
+    # * is a special character in YAML so you have to quote this string
+    #          ┌───────────── minute (0 - 59)
+    #          │ ┌───────────── hour (0 - 23)
+    #          │ │ ┌───────────── day of the month (1 - 31)
+    #          │ │ │ ┌───────────── month (1 - 12 or JAN-DEC)
+    #          │ │ │ │ ┌───────────── day of the week (0 - 6 or SUN-SAT)
+    - cron:   '0 6 * * *' # run once a day at 6 AM UTC
+  workflow_dispatch: # adds ability to run this manually
+
+defaults:
+  run:
+    shell: bash -euxo pipefail {0}
+
+concurrency:
+  # Allow only one workflow globally because we do not want to be too noisy in production environment
+  group: benchbase-tpcc-workflow
+  cancel-in-progress: false
+
+permissions:
+  contents: read
+
+jobs:
+  benchbase-tpcc:
+    strategy:
+      fail-fast: false # allow other variants to continue even if one fails
+      matrix:
+        include:
+          - warehouses: 50 # defines number of warehouses and is used to compute number of terminals
+            max_rate: 800  # measured max TPS at scale factor based on experiments. Adjust if performance is better/worse
+            min_cu: 0.25   # simulate free tier plan (0.25 -2 CU)
+            max_cu: 2
+          - warehouses: 500 # serverless plan (2-8 CU)
+            max_rate: 2000
+            min_cu: 2
+            max_cu: 8
+          - warehouses: 1000 # business plan (2-16 CU)
+            max_rate: 2900
+            min_cu: 2
+            max_cu: 16
+      max-parallel: 1 # we want to run each workload size sequentially to avoid noisy neighbors
+    permissions:
+      contents: write
+      statuses: write
+      id-token: write # aws-actions/configure-aws-credentials
+    env:
+      PG_CONFIG: /tmp/neon/pg_install/v17/bin/pg_config
+      PSQL: /tmp/neon/pg_install/v17/bin/psql
+      PG_17_LIB_PATH: /tmp/neon/pg_install/v17/lib
+      POSTGRES_VERSION: 17
+    runs-on: [ self-hosted, us-east-2, x64 ]
+    timeout-minutes: 1440
+
+    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
+
+    - name: Configure AWS credentials # necessary to download artefacts
+      uses: aws-actions/configure-aws-credentials@e3dd6a429d7300a6a4c196c26e071d42e0343502 # v4.0.2
+      with:
+        aws-region: eu-central-1
+        role-to-assume: ${{ vars.DEV_AWS_OIDC_ROLE_ARN }}
+        role-duration-seconds: 18000 # 5 hours is currently max associated with IAM role
+
+    - name: Download Neon artifact
+      uses: ./.github/actions/download
+      with:
+        name: neon-${{ runner.os }}-${{ runner.arch }}-release-artifact
+        path: /tmp/neon/
+        prefix: latest
+        aws-oidc-role-arn: ${{ vars.DEV_AWS_OIDC_ROLE_ARN }}
+
+    - name: Create Neon Project
+      id: create-neon-project-tpcc
+      uses: ./.github/actions/neon-project-create
+      with:
+        region_id: aws-us-east-2
+        postgres_version: ${{ env.POSTGRES_VERSION }}
+        compute_units: '[${{ matrix.min_cu }}, ${{ matrix.max_cu }}]'
+        api_key: ${{ secrets.NEON_PRODUCTION_API_KEY_4_BENCHMARKS }}
+        api_host: console.neon.tech  # production (!)
+
+    - name: Initialize Neon project
+      env:
+          BENCHMARK_TPCC_CONNSTR: ${{ steps.create-neon-project-tpcc.outputs.dsn }}
+          PROJECT_ID: ${{ steps.create-neon-project-tpcc.outputs.project_id }}
+      run: |
+        echo "Initializing Neon project with project_id: ${PROJECT_ID}"
+        export LD_LIBRARY_PATH=${PG_17_LIB_PATH}
+        
+        # Retry logic for psql connection with 1 minute sleep between attempts
+        for attempt in {1..3}; do
+          echo "Attempt ${attempt}/3: Creating extensions in Neon project"
+          if ${PSQL} "${BENCHMARK_TPCC_CONNSTR}" -c "CREATE EXTENSION IF NOT EXISTS neon; CREATE EXTENSION IF NOT EXISTS neon_utils;"; then
+            echo "Successfully created extensions"
+            break
+          else
+            echo "Failed to create extensions on attempt ${attempt}"
+            if [ ${attempt} -lt 3 ]; then
+              echo "Waiting 60 seconds before retry..."
+              sleep 60
+            else
+              echo "All attempts failed, exiting"
+              exit 1
+            fi
+          fi
+        done
+        
+        echo "BENCHMARK_TPCC_CONNSTR=${BENCHMARK_TPCC_CONNSTR}" >> $GITHUB_ENV
+
+    - name: Generate BenchBase workload configuration
+      env:
+        WAREHOUSES: ${{ matrix.warehouses }}
+        MAX_RATE: ${{ matrix.max_rate }}
+      run: |
+        echo "Generating BenchBase configs for warehouses: ${WAREHOUSES}, max_rate: ${MAX_RATE}"
+        
+        # Extract hostname and password from connection string
+        # Format: postgresql://username:password@hostname/database?params (no port for Neon)
+        HOSTNAME=$(echo "${BENCHMARK_TPCC_CONNSTR}" | sed -n 's|.*://[^:]*:[^@]*@\([^/]*\)/.*|\1|p')
+        PASSWORD=$(echo "${BENCHMARK_TPCC_CONNSTR}" | sed -n 's|.*://[^:]*:\([^@]*\)@.*|\1|p')
+        
+        echo "Extracted hostname: ${HOSTNAME}"
+        
+        # Use runner temp (NVMe) as working directory
+        cd "${RUNNER_TEMP}"
+        
+        # Copy the generator script
+        cp "${GITHUB_WORKSPACE}/test_runner/performance/benchbase_tpc_c_helpers/generate_workload_size.py" .
+        
+        # Generate configs and scripts
+        python3 generate_workload_size.py \
+          --warehouses ${WAREHOUSES} \
+          --max-rate ${MAX_RATE} \
+          --hostname ${HOSTNAME} \
+          --password ${PASSWORD} \
+          --runner-arch ${{ runner.arch }}
+        
+        # Fix path mismatch: move generated configs and scripts to expected locations
+        mv ../configs ./configs
+        mv ../scripts ./scripts
+
+    - name: Prepare database (load data)
+      env:
+        WAREHOUSES: ${{ matrix.warehouses }}
+      run: |
+        cd "${RUNNER_TEMP}"
+        
+        echo "Loading ${WAREHOUSES} warehouses into database..."
+        
+        # Run the loader script and capture output to log file while preserving stdout/stderr
+        ./scripts/load_${WAREHOUSES}_warehouses.sh 2>&1 | tee "load_${WAREHOUSES}_warehouses.log"
+        
+        echo "Database loading completed"
+
+    - name: Run TPC-C benchmark (warmup phase, then benchmark at 70% of configuredmax TPS)
+      env:
+        WAREHOUSES: ${{ matrix.warehouses }}
+      run: |
+        cd "${RUNNER_TEMP}"
+        
+        echo "Running TPC-C benchmark with ${WAREHOUSES} warehouses..."
+        
+        # Run the optimal rate benchmark
+        ./scripts/execute_${WAREHOUSES}_warehouses_opt_rate.sh
+        
+        echo "Benchmark execution completed"
+
+    - name: Run TPC-C benchmark (warmup phase, then ramp down TPS and up again in 5 minute intervals)
+
+      env:
+          WAREHOUSES: ${{ matrix.warehouses }}
+      run: |
+        cd "${RUNNER_TEMP}"
+        
+        echo "Running TPC-C ramp-down-up with ${WAREHOUSES} warehouses..."
+        
+        # Run the optimal rate benchmark
+        ./scripts/execute_${WAREHOUSES}_warehouses_ramp_up.sh
+        
+        echo "Benchmark execution completed"
+
+    - name: Process results (upload to test results database and generate diagrams)
+      env:
+        WAREHOUSES: ${{ matrix.warehouses }}
+        MIN_CU: ${{ matrix.min_cu }}
+        MAX_CU: ${{ matrix.max_cu }}
+        PROJECT_ID: ${{ steps.create-neon-project-tpcc.outputs.project_id }}
+        REVISION: ${{ github.sha }}
+        PERF_DB_CONNSTR: ${{ secrets.PERF_TEST_RESULT_CONNSTR }}
+      run: |
+        cd "${RUNNER_TEMP}"
+        
+        echo "Creating temporary Python environment for results processing..."
+        
+        # Create temporary virtual environment
+        python3 -m venv temp_results_env
+        source temp_results_env/bin/activate
+        
+        # Install required packages in virtual environment
+        pip install matplotlib pandas psycopg2-binary
+        
+        echo "Copying results processing scripts..."
+        
+        # Copy both processing scripts
+        cp "${GITHUB_WORKSPACE}/test_runner/performance/benchbase_tpc_c_helpers/generate_diagrams.py" .
+        cp "${GITHUB_WORKSPACE}/test_runner/performance/benchbase_tpc_c_helpers/upload_results_to_perf_test_results.py" .
+        
+        echo "Processing load phase metrics..."
+        
+        # Find and process load log
+        LOAD_LOG=$(find . -name "load_${WAREHOUSES}_warehouses.log" -type f | head -1)
+        if [ -n "$LOAD_LOG" ]; then
+          echo "Processing load metrics from: $LOAD_LOG"
+          python upload_results_to_perf_test_results.py \
+            --load-log "$LOAD_LOG" \
+            --run-type "load" \
+            --warehouses "${WAREHOUSES}" \
+            --min-cu "${MIN_CU}" \
+            --max-cu "${MAX_CU}" \
+            --project-id "${PROJECT_ID}" \
+            --revision "${REVISION}" \
+            --connection-string "${PERF_DB_CONNSTR}"
+        else
+          echo "Warning: Load log file not found: load_${WAREHOUSES}_warehouses.log"
+        fi
+        
+        echo "Processing warmup results for optimal rate..."
+        
+        # Find and process warmup results
+        WARMUP_CSV=$(find results_warmup -name "*.results.csv" -type f | head -1)
+        WARMUP_JSON=$(find results_warmup -name "*.summary.json" -type f | head -1)
+        
+        if [ -n "$WARMUP_CSV" ] && [ -n "$WARMUP_JSON" ]; then
+          echo "Generating warmup diagram from: $WARMUP_CSV"
+          python generate_diagrams.py \
+            --input-csv "$WARMUP_CSV" \
+            --output-svg "warmup_${WAREHOUSES}_warehouses_performance.svg" \
+            --title-suffix "Warmup at max TPS"
+            
+          echo "Uploading warmup metrics from: $WARMUP_JSON"
+          python upload_results_to_perf_test_results.py \
+            --summary-json "$WARMUP_JSON" \
+            --results-csv "$WARMUP_CSV" \
+            --run-type "warmup" \
+            --min-cu "${MIN_CU}" \
+            --max-cu "${MAX_CU}" \
+            --project-id "${PROJECT_ID}" \
+            --revision "${REVISION}" \
+            --connection-string "${PERF_DB_CONNSTR}"
+        else
+          echo "Warning: Missing warmup results files (CSV: $WARMUP_CSV, JSON: $WARMUP_JSON)"
+        fi
+        
+        echo "Processing optimal rate results..."
+        
+        # Find and process optimal rate results  
+        OPTRATE_CSV=$(find results_opt_rate -name "*.results.csv" -type f | head -1)
+        OPTRATE_JSON=$(find results_opt_rate -name "*.summary.json" -type f | head -1)
+        
+        if [ -n "$OPTRATE_CSV" ] && [ -n "$OPTRATE_JSON" ]; then
+          echo "Generating optimal rate diagram from: $OPTRATE_CSV"
+          python generate_diagrams.py \
+            --input-csv "$OPTRATE_CSV" \
+            --output-svg "benchmark_${WAREHOUSES}_warehouses_performance.svg" \
+            --title-suffix "70% of max TPS"
+            
+          echo "Uploading optimal rate metrics from: $OPTRATE_JSON"
+          python upload_results_to_perf_test_results.py \
+            --summary-json "$OPTRATE_JSON" \
+            --results-csv "$OPTRATE_CSV" \
+            --run-type "opt-rate" \
+            --min-cu "${MIN_CU}" \
+            --max-cu "${MAX_CU}" \
+            --project-id "${PROJECT_ID}" \
+            --revision "${REVISION}" \
+            --connection-string "${PERF_DB_CONNSTR}"
+        else
+          echo "Warning: Missing optimal rate results files (CSV: $OPTRATE_CSV, JSON: $OPTRATE_JSON)"
+        fi
+
+        echo "Processing warmup 2 results for ramp down/up phase..."
+        
+        # Find and process warmup results
+        WARMUP_CSV=$(find results_warmup -name "*.results.csv" -type f | tail -1)
+        WARMUP_JSON=$(find results_warmup -name "*.summary.json" -type f | tail -1)
+        
+        if [ -n "$WARMUP_CSV" ] && [ -n "$WARMUP_JSON" ]; then
+          echo "Generating warmup diagram from: $WARMUP_CSV"
+          python generate_diagrams.py \
+            --input-csv "$WARMUP_CSV" \
+            --output-svg "warmup_2_${WAREHOUSES}_warehouses_performance.svg" \
+            --title-suffix "Warmup at max TPS"
+            
+          echo "Uploading warmup metrics from: $WARMUP_JSON"
+          python upload_results_to_perf_test_results.py \
+            --summary-json "$WARMUP_JSON" \
+            --results-csv "$WARMUP_CSV" \
+            --run-type "warmup" \
+            --min-cu "${MIN_CU}" \
+            --max-cu "${MAX_CU}" \
+            --project-id "${PROJECT_ID}" \
+            --revision "${REVISION}" \
+            --connection-string "${PERF_DB_CONNSTR}"
+        else
+          echo "Warning: Missing warmup results files (CSV: $WARMUP_CSV, JSON: $WARMUP_JSON)"
+        fi
+        
+        echo "Processing ramp results..."
+        
+        # Find and process ramp results  
+        RAMPUP_CSV=$(find results_ramp_up -name "*.results.csv" -type f | head -1)
+        RAMPUP_JSON=$(find results_ramp_up -name "*.summary.json" -type f | head -1)
+        
+        if [ -n "$RAMPUP_CSV" ] && [ -n "$RAMPUP_JSON" ]; then
+          echo "Generating ramp diagram from: $RAMPUP_CSV"
+          python generate_diagrams.py \
+            --input-csv "$RAMPUP_CSV" \
+            --output-svg "ramp_${WAREHOUSES}_warehouses_performance.svg" \
+            --title-suffix "ramp TPS down and up in 5 minute intervals"
+            
+          echo "Uploading ramp metrics from: $RAMPUP_JSON"
+          python upload_results_to_perf_test_results.py \
+            --summary-json "$RAMPUP_JSON" \
+            --results-csv "$RAMPUP_CSV" \
+            --run-type "ramp-up" \
+            --min-cu "${MIN_CU}" \
+            --max-cu "${MAX_CU}" \
+            --project-id "${PROJECT_ID}" \
+            --revision "${REVISION}" \
+            --connection-string "${PERF_DB_CONNSTR}"
+        else
+          echo "Warning: Missing ramp results files (CSV: $RAMPUP_CSV, JSON: $RAMPUP_JSON)"
+        fi
+        
+        # Deactivate and clean up virtual environment
+        deactivate
+        rm -rf temp_results_env
+        rm upload_results_to_perf_test_results.py
+        
+        echo "Results processing completed and environment cleaned up"
+
+    - name: Set date for upload
+      id: set-date
+      run: echo "date=$(date +%Y-%m-%d)" >> $GITHUB_OUTPUT
+
+    - name: Configure AWS credentials # necessary to upload results
+      uses: aws-actions/configure-aws-credentials@e3dd6a429d7300a6a4c196c26e071d42e0343502 # v4.0.2
+      with:
+        aws-region: us-east-2
+        role-to-assume: ${{ vars.DEV_AWS_OIDC_ROLE_ARN }}
+        role-duration-seconds: 900 # 900 is minimum value 
+        
+    - name: Upload benchmark results to S3
+      env:
+        S3_BUCKET: neon-public-benchmark-results
+        S3_PREFIX: benchbase-tpc-c/${{ steps.set-date.outputs.date }}/${{ github.run_id }}/${{ matrix.warehouses }}-warehouses
+      run: |
+        echo "Redacting passwords from configuration files before upload..."
+        
+        # Mask all passwords in XML config files
+        find "${RUNNER_TEMP}/configs" -name "*.xml" -type f -exec sed -i 's|<password>[^<]*</password>|<password>redacted</password>|g' {} \;
+        
+        echo "Uploading benchmark results to s3://${S3_BUCKET}/${S3_PREFIX}/"
+        
+        # Upload the entire benchmark directory recursively
+        aws s3 cp --only-show-errors --recursive "${RUNNER_TEMP}" s3://${S3_BUCKET}/${S3_PREFIX}/
+        
+        echo "Upload completed"
+        
+    - name: Delete Neon Project
+      if: ${{ always() }}
+      uses: ./.github/actions/neon-project-delete
+      with:
+        project_id: ${{ steps.create-neon-project-tpcc.outputs.project_id }}
+        api_key: ${{ secrets.NEON_PRODUCTION_API_KEY_4_BENCHMARKS }} 
+        api_host: console.neon.tech  # production (!)

.github/workflows/benchmarking.yml

@@ -219,6 +219,7 @@ jobs:
           --ignore test_runner/performance/test_cumulative_statistics_persistence.py
           --ignore test_runner/performance/test_perf_many_relations.py
           --ignore test_runner/performance/test_perf_oltp_large_tenant.py
+          --ignore test_runner/performance/test_lfc_prewarm.py
       env:
         BENCHMARK_CONNSTR: ${{ steps.create-neon-project.outputs.dsn }}
         VIP_VAP_ACCESS_TOKEN: "${{ secrets.VIP_VAP_ACCESS_TOKEN }}"
@@ -410,6 +411,77 @@ jobs:
       env:
         SLACK_BOT_TOKEN: ${{ secrets.SLACK_BOT_TOKEN }}
 
+  prewarm-test:
+    if: ${{ github.event.inputs.run_only_pgvector_tests == 'false' || github.event.inputs.run_only_pgvector_tests == null }}
+    permissions:
+      contents: write
+      statuses: write
+      id-token: write # aws-actions/configure-aws-credentials
+    env:
+      PROJECT_ID: ${{ vars.PREWARM_PROJECT_ID }}
+      POSTGRES_DISTRIB_DIR: /tmp/neon/pg_install
+      DEFAULT_PG_VERSION: 17
+      TEST_OUTPUT: /tmp/test_output
+      BUILD_TYPE: remote
+      SAVE_PERF_REPORT: ${{ github.event.inputs.save_perf_report || ( github.ref_name == 'main' ) }}
+      PLATFORM: "neon-staging"
+
+    runs-on: [ self-hosted, us-east-2, x64 ]
+    container:
+      image: ghcr.io/neondatabase/build-tools:pinned-bookworm
+      credentials:
+        username: ${{ github.actor }}
+        password: ${{ secrets.GITHUB_TOKEN }}
+      options: --init
+
+    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
+
+    - name: Configure AWS credentials
+      uses: aws-actions/configure-aws-credentials@e3dd6a429d7300a6a4c196c26e071d42e0343502 # v4.0.2
+      with:
+        aws-region: eu-central-1
+        role-to-assume: ${{ vars.DEV_AWS_OIDC_ROLE_ARN }}
+        role-duration-seconds: 18000 # 5 hours
+
+    - name: Download Neon artifact
+      uses: ./.github/actions/download
+      with:
+        name: neon-${{ runner.os }}-${{ runner.arch }}-release-artifact
+        path: /tmp/neon/
+        prefix: latest
+        aws-oidc-role-arn: ${{ vars.DEV_AWS_OIDC_ROLE_ARN }}
+
+    - name: Run prewarm benchmark
+      uses: ./.github/actions/run-python-test-set
+      with:
+        build_type: ${{ env.BUILD_TYPE }}
+        test_selection: performance/test_lfc_prewarm.py
+        run_in_parallel: false
+        save_perf_report: ${{ env.SAVE_PERF_REPORT }}
+        extra_params: -m remote_cluster --timeout 5400
+        pg_version: ${{ env.DEFAULT_PG_VERSION }}
+        aws-oidc-role-arn: ${{ vars.DEV_AWS_OIDC_ROLE_ARN }}
+      env:
+        VIP_VAP_ACCESS_TOKEN: "${{ secrets.VIP_VAP_ACCESS_TOKEN }}"
+        PERF_TEST_RESULT_CONNSTR: "${{ secrets.PERF_TEST_RESULT_CONNSTR }}"
+        NEON_API_KEY: ${{ secrets.NEON_STAGING_API_KEY }}
+
+    - name: Create Allure report
+      id: create-allure-report
+      if: ${{ !cancelled() }}
+      uses: ./.github/actions/allure-report-generate
+      with:
+        store-test-results-into-db: true
+        aws-oidc-role-arn: ${{ vars.DEV_AWS_OIDC_ROLE_ARN }}
+      env:
+        REGRESS_TEST_RESULT_CONNSTR_NEW: ${{ secrets.REGRESS_TEST_RESULT_CONNSTR_NEW }}
+
   generate-matrices:
     if: ${{ github.event.inputs.run_only_pgvector_tests == 'false' || github.event.inputs.run_only_pgvector_tests == null }}
     # Create matrices for the benchmarking jobs, so we run benchmarks on rds only once a week (on Saturday)

.github/workflows/build-build-tools-image.yml

@@ -72,7 +72,7 @@ jobs:
           ARCHS: ${{ inputs.archs || '["x64","arm64"]' }}
           DEBIANS: ${{ inputs.debians || '["bullseye","bookworm"]' }}
           IMAGE_TAG: |
-            ${{ hashFiles('build-tools.Dockerfile',
+            ${{ hashFiles('build-tools/Dockerfile',
                           '.github/workflows/build-build-tools-image.yml') }}
         run: |
           echo "archs=${ARCHS}"           | tee -a ${GITHUB_OUTPUT}
@@ -144,9 +144,11 @@ jobs:
 
       - uses: docker/build-push-action@471d1dc4e07e5cdedd4c2171150001c434f0b7a4 # v6.15.0
         with:
-          file: build-tools.Dockerfile
+          file: build-tools/Dockerfile
           context: .
-          provenance: false
+          attests: |
+            type=provenance,mode=max
+            type=sbom,generator=docker.io/docker/buildkit-syft-scanner:1
           push: true
           pull: true
           build-args: |

.github/workflows/build-macos.yml

@@ -54,6 +54,10 @@ jobs:
         uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
         with:
           submodules: true
+      
+      - uses: ./.github/actions/prepare-for-subzero
+        with:
+          token: ${{ secrets.CI_ACCESS_TOKEN }}
 
       - name: Install build dependencies
         run: |

.github/workflows/build_and_test.yml

@@ -87,22 +87,27 @@ jobs:
     uses: ./.github/workflows/build-build-tools-image.yml
     secrets: inherit
 
-  lint-openapi-spec:
-    runs-on: ubuntu-22.04
-    needs: [ meta, check-permissions ]
+  lint-yamls:
+    needs: [ meta, check-permissions, build-build-tools-image ]
     # We do need to run this in `.*-rc-pr` because of hotfixes.
     if: ${{ contains(fromJSON('["pr", "push-main", "storage-rc-pr", "proxy-rc-pr", "compute-rc-pr"]'), needs.meta.outputs.run-kind) }}
+    runs-on: [ self-hosted, small ]
+    container:
+      image: ${{ needs.build-build-tools-image.outputs.image }}
+      credentials:
+        username: ${{ github.actor }}
+        password: ${{ secrets.GITHUB_TOKEN }}
+      options: --init
+
     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
-      - uses: docker/login-action@74a5d142397b4f367a81961eba4e8cd7edddf772 # v3.4.0
-        with:
-          registry: ghcr.io
-          username: ${{ github.actor }}
-          password: ${{ secrets.GITHUB_TOKEN }}
+
+      - run: make -C compute manifest-schema-validation
       - run: make lint-openapi-spec
 
   check-codestyle-python:
@@ -217,28 +222,6 @@ jobs:
       build-tools-image: ${{ needs.build-build-tools-image.outputs.image }}-bookworm
     secrets: inherit
 
-  validate-compute-manifest:
-    runs-on: ubuntu-22.04
-    needs: [ meta, check-permissions ]
-    # We do need to run this in `.*-rc-pr` because of hotfixes.
-    if: ${{ contains(fromJSON('["pr", "push-main", "storage-rc-pr", "proxy-rc-pr", "compute-rc-pr"]'), needs.meta.outputs.run-kind) }}
-    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
-
-      - name: Set up Node.js
-        uses: actions/setup-node@49933ea5288caeca8642d1e84afbd3f7d6820020 # v4.4.0
-        with:
-          node-version: '24'
-
-      - name: Validate manifest against schema
-        run: |
-          make -C compute manifest-schema-validation
-
   build-and-test-locally:
     needs: [ meta, build-build-tools-image ]
     # We do need to run this in `.*-rc-pr` because of hotfixes.
@@ -649,7 +632,11 @@ jobs:
             BUILD_TAG=${{ needs.meta.outputs.release-tag || needs.meta.outputs.build-tag }}
             TAG=${{ needs.build-build-tools-image.outputs.image-tag }}-bookworm
             DEBIAN_VERSION=bookworm
-          provenance: false
+          secrets: |
+            SUBZERO_ACCESS_TOKEN=${{ secrets.CI_ACCESS_TOKEN }}
+          attests: |
+            type=provenance,mode=max
+            type=sbom,generator=docker.io/docker/buildkit-syft-scanner:1
           push: true
           pull: true
           file: Dockerfile
@@ -762,7 +749,9 @@ jobs:
             PG_VERSION=${{ matrix.version.pg }}
             BUILD_TAG=${{ needs.meta.outputs.release-tag || needs.meta.outputs.build-tag }}
             DEBIAN_VERSION=${{ matrix.version.debian }}
-          provenance: false
+          attests: |
+            type=provenance,mode=max
+            type=sbom,generator=docker.io/docker/buildkit-syft-scanner:1
           push: true
           pull: true
           file: compute/compute-node.Dockerfile
@@ -781,7 +770,9 @@ jobs:
             PG_VERSION=${{ matrix.version.pg }}
             BUILD_TAG=${{ needs.meta.outputs.release-tag || needs.meta.outputs.build-tag }}
             DEBIAN_VERSION=${{ matrix.version.debian }}
-          provenance: false
+          attests: |
+            type=provenance,mode=max
+            type=sbom,generator=docker.io/docker/buildkit-syft-scanner:1
           push: true
           pull: true
           file: compute/compute-node.Dockerfile

.github/workflows/neon_extra_builds.yml

@@ -72,6 +72,7 @@ jobs:
   check-macos-build:
     needs: [ check-permissions, files-changed ]
     uses: ./.github/workflows/build-macos.yml
+    secrets: inherit
     with:
       pg_versions: ${{ needs.files-changed.outputs.postgres_changes }}
       rebuild_rust_code: ${{ fromJSON(needs.files-changed.outputs.rebuild_rust_code) }}

.github/workflows/pg-clients.yml

@@ -48,8 +48,20 @@ jobs:
     uses: ./.github/workflows/build-build-tools-image.yml
     secrets: inherit
 
+  generate-ch-tmppw:
+    runs-on: ubuntu-22.04
+    outputs:
+      tmp_val: ${{ steps.pwgen.outputs.tmp_val }}
+    steps:
+      - name: Generate a random password
+        id: pwgen
+        run: |
+          set +x
+          p=$(dd if=/dev/random bs=14 count=1 2>/dev/null | base64)
+          echo tmp_val="${p//\//}" >> "${GITHUB_OUTPUT}"
+
   test-logical-replication:
-    needs: [ build-build-tools-image ]
+    needs: [ build-build-tools-image, generate-ch-tmppw ]
     runs-on: ubuntu-22.04
 
     container:
@@ -60,16 +72,21 @@ jobs:
       options: --init --user root
     services:
       clickhouse:
-        image: clickhouse/clickhouse-server:24.6.3.64
+        image: clickhouse/clickhouse-server:25.6
+        env:
+          CLICKHOUSE_PASSWORD: ${{ needs.generate-ch-tmppw.outputs.tmp_val }}
+          PGSSLCERT: /tmp/postgresql.crt
         ports:
           - 9000:9000
           - 8123:8123
       zookeeper:
-        image: quay.io/debezium/zookeeper:2.7
+        image: quay.io/debezium/zookeeper:3.1.3.Final
         ports:
           - 2181:2181
+          - 2888:2888
+          - 3888:3888
       kafka:
-        image: quay.io/debezium/kafka:2.7
+        image: quay.io/debezium/kafka:3.1.3.Final
         env:
           ZOOKEEPER_CONNECT: "zookeeper:2181"
           KAFKA_ADVERTISED_LISTENERS: PLAINTEXT://kafka:9092
@@ -79,7 +96,7 @@ jobs:
         ports:
           - 9092:9092
       debezium:
-        image: quay.io/debezium/connect:2.7
+        image: quay.io/debezium/connect:3.1.3.Final
         env:
           BOOTSTRAP_SERVERS: kafka:9092
           GROUP_ID: 1
@@ -125,6 +142,7 @@ jobs:
           aws-oidc-role-arn: ${{ vars.DEV_AWS_OIDC_ROLE_ARN }}
         env:
           BENCHMARK_CONNSTR: ${{ steps.create-neon-project.outputs.dsn }}
+          CLICKHOUSE_PASSWORD: ${{ needs.generate-ch-tmppw.outputs.tmp_val }}
 
       - name: Delete Neon Project
         if: always()

.github/workflows/proxy-benchmark.yml

@@ -3,7 +3,7 @@ name: Periodic proxy performance test on unit-perf-aws-arm runners
 on:
   push: # TODO: remove after testing
     branches:
-      - test-proxy-bench # Runs on pushes to branches starting with test-proxy-bench
+      - test-proxy-bench # Runs on pushes to test-proxy-bench branch
   # schedule:
     # * is a special character in YAML so you have to quote this string
     #        ┌───────────── minute (0 - 59)
@@ -32,7 +32,7 @@ jobs:
       statuses: write
       contents: write
       pull-requests: write
-    runs-on: [self-hosted, unit-perf-aws-arm]
+    runs-on: [ self-hosted, unit-perf-aws-arm ]
     timeout-minutes: 60  # 1h timeout
     container:
       image: ghcr.io/neondatabase/build-tools:pinned-bookworm
@@ -55,30 +55,58 @@ jobs:
         {
           echo "PROXY_BENCH_PATH=$PROXY_BENCH_PATH"
           echo "NEON_DIR=${RUNNER_TEMP}/neon"
+          echo "NEON_PROXY_PATH=${RUNNER_TEMP}/neon/bin/proxy"
           echo "TEST_OUTPUT=${PROXY_BENCH_PATH}/test_output"
           echo ""
         } >> "$GITHUB_ENV"
 
-    - name: Run proxy-bench
-      run: ${PROXY_BENCH_PATH}/run.sh
+    - 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: Ingest Bench Results # neon repo script
+    - name: Install Python deps
+      shell: bash -euxo pipefail {0}
+      run: ./scripts/pysync
+
+    - name: show ulimits
+      shell: bash -euxo pipefail {0}
+      run: |
+        ulimit -a
+
+    - name: Run proxy-bench
+      working-directory: ${{ env.PROXY_BENCH_PATH }}
+      run: ./run.sh --with-grafana --bare-metal
+
+    - name: Ingest Bench Results
       if: always()
+      working-directory: ${{ env.NEON_DIR }}
       run: |
         mkdir -p $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}
       if: always()
       env:
         PERF_TEST_RESULT_CONNSTR: "${{ secrets.PROXY_TEST_RESULT_CONNSTR }}"
         REPORT_FROM: $TEST_OUTPUT
+      working-directory: ${{ env.NEON_DIR }}
       run: $NEON_DIR/scripts/generate_and_push_perf_report.sh
 
-    - name: Docker cleanup
-      if: always()
-      run: docker compose down
-
     - name: Notify Failure
       if: failure()
-      run: echo "Proxy bench job failed" && exit 1
+      run: echo "Proxy bench job failed" && exit 1
+
+    - name: Cleanup Test Resources
+      if: always()
+      shell: bash -euxo pipefail {0}
+      run: |
+        # Cleanup the test resources
+        if [[ -d "${TEST_OUTPUT}" ]]; then
+          rm -rf ${TEST_OUTPUT}
+        fi
+        if [[ -d "${PROXY_BENCH_PATH}/test_output" ]]; then
+          rm -rf ${PROXY_BENCH_PATH}/test_output
+        fi

.gitignore

@@ -26,6 +26,14 @@ docker-compose/docker-compose-parallel.yml
 *.o
 *.so
 *.Po
+*.pid
 
 # pgindent typedef lists
 *.list
+
+# Node
+**/node_modules/
+
+# various files for local testing
+/proxy/.subzero
+local_proxy.json

.gitmodules

@@ -1,16 +1,16 @@
 [submodule "vendor/postgres-v14"]
 	path = vendor/postgres-v14
-	url = https://github.com/neondatabase/postgres.git
+	url = ../postgres.git
 	branch = REL_14_STABLE_neon
 [submodule "vendor/postgres-v15"]
 	path = vendor/postgres-v15
-	url = https://github.com/neondatabase/postgres.git
+	url = ../postgres.git
 	branch = REL_15_STABLE_neon
 [submodule "vendor/postgres-v16"]
 	path = vendor/postgres-v16
-	url = https://github.com/neondatabase/postgres.git
+	url = ../postgres.git
 	branch = REL_16_STABLE_neon
 [submodule "vendor/postgres-v17"]
 	path = vendor/postgres-v17
-	url = https://github.com/neondatabase/postgres.git
+	url = ../postgres.git
 	branch = REL_17_STABLE_neon

Cargo.toml

@@ -21,7 +21,6 @@ members = [
     "workspace_hack",
     "libs/compute_api",
     "libs/http-utils",
-    "libs/neon_failpoint",
     "libs/pageserver_api",
     "libs/postgres_ffi",
     "libs/postgres_ffi_types",
@@ -47,6 +46,7 @@ members = [
     "libs/proxy/json",
     "libs/proxy/postgres-protocol2",
     "libs/proxy/postgres-types2",
+    "libs/proxy/subzero_core",
     "libs/proxy/tokio-postgres2",
     "endpoint_storage",
     "pgxn/neon/communicator",
@@ -98,6 +98,7 @@ diatomic-waker = { version = "0.2.3" }
 either = "1.8"
 enum-map = "2.4.2"
 enumset = "1.0.12"
+fail = "0.5.0"
 fallible-iterator = "0.2"
 framed-websockets = { version = "0.1.0", git = "https://github.com/neondatabase/framed-websockets" }
 futures = "0.3"
@@ -130,10 +131,11 @@ jemalloc_pprof = { version = "0.7", features = ["symbolize", "flamegraph"] }
 jsonwebtoken = "9"
 lasso = "0.7"
 libc = "0.2"
+lock_api = "0.4.13"
 md5 = "0.7.0"
 measured = { version = "0.0.22", features=["lasso"] }
 measured-process = { version = "0.0.22" }
-memoffset = "0.9"
+moka = { version = "0.12", features = ["sync"] }
 nix = { version = "0.30.1", features = ["dir", "fs", "mman", "process", "socket", "signal", "poll"] }
 # Do not update to >= 7.0.0, at least. The update will have a significant impact
 # on compute startup metrics (start_postgres_ms), >= 25% degradation.
@@ -141,10 +143,10 @@ notify = "6.0.0"
 num_cpus = "1.15"
 num-traits = "0.2.19"
 once_cell = "1.13"
-opentelemetry = "0.27"
-opentelemetry_sdk = "0.27"
-opentelemetry-otlp = { version = "0.27", default-features = false, features = ["http-proto", "trace", "http", "reqwest-client"] }
-opentelemetry-semantic-conventions = "0.27"
+opentelemetry = "0.30"
+opentelemetry_sdk = "0.30"
+opentelemetry-otlp = { version = "0.30", default-features = false, features = ["http-proto", "trace", "http", "reqwest-blocking-client"] }
+opentelemetry-semantic-conventions = "0.30"
 parking_lot = "0.12"
 parquet = { version = "53", default-features = false, features = ["zstd"] }
 parquet_derive = "53"
@@ -156,16 +158,18 @@ procfs = "0.16"
 prometheus = {version = "0.13", default-features=false, features = ["process"]} # removes protobuf dependency
 prost = "0.13.5"
 prost-types = "0.13.5"
-rand = "0.8"
+rand = "0.9"
+# Remove after p256 is updated to 0.14.
+rand_core = "=0.6"
 redis = { version = "0.29.2", features = ["tokio-rustls-comp", "keep-alive"] }
 regex = "1.10.2"
 reqwest = { version = "0.12", default-features = false, features = ["rustls-tls"] }
-reqwest-tracing = { version = "0.5", features = ["opentelemetry_0_27"] }
+reqwest-tracing = { version = "0.5", features = ["opentelemetry_0_30"] }
 reqwest-middleware = "0.4"
 reqwest-retry = "0.7"
 routerify = "3"
 rpds = "0.13"
-rustc-hash = "1.1.0"
+rustc-hash = "2.1.1"
 rustls = { version = "0.23.16", default-features = false }
 rustls-pemfile = "2"
 rustls-pki-types = "1.11"
@@ -201,7 +205,7 @@ tokio-epoll-uring = { git = "https://github.com/neondatabase/tokio-epoll-uring.g
 tokio-io-timeout = "1.2.0"
 tokio-postgres-rustls = "0.12.0"
 tokio-rustls = { version = "0.26.0", default-features = false, features = ["tls12", "ring"]}
-tokio-stream = "0.1"
+tokio-stream = { version = "0.1", features = ["sync"] }
 tokio-tar = "0.3"
 tokio-util = { version = "0.7.10", features = ["io", "io-util", "rt"] }
 toml = "0.8"
@@ -210,17 +214,15 @@ tonic = { version = "0.13.1", default-features = false, features = ["channel", "
 tonic-reflection = { version = "0.13.1", features = ["server"] }
 tower = { version = "0.5.2", default-features = false }
 tower-http = { version = "0.6.2", features = ["auth", "request-id", "trace"] }
-
-# This revision uses opentelemetry 0.27. There's no tag for it.
-tower-otel = { git = "https://github.com/mattiapenati/tower-otel", rev = "56a7321053bcb72443888257b622ba0d43a11fcd" }
-
+tower-otel = { version = "0.6", features = ["axum"] }
 tower-service = "0.3.3"
 tracing = "0.1"
 tracing-error = "0.2"
 tracing-log = "0.2"
-tracing-opentelemetry = "0.28"
+tracing-opentelemetry = "0.31"
 tracing-serde = "0.2.0"
 tracing-subscriber = { version = "0.3", default-features = false, features = ["smallvec", "fmt", "tracing-log", "std", "env-filter", "json"] }
+tracing-appender = "0.2.3"
 try-lock = "0.2.5"
 test-log = { version = "0.2.17", default-features = false, features = ["log"] }
 twox-hash = { version = "1.6.3", default-features = false }
@@ -231,9 +233,10 @@ uuid = { version = "1.6.1", features = ["v4", "v7", "serde"] }
 walkdir = "2.3.2"
 rustls-native-certs = "0.8"
 whoami = "1.5.1"
-zerocopy = { version = "0.8", features = ["derive", "simd"] }
 json-structural-diff = { version = "0.2.0" }
 x509-cert = { version = "0.2.5" }
+zerocopy = { version = "0.8", features = ["derive", "simd"] }
+zeroize = "1.8"
 
 ## TODO replace this with tracing
 env_logger = "0.11"
@@ -258,7 +261,6 @@ desim = { version = "0.1", path = "./libs/desim" }
 endpoint_storage = { version = "0.0.1", path = "./endpoint_storage/" }
 http-utils = { version = "0.1", path = "./libs/http-utils/" }
 metrics = { version = "0.1", path = "./libs/metrics/" }
-neon_failpoint = { version = "0.1", path = "./libs/neon_failpoint/" }
 neon-shmem = { version = "0.1", path = "./libs/neon-shmem/" }
 pageserver = { path = "./pageserver" }
 pageserver_api = { version = "0.1", path = "./libs/pageserver_api/" }

Dockerfile

@@ -63,7 +63,14 @@ WORKDIR /home/nonroot
 
 COPY --chown=nonroot . .
 
-RUN cargo chef prepare --recipe-path recipe.json
+RUN --mount=type=secret,uid=1000,id=SUBZERO_ACCESS_TOKEN \
+    set -e \
+    && if [ -s /run/secrets/SUBZERO_ACCESS_TOKEN ]; then \
+        export CARGO_NET_GIT_FETCH_WITH_CLI=true && \
+        git config --global url."https://$(cat /run/secrets/SUBZERO_ACCESS_TOKEN)@github.com/neondatabase/subzero".insteadOf "https://github.com/neondatabase/subzero" && \
+        cargo add -p proxy subzero-core --git https://github.com/neondatabase/subzero --rev 396264617e78e8be428682f87469bb25429af88a; \
+    fi \
+    && cargo chef prepare --recipe-path recipe.json
 
 # Main build image
 FROM $REPOSITORY/$IMAGE:$TAG AS build
@@ -71,20 +78,33 @@ WORKDIR /home/nonroot
 ARG GIT_VERSION=local
 ARG BUILD_TAG
 ARG ADDITIONAL_RUSTFLAGS=""
+ENV CARGO_FEATURES="default"
 
 # 3. Build cargo dependencies. Note that this step doesn't depend on anything else than
 # `recipe.json`, so the layer can be reused as long as none of the dependencies change.
 COPY --from=plan     /home/nonroot/recipe.json                              recipe.json
-RUN set -e \
+RUN --mount=type=secret,uid=1000,id=SUBZERO_ACCESS_TOKEN \
+    set -e \
+    && if [ -s /run/secrets/SUBZERO_ACCESS_TOKEN ]; then \
+        export CARGO_NET_GIT_FETCH_WITH_CLI=true && \
+        git config --global url."https://$(cat /run/secrets/SUBZERO_ACCESS_TOKEN)@github.com/neondatabase/subzero".insteadOf "https://github.com/neondatabase/subzero"; \
+    fi \
     && RUSTFLAGS="-Clinker=clang -Clink-arg=-fuse-ld=mold -Clink-arg=-Wl,--no-rosegment -Cforce-frame-pointers=yes ${ADDITIONAL_RUSTFLAGS}" cargo chef cook --locked --release --recipe-path recipe.json
 
 # Perform the main build. We reuse the Postgres build artifacts from the intermediate 'pg-build'
 # layer, and the cargo dependencies built in the previous step.
 COPY --chown=nonroot --from=pg-build /home/nonroot/pg_install/ pg_install
 COPY --chown=nonroot . .
+COPY --chown=nonroot --from=plan     /home/nonroot/proxy/Cargo.toml         proxy/Cargo.toml
+COPY --chown=nonroot --from=plan     /home/nonroot/Cargo.lock               Cargo.lock
 
-RUN set -e \
-    && RUSTFLAGS="-Clinker=clang -Clink-arg=-fuse-ld=mold -Clink-arg=-Wl,--no-rosegment -Cforce-frame-pointers=yes ${ADDITIONAL_RUSTFLAGS}" cargo build \
+RUN  --mount=type=secret,uid=1000,id=SUBZERO_ACCESS_TOKEN \
+    set -e \
+    && if [ -s /run/secrets/SUBZERO_ACCESS_TOKEN ]; then \
+        export CARGO_FEATURES="rest_broker"; \
+    fi \
+    && RUSTFLAGS="-Clinker=clang -Clink-arg=-fuse-ld=mold -Clink-arg=-Wl,--no-rosegment -Cforce-frame-pointers=yes ${ADDITIONAL_RUSTFLAGS}" cargo auditable build \
+      --features $CARGO_FEATURES \
       --bin pg_sni_router  \
       --bin pageserver  \
       --bin pagectl  \

Makefile

@@ -2,7 +2,7 @@ ROOT_PROJECT_DIR := $(dir $(abspath $(lastword $(MAKEFILE_LIST))))
 
 # Where to install Postgres, default is ./pg_install, maybe useful for package
 # managers.
-POSTGRES_INSTALL_DIR ?= $(ROOT_PROJECT_DIR)/pg_install/
+POSTGRES_INSTALL_DIR ?= $(ROOT_PROJECT_DIR)/pg_install
 
 # Supported PostgreSQL versions
 POSTGRES_VERSIONS = v17 v16 v15 v14
@@ -14,7 +14,7 @@ POSTGRES_VERSIONS = v17 v16 v15 v14
 # it is derived from BUILD_TYPE.
 
 # All intermediate build artifacts are stored here.
-BUILD_DIR := build
+BUILD_DIR := $(ROOT_PROJECT_DIR)/build
 
 ICU_PREFIX_DIR := /usr/local/icu
 
@@ -212,7 +212,7 @@ neon-pgindent: postgres-v17-pg-bsd-indent neon-pg-ext-v17
 		FIND_TYPEDEF=$(ROOT_PROJECT_DIR)/vendor/postgres-v17/src/tools/find_typedef \
 		INDENT=$(BUILD_DIR)/v17/src/tools/pg_bsd_indent/pg_bsd_indent \
 		PGINDENT_SCRIPT=$(ROOT_PROJECT_DIR)/vendor/postgres-v17/src/tools/pgindent/pgindent \
-		-C $(BUILD_DIR)/neon-v17 \
+		-C $(BUILD_DIR)/pgxn-v17/neon \
 		-f $(ROOT_PROJECT_DIR)/pgxn/neon/Makefile pgindent
 
 
@@ -220,11 +220,15 @@ neon-pgindent: postgres-v17-pg-bsd-indent neon-pg-ext-v17
 setup-pre-commit-hook:
 	ln -s -f $(ROOT_PROJECT_DIR)/pre-commit.py .git/hooks/pre-commit
 
+build-tools/node_modules: build-tools/package.json
+	cd build-tools && $(if $(CI),npm ci,npm install)
+	touch build-tools/node_modules
+
 .PHONY: lint-openapi-spec
-lint-openapi-spec:
+lint-openapi-spec: build-tools/node_modules
 	# operation-2xx-response: pageserver timeline delete returns 404 on success
 	find . -iname "openapi_spec.y*ml" -exec\
-		docker run --rm -v ${PWD}:/spec ghcr.io/redocly/cli:1.34.4\
+		npx --prefix=build-tools/ redocly\
 			--skip-rule=operation-operationId --skip-rule=operation-summary --extends=minimal\
 			--skip-rule=no-server-example.com --skip-rule=operation-2xx-response\
 			lint {} \+

build-tools/Dockerfile

@@ -35,17 +35,17 @@ RUN echo 'Acquire::Retries "5";' > /etc/apt/apt.conf.d/80-retries && \
     echo -e "retry_connrefused=on\ntimeout=15\ntries=5\nretry-on-host-error=on\n" > /root/.wgetrc && \
     echo -e "--retry-connrefused\n--connect-timeout 15\n--retry 5\n--max-time 300\n" > /root/.curlrc
 
-COPY build_tools/patches/pgcopydbv017.patch /pgcopydbv017.patch
+COPY build-tools/patches/pgcopydbv017.patch /pgcopydbv017.patch
 
 RUN if [ "${DEBIAN_VERSION}" = "bookworm" ]; then \
         set -e && \
-        apt update && \
-        apt install -y --no-install-recommends \
+        apt-get update && \
+        apt-get install -y --no-install-recommends \
         ca-certificates wget gpg && \
         wget -qO - https://www.postgresql.org/media/keys/ACCC4CF8.asc | gpg --dearmor -o /usr/share/keyrings/postgresql-keyring.gpg && \
         echo "deb [signed-by=/usr/share/keyrings/postgresql-keyring.gpg] http://apt.postgresql.org/pub/repos/apt bookworm-pgdg main" > /etc/apt/sources.list.d/pgdg.list && \
         apt-get update && \
-        apt install -y --no-install-recommends \
+        apt-get install -y --no-install-recommends \
         build-essential \
         autotools-dev \
         libedit-dev \
@@ -89,8 +89,7 @@ RUN useradd -ms /bin/bash nonroot -b /home
 # Use strict mode for bash to catch errors early
 SHELL ["/bin/bash", "-euo", "pipefail", "-c"]
 
-RUN mkdir -p /pgcopydb/bin && \
-    mkdir -p /pgcopydb/lib && \
+RUN mkdir -p /pgcopydb/{bin,lib} && \    
     chmod -R 755 /pgcopydb && \
     chown -R nonroot:nonroot /pgcopydb
 
@@ -106,8 +105,8 @@ RUN echo 'Acquire::Retries "5";' > /etc/apt/apt.conf.d/80-retries && \
 # 'gdb' is included so that we get backtraces of core dumps produced in
 # regression tests
 RUN set -e \
-    && apt update \
-    && apt install -y \
+    && apt-get update \
+    && apt-get install -y --no-install-recommends \
         autoconf \
         automake \
         bison \
@@ -183,16 +182,22 @@ RUN curl -sL "https://github.com/peak/s5cmd/releases/download/v${S5CMD_VERSION}/
 ENV LLVM_VERSION=20
 RUN curl -fsSL 'https://apt.llvm.org/llvm-snapshot.gpg.key' | apt-key add - \
     && echo "deb http://apt.llvm.org/${DEBIAN_VERSION}/ llvm-toolchain-${DEBIAN_VERSION}-${LLVM_VERSION} main" > /etc/apt/sources.list.d/llvm.stable.list \
-    && apt update \
-    && apt install -y clang-${LLVM_VERSION} llvm-${LLVM_VERSION} \
+    && apt-get update \
+    && apt-get install -y --no-install-recommends clang-${LLVM_VERSION} llvm-${LLVM_VERSION} \
     && bash -c 'for f in /usr/bin/clang*-${LLVM_VERSION} /usr/bin/llvm*-${LLVM_VERSION}; do ln -s "${f}" "${f%-${LLVM_VERSION}}"; done' \
     && rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/*
 
+# Install node
+ENV NODE_VERSION=24
+RUN curl -fsSL https://deb.nodesource.com/setup_${NODE_VERSION}.x | bash - \
+    && apt-get install -y --no-install-recommends nodejs \
+    && rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/*
+
 # Install docker
 RUN curl -fsSL https://download.docker.com/linux/ubuntu/gpg | gpg --dearmor -o /usr/share/keyrings/docker-archive-keyring.gpg \
     && echo "deb [arch=$(dpkg --print-architecture) signed-by=/usr/share/keyrings/docker-archive-keyring.gpg] https://download.docker.com/linux/debian ${DEBIAN_VERSION} stable" > /etc/apt/sources.list.d/docker.list \
-    && apt update \
-    && apt install -y docker-ce docker-ce-cli \
+    && apt-get update \
+    && apt-get install -y --no-install-recommends docker-ce docker-ce-cli \
     && rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/*
 
 # Configure sudo & docker
@@ -209,12 +214,11 @@ RUN curl "https://awscli.amazonaws.com/awscli-exe-linux-$(uname -m).zip" -o "aws
 # Mold: A Modern Linker
 ENV MOLD_VERSION=v2.37.1
 RUN set -e \
-    && git clone https://github.com/rui314/mold.git \
+    && git clone -b "${MOLD_VERSION}" --depth 1 https://github.com/rui314/mold.git \
     && mkdir mold/build \
-    && cd mold/build \
-    && git checkout ${MOLD_VERSION} \
+    && cd mold/build \    
     && cmake -DCMAKE_BUILD_TYPE=Release -DCMAKE_CXX_COMPILER=clang++ .. \
-    && cmake --build . -j $(nproc) \
+    && cmake --build . -j "$(nproc)" \
     && cmake --install . \
     && cd .. \
     && rm -rf mold
@@ -248,7 +252,7 @@ ENV ICU_VERSION=67.1
 ENV ICU_PREFIX=/usr/local/icu
 
 # Download and build static ICU
-RUN wget -O /tmp/libicu-${ICU_VERSION}.tgz https://github.com/unicode-org/icu/releases/download/release-${ICU_VERSION//./-}/icu4c-${ICU_VERSION//./_}-src.tgz && \
+RUN wget -O "/tmp/libicu-${ICU_VERSION}.tgz" https://github.com/unicode-org/icu/releases/download/release-${ICU_VERSION//./-}/icu4c-${ICU_VERSION//./_}-src.tgz && \
     echo "94a80cd6f251a53bd2a997f6f1b5ac6653fe791dfab66e1eb0227740fb86d5dc /tmp/libicu-${ICU_VERSION}.tgz" | sha256sum --check && \
     mkdir /tmp/icu && \
     pushd /tmp/icu && \
@@ -259,8 +263,7 @@ RUN wget -O /tmp/libicu-${ICU_VERSION}.tgz https://github.com/unicode-org/icu/re
     make install && \
     popd && \
     rm -rf icu && \
-    rm -f /tmp/libicu-${ICU_VERSION}.tgz && \
-    popd
+    rm -f /tmp/libicu-${ICU_VERSION}.tgz
 
 # Switch to nonroot user
 USER nonroot:nonroot
@@ -273,19 +276,19 @@ ENV PYTHON_VERSION=3.11.12 \
     PYENV_ROOT=/home/nonroot/.pyenv \
     PATH=/home/nonroot/.pyenv/shims:/home/nonroot/.pyenv/bin:/home/nonroot/.poetry/bin:$PATH
 RUN set -e \
-    && cd $HOME \
+    && cd "$HOME" \
     && curl -sSO https://raw.githubusercontent.com/pyenv/pyenv-installer/master/bin/pyenv-installer \
     && chmod +x pyenv-installer \
     && ./pyenv-installer \
     && export PYENV_ROOT=/home/nonroot/.pyenv \
     && export PATH="$PYENV_ROOT/bin:$PATH" \
     && export PATH="$PYENV_ROOT/shims:$PATH" \
-    && pyenv install ${PYTHON_VERSION} \
-    && pyenv global ${PYTHON_VERSION} \
+    && pyenv install "${PYTHON_VERSION}" \
+    && pyenv global "${PYTHON_VERSION}" \
     && python --version \
-    && pip install --upgrade pip \
+    && pip install --no-cache-dir --upgrade pip \
     && pip --version \
-    && pip install pipenv wheel poetry
+    && pip install --no-cache-dir pipenv wheel poetry
 
 # Switch to nonroot user (again)
 USER nonroot:nonroot
@@ -296,6 +299,7 @@ WORKDIR /home/nonroot
 ENV RUSTC_VERSION=1.88.0
 ENV RUSTUP_HOME="/home/nonroot/.rustup"
 ENV PATH="/home/nonroot/.cargo/bin:${PATH}"
+ARG CARGO_AUDITABLE_VERSION=0.7.0
 ARG RUSTFILT_VERSION=0.2.1
 ARG CARGO_HAKARI_VERSION=0.9.36
 ARG CARGO_DENY_VERSION=0.18.2
@@ -311,14 +315,16 @@ RUN curl -sSO https://static.rust-lang.org/rustup/dist/$(uname -m)-unknown-linux
     . "$HOME/.cargo/env" && \
     cargo --version && rustup --version && \
     rustup component add llvm-tools rustfmt clippy && \
-    cargo install rustfilt            --version ${RUSTFILT_VERSION} --locked && \
-    cargo install cargo-hakari        --version ${CARGO_HAKARI_VERSION} --locked && \
-    cargo install cargo-deny          --version ${CARGO_DENY_VERSION} --locked && \
-    cargo install cargo-hack          --version ${CARGO_HACK_VERSION} --locked && \
-    cargo install cargo-nextest       --version ${CARGO_NEXTEST_VERSION} --locked && \
-    cargo install cargo-chef          --version ${CARGO_CHEF_VERSION} --locked && \
-    cargo install diesel_cli          --version ${CARGO_DIESEL_CLI_VERSION} --locked \
-                                      --features postgres-bundled --no-default-features && \
+    cargo install cargo-auditable           --locked --version "${CARGO_AUDITABLE_VERSION}" && \
+    cargo auditable install cargo-auditable --locked --version "${CARGO_AUDITABLE_VERSION}" --force && \
+    cargo auditable install rustfilt                 --version "${RUSTFILT_VERSION}" && \
+    cargo auditable install cargo-hakari    --locked --version "${CARGO_HAKARI_VERSION}" && \
+    cargo auditable install cargo-deny      --locked --version "${CARGO_DENY_VERSION}" && \
+    cargo auditable install cargo-hack      --locked --version "${CARGO_HACK_VERSION}" && \
+    cargo auditable install cargo-nextest   --locked --version "${CARGO_NEXTEST_VERSION}" && \
+    cargo auditable install cargo-chef      --locked --version "${CARGO_CHEF_VERSION}" && \
+    cargo auditable install diesel_cli      --locked --version "${CARGO_DIESEL_CLI_VERSION}" \
+                                            --features postgres-bundled --no-default-features && \
     rm -rf /home/nonroot/.cargo/registry && \
     rm -rf /home/nonroot/.cargo/git
 

build-tools/package.json

@@ -0,0 +1,8 @@
+{
+  "name": "build-tools",
+  "private": true,
+  "devDependencies": {
+    "@redocly/cli": "1.34.5",
+    "@sourcemeta/jsonschema": "10.0.0"
+  }
+}

compute/Makefile

@@ -50,9 +50,9 @@ jsonnetfmt-format:
 	jsonnetfmt --in-place $(jsonnet_files)
 
 .PHONY: manifest-schema-validation
-manifest-schema-validation: node_modules
-	node_modules/.bin/jsonschema validate -d https://json-schema.org/draft/2020-12/schema manifest.schema.json manifest.yaml
+manifest-schema-validation: ../build-tools/node_modules
+	npx --prefix=../build-tools/ jsonschema validate -d https://json-schema.org/draft/2020-12/schema manifest.schema.json manifest.yaml
 
-node_modules: package.json
-	npm install
-	touch node_modules
+../build-tools/node_modules: ../build-tools/package.json
+	cd ../build-tools && $(if $(CI),npm ci,npm install)
+	touch ../build-tools/node_modules

compute/compute-node.Dockerfile

@@ -9,7 +9,7 @@
 #
 # build-tools:   This contains Rust compiler toolchain and other tools needed at compile
 #                time. This is also used for the storage builds. This image is defined in
-#                build-tools.Dockerfile.
+#                build-tools/Dockerfile.
 #
 # build-deps:    Contains C compiler, other build tools, and compile-time dependencies
 #                needed to compile PostgreSQL and most extensions. (Some extensions need
@@ -115,7 +115,7 @@ ARG EXTENSIONS=all
 FROM $BASE_IMAGE_SHA AS build-deps
 ARG DEBIAN_VERSION
 
-# Keep in sync with build-tools.Dockerfile
+# Keep in sync with build-tools/Dockerfile
 ENV PROTOC_VERSION=25.1
 
 # Use strict mode for bash to catch errors early
@@ -133,7 +133,7 @@ RUN case $DEBIAN_VERSION in \
       # Install newer version (3.25) from backports.
       # libstdc++-10-dev is required for plv8
       bullseye) \
-        echo "deb http://deb.debian.org/debian bullseye-backports main" > /etc/apt/sources.list.d/bullseye-backports.list; \
+        echo "deb http://archive.debian.org/debian bullseye-backports main" > /etc/apt/sources.list.d/bullseye-backports.list; \
         VERSION_INSTALLS="cmake/bullseye-backports cmake-data/bullseye-backports libstdc++-10-dev"; \
       ;; \
       # Version-specific installs for Bookworm (PG17):
@@ -170,7 +170,29 @@ RUN case $DEBIAN_VERSION in \
 FROM build-deps AS pg-build
 ARG PG_VERSION
 COPY vendor/postgres-${PG_VERSION:?} postgres
+COPY compute/patches/postgres_fdw.patch .
+COPY compute/patches/pg_stat_statements_pg14-16.patch .
+COPY compute/patches/pg_stat_statements_pg17.patch .
 RUN cd postgres && \
+    # Apply patches to some contrib extensions
+    # For example, we need to grant EXECUTE on pg_stat_statements_reset() to {privileged_role_name}.
+    # In vanilla Postgres this function is limited to Postgres role superuser.
+    # In Neon we have {privileged_role_name} role that is not a superuser but replaces superuser in some cases.
+    # We could add the additional grant statements to the Postgres repository but it would be hard to maintain,
+    # whenever we need to pick up a new Postgres version and we want to limit the changes in our Postgres fork,
+    # so we do it here.
+    case "${PG_VERSION}" in \
+    "v14" | "v15" | "v16") \
+    patch -p1 < /pg_stat_statements_pg14-16.patch; \
+    ;; \
+    "v17") \
+    patch -p1 < /pg_stat_statements_pg17.patch; \
+    ;; \
+    *) \
+    # To do not forget to migrate patches to the next major version
+    echo "No contrib patches for this PostgreSQL version" && exit 1;; \
+    esac && \
+    patch -p1 < /postgres_fdw.patch && \
     export CONFIGURE_CMD="./configure CFLAGS='-O2 -g3 -fsigned-char' --enable-debug --with-openssl --with-uuid=ossp \
     --with-icu --with-libxml --with-libxslt --with-lz4" && \
     if [ "${PG_VERSION:?}" != "v14" ]; then \
@@ -184,8 +206,6 @@ RUN cd postgres && \
     echo 'trusted = true' >> /usr/local/pgsql/share/extension/autoinc.control && \
     echo 'trusted = true' >> /usr/local/pgsql/share/extension/dblink.control && \
     echo 'trusted = true' >> /usr/local/pgsql/share/extension/postgres_fdw.control && \
-    file=/usr/local/pgsql/share/extension/postgres_fdw--1.0.sql && [ -e $file ] && \
-    echo 'GRANT USAGE ON FOREIGN DATA WRAPPER postgres_fdw TO neon_superuser;' >> $file && \
     echo 'trusted = true' >> /usr/local/pgsql/share/extension/bloom.control && \
     echo 'trusted = true' >> /usr/local/pgsql/share/extension/earthdistance.control && \
     echo 'trusted = true' >> /usr/local/pgsql/share/extension/insert_username.control && \
@@ -195,34 +215,7 @@ RUN cd postgres && \
     echo 'trusted = true' >> /usr/local/pgsql/share/extension/pgrowlocks.control && \
     echo 'trusted = true' >> /usr/local/pgsql/share/extension/pgstattuple.control && \
     echo 'trusted = true' >> /usr/local/pgsql/share/extension/refint.control && \
-    echo 'trusted = true' >> /usr/local/pgsql/share/extension/xml2.control && \
-    # We need to grant EXECUTE on pg_stat_statements_reset() to neon_superuser.
-    # In vanilla postgres this function is limited to Postgres role superuser.
-    # In neon we have neon_superuser role that is not a superuser but replaces superuser in some cases.
-    # We could add the additional grant statements to the postgres repository but it would be hard to maintain,
-    # whenever we need to pick up a new postgres version and we want to limit the changes in our postgres fork,
-    # so we do it here.
-    for file in /usr/local/pgsql/share/extension/pg_stat_statements--*.sql; do \
-        filename=$(basename "$file"); \
-        # Note that there are no downgrade scripts for pg_stat_statements, so we \
-        # don't have to modify any downgrade paths or (much) older versions: we only \
-        # have to make sure every creation of the pg_stat_statements_reset function \
-        # also adds execute permissions to the neon_superuser.
-        case $filename in \
-          pg_stat_statements--1.4.sql) \
-            # pg_stat_statements_reset is first created with 1.4
-            echo 'GRANT EXECUTE ON FUNCTION pg_stat_statements_reset() TO neon_superuser;' >> $file; \
-            ;; \
-          pg_stat_statements--1.6--1.7.sql) \
-            # Then with the 1.6-1.7 migration it is re-created with a new signature, thus add the permissions back
-            echo 'GRANT EXECUTE ON FUNCTION pg_stat_statements_reset(Oid, Oid, bigint) TO neon_superuser;' >> $file; \
-            ;; \
-          pg_stat_statements--1.10--1.11.sql) \
-            # Then with the 1.10-1.11 migration it is re-created with a new signature again, thus add the permissions back
-            echo 'GRANT EXECUTE ON FUNCTION pg_stat_statements_reset(Oid, Oid, bigint, boolean) TO neon_superuser;' >> $file; \
-            ;; \
-        esac; \
-    done;
+    echo 'trusted = true' >> /usr/local/pgsql/share/extension/xml2.control
 
 # Set PATH for all the subsequent build steps
 ENV PATH="/usr/local/pgsql/bin:$PATH"
@@ -1524,7 +1517,7 @@ WORKDIR /ext-src
 COPY compute/patches/pg_duckdb_v031.patch .
 COPY compute/patches/duckdb_v120.patch .
 # pg_duckdb build requires source dir to be a git repo to get submodules
-# allow neon_superuser to execute some functions that in pg_duckdb are available to superuser only:
+# allow {privileged_role_name} to execute some functions that in pg_duckdb are available to superuser only:
 # - extension management function duckdb.install_extension()
 # - access to duckdb.extensions table and its sequence
 RUN git clone --depth 1 --branch v0.3.1 https://github.com/duckdb/pg_duckdb.git pg_duckdb-src && \
@@ -1790,7 +1783,7 @@ RUN set -e \
 #########################################################################################
 FROM build-deps AS exporters
 ARG TARGETARCH
-# Keep sql_exporter version same as in build-tools.Dockerfile and
+# Keep sql_exporter version same as in build-tools/Dockerfile and
 # test_runner/regress/test_compute_metrics.py
 # See comment on the top of the file regading `echo`, `-e` and `\n`
 RUN if [ "$TARGETARCH" = "amd64" ]; then\

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/package.json

@@ -1,7 +0,0 @@
-{
-  "name": "neon-compute",
-  "private": true,
-  "dependencies": {
-    "@sourcemeta/jsonschema": "9.3.4"
-  }
-} 

compute/patches/anon_v2.patch

@@ -1,22 +1,26 @@
 diff --git a/sql/anon.sql b/sql/anon.sql
-index 0cdc769..b450327 100644
+index 0cdc769..5eab1d6 100644
 --- a/sql/anon.sql
 +++ b/sql/anon.sql
-@@ -1141,3 +1141,15 @@ $$
+@@ -1141,3 +1141,19 @@ $$
  -- TODO : https://en.wikipedia.org/wiki/L-diversity
  
  -- TODO : https://en.wikipedia.org/wiki/T-closeness
 +
 +-- NEON Patches
 +
-+GRANT ALL ON SCHEMA anon to neon_superuser;
-+GRANT ALL ON ALL TABLES IN SCHEMA anon TO neon_superuser;
-+
 +DO $$
++DECLARE
++  privileged_role_name text;
 +BEGIN
-+    IF current_setting('server_version_num')::int >= 150000 THEN
-+        GRANT SET ON PARAMETER anon.transparent_dynamic_masking TO neon_superuser;
-+    END IF;
++  privileged_role_name := current_setting('neon.privileged_role_name');
++
++  EXECUTE format('GRANT ALL ON SCHEMA anon to %I', privileged_role_name);
++  EXECUTE format('GRANT ALL ON ALL TABLES IN SCHEMA anon TO %I', privileged_role_name);
++
++  IF current_setting('server_version_num')::int >= 150000 THEN
++    EXECUTE format('GRANT SET ON PARAMETER anon.transparent_dynamic_masking TO %I', privileged_role_name);
++  END IF;
 +END $$;
 diff --git a/sql/init.sql b/sql/init.sql
 index 7da6553..9b6164b 100644

compute/patches/pg_duckdb_v031.patch

@@ -21,13 +21,21 @@ index 3235cc8..6b892bc 100644
  include Makefile.global
  
 diff --git a/sql/pg_duckdb--0.2.0--0.3.0.sql b/sql/pg_duckdb--0.2.0--0.3.0.sql
-index d777d76..af60106 100644
+index d777d76..3b54396 100644
 --- a/sql/pg_duckdb--0.2.0--0.3.0.sql
 +++ b/sql/pg_duckdb--0.2.0--0.3.0.sql
-@@ -1056,3 +1056,6 @@ GRANT ALL ON FUNCTION duckdb.cache(TEXT, TEXT) TO PUBLIC;
+@@ -1056,3 +1056,14 @@ GRANT ALL ON FUNCTION duckdb.cache(TEXT, TEXT) TO PUBLIC;
  GRANT ALL ON FUNCTION duckdb.cache_info() TO PUBLIC;
  GRANT ALL ON FUNCTION duckdb.cache_delete(TEXT) TO PUBLIC;
  GRANT ALL ON PROCEDURE duckdb.recycle_ddb() TO PUBLIC;
-+GRANT ALL ON FUNCTION duckdb.install_extension(TEXT) TO neon_superuser;
-+GRANT ALL ON TABLE duckdb.extensions TO neon_superuser;
-+GRANT ALL ON SEQUENCE duckdb.extensions_table_seq TO neon_superuser;
++
++DO $$
++DECLARE
++  privileged_role_name text;
++BEGIN
++  privileged_role_name := current_setting('neon.privileged_role_name');
++
++  EXECUTE format('GRANT ALL ON FUNCTION duckdb.install_extension(TEXT) TO %I', privileged_role_name);
++  EXECUTE format('GRANT ALL ON TABLE duckdb.extensions TO %I', privileged_role_name);
++  EXECUTE format('GRANT ALL ON SEQUENCE duckdb.extensions_table_seq TO %I', privileged_role_name);
++END $$;

compute/patches/pg_repack.patch

@@ -1,5 +1,11 @@
+commit 5eb393810cf7c7bafa4e394dad2e349e2a8cb2cb
+Author: Alexey Masterov <alexey.masterov@databricks.com>
+Date:   Mon Jul 28 18:11:02 2025 +0200
+
+    Patch for pg_repack
+
 diff --git a/regress/Makefile b/regress/Makefile
-index bf6edcb..89b4c7f 100644
+index bf6edcb..110e734 100644
 --- a/regress/Makefile
 +++ b/regress/Makefile
 @@ -17,7 +17,7 @@ INTVERSION := $(shell echo $$(($$(echo $(VERSION).0 | sed 's/\([[:digit:]]\{1,\}
@@ -7,18 +13,36 @@ index bf6edcb..89b4c7f 100644
  #
  
 -REGRESS := init-extension repack-setup repack-run error-on-invalid-idx no-error-on-invalid-idx after-schema repack-check nosuper tablespace get_order_by trigger
-+REGRESS := init-extension repack-setup repack-run error-on-invalid-idx no-error-on-invalid-idx after-schema repack-check nosuper get_order_by trigger
++REGRESS := init-extension noautovacuum repack-setup repack-run error-on-invalid-idx no-error-on-invalid-idx after-schema repack-check nosuper get_order_by trigger autovacuum
  
  USE_PGXS = 1	# use pgxs if not in contrib directory
  PGXS := $(shell $(PG_CONFIG) --pgxs)
-diff --git a/regress/expected/init-extension.out b/regress/expected/init-extension.out
-index 9f2e171..f6e4f8d 100644
---- a/regress/expected/init-extension.out
-+++ b/regress/expected/init-extension.out
-@@ -1,3 +1,2 @@
- SET client_min_messages = warning;
- CREATE EXTENSION pg_repack;
--RESET client_min_messages;
+diff --git a/regress/expected/autovacuum.out b/regress/expected/autovacuum.out
+new file mode 100644
+index 0000000..e7f2363
+--- /dev/null
++++ b/regress/expected/autovacuum.out
+@@ -0,0 +1,7 @@
++ALTER SYSTEM SET autovacuum='on';
++SELECT pg_reload_conf();
++ pg_reload_conf 
++----------------
++ t
++(1 row)
++
+diff --git a/regress/expected/noautovacuum.out b/regress/expected/noautovacuum.out
+new file mode 100644
+index 0000000..fc7978e
+--- /dev/null
++++ b/regress/expected/noautovacuum.out
+@@ -0,0 +1,7 @@
++ALTER SYSTEM SET autovacuum='off';
++SELECT pg_reload_conf();
++ pg_reload_conf 
++----------------
++ t
++(1 row)
++
 diff --git a/regress/expected/nosuper.out b/regress/expected/nosuper.out
 index 8d0a94e..63b68bf 100644
 --- a/regress/expected/nosuper.out
@@ -50,14 +74,22 @@ index 8d0a94e..63b68bf 100644
  INFO: repacking table "public.tbl_cluster"
  ERROR: query failed: ERROR:  current transaction is aborted, commands ignored until end of transaction block
  DETAIL: query was: RESET lock_timeout
-diff --git a/regress/sql/init-extension.sql b/regress/sql/init-extension.sql
-index 9f2e171..f6e4f8d 100644
---- a/regress/sql/init-extension.sql
-+++ b/regress/sql/init-extension.sql
-@@ -1,3 +1,2 @@
- SET client_min_messages = warning;
- CREATE EXTENSION pg_repack;
--RESET client_min_messages;
+diff --git a/regress/sql/autovacuum.sql b/regress/sql/autovacuum.sql
+new file mode 100644
+index 0000000..a8eda63
+--- /dev/null
++++ b/regress/sql/autovacuum.sql
+@@ -0,0 +1,2 @@
++ALTER SYSTEM SET autovacuum='on';
++SELECT pg_reload_conf();
+diff --git a/regress/sql/noautovacuum.sql b/regress/sql/noautovacuum.sql
+new file mode 100644
+index 0000000..13d4836
+--- /dev/null
++++ b/regress/sql/noautovacuum.sql
+@@ -0,0 +1,2 @@
++ALTER SYSTEM SET autovacuum='off';
++SELECT pg_reload_conf();
 diff --git a/regress/sql/nosuper.sql b/regress/sql/nosuper.sql
 index 072f0fa..dbe60f8 100644
 --- a/regress/sql/nosuper.sql

compute/patches/pg_stat_statements_pg14-16.patch

@@ -0,0 +1,34 @@
+diff --git a/contrib/pg_stat_statements/pg_stat_statements--1.4.sql b/contrib/pg_stat_statements/pg_stat_statements--1.4.sql
+index 58cdf600fce..8be57a996f6 100644
+--- a/contrib/pg_stat_statements/pg_stat_statements--1.4.sql
++++ b/contrib/pg_stat_statements/pg_stat_statements--1.4.sql
+@@ -46,3 +46,12 @@ GRANT SELECT ON pg_stat_statements TO PUBLIC;
+ 
+ -- Don't want this to be available to non-superusers.
+ REVOKE ALL ON FUNCTION pg_stat_statements_reset() FROM PUBLIC;
++
++DO $$
++DECLARE
++  privileged_role_name text;
++BEGIN
++  privileged_role_name := current_setting('neon.privileged_role_name');
++
++  EXECUTE format('GRANT EXECUTE ON FUNCTION pg_stat_statements_reset() TO %I', privileged_role_name);
++END $$;
+diff --git a/contrib/pg_stat_statements/pg_stat_statements--1.6--1.7.sql b/contrib/pg_stat_statements/pg_stat_statements--1.6--1.7.sql
+index 6fc3fed4c93..256345a8f79 100644
+--- a/contrib/pg_stat_statements/pg_stat_statements--1.6--1.7.sql
++++ b/contrib/pg_stat_statements/pg_stat_statements--1.6--1.7.sql
+@@ -20,3 +20,12 @@ LANGUAGE C STRICT PARALLEL SAFE;
+ 
+ -- Don't want this to be available to non-superusers.
+ REVOKE ALL ON FUNCTION pg_stat_statements_reset(Oid, Oid, bigint) FROM PUBLIC;
++
++DO $$
++DECLARE
++  privileged_role_name text;
++BEGIN
++  privileged_role_name := current_setting('neon.privileged_role_name');
++
++  EXECUTE format('GRANT EXECUTE ON FUNCTION pg_stat_statements_reset(Oid, Oid, bigint) TO %I', privileged_role_name);
++END $$;

compute/patches/pg_stat_statements_pg17.patch

@@ -0,0 +1,52 @@
+diff --git a/contrib/pg_stat_statements/pg_stat_statements--1.10--1.11.sql b/contrib/pg_stat_statements/pg_stat_statements--1.10--1.11.sql
+index 0bb2c397711..32764db1d8b 100644
+--- a/contrib/pg_stat_statements/pg_stat_statements--1.10--1.11.sql
++++ b/contrib/pg_stat_statements/pg_stat_statements--1.10--1.11.sql
+@@ -80,3 +80,12 @@ LANGUAGE C STRICT PARALLEL SAFE;
+ 
+ -- Don't want this to be available to non-superusers.
+ REVOKE ALL ON FUNCTION pg_stat_statements_reset(Oid, Oid, bigint, boolean) FROM PUBLIC;
++
++DO $$
++DECLARE
++  privileged_role_name text;
++BEGIN
++  privileged_role_name := current_setting('neon.privileged_role_name');
++
++  EXECUTE format('GRANT EXECUTE ON FUNCTION pg_stat_statements_reset(Oid, Oid, bigint, boolean) TO %I', privileged_role_name);
++END $$;
+\ No newline at end of file
+diff --git a/contrib/pg_stat_statements/pg_stat_statements--1.4.sql b/contrib/pg_stat_statements/pg_stat_statements--1.4.sql
+index 58cdf600fce..8be57a996f6 100644
+--- a/contrib/pg_stat_statements/pg_stat_statements--1.4.sql
++++ b/contrib/pg_stat_statements/pg_stat_statements--1.4.sql
+@@ -46,3 +46,12 @@ GRANT SELECT ON pg_stat_statements TO PUBLIC;
+ 
+ -- Don't want this to be available to non-superusers.
+ REVOKE ALL ON FUNCTION pg_stat_statements_reset() FROM PUBLIC;
++
++DO $$
++DECLARE
++  privileged_role_name text;
++BEGIN
++  privileged_role_name := current_setting('neon.privileged_role_name');
++
++  EXECUTE format('GRANT EXECUTE ON FUNCTION pg_stat_statements_reset() TO %I', privileged_role_name);
++END $$;
+diff --git a/contrib/pg_stat_statements/pg_stat_statements--1.6--1.7.sql b/contrib/pg_stat_statements/pg_stat_statements--1.6--1.7.sql
+index 6fc3fed4c93..256345a8f79 100644
+--- a/contrib/pg_stat_statements/pg_stat_statements--1.6--1.7.sql
++++ b/contrib/pg_stat_statements/pg_stat_statements--1.6--1.7.sql
+@@ -20,3 +20,12 @@ LANGUAGE C STRICT PARALLEL SAFE;
+ 
+ -- Don't want this to be available to non-superusers.
+ REVOKE ALL ON FUNCTION pg_stat_statements_reset(Oid, Oid, bigint) FROM PUBLIC;
++
++DO $$
++DECLARE
++  privileged_role_name text;
++BEGIN
++  privileged_role_name := current_setting('neon.privileged_role_name');
++
++  EXECUTE format('GRANT EXECUTE ON FUNCTION pg_stat_statements_reset(Oid, Oid, bigint) TO %I', privileged_role_name);
++END $$;

compute/patches/postgres_fdw.patch

@@ -0,0 +1,17 @@
+diff --git a/contrib/postgres_fdw/postgres_fdw--1.0.sql b/contrib/postgres_fdw/postgres_fdw--1.0.sql
+index a0f0fc1bf45..ee077f2eea6 100644
+--- a/contrib/postgres_fdw/postgres_fdw--1.0.sql
++++ b/contrib/postgres_fdw/postgres_fdw--1.0.sql
+@@ -16,3 +16,12 @@ LANGUAGE C STRICT;
+ CREATE FOREIGN DATA WRAPPER postgres_fdw
+   HANDLER postgres_fdw_handler
+   VALIDATOR postgres_fdw_validator;
++
++DO $$
++DECLARE
++  privileged_role_name text;
++BEGIN
++  privileged_role_name := current_setting('neon.privileged_role_name');
++
++  EXECUTE format('GRANT USAGE ON FOREIGN DATA WRAPPER postgres_fdw TO %I', privileged_role_name);
++END $$;

compute/vm-image-spec-bookworm.yaml

@@ -26,7 +26,13 @@ commands:
   - name: postgres-exporter
     user: nobody
     sysvInitAction: respawn
-    shell: 'DATA_SOURCE_NAME="user=cloud_admin sslmode=disable dbname=postgres application_name=postgres-exporter pgaudit.log=none" /bin/postgres_exporter --config.file=/etc/postgres_exporter.yml'
+    # Turn off database collector (`--no-collector.database`), we don't use `pg_database_size_bytes` metric anyway, see
+    # https://github.com/neondatabase/flux-fleet/blob/5e19b3fd897667b70d9a7ad4aa06df0ca22b49ff/apps/base/compute-metrics/scrape-compute-pg-exporter-neon.yaml#L29
+    # but it's enabled by default and it doesn't filter out invalid databases, see
+    # https://github.com/prometheus-community/postgres_exporter/blob/06a553c8166512c9d9c5ccf257b0f9bba8751dbc/collector/pg_database.go#L67
+    # so if it hits one, it starts spamming logs
+    #   ERROR:  [NEON_SMGR] [reqid d9700000018] could not read db size of db 705302 from page server at lsn 5/A2457EB0
+    shell: 'DATA_SOURCE_NAME="user=cloud_admin sslmode=disable dbname=postgres application_name=postgres-exporter pgaudit.log=none" /bin/postgres_exporter --no-collector.database --config.file=/etc/postgres_exporter.yml'
   - name: pgbouncer-exporter
     user: postgres
     sysvInitAction: respawn

compute/vm-image-spec-bullseye.yaml

@@ -26,7 +26,13 @@ commands:
   - name: postgres-exporter
     user: nobody
     sysvInitAction: respawn
-    shell: 'DATA_SOURCE_NAME="user=cloud_admin sslmode=disable dbname=postgres application_name=postgres-exporter pgaudit.log=none" /bin/postgres_exporter --config.file=/etc/postgres_exporter.yml'
+    # Turn off database collector (`--no-collector.database`), we don't use `pg_database_size_bytes` metric anyway, see
+    # https://github.com/neondatabase/flux-fleet/blob/5e19b3fd897667b70d9a7ad4aa06df0ca22b49ff/apps/base/compute-metrics/scrape-compute-pg-exporter-neon.yaml#L29
+    # but it's enabled by default and it doesn't filter out invalid databases, see
+    # https://github.com/prometheus-community/postgres_exporter/blob/06a553c8166512c9d9c5ccf257b0f9bba8751dbc/collector/pg_database.go#L67
+    # so if it hits one, it starts spamming logs
+    #   ERROR:  [NEON_SMGR] [reqid d9700000018] could not read db size of db 705302 from page server at lsn 5/A2457EB0
+    shell: 'DATA_SOURCE_NAME="user=cloud_admin sslmode=disable dbname=postgres application_name=postgres-exporter pgaudit.log=none" /bin/postgres_exporter --no-collector.database --config.file=/etc/postgres_exporter.yml'
   - name: pgbouncer-exporter
     user: postgres
     sysvInitAction: respawn

compute_tools/Cargo.toml

@@ -7,7 +7,7 @@ license.workspace = true
 [features]
 default = []
 # Enables test specific features.
-testing = ["neon_failpoint/testing"]
+testing = ["fail/failpoints"]
 
 [dependencies]
 async-compression.workspace = true
@@ -23,11 +23,14 @@ camino.workspace = true
 chrono.workspace = true
 cfg-if.workspace = true
 clap.workspace = true
-neon_failpoint.workspace = true
+fail.workspace = true
 flate2.workspace = true
 futures.workspace = true
 http.workspace = true
+http-body-util.workspace = true
 hostname-validator = "1.1"
+hyper.workspace = true
+hyper-util.workspace = true
 indexmap.workspace = true
 itertools.workspace = true
 jsonwebtoken.workspace = true
@@ -44,6 +47,7 @@ postgres.workspace = true
 regex.workspace = true
 reqwest = { workspace = true, features = ["json"] }
 ring = "0.17"
+scopeguard.workspace = true
 serde.workspace = true
 serde_with.workspace = true
 serde_json.workspace = true
@@ -58,6 +62,7 @@ tokio-stream.workspace = true
 tonic.workspace = true
 tower-otel.workspace = true
 tracing.workspace = true
+tracing-appender.workspace = true
 tracing-opentelemetry.workspace = true
 tracing-subscriber.workspace = true
 tracing-utils.workspace = true

compute_tools/README.md

@@ -52,8 +52,14 @@ stateDiagram-v2
   Init --> Running : Started Postgres
   Running --> TerminationPendingFast : Requested termination
   Running --> TerminationPendingImmediate : Requested termination
+  Running --> ConfigurationPending : Received a /configure request with spec
+  Running --> RefreshConfigurationPending : Received a /refresh_configuration request, compute node will pull a new spec and reconfigure
+  RefreshConfigurationPending --> RefreshConfiguration: Received compute spec and started configuration
+  RefreshConfiguration --> Running : Compute has been re-configured
+  RefreshConfiguration --> RefreshConfigurationPending : Configuration failed and to be retried
   TerminationPendingFast --> Terminated compute with 30s delay for cplane to inspect status
   TerminationPendingImmediate --> Terminated : Terminated compute immediately
+  Failed --> RefreshConfigurationPending : Received a /refresh_configuration request
   Failed --> [*] : Compute exited
   Terminated --> [*] : Compute exited
 ```

compute_tools/src/bin/compute_ctl.rs

@@ -49,9 +49,10 @@ use compute_tools::compute::{
     BUILD_TAG, ComputeNode, ComputeNodeParams, forward_termination_signal,
 };
 use compute_tools::extension_server::get_pg_version_string;
-use compute_tools::logger::*;
 use compute_tools::params::*;
+use compute_tools::pg_isready::get_pg_isready_bin;
 use compute_tools::spec::*;
+use compute_tools::{hadron_metrics, installed_extensions, logger::*};
 use rlimit::{Resource, setrlimit};
 use signal_hook::consts::{SIGINT, SIGQUIT, SIGTERM};
 use signal_hook::iterator::Signals;
@@ -81,12 +82,29 @@ struct Cli {
     #[arg(long, default_value_t = 3081)]
     pub internal_http_port: u16,
 
+    /// Backwards-compatible --http-port for Hadron deployments. Functionally the
+    /// same as --external-http-port.
+    #[arg(
+        long,
+        conflicts_with = "external_http_port",
+        conflicts_with = "internal_http_port"
+    )]
+    pub http_port: Option<u16>,
+
     #[arg(short = 'D', long, value_name = "DATADIR")]
     pub pgdata: String,
 
     #[arg(short = 'C', long, value_name = "DATABASE_URL")]
     pub connstr: String,
 
+    #[arg(
+        long,
+        default_value = "neon_superuser",
+        value_name = "PRIVILEGED_ROLE_NAME",
+        value_parser = Self::parse_privileged_role_name
+    )]
+    pub privileged_role_name: String,
+
     #[cfg(target_os = "linux")]
     #[arg(long, default_value = "neon-postgres")]
     pub cgroup: String,
@@ -130,6 +148,12 @@ struct Cli {
     /// Run in development mode, skipping VM-specific operations like process termination
     #[arg(long, action = clap::ArgAction::SetTrue)]
     pub dev: bool,
+
+    #[arg(long)]
+    pub pg_init_timeout: Option<u64>,
+
+    #[arg(long, default_value_t = false, action = clap::ArgAction::Set)]
+    pub lakebase_mode: bool,
 }
 
 impl Cli {
@@ -149,12 +173,47 @@ impl Cli {
 
         Ok(url)
     }
+
+    /// For simplicity, we do not escape `privileged_role_name` anywhere in the code.
+    /// Since it's a system role, which we fully control, that's fine. Still, let's
+    /// validate it to avoid any surprises.
+    fn parse_privileged_role_name(value: &str) -> Result<String> {
+        use regex::Regex;
+
+        let pattern = Regex::new(r"^[a-z_]+$").unwrap();
+
+        if !pattern.is_match(value) {
+            bail!("--privileged-role-name can only contain lowercase letters and underscores")
+        }
+
+        Ok(value.to_string())
+    }
+}
+
+// Hadron helpers to get compatible compute_ctl http ports from Cli. The old `--http-port`
+// arg is used and acts the same as `--external-http-port`. The internal http port is defined
+// to be http_port + 1. Hadron runs in the dblet environment which uses the host network, so
+// we need to be careful with the ports to choose.
+fn get_external_http_port(cli: &Cli) -> u16 {
+    if cli.lakebase_mode {
+        return cli.http_port.unwrap_or(cli.external_http_port);
+    }
+    cli.external_http_port
+}
+fn get_internal_http_port(cli: &Cli) -> u16 {
+    if cli.lakebase_mode {
+        return cli
+            .http_port
+            .map(|p| p + 1)
+            .unwrap_or(cli.internal_http_port);
+    }
+    cli.internal_http_port
 }
 
 fn main() -> Result<()> {
     let cli = Cli::parse();
 
-    failpoint_support::init().unwrap();
+    let scenario = failpoint_support::init();
 
     // For historical reasons, the main thread that processes the config and launches postgres
     // is synchronous, but we always have this tokio runtime available and we "enter" it so
@@ -165,24 +224,38 @@ fn main() -> Result<()> {
         .build()?;
     let _rt_guard = runtime.enter();
 
-    runtime.block_on(init(cli.dev))?;
+    let mut log_dir = None;
+    if cli.lakebase_mode {
+        log_dir = std::env::var("COMPUTE_CTL_LOG_DIRECTORY").ok();
+    }
+
+    let (tracing_provider, _file_logs_guard) = init(cli.dev, log_dir)?;
 
     // enable core dumping for all child processes
     setrlimit(Resource::CORE, rlimit::INFINITY, rlimit::INFINITY)?;
 
+    if cli.lakebase_mode {
+        installed_extensions::initialize_metrics();
+        hadron_metrics::initialize_metrics();
+    }
+
     let connstr = Url::parse(&cli.connstr).context("cannot parse connstr as a URL")?;
 
     let config = get_config(&cli)?;
 
+    let external_http_port = get_external_http_port(&cli);
+    let internal_http_port = get_internal_http_port(&cli);
+
     let compute_node = ComputeNode::new(
         ComputeNodeParams {
             compute_id: cli.compute_id,
             connstr,
+            privileged_role_name: cli.privileged_role_name.clone(),
             pgdata: cli.pgdata.clone(),
             pgbin: cli.pgbin.clone(),
             pgversion: get_pg_version_string(&cli.pgbin),
-            external_http_port: cli.external_http_port,
-            internal_http_port: cli.internal_http_port,
+            external_http_port,
+            internal_http_port,
             remote_ext_base_url: cli.remote_ext_base_url.clone(),
             resize_swap_on_bind: cli.resize_swap_on_bind,
             set_disk_quota_for_fs: cli.set_disk_quota_for_fs,
@@ -195,17 +268,32 @@ fn main() -> Result<()> {
             installed_extensions_collection_interval: Arc::new(AtomicU64::new(
                 cli.installed_extensions_collection_interval,
             )),
+            pg_init_timeout: cli.pg_init_timeout.map(Duration::from_secs),
+            pg_isready_bin: get_pg_isready_bin(&cli.pgbin),
+            instance_id: std::env::var("INSTANCE_ID").ok(),
+            lakebase_mode: cli.lakebase_mode,
+            build_tag: BUILD_TAG.to_string(),
+            control_plane_uri: cli.control_plane_uri,
+            config_path_test_only: cli.config,
         },
         config,
     )?;
 
-    let exit_code = compute_node.run()?;
+    let exit_code = compute_node.run().context("running compute node")?;
 
-    deinit_and_exit(exit_code);
+    scenario.teardown();
+
+    deinit_and_exit(tracing_provider, exit_code);
 }
 
-async fn init(dev_mode: bool) -> Result<()> {
-    init_tracing_and_logging(DEFAULT_LOG_LEVEL).await?;
+fn init(
+    dev_mode: bool,
+    log_dir: Option<String>,
+) -> Result<(
+    Option<tracing_utils::Provider>,
+    Option<tracing_appender::non_blocking::WorkerGuard>,
+)> {
+    let (provider, file_logs_guard) = init_tracing_and_logging(DEFAULT_LOG_LEVEL, &log_dir)?;
 
     let mut signals = Signals::new([SIGINT, SIGTERM, SIGQUIT])?;
     thread::spawn(move || {
@@ -216,7 +304,7 @@ async fn init(dev_mode: bool) -> Result<()> {
 
     info!("compute build_tag: {}", &BUILD_TAG.to_string());
 
-    Ok(())
+    Ok((provider, file_logs_guard))
 }
 
 fn get_config(cli: &Cli) -> Result<ComputeConfig> {
@@ -241,25 +329,27 @@ fn get_config(cli: &Cli) -> Result<ComputeConfig> {
     }
 }
 
-fn deinit_and_exit(exit_code: Option<i32>) -> ! {
-    // Shutdown trace pipeline gracefully, so that it has a chance to send any
-    // pending traces before we exit. Shutting down OTEL tracing provider may
-    // hang for quite some time, see, for example:
-    // - https://github.com/open-telemetry/opentelemetry-rust/issues/868
-    // - and our problems with staging https://github.com/neondatabase/cloud/issues/3707#issuecomment-1493983636
-    //
-    // Yet, we want computes to shut down fast enough, as we may need a new one
-    // for the same timeline ASAP. So wait no longer than 2s for the shutdown to
-    // complete, then just error out and exit the main thread.
-    info!("shutting down tracing");
-    let (sender, receiver) = mpsc::channel();
-    let _ = thread::spawn(move || {
-        tracing_utils::shutdown_tracing();
-        sender.send(()).ok()
-    });
-    let shutdown_res = receiver.recv_timeout(Duration::from_millis(2000));
-    if shutdown_res.is_err() {
-        error!("timed out while shutting down tracing, exiting anyway");
+fn deinit_and_exit(tracing_provider: Option<tracing_utils::Provider>, exit_code: Option<i32>) -> ! {
+    if let Some(p) = tracing_provider {
+        // Shutdown trace pipeline gracefully, so that it has a chance to send any
+        // pending traces before we exit. Shutting down OTEL tracing provider may
+        // hang for quite some time, see, for example:
+        // - https://github.com/open-telemetry/opentelemetry-rust/issues/868
+        // - and our problems with staging https://github.com/neondatabase/cloud/issues/3707#issuecomment-1493983636
+        //
+        // Yet, we want computes to shut down fast enough, as we may need a new one
+        // for the same timeline ASAP. So wait no longer than 2s for the shutdown to
+        // complete, then just error out and exit the main thread.
+        info!("shutting down tracing");
+        let (sender, receiver) = mpsc::channel();
+        let _ = thread::spawn(move || {
+            _ = p.shutdown();
+            sender.send(()).ok()
+        });
+        let shutdown_res = receiver.recv_timeout(Duration::from_millis(2000));
+        if shutdown_res.is_err() {
+            error!("timed out while shutting down tracing, exiting anyway");
+        }
     }
 
     info!("shutting down");
@@ -325,4 +415,49 @@ mod test {
         ])
         .expect_err("URL parameters are not allowed");
     }
+
+    #[test]
+    fn verify_privileged_role_name() {
+        // Valid name
+        let cli = Cli::parse_from([
+            "compute_ctl",
+            "--pgdata=test",
+            "--connstr=test",
+            "--compute-id=test",
+            "--privileged-role-name",
+            "my_superuser",
+        ]);
+        assert_eq!(cli.privileged_role_name, "my_superuser");
+
+        // Invalid names
+        Cli::try_parse_from([
+            "compute_ctl",
+            "--pgdata=test",
+            "--connstr=test",
+            "--compute-id=test",
+            "--privileged-role-name",
+            "NeonSuperuser",
+        ])
+        .expect_err("uppercase letters are not allowed");
+
+        Cli::try_parse_from([
+            "compute_ctl",
+            "--pgdata=test",
+            "--connstr=test",
+            "--compute-id=test",
+            "--privileged-role-name",
+            "$'neon_superuser",
+        ])
+        .expect_err("special characters are not allowed");
+
+        Cli::try_parse_from([
+            "compute_ctl",
+            "--pgdata=test",
+            "--connstr=test",
+            "--compute-id=test",
+            "--privileged-role-name",
+            "",
+        ])
+        .expect_err("empty name is not allowed");
+    }
 }

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/communicator_socket_client.rs

@@ -0,0 +1,98 @@
+//! Client for making request to a running Postgres server's communicator control socket.
+//!
+//! The storage communicator process that runs inside Postgres exposes an HTTP endpoint in
+//! a Unix Domain Socket in the Postgres data directory. This provides access to it.
+
+use std::path::Path;
+
+use anyhow::Context;
+use hyper::client::conn::http1::SendRequest;
+use hyper_util::rt::TokioIo;
+
+/// Name of the socket within the Postgres data directory. This better match that in
+/// `pgxn/neon/communicator/src/lib.rs`.
+const NEON_COMMUNICATOR_SOCKET_NAME: &str = "neon-communicator.socket";
+
+/// Open a connection to the communicator's control socket, prepare to send requests to it
+/// with hyper.
+pub async fn connect_communicator_socket<B>(pgdata: &Path) -> anyhow::Result<SendRequest<B>>
+where
+    B: hyper::body::Body + 'static + Send,
+    B::Data: Send,
+    B::Error: Into<Box<dyn std::error::Error + Send + Sync>>,
+{
+    let socket_path = pgdata.join(NEON_COMMUNICATOR_SOCKET_NAME);
+    let socket_path_len = socket_path.display().to_string().len();
+
+    // There is a limit of around 100 bytes (108 on Linux?) on the length of the path to a
+    // Unix Domain socket. The limit is on the connect(2) function used to open the
+    // socket, not on the absolute path itself. Postgres changes the current directory to
+    // the data directory and uses a relative path to bind to the socket, and the relative
+    // path "./neon-communicator.socket" is always short, but when compute_ctl needs to
+    // open the socket, we need to use a full path, which can be arbitrarily long.
+    //
+    // There are a few ways we could work around this:
+    //
+    // 1. Change the current directory to the Postgres data directory and use a relative
+    //    path in the connect(2) call. That's problematic because the current directory
+    //    applies to the whole process. We could change the current directory early in
+    //    compute_ctl startup, and that might be a good idea anyway for other reasons too:
+    //    it would be more robust if the data directory is moved around or unlinked for
+    //    some reason, and you would be less likely to accidentally litter other parts of
+    //    the filesystem with e.g. temporary files. However, that's a pretty invasive
+    //    change.
+    //
+    // 2. On Linux, you could open() the data directory, and refer to the the socket
+    //    inside it as "/proc/self/fd/<fd>/neon-communicator.socket". But that's
+    //    Linux-only.
+    //
+    // 3. Create a symbolic link to the socket with a shorter path, and use that.
+    //
+    // We use the symbolic link approach here. Hopefully the paths we use in production
+    // are shorter, so that we can open the socket directly, so that this hack is needed
+    // only in development.
+    let connect_result = if socket_path_len < 100 {
+        // We can open the path directly with no hacks.
+        tokio::net::UnixStream::connect(socket_path).await
+    } else {
+        // The path to the socket is too long. Create a symlink to it with a shorter path.
+        let short_path = std::env::temp_dir().join(format!(
+            "compute_ctl.short-socket.{}.{}",
+            std::process::id(),
+            tokio::task::id()
+        ));
+        std::os::unix::fs::symlink(&socket_path, &short_path)?;
+
+        // Delete the symlink as soon as we have connected to it. There's a small chance
+        // of leaking if the process dies before we remove it, so try to keep that window
+        // as small as possible.
+        scopeguard::defer! {
+            if let Err(err) = std::fs::remove_file(&short_path) {
+                tracing::warn!("could not remove symlink \"{}\" created for socket: {}",
+                               short_path.display(), err);
+            }
+        }
+
+        tracing::info!(
+            "created symlink \"{}\" for socket \"{}\", opening it now",
+            short_path.display(),
+            socket_path.display()
+        );
+
+        tokio::net::UnixStream::connect(&short_path).await
+    };
+
+    let stream = connect_result.context("connecting to communicator control socket")?;
+
+    let io = TokioIo::new(stream);
+    let (request_sender, connection) = hyper::client::conn::http1::handshake(io).await?;
+
+    // spawn a task to poll the connection and drive the HTTP state
+    tokio::spawn(async move {
+        if let Err(err) = connection.await {
+            eprintln!("Error in connection: {err}");
+        }
+    });
+
+    Ok(request_sender)
+}

compute_tools/src/compute.rs

@@ -6,7 +6,8 @@ use compute_api::responses::{
     LfcPrewarmState, PromoteState, TlsConfig,
 };
 use compute_api::spec::{
-    ComputeAudit, ComputeFeature, ComputeMode, ComputeSpec, ExtVersion, PageserverProtocol, PgIdent,
+    ComputeAudit, ComputeFeature, ComputeMode, ComputeSpec, ExtVersion, GenericOption,
+    PageserverConnectionInfo, PageserverProtocol, PgIdent, Role,
 };
 use futures::StreamExt;
 use futures::future::join_all;
@@ -21,6 +22,7 @@ use postgres::NoTls;
 use postgres::error::SqlState;
 use remote_storage::{DownloadError, RemotePath};
 use std::collections::{HashMap, HashSet};
+use std::ffi::OsString;
 use std::os::unix::fs::{PermissionsExt, symlink};
 use std::path::Path;
 use std::process::{Command, Stdio};
@@ -30,18 +32,23 @@ 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::backoff::{
+    DEFAULT_BASE_BACKOFF_SECONDS, DEFAULT_MAX_BACKOFF_SECONDS, exponential_backoff_duration,
+};
 use utils::id::{TenantId, TimelineId};
 use utils::lsn::Lsn;
 use utils::measured_stream::MeasuredReader;
 use utils::pid_file;
-use utils::shard::{ShardCount, ShardIndex, ShardNumber};
+use utils::shard::{ShardIndex, ShardNumber, ShardStripeSize};
 
 use crate::configurator::launch_configurator;
 use crate::disk_quota::set_disk_quota;
+use crate::hadron_metrics::COMPUTE_ATTACHED;
 use crate::installed_extensions::get_installed_extensions;
-use crate::logger::startup_context_from_env;
+use crate::logger::{self, startup_context_from_env};
 use crate::lsn_lease::launch_lsn_lease_bg_task_for_static;
 use crate::metrics::COMPUTE_CTL_UP;
 use crate::monitor::launch_monitor;
@@ -74,12 +81,20 @@ const DEFAULT_INSTALLED_EXTENSIONS_COLLECTION_INTERVAL: u64 = 3600;
 
 /// Static configuration params that don't change after startup. These mostly
 /// come from the CLI args, or are derived from them.
+#[derive(Clone, Debug)]
 pub struct ComputeNodeParams {
     /// The ID of the compute
     pub compute_id: String,
-    // Url type maintains proper escaping
+
+    /// Url type maintains proper escaping
     pub connstr: url::Url,
 
+    /// The name of the 'weak' superuser role, which we give to the users.
+    /// It follows the allow list approach, i.e., we take a standard role
+    /// and grant it extra permissions with explicit GRANTs here and there,
+    /// and core patches.
+    pub privileged_role_name: String,
+
     pub resize_swap_on_bind: bool,
     pub set_disk_quota_for_fs: Option<String>,
 
@@ -105,6 +120,17 @@ pub struct ComputeNodeParams {
 
     /// Interval for installed extensions collection
     pub installed_extensions_collection_interval: Arc<AtomicU64>,
+    /// Hadron instance ID of the compute node.
+    pub instance_id: Option<String>,
+    /// Timeout of PG compute startup in the Init state.
+    pub pg_init_timeout: Option<Duration>,
+    // Path to the `pg_isready` binary.
+    pub pg_isready_bin: String,
+    pub lakebase_mode: bool,
+
+    pub build_tag: String,
+    pub control_plane_uri: Option<String>,
+    pub config_path_test_only: Option<OsString>,
 }
 
 type TaskHandle = Mutex<Option<JoinHandle<()>>>;
@@ -146,6 +172,7 @@ pub struct RemoteExtensionMetrics {
 #[derive(Clone, Debug)]
 pub struct ComputeState {
     pub start_time: DateTime<Utc>,
+    pub pg_start_time: Option<DateTime<Utc>>,
     pub status: ComputeStatus,
     /// Timestamp of the last Postgres activity. It could be `None` if
     /// compute wasn't used since start.
@@ -169,6 +196,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
@@ -183,6 +211,7 @@ impl ComputeState {
     pub fn new() -> Self {
         Self {
             start_time: Utc::now(),
+            pg_start_time: None,
             status: ComputeStatus::Empty,
             last_active: None,
             error: None,
@@ -193,6 +222,7 @@ impl ComputeState {
             lfc_offload_state: LfcOffloadState::default(),
             terminate_flush_lsn: None,
             promote_state: None,
+            lfc_prewarm_token: CancellationToken::new(),
         }
     }
 
@@ -225,7 +255,7 @@ pub struct ParsedSpec {
     pub spec: ComputeSpec,
     pub tenant_id: TenantId,
     pub timeline_id: TimelineId,
-    pub pageserver_connstr: String,
+    pub pageserver_conninfo: PageserverConnectionInfo,
     pub safekeeper_connstrings: Vec<String>,
     pub storage_auth_token: Option<String>,
     /// k8s dns name and port
@@ -273,25 +303,47 @@ impl ParsedSpec {
 }
 
 impl TryFrom<ComputeSpec> for ParsedSpec {
-    type Error = String;
-    fn try_from(spec: ComputeSpec) -> Result<Self, String> {
+    type Error = anyhow::Error;
+    fn try_from(spec: ComputeSpec) -> Result<Self, anyhow::Error> {
         // Extract the options from the spec file that are needed to connect to
         // the storage system.
         //
-        // For backwards-compatibility, the top-level fields in the spec file
-        // may be empty. In that case, we need to dig them from the GUCs in the
-        // cluster.settings field.
-        let pageserver_connstr = spec
-            .pageserver_connstring
-            .clone()
-            .or_else(|| spec.cluster.settings.find("neon.pageserver_connstring"))
-            .ok_or("pageserver connstr should be provided")?;
+        // In compute specs generated by old control plane versions, the spec file might
+        // be missing the `pageserver_connection_info` field. In that case, we need to dig
+        // the pageserver connection info from the `pageserver_connstr` field instead, or
+        // if that's missing too, from the GUC in the cluster.settings field.
+        let mut pageserver_conninfo = spec.pageserver_connection_info.clone();
+        if pageserver_conninfo.is_none() {
+            if let Some(pageserver_connstr_field) = &spec.pageserver_connstring {
+                pageserver_conninfo = Some(PageserverConnectionInfo::from_connstr(
+                    pageserver_connstr_field,
+                    spec.shard_stripe_size,
+                )?);
+            }
+        }
+        if pageserver_conninfo.is_none() {
+            if let Some(guc) = spec.cluster.settings.find("neon.pageserver_connstring") {
+                let stripe_size = if let Some(guc) = spec.cluster.settings.find("neon.stripe_size")
+                {
+                    Some(ShardStripeSize(u32::from_str(&guc)?))
+                } else {
+                    None
+                };
+                pageserver_conninfo =
+                    Some(PageserverConnectionInfo::from_connstr(&guc, stripe_size)?);
+            }
+        }
+        let pageserver_conninfo = pageserver_conninfo.ok_or(anyhow::anyhow!(
+            "pageserver connection information should be provided"
+        ))?;
+
+        // Similarly for safekeeper connection strings
         let safekeeper_connstrings = if spec.safekeeper_connstrings.is_empty() {
             if matches!(spec.mode, ComputeMode::Primary) {
                 spec.cluster
                     .settings
                     .find("neon.safekeepers")
-                    .ok_or("safekeeper connstrings should be provided")?
+                    .ok_or(anyhow::anyhow!("safekeeper connstrings should be provided"))?
                     .split(',')
                     .map(|str| str.to_string())
                     .collect()
@@ -306,22 +358,22 @@ impl TryFrom<ComputeSpec> for ParsedSpec {
         let tenant_id: TenantId = if let Some(tenant_id) = spec.tenant_id {
             tenant_id
         } else {
-            spec.cluster
+            let guc = spec
+                .cluster
                 .settings
                 .find("neon.tenant_id")
-                .ok_or("tenant id should be provided")
-                .map(|s| TenantId::from_str(&s))?
-                .or(Err("invalid tenant id"))?
+                .ok_or(anyhow::anyhow!("tenant id should be provided"))?;
+            TenantId::from_str(&guc).context("invalid tenant id")?
         };
         let timeline_id: TimelineId = if let Some(timeline_id) = spec.timeline_id {
             timeline_id
         } else {
-            spec.cluster
+            let guc = spec
+                .cluster
                 .settings
                 .find("neon.timeline_id")
-                .ok_or("timeline id should be provided")
-                .map(|s| TimelineId::from_str(&s))?
-                .or(Err("invalid timeline id"))?
+                .ok_or(anyhow::anyhow!("timeline id should be provided"))?;
+            TimelineId::from_str(&guc).context(anyhow::anyhow!("invalid timeline id"))?
         };
 
         let endpoint_storage_addr: Option<String> = spec
@@ -335,7 +387,7 @@ impl TryFrom<ComputeSpec> for ParsedSpec {
 
         let res = ParsedSpec {
             spec,
-            pageserver_connstr,
+            pageserver_conninfo,
             safekeeper_connstrings,
             storage_auth_token,
             tenant_id,
@@ -345,7 +397,7 @@ impl TryFrom<ComputeSpec> for ParsedSpec {
         };
 
         // Now check validity of the parsed specification
-        res.validate()?;
+        res.validate().map_err(anyhow::Error::msg)?;
         Ok(res)
     }
 }
@@ -390,6 +442,130 @@ struct StartVmMonitorResult {
     vm_monitor: Option<JoinHandle<Result<()>>>,
 }
 
+// BEGIN_HADRON
+/// This function creates roles that are used by Databricks.
+/// These roles are not needs to be botostrapped at PG Compute provisioning time.
+/// The auth method for these roles are configured in databricks_pg_hba.conf in universe repository.
+pub(crate) fn create_databricks_roles() -> Vec<String> {
+    let roles = vec![
+        // Role for prometheus_stats_exporter
+        Role {
+            name: "databricks_monitor".to_string(),
+            // This uses "local" connection and auth method for that is "trust", so no password is needed.
+            encrypted_password: None,
+            options: Some(vec![GenericOption {
+                name: "IN ROLE pg_monitor".to_string(),
+                value: None,
+                vartype: "string".to_string(),
+            }]),
+        },
+        // Role for brickstore control plane
+        Role {
+            name: "databricks_control_plane".to_string(),
+            // Certificate user does not need password.
+            encrypted_password: None,
+            options: Some(vec![GenericOption {
+                name: "SUPERUSER".to_string(),
+                value: None,
+                vartype: "string".to_string(),
+            }]),
+        },
+        // Role for brickstore httpgateway.
+        Role {
+            name: "databricks_gateway".to_string(),
+            // Certificate user does not need password.
+            encrypted_password: None,
+            options: None,
+        },
+    ];
+
+    roles
+        .into_iter()
+        .map(|role| {
+            let query = format!(
+                r#"
+                DO $$
+                    BEGIN
+                        IF NOT EXISTS (
+                            SELECT FROM pg_catalog.pg_roles WHERE rolname = '{}')
+                        THEN
+                            CREATE ROLE {} {};
+                        END IF;
+                    END
+                $$;"#,
+                role.name,
+                role.name.pg_quote(),
+                role.to_pg_options(),
+            );
+            query
+        })
+        .collect()
+}
+
+/// Databricks-specific environment variables to be passed to the `postgres` sub-process.
+pub struct DatabricksEnvVars {
+    /// The Databricks "endpoint ID" of the compute instance. Used by `postgres` to check
+    /// the token scopes of internal auth tokens.
+    pub endpoint_id: String,
+    /// Hostname of the Databricks workspace URL this compute instance belongs to.
+    /// Used by postgres to verify Databricks PAT tokens.
+    pub workspace_host: String,
+
+    pub lakebase_mode: bool,
+}
+
+impl DatabricksEnvVars {
+    pub fn new(
+        compute_spec: &ComputeSpec,
+        compute_id: Option<&String>,
+        instance_id: Option<String>,
+        lakebase_mode: bool,
+    ) -> Self {
+        let endpoint_id = if let Some(instance_id) = instance_id {
+            // Use instance_id as endpoint_id if it is set. This code path is for PuPr model.
+            instance_id
+        } else {
+            // Use compute_id as endpoint_id if instance_id is not set. The code path is for PrPr model.
+            // compute_id is a string format of "{endpoint_id}/{compute_idx}"
+            // endpoint_id is a uuid. We only need to pass down endpoint_id to postgres.
+            // Panics if compute_id is not set or not in the expected format.
+            compute_id.unwrap().split('/').next().unwrap().to_string()
+        };
+        let workspace_host = compute_spec
+            .databricks_settings
+            .as_ref()
+            .map(|s| s.databricks_workspace_host.clone())
+            .unwrap_or("".to_string());
+        Self {
+            endpoint_id,
+            workspace_host,
+            lakebase_mode,
+        }
+    }
+
+    /// Constants for the names of Databricks-specific postgres environment variables.
+    const DATABRICKS_ENDPOINT_ID_ENVVAR: &'static str = "DATABRICKS_ENDPOINT_ID";
+    const DATABRICKS_WORKSPACE_HOST_ENVVAR: &'static str = "DATABRICKS_WORKSPACE_HOST";
+
+    /// Convert DatabricksEnvVars to a list of string pairs that can be passed as env vars. Consumes `self`.
+    pub fn to_env_var_list(self) -> Vec<(String, String)> {
+        if !self.lakebase_mode {
+            // In neon env, we don't need to pass down the env vars to postgres.
+            return vec![];
+        }
+        vec![
+            (
+                Self::DATABRICKS_ENDPOINT_ID_ENVVAR.to_string(),
+                self.endpoint_id.clone(),
+            ),
+            (
+                Self::DATABRICKS_WORKSPACE_HOST_ENVVAR.to_string(),
+                self.workspace_host.clone(),
+            ),
+        ]
+    }
+}
+
 impl ComputeNode {
     pub fn new(params: ComputeNodeParams, config: ComputeConfig) -> Result<Self> {
         let connstr = params.connstr.as_str();
@@ -413,7 +589,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
@@ -426,7 +602,11 @@ impl ComputeNode {
         let mut new_state = ComputeState::new();
         if let Some(spec) = config.spec {
             let pspec = ParsedSpec::try_from(spec).map_err(|msg| anyhow::anyhow!(msg))?;
-            new_state.pspec = Some(pspec);
+            if params.lakebase_mode {
+                ComputeNode::set_spec(&params, &mut new_state, pspec);
+            } else {
+                new_state.pspec = Some(pspec);
+            }
         }
 
         Ok(ComputeNode {
@@ -471,6 +651,7 @@ impl ComputeNode {
             port: this.params.external_http_port,
             config: this.compute_ctl_config.clone(),
             compute_id: this.params.compute_id.clone(),
+            instance_id: this.params.instance_id.clone(),
         }
         .launch(&this);
 
@@ -640,6 +821,9 @@ impl ComputeNode {
             };
             _this_entered = start_compute_span.enter();
 
+            // Hadron: Record postgres start time (used to enforce pg_init_timeout).
+            state_guard.pg_start_time.replace(Utc::now());
+
             state_guard.set_status(ComputeStatus::Init, &self.state_changed);
             compute_state = state_guard.clone()
         }
@@ -1020,7 +1204,14 @@ impl ComputeNode {
         // If it is something different then create_dir() will error out anyway.
         let pgdata = &self.params.pgdata;
         let _ok = fs::remove_dir_all(pgdata);
-        fs::create_dir(pgdata)?;
+        if self.params.lakebase_mode {
+            // Ignore creation errors if the directory already exists (e.g. mounting it ahead of time).
+            // If it is something different then PG startup will error out anyway.
+            let _ok = fs::create_dir(pgdata);
+        } else {
+            fs::create_dir(pgdata)?;
+        }
+
         fs::set_permissions(pgdata, fs::Permissions::from_mode(0o700))?;
 
         Ok(())
@@ -1032,14 +1223,14 @@ impl ComputeNode {
     fn try_get_basebackup(&self, compute_state: &ComputeState, lsn: Lsn) -> Result<()> {
         let spec = compute_state.pspec.as_ref().expect("spec must be set");
 
-        let shard0_connstr = spec.pageserver_connstr.split(',').next().unwrap();
         let started = Instant::now();
-
-        let (connected, size) = match PageserverProtocol::from_connstring(shard0_connstr)? {
-            PageserverProtocol::Libpq => self.try_get_basebackup_libpq(spec, lsn)?,
+        let (connected, size) = match spec.pageserver_conninfo.prefer_protocol {
             PageserverProtocol::Grpc => self.try_get_basebackup_grpc(spec, lsn)?,
+            PageserverProtocol::Libpq => self.try_get_basebackup_libpq(spec, lsn)?,
         };
 
+        self.fix_zenith_signal_neon_signal()?;
+
         let mut state = self.state.lock().unwrap();
         state.metrics.pageserver_connect_micros =
             connected.duration_since(started).as_micros() as u64;
@@ -1049,26 +1240,44 @@ impl ComputeNode {
         Ok(())
     }
 
+    /// Move the Zenith signal file to Neon signal file location.
+    /// This makes Compute compatible with older PageServers that don't yet
+    /// know about the Zenith->Neon rename.
+    fn fix_zenith_signal_neon_signal(&self) -> Result<()> {
+        let datadir = Path::new(&self.params.pgdata);
+
+        let neonsig = datadir.join("neon.signal");
+
+        if neonsig.is_file() {
+            return Ok(());
+        }
+
+        let zenithsig = datadir.join("zenith.signal");
+
+        if zenithsig.is_file() {
+            fs::copy(zenithsig, neonsig)?;
+        }
+
+        Ok(())
+    }
+
     /// Fetches a basebackup via gRPC. The connstring must use grpc://. Returns the timestamp when
     /// the connection was established, and the (compressed) size of the basebackup.
     fn try_get_basebackup_grpc(&self, spec: &ParsedSpec, lsn: Lsn) -> Result<(Instant, usize)> {
-        let shard0_connstr = spec
-            .pageserver_connstr
-            .split(',')
-            .next()
-            .unwrap()
-            .to_string();
-        let shard_index = match spec.pageserver_connstr.split(',').count() as u8 {
-            0 | 1 => ShardIndex::unsharded(),
-            count => ShardIndex::new(ShardNumber(0), ShardCount(count)),
+        let shard0_index = ShardIndex {
+            shard_number: ShardNumber(0),
+            shard_count: spec.pageserver_conninfo.shard_count,
         };
-
+        let shard0_url = spec
+            .pageserver_conninfo
+            .shard_url(ShardNumber(0), PageserverProtocol::Grpc)?
+            .to_owned();
         let (reader, connected) = tokio::runtime::Handle::current().block_on(async move {
             let mut client = page_api::Client::connect(
-                shard0_connstr,
+                shard0_url,
                 spec.tenant_id,
                 spec.timeline_id,
-                shard_index,
+                shard0_index,
                 spec.storage_auth_token.clone(),
                 None, // NB: base backups use payload compression
             )
@@ -1100,7 +1309,9 @@ impl ComputeNode {
     /// Fetches a basebackup via libpq. The connstring must use postgresql://. Returns the timestamp
     /// when the connection was established, and the (compressed) size of the basebackup.
     fn try_get_basebackup_libpq(&self, spec: &ParsedSpec, lsn: Lsn) -> Result<(Instant, usize)> {
-        let shard0_connstr = spec.pageserver_connstr.split(',').next().unwrap();
+        let shard0_connstr = spec
+            .pageserver_conninfo
+            .shard_url(ShardNumber(0), PageserverProtocol::Libpq)?;
         let mut config = postgres::Config::from_str(shard0_connstr)?;
 
         // Use the storage auth token from the config file, if given.
@@ -1187,10 +1398,7 @@ impl ComputeNode {
                     return result;
                 }
                 Err(ref e) if attempts < max_attempts => {
-                    warn!(
-                        "Failed to get basebackup: {} (attempt {}/{})",
-                        e, attempts, max_attempts
-                    );
+                    warn!("Failed to get basebackup: {e:?} (attempt {attempts}/{max_attempts})");
                     std::thread::sleep(std::time::Duration::from_millis(retry_period_ms as u64));
                     retry_period_ms *= 1.5;
                 }
@@ -1263,9 +1471,7 @@ impl ComputeNode {
 
         // In case of error, log and fail the check, but don't crash.
         // We're playing it safe because these errors could be transient
-        // and we don't yet retry. Also being careful here allows us to
-        // be backwards compatible with safekeepers that don't have the
-        // TIMELINE_STATUS API yet.
+        // and we don't yet retry.
         if responses.len() < quorum {
             error!(
                 "failed sync safekeepers check {:?} {:?} {:?}",
@@ -1354,6 +1560,41 @@ impl ComputeNode {
         Ok(lsn)
     }
 
+    fn sync_safekeepers_with_retries(&self, storage_auth_token: Option<String>) -> Result<Lsn> {
+        let max_retries = 5;
+        let mut attempts = 0;
+        loop {
+            let result = self.sync_safekeepers(storage_auth_token.clone());
+            match &result {
+                Ok(_) => {
+                    if attempts > 0 {
+                        tracing::info!("sync_safekeepers succeeded after {attempts} retries");
+                    }
+                    return result;
+                }
+                Err(e) if attempts < max_retries => {
+                    tracing::info!(
+                        "sync_safekeepers failed, will retry (attempt {attempts}): {e:#}"
+                    );
+                }
+                Err(err) => {
+                    tracing::warn!(
+                        "sync_safekeepers still failed after {attempts} retries, giving up: {err:?}"
+                    );
+                    return result;
+                }
+            }
+            // sleep and retry
+            let backoff = exponential_backoff_duration(
+                attempts,
+                DEFAULT_BASE_BACKOFF_SECONDS,
+                DEFAULT_MAX_BACKOFF_SECONDS,
+            );
+            std::thread::sleep(backoff);
+            attempts += 1;
+        }
+    }
+
     /// Do all the preparations like PGDATA directory creation, configuration,
     /// safekeepers sync, basebackup, etc.
     #[instrument(skip_all)]
@@ -1363,14 +1604,20 @@ impl ComputeNode {
         let pgdata_path = Path::new(&self.params.pgdata);
 
         let tls_config = self.tls_config(&pspec.spec);
+        let databricks_settings = spec.databricks_settings.as_ref();
+        let postgres_port = self.params.connstr.port();
 
         // Remove/create an empty pgdata directory and put configuration there.
         self.create_pgdata()?;
         config::write_postgres_conf(
             pgdata_path,
+            &self.params,
             &pspec.spec,
+            postgres_port,
             self.params.internal_http_port,
             tls_config,
+            databricks_settings,
+            self.params.lakebase_mode,
         )?;
 
         // Syncing safekeepers is only safe with primary nodes: if a primary
@@ -1383,7 +1630,7 @@ impl ComputeNode {
                     lsn
                 } else {
                     info!("starting safekeepers syncing");
-                    self.sync_safekeepers(pspec.storage_auth_token.clone())
+                    self.sync_safekeepers_with_retries(pspec.storage_auth_token.clone())
                         .with_context(|| "failed to sync safekeepers")?
                 };
                 info!("safekeepers synced at LSN {}", lsn);
@@ -1399,19 +1646,31 @@ impl ComputeNode {
             }
         };
 
-        info!(
-            "getting basebackup@{} from pageserver {}",
-            lsn, &pspec.pageserver_connstr
-        );
-        self.get_basebackup(compute_state, lsn).with_context(|| {
-            format!(
-                "failed to get basebackup@{} from pageserver {}",
-                lsn, &pspec.pageserver_connstr
-            )
-        })?;
+        self.get_basebackup(compute_state, lsn)
+            .with_context(|| format!("failed to get basebackup@{lsn}"))?;
 
-        // Update pg_hba.conf received with basebackup.
-        update_pg_hba(pgdata_path)?;
+        if let Some(settings) = databricks_settings {
+            copy_tls_certificates(
+                &settings.pg_compute_tls_settings.key_file,
+                &settings.pg_compute_tls_settings.cert_file,
+                pgdata_path,
+            )?;
+
+            // Update pg_hba.conf received with basebackup including additional databricks settings.
+            update_pg_hba(pgdata_path, Some(&settings.databricks_pg_hba))?;
+            update_pg_ident(pgdata_path, Some(&settings.databricks_pg_ident))?;
+        } else {
+            // Update pg_hba.conf received with basebackup.
+            update_pg_hba(pgdata_path, None)?;
+        }
+
+        if let Some(databricks_settings) = spec.databricks_settings.as_ref() {
+            copy_tls_certificates(
+                &databricks_settings.pg_compute_tls_settings.key_file,
+                &databricks_settings.pg_compute_tls_settings.cert_file,
+                pgdata_path,
+            )?;
+        }
 
         // Place pg_dynshmem under /dev/shm. This allows us to use
         // 'dynamic_shared_memory_type = mmap' so that the files are placed in
@@ -1452,7 +1711,7 @@ impl ComputeNode {
         // symlink doesn't affect anything.
         //
         // See https://github.com/neondatabase/autoscaling/issues/800
-        std::fs::remove_dir(pgdata_path.join("pg_dynshmem"))?;
+        std::fs::remove_dir_all(pgdata_path.join("pg_dynshmem"))?;
         symlink("/dev/shm/", pgdata_path.join("pg_dynshmem"))?;
 
         match spec.mode {
@@ -1467,6 +1726,12 @@ impl ComputeNode {
 
     /// Start and stop a postgres process to warm up the VM for startup.
     pub fn prewarm_postgres_vm_memory(&self) -> Result<()> {
+        if self.params.lakebase_mode {
+            // We are running in Hadron mode. Disabling this prewarming step for now as it could run
+            // into dblet port conflicts and also doesn't add much value with our current infra.
+            info!("Skipping postgres prewarming in Hadron mode");
+            return Ok(());
+        }
         info!("prewarming VM memory");
 
         // Create pgdata
@@ -1524,14 +1789,36 @@ impl ComputeNode {
     pub fn start_postgres(&self, storage_auth_token: Option<String>) -> Result<PostgresHandle> {
         let pgdata_path = Path::new(&self.params.pgdata);
 
+        let env_vars: Vec<(String, String)> = if self.params.lakebase_mode {
+            let databricks_env_vars = {
+                let state = self.state.lock().unwrap();
+                let spec = &state.pspec.as_ref().unwrap().spec;
+                DatabricksEnvVars::new(
+                    spec,
+                    Some(&self.params.compute_id),
+                    self.params.instance_id.clone(),
+                    self.params.lakebase_mode,
+                )
+            };
+
+            info!(
+                "Starting Postgres for databricks endpoint id: {}",
+                &databricks_env_vars.endpoint_id
+            );
+
+            let mut env_vars = databricks_env_vars.to_env_var_list();
+            env_vars.extend(storage_auth_token.map(|t| ("NEON_AUTH_TOKEN".to_string(), t)));
+            env_vars
+        } else if let Some(storage_auth_token) = &storage_auth_token {
+            vec![("NEON_AUTH_TOKEN".to_owned(), storage_auth_token.to_owned())]
+        } else {
+            vec![]
+        };
+
         // Run postgres as a child process.
         let mut pg = maybe_cgexec(&self.params.pgbin)
             .args(["-D", &self.params.pgdata])
-            .envs(if let Some(storage_auth_token) = &storage_auth_token {
-                vec![("NEON_AUTH_TOKEN", storage_auth_token)]
-            } else {
-                vec![]
-            })
+            .envs(env_vars)
             .stderr(Stdio::piped())
             .spawn()
             .expect("cannot start postgres process");
@@ -1638,7 +1925,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 =
@@ -1683,7 +1970,15 @@ impl ComputeNode {
     /// Do initial configuration of the already started Postgres.
     #[instrument(skip_all)]
     pub fn apply_config(&self, compute_state: &ComputeState) -> Result<()> {
-        let conf = self.get_tokio_conn_conf(Some("compute_ctl:apply_config"));
+        let mut conf = self.get_tokio_conn_conf(Some("compute_ctl:apply_config"));
+
+        if self.params.lakebase_mode {
+            // Set a 2-minute statement_timeout for the session applying config. The individual SQL statements
+            // used in apply_spec_sql() should not take long (they are just creating users and installing
+            // extensions). If any of them are stuck for an extended period of time it usually indicates a
+            // pageserver connectivity problem and we should bail out.
+            conf.options("-c statement_timeout=2min");
+        }
 
         let conf = Arc::new(conf);
         let spec = Arc::new(
@@ -1716,6 +2011,8 @@ impl ComputeNode {
         }
 
         // Run migrations separately to not hold up cold starts
+        let lakebase_mode = self.params.lakebase_mode;
+        let params = self.params.clone();
         tokio::spawn(async move {
             let mut conf = conf.as_ref().clone();
             conf.application_name("compute_ctl:migrations");
@@ -1727,7 +2024,7 @@ impl ComputeNode {
                             eprintln!("connection error: {e}");
                         }
                     });
-                    if let Err(e) = handle_migrations(&mut client).await {
+                    if let Err(e) = handle_migrations(params, &mut client, lakebase_mode).await {
                         error!("Failed to run migrations: {}", e);
                     }
                 }
@@ -1743,6 +2040,34 @@ impl ComputeNode {
         Ok::<(), anyhow::Error>(())
     }
 
+    // Signal to the configurator to refresh the configuration by pulling a new spec from the HCC.
+    // Note that this merely triggers a notification on a condition variable the configurator thread
+    // waits on. The configurator thread (in configurator.rs) pulls the new spec from the HCC and
+    // applies it.
+    pub async fn signal_refresh_configuration(&self) -> Result<()> {
+        let states_allowing_configuration_refresh = [
+            ComputeStatus::Running,
+            ComputeStatus::Failed,
+            ComputeStatus::RefreshConfigurationPending,
+        ];
+
+        let mut state = self.state.lock().expect("state lock poisoned");
+        if states_allowing_configuration_refresh.contains(&state.status) {
+            state.status = ComputeStatus::RefreshConfigurationPending;
+            self.state_changed.notify_all();
+            Ok(())
+        } else if state.status == ComputeStatus::Init {
+            // If the compute is in Init state, we can't refresh the configuration immediately,
+            // but we should be able to do that soon.
+            Ok(())
+        } else {
+            Err(anyhow::anyhow!(
+                "Cannot refresh compute configuration in state {:?}",
+                state.status
+            ))
+        }
+    }
+
     // Wrapped this around `pg_ctl reload`, but right now we don't use
     // `pg_ctl` for start / stop.
     #[instrument(skip_all)]
@@ -1804,11 +2129,16 @@ impl ComputeNode {
 
         // Write new config
         let pgdata_path = Path::new(&self.params.pgdata);
+        let postgres_port = self.params.connstr.port();
         config::write_postgres_conf(
             pgdata_path,
+            &self.params,
             &spec,
+            postgres_port,
             self.params.internal_http_port,
             tls_config,
+            spec.databricks_settings.as_ref(),
+            self.params.lakebase_mode,
         )?;
 
         self.pg_reload_conf()?;
@@ -1914,6 +2244,8 @@ impl ComputeNode {
                             // wait
                             ComputeStatus::Init
                             | ComputeStatus::Configuration
+                            | ComputeStatus::RefreshConfiguration
+                            | ComputeStatus::RefreshConfigurationPending
                             | ComputeStatus::Empty => {
                                 state = self.state_changed.wait(state).unwrap();
                             }
@@ -1964,7 +2296,17 @@ impl ComputeNode {
     pub fn check_for_core_dumps(&self) -> Result<()> {
         let core_dump_dir = match std::env::consts::OS {
             "macos" => Path::new("/cores/"),
-            _ => Path::new(&self.params.pgdata),
+            // BEGIN HADRON
+            // NB: Read core dump files from a fixed location outside of
+            // the data directory since `compute_ctl` wipes the data directory
+            // across container restarts.
+            _ => {
+                if self.params.lakebase_mode {
+                    Path::new("/databricks/logs/brickstore")
+                } else {
+                    Path::new(&self.params.pgdata)
+                }
+            } // END HADRON
         };
 
         // Collect core dump paths if any
@@ -2038,13 +2380,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;
@@ -2172,11 +2514,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
@@ -2245,8 +2587,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
@@ -2277,7 +2620,7 @@ LIMIT 100",
         if let Some(libs) = spec.cluster.settings.find("shared_preload_libraries") {
             libs_vec = libs
                 .split(&[',', '\'', ' '])
-                .filter(|s| *s != "neon" && !s.is_empty())
+                .filter(|s| *s != "neon" && *s != "databricks_auth" && !s.is_empty())
                 .map(str::to_string)
                 .collect();
         }
@@ -2296,7 +2639,7 @@ LIMIT 100",
             if let Some(libs) = shared_preload_libraries_line.split("='").nth(1) {
                 preload_libs_vec = libs
                     .split(&[',', '\'', ' '])
-                    .filter(|s| *s != "neon" && !s.is_empty())
+                    .filter(|s| *s != "neon" && *s != "databricks_auth" && !s.is_empty())
                     .map(str::to_string)
                     .collect();
             }
@@ -2349,22 +2692,22 @@ LIMIT 100",
     /// The operation will time out after a specified duration.
     pub fn wait_timeout_while_pageserver_connstr_unchanged(&self, duration: Duration) {
         let state = self.state.lock().unwrap();
-        let old_pageserver_connstr = state
+        let old_pageserver_conninfo = state
             .pspec
             .as_ref()
             .expect("spec must be set")
-            .pageserver_connstr
+            .pageserver_conninfo
             .clone();
         let mut unchanged = true;
         let _ = self
             .state_changed
             .wait_timeout_while(state, duration, |s| {
-                let pageserver_connstr = &s
+                let pageserver_conninfo = &s
                     .pspec
                     .as_ref()
                     .expect("spec must be set")
-                    .pageserver_connstr;
-                unchanged = pageserver_connstr == &old_pageserver_connstr;
+                    .pageserver_conninfo;
+                unchanged = pageserver_conninfo == &old_pageserver_conninfo;
                 unchanged
             })
             .unwrap();
@@ -2418,14 +2761,31 @@ LIMIT 100",
     pub fn spawn_lfc_offload_task(self: &Arc<Self>, interval: Duration) {
         self.terminate_lfc_offload_task();
         let secs = interval.as_secs();
-        info!("spawning lfc offload worker with {secs}s interval");
         let this = self.clone();
+
+        info!("spawning LFC offload worker with {secs}s interval");
         let handle = spawn(async move {
             let mut interval = time::interval(interval);
             interval.tick().await; // returns immediately
             loop {
                 interval.tick().await;
-                this.offload_lfc_async().await;
+
+                let prewarm_state = this.state.lock().unwrap().lfc_prewarm_state.clone();
+                // Do not offload LFC state if we are currently prewarming or any issue occurred.
+                // If we'd do that, we might override the LFC state in endpoint storage with some
+                // incomplete state. Imagine a situation:
+                // 1. Endpoint started with `autoprewarm: true`
+                // 2. While prewarming is not completed, we upload the new incomplete state
+                // 3. Compute gets interrupted and restarts
+                // 4. We start again and try to prewarm with the state from 2. instead of the previous complete state
+                if matches!(
+                    prewarm_state,
+                    LfcPrewarmState::Completed { .. }
+                        | LfcPrewarmState::NotPrewarmed
+                        | LfcPrewarmState::Skipped
+                ) {
+                    this.offload_lfc_async().await;
+                }
             }
         });
         *self.lfc_offload_task.lock().unwrap() = Some(handle);
@@ -2453,6 +2813,34 @@ LIMIT 100",
             );
         }
     }
+
+    /// Set the compute spec and update related metrics.
+    /// This is the central place where pspec is updated.
+    pub fn set_spec(params: &ComputeNodeParams, state: &mut ComputeState, pspec: ParsedSpec) {
+        state.pspec = Some(pspec);
+        ComputeNode::update_attached_metric(params, state);
+        let _ = logger::update_ids(&params.instance_id, &Some(params.compute_id.clone()));
+    }
+
+    pub fn update_attached_metric(params: &ComputeNodeParams, state: &mut ComputeState) {
+        // Update the pg_cctl_attached gauge when all identifiers are available.
+        if let Some(instance_id) = &params.instance_id {
+            if let Some(pspec) = &state.pspec {
+                // Clear all values in the metric
+                COMPUTE_ATTACHED.reset();
+
+                // Set new metric value
+                COMPUTE_ATTACHED
+                    .with_label_values(&[
+                        &params.compute_id,
+                        instance_id,
+                        &pspec.tenant_id.to_string(),
+                        &pspec.timeline_id.to_string(),
+                    ])
+                    .set(1);
+            }
+        }
+    }
 }
 
 pub async fn installed_extensions(conf: tokio_postgres::Config) -> Result<()> {
@@ -2464,7 +2852,7 @@ pub async fn installed_extensions(conf: tokio_postgres::Config) -> Result<()> {
                 serde_json::to_string(&extensions).expect("failed to serialize extensions list")
             );
         }
-        Err(err) => error!("could not get installed extensions: {err:?}"),
+        Err(err) => error!("could not get installed extensions: {err}"),
     }
     Ok(())
 }
@@ -2577,7 +2965,10 @@ mod tests {
 
         match ParsedSpec::try_from(spec.clone()) {
             Ok(_p) => panic!("Failed to detect duplicate entry"),
-            Err(e) => assert!(e.starts_with("duplicate entry in safekeeper_connstrings:")),
+            Err(e) => assert!(
+                e.to_string()
+                    .starts_with("duplicate entry in safekeeper_connstrings:")
+            ),
         };
     }
 }

compute_tools/src/compute_prewarm.rs

@@ -7,18 +7,11 @@ use http::StatusCode;
 use reqwest::Client;
 use std::mem::replace;
 use std::sync::Arc;
-use tokio::{io::AsyncReadExt, spawn};
+use std::time::Instant;
+use tokio::{io::AsyncReadExt, select, spawn};
+use tokio_util::sync::CancellationToken;
 use tracing::{error, info};
 
-#[derive(serde::Serialize, Default)]
-pub struct LfcPrewarmStateWithProgress {
-    #[serde(flatten)]
-    base: LfcPrewarmState,
-    total: i32,
-    prewarmed: i32,
-    skipped: i32,
-}
-
 /// A pair of url and a token to query endpoint storage for LFC prewarm-related tasks
 struct EndpointStoragePair {
     url: String,
@@ -27,7 +20,7 @@ struct EndpointStoragePair {
 
 const KEY: &str = "lfc_state";
 impl EndpointStoragePair {
-    /// endpoint_id is set to None while prewarming from other endpoint, see replica promotion
+    /// endpoint_id is set to None while prewarming from other endpoint, see compute_promote.rs
     /// If not None, takes precedence over pspec.spec.endpoint_id
     fn from_spec_and_endpoint(
         pspec: &crate::compute::ParsedSpec,
@@ -53,63 +46,46 @@ impl EndpointStoragePair {
 }
 
 impl ComputeNode {
-    // If prewarm failed, we want to get overall number of segments as well as done ones.
-    // However, this function should be reliable even if querying postgres failed.
-    pub async fn lfc_prewarm_state(&self) -> LfcPrewarmStateWithProgress {
-        info!("requesting LFC prewarm state from postgres");
-        let mut state = LfcPrewarmStateWithProgress::default();
-        {
-            state.base = self.state.lock().unwrap().lfc_prewarm_state.clone();
-        }
-
-        let client = match ComputeNode::get_maintenance_client(&self.tokio_conn_conf).await {
-            Ok(client) => client,
-            Err(err) => {
-                error!(%err, "connecting to postgres");
-                return state;
-            }
-        };
-        let row = match client
-            .query_one("select * from neon.get_prewarm_info()", &[])
-            .await
-        {
-            Ok(row) => row,
-            Err(err) => {
-                error!(%err, "querying LFC prewarm status");
-                return state;
-            }
-        };
-        state.total = row.try_get(0).unwrap_or_default();
-        state.prewarmed = row.try_get(1).unwrap_or_default();
-        state.skipped = row.try_get(2).unwrap_or_default();
-        state
+    pub async fn lfc_prewarm_state(&self) -> LfcPrewarmState {
+        self.state.lock().unwrap().lfc_prewarm_state.clone()
     }
 
     pub fn lfc_offload_state(&self) -> LfcOffloadState {
         self.state.lock().unwrap().lfc_offload_state.clone()
     }
 
-    /// If there is a prewarm request ongoing, return false, true otherwise
+    /// 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 Err(err) = cloned.prewarm_impl(from_endpoint).await else {
-                cloned.state.lock().unwrap().lfc_prewarm_state = LfcPrewarmState::Completed;
-                return;
-            };
-            crate::metrics::LFC_PREWARM_ERRORS.inc();
-            error!(%err, "prewarming lfc");
-            cloned.state.lock().unwrap().lfc_prewarm_state = LfcPrewarmState::Failed {
-                error: err.to_string(),
+            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");
+                    let error = format!("{err:#}");
+                    LfcPrewarmState::Failed { error }
+                }
             };
+
+            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
     }
@@ -120,43 +96,102 @@ impl ComputeNode {
         EndpointStoragePair::from_spec_and_endpoint(state.pspec.as_ref().unwrap(), from_endpoint)
     }
 
-    async fn prewarm_impl(&self, from_endpoint: Option<String>) -> Result<()> {
-        let EndpointStoragePair { url, token } = self.endpoint_storage_pair(from_endpoint)?;
-        info!(%url, "requesting LFC state from endpoint storage");
+    /// Request LFC state from endpoint storage and load corresponding pages into Postgres.
+    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)?;
 
-        let request = Client::new().get(&url).bearer_auth(token);
-        let res = request.send().await.context("querying endpoint storage")?;
-        let status = res.status();
-        if status != StatusCode::OK {
-            bail!("{status} querying endpoint storage")
+        #[cfg(feature = "testing")]
+        fail::fail_point!("compute-prewarm", |_| bail!("compute-prewarm failpoint"));
+
+        info!(%url, "requesting LFC state from endpoint storage");
+        let mut now = Instant::now();
+        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(LfcPrewarmState::Skipped),
+            status => bail!("{status} querying endpoint storage"),
+        }
+        let state_download_time_ms = now.elapsed().as_millis() as u32;
+        now = Instant::now();
 
         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 uncompressed_len = uncompressed.len();
-        info!(%url, "downloaded LFC state, uncompressed size {uncompressed_len}, loading into postgres");
+        let lfc_state = select! {
+            _ = token.cancelled() => return Ok(LfcPrewarmState::Cancelled),
+            lfc_state = response.bytes() => lfc_state
+        }
+        .context("getting request body from endpoint storage")?;
 
-        ComputeNode::get_maintenance_client(&self.tokio_conn_conf)
+        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 uncompress_time_ms = now.elapsed().as_millis() as u32;
+        now = Instant::now();
+
+        let uncompressed_len = uncompressed.len();
+        info!(%url, "downloaded LFC state, uncompressed size {uncompressed_len}");
+
+        // 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])
+            .context("connecting to postgres")?;
+        let pg_token = client.cancel_token();
+
+        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(|_| ())?;
+        let prewarm_time_ms = now.elapsed().as_millis() as u32;
+
+        let row = client
+            .query_one("select * from neon.get_prewarm_info()", &[])
             .await
-            .context("loading LFC state into postgres")
-            .map(|_| ())
+            .context("querying prewarm info")?;
+        let total = row.try_get(0).unwrap_or_default();
+        let prewarmed = row.try_get(1).unwrap_or_default();
+        let skipped = row.try_get(2).unwrap_or_default();
+
+        Ok(LfcPrewarmState::Completed {
+            total,
+            prewarmed,
+            skipped,
+            state_download_time_ms,
+            uncompress_time_ms,
+            prewarm_time_ms,
+        })
     }
 
     /// If offload request is ongoing, return false, true otherwise
     pub fn offload_lfc(self: &Arc<Self>) -> bool {
         {
             let state = &mut self.state.lock().unwrap().lfc_offload_state;
-            if replace(state, LfcOffloadState::Offloading) == LfcOffloadState::Offloading {
+            if matches!(
+                replace(state, LfcOffloadState::Offloading),
+                LfcOffloadState::Offloading
+            ) {
                 return false;
             }
         }
@@ -168,7 +203,10 @@ impl ComputeNode {
     pub async fn offload_lfc_async(self: &Arc<Self>) {
         {
             let state = &mut self.state.lock().unwrap().lfc_offload_state;
-            if replace(state, LfcOffloadState::Offloading) == LfcOffloadState::Offloading {
+            if matches!(
+                replace(state, LfcOffloadState::Offloading),
+                LfcOffloadState::Offloading
+            ) {
                 return;
             }
         }
@@ -177,42 +215,69 @@ 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;
-        };
-        crate::metrics::LFC_OFFLOAD_ERRORS.inc();
-        error!(%err, "offloading lfc");
-        self.state.lock().unwrap().lfc_offload_state = LfcOffloadState::Failed {
-            error: err.to_string(),
+        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 }
+            }
         };
+        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");
+        info!(%url, "requesting LFC state from Postgres");
 
-        let mut compressed = Vec::new();
-        ComputeNode::get_maintenance_client(&self.tokio_conn_conf)
+        let mut now = Instant::now();
+        let row = ComputeNode::get_maintenance_client(&self.tokio_conn_conf)
             .await
             .context("connecting to postgres")?
             .query_one("select neon.get_local_cache_state()", &[])
             .await
-            .context("querying LFC state")?
-            .try_get::<usize, &[u8]>(0)
-            .context("deserializing LFC state")
-            .map(ZstdEncoder::new)?
+            .context("querying LFC state")?;
+        let state = row
+            .try_get::<usize, Option<&[u8]>>(0)
+            .context("deserializing LFC state")?;
+        let Some(state) = state else {
+            info!(%url, "empty LFC state, not exporting");
+            return Ok(LfcOffloadState::Skipped);
+        };
+        let state_query_time_ms = now.elapsed().as_millis() as u32;
+        now = Instant::now();
+
+        let mut compressed = Vec::new();
+        ZstdEncoder::new(state)
             .read_to_end(&mut compressed)
             .await
             .context("compressing LFC state")?;
+        let compress_time_ms = now.elapsed().as_millis() as u32;
+        now = Instant::now();
+
         let compressed_len = compressed.len();
-        info!(%url, "downloaded LFC state, compressed size {compressed_len}, writing to endpoint storage");
+        info!(%url, "downloaded LFC state, compressed size {compressed_len}");
 
         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) => bail!("Error writing to endpoint storage: {}", res.status()),
-            Err(err) => Err(err).context("writing to endpoint storage"),
+        let response = request
+            .send()
+            .await
+            .context("writing to endpoint storage")?;
+        let state_upload_time_ms = now.elapsed().as_millis() as u32;
+        let status = response.status();
+        if status != StatusCode::OK {
+            bail!("request to endpoint storage failed: {status}");
         }
+
+        Ok(LfcOffloadState::Completed {
+            compress_time_ms,
+            state_query_time_ms,
+            state_upload_time_ms,
+        })
+    }
+
+    pub fn cancel_prewarm(self: &Arc<Self>) {
+        self.state.lock().unwrap().lfc_prewarm_token.cancel();
     }
 }

compute_tools/src/compute_promote.rs

@@ -1,69 +1,56 @@
 use crate::compute::ComputeNode;
-use anyhow::{Context, Result, bail};
-use compute_api::{
-    responses::{LfcPrewarmState, PromoteState, SafekeepersLsn},
-    spec::ComputeMode,
-};
-use std::{sync::Arc, time::Duration};
-use tokio::time::sleep;
-use utils::lsn::Lsn;
+use anyhow::{Context, bail};
+use compute_api::responses::{LfcPrewarmState, PromoteConfig, PromoteState};
+use std::time::Instant;
+use tracing::info;
 
 impl ComputeNode {
-    /// Returns only when promote fails or succeeds. If a network error occurs
-    /// and http client disconnects, this does not stop promotion, and subsequent
-    /// calls block until promote finishes.
+    /// Returns only when promote fails or succeeds. If http client calling this function
+    /// disconnects, this does not stop promotion, and subsequent calls block until promote finishes.
     /// Called by control plane on secondary after primary endpoint is terminated
-    pub async fn promote(self: &Arc<Self>, safekeepers_lsn: SafekeepersLsn) -> PromoteState {
-        let cloned = self.clone();
+    /// Has a failpoint "compute-promotion"
+    pub async fn promote(self: &std::sync::Arc<Self>, cfg: PromoteConfig) -> PromoteState {
+        let this = self.clone();
+        let promote_fn = async move || match this.promote_impl(cfg).await {
+            Ok(state) => state,
+            Err(err) => {
+                tracing::error!(%err, "promoting replica");
+                let error = format!("{err:#}");
+                PromoteState::Failed { error }
+            }
+        };
         let start_promotion = || {
             let (tx, rx) = tokio::sync::watch::channel(PromoteState::NotPromoted);
-            tokio::spawn(async move {
-                tx.send(match cloned.promote_impl(safekeepers_lsn).await {
-                    Ok(_) => PromoteState::Completed,
-                    Err(err) => {
-                        tracing::error!(%err, "promoting");
-                        PromoteState::Failed {
-                            error: err.to_string(),
-                        }
-                    }
-                })
-            });
+            tokio::spawn(async move { tx.send(promote_fn().await) });
             rx
         };
 
         let mut task;
-        // self.state is unlocked after block ends so we lock it in promote_impl
-        // and task.changed() is reached
+        // promote_impl locks self.state so we need to unlock it before calling task.changed()
         {
-            task = self
-                .state
-                .lock()
-                .unwrap()
-                .promote_state
-                .get_or_insert_with(start_promotion)
-                .clone()
+            let promote_state = &mut self.state.lock().unwrap().promote_state;
+            task = promote_state.get_or_insert_with(start_promotion).clone()
+        }
+        if task.changed().await.is_err() {
+            let error = "promote sender dropped".to_string();
+            return PromoteState::Failed { error };
         }
-        task.changed().await.expect("promote sender dropped");
         task.borrow().clone()
     }
 
-    // Why do we have to supply safekeepers?
-    // For secondary we use primary_connection_conninfo so safekeepers field is empty
-    async fn promote_impl(&self, safekeepers_lsn: SafekeepersLsn) -> Result<()> {
+    async fn promote_impl(&self, cfg: PromoteConfig) -> anyhow::Result<PromoteState> {
         {
             let state = self.state.lock().unwrap();
             let mode = &state.pspec.as_ref().unwrap().spec.mode;
-            if *mode != ComputeMode::Replica {
-                bail!("{} is not replica", mode.to_type_str());
+            if *mode != compute_api::spec::ComputeMode::Replica {
+                bail!("compute mode \"{}\" is not replica", mode.to_type_str());
             }
-
-            // we don't need to query Postgres so not self.lfc_prewarm_state()
             match &state.lfc_prewarm_state {
-                LfcPrewarmState::NotPrewarmed | LfcPrewarmState::Prewarming => {
-                    bail!("prewarm not requested or pending")
+                status @ (LfcPrewarmState::NotPrewarmed | LfcPrewarmState::Prewarming) => {
+                    bail!("compute {status}")
                 }
                 LfcPrewarmState::Failed { error } => {
-                    tracing::warn!(%error, "replica prewarm failed")
+                    tracing::warn!(%error, "compute prewarm failed")
                 }
                 _ => {}
             }
@@ -72,51 +59,66 @@ impl ComputeNode {
         let client = ComputeNode::get_maintenance_client(&self.tokio_conn_conf)
             .await
             .context("connecting to postgres")?;
+        let mut now = Instant::now();
 
-        let primary_lsn = safekeepers_lsn.wal_flush_lsn;
-        let mut last_wal_replay_lsn: Lsn = Lsn::INVALID;
+        let primary_lsn = cfg.wal_flush_lsn;
+        let mut standby_lsn = utils::lsn::Lsn::INVALID;
         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();
-            last_wal_replay_lsn = lsn.into();
-            if last_wal_replay_lsn >= primary_lsn {
+            standby_lsn = lsn.into();
+            if standby_lsn >= primary_lsn {
                 break;
             }
-            tracing::info!("Try {i}, replica lsn {last_wal_replay_lsn}, primary lsn {primary_lsn}");
-            sleep(Duration::from_secs(1)).await;
+            info!(%standby_lsn, %primary_lsn, "catching up, try {i}");
+            tokio::time::sleep(std::time::Duration::from_secs(1)).await;
         }
-        if last_wal_replay_lsn < primary_lsn {
+        if standby_lsn < primary_lsn {
             bail!("didn't catch up with primary in {RETRIES} retries");
         }
+        let lsn_wait_time_ms = now.elapsed().as_millis() as u32;
+        now = Instant::now();
 
         // using $1 doesn't work with ALTER SYSTEM SET
         let safekeepers_sql = format!(
             "ALTER SYSTEM SET neon.safekeepers='{}'",
-            safekeepers_lsn.safekeepers
+            cfg.spec.safekeeper_connstrings.join(",")
         );
         client
             .query(&safekeepers_sql, &[])
             .await
             .context("setting safekeepers")?;
         client
-            .query("SELECT pg_reload_conf()", &[])
+            .query(
+                "ALTER SYSTEM SET synchronous_standby_names=walproposer",
+                &[],
+            )
+            .await
+            .context("setting synchronous_standby_names")?;
+        client
+            .query("SELECT pg_catalog.pg_reload_conf()", &[])
             .await
             .context("reloading postgres config")?;
+
+        #[cfg(feature = "testing")]
+        fail::fail_point!("compute-promotion", |_| bail!(
+            "compute-promotion failpoint"
+        ));
+
         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) {
-            bail!("pg_promote() returned false");
+            bail!("pg_promote() failed");
         }
+        let pg_promote_time_ms = now.elapsed().as_millis() as u32;
+        let now = Instant::now();
 
-        let client = ComputeNode::get_maintenance_client(&self.tokio_conn_conf)
-            .await
-            .context("connecting to postgres")?;
         let row = client
             .query_one("SHOW transaction_read_only", &[])
             .await
@@ -125,8 +127,47 @@ impl ComputeNode {
             bail!("replica in read only mode after promotion");
         }
 
-        let mut state = self.state.lock().unwrap();
-        state.pspec.as_mut().unwrap().spec.mode = ComputeMode::Primary;
-        Ok(())
+        // Already checked validity in http handler
+        #[allow(unused_mut)]
+        let mut new_pspec = crate::compute::ParsedSpec::try_from(cfg.spec).expect("invalid spec");
+        {
+            let mut state = self.state.lock().unwrap();
+
+            // Local setup has different ports for pg process (port=) for primary and secondary.
+            // Primary is stopped so we need secondary's "port" value
+            #[cfg(feature = "testing")]
+            {
+                let old_spec = &state.pspec.as_ref().unwrap().spec;
+                let Some(old_conf) = old_spec.cluster.postgresql_conf.as_ref() else {
+                    bail!("pspec.spec.cluster.postgresql_conf missing for endpoint");
+                };
+                let set: std::collections::HashMap<&str, &str> = old_conf
+                    .split_terminator('\n')
+                    .map(|e| e.split_once("=").expect("invalid item"))
+                    .collect();
+
+                let Some(new_conf) = new_pspec.spec.cluster.postgresql_conf.as_mut() else {
+                    bail!("pspec.spec.cluster.postgresql_conf missing for supplied config");
+                };
+                new_conf.push_str(&format!("port={}\n", set["port"]));
+            }
+
+            tracing::debug!("applied spec: {:#?}", new_pspec.spec);
+            if self.params.lakebase_mode {
+                ComputeNode::set_spec(&self.params, &mut state, new_pspec);
+            } else {
+                state.pspec = Some(new_pspec);
+            }
+        }
+
+        info!("applied new spec, reconfiguring as primary");
+        self.reconfigure()?;
+        let reconfigure_time_ms = now.elapsed().as_millis() as u32;
+
+        Ok(PromoteState::Completed {
+            lsn_wait_time_ms,
+            pg_promote_time_ms,
+            reconfigure_time_ms,
+        })
     }
 }

compute_tools/src/config.rs

@@ -7,13 +7,19 @@ use std::io::prelude::*;
 use std::path::Path;
 
 use compute_api::responses::TlsConfig;
-use compute_api::spec::{ComputeAudit, ComputeMode, ComputeSpec, GenericOption};
+use compute_api::spec::{
+    ComputeAudit, ComputeMode, ComputeSpec, DatabricksSettings, GenericOption,
+};
 
+use crate::compute::ComputeNodeParams;
 use crate::pg_helpers::{
-    GenericOptionExt, GenericOptionsSearch, PgOptionsSerialize, escape_conf_value,
+    DatabricksSettingsExt as _, GenericOptionExt, GenericOptionsSearch, PgOptionsSerialize,
+    escape_conf_value,
 };
 use crate::tls::{self, SERVER_CRT, SERVER_KEY};
 
+use utils::shard::{ShardIndex, ShardNumber};
+
 /// Check that `line` is inside a text file and put it there if it is not.
 /// Create file if it doesn't exist.
 pub fn line_in_file(path: &Path, line: &str) -> Result<bool> {
@@ -39,11 +45,16 @@ pub fn line_in_file(path: &Path, line: &str) -> Result<bool> {
 }
 
 /// Create or completely rewrite configuration file specified by `path`
+#[allow(clippy::too_many_arguments)]
 pub fn write_postgres_conf(
     pgdata_path: &Path,
+    params: &ComputeNodeParams,
     spec: &ComputeSpec,
+    postgres_port: Option<u16>,
     extension_server_port: u16,
     tls_config: &Option<TlsConfig>,
+    databricks_settings: Option<&DatabricksSettings>,
+    lakebase_mode: bool,
 ) -> Result<()> {
     let path = pgdata_path.join("postgresql.conf");
     // File::create() destroys the file content if it exists.
@@ -56,12 +67,76 @@ pub fn write_postgres_conf(
 
     // Add options for connecting to storage
     writeln!(file, "# Neon storage settings")?;
-    if let Some(s) = &spec.pageserver_connstring {
-        writeln!(file, "neon.pageserver_connstring={}", escape_conf_value(s))?;
-    }
-    if let Some(stripe_size) = spec.shard_stripe_size {
-        writeln!(file, "neon.stripe_size={stripe_size}")?;
+    writeln!(file)?;
+    if let Some(conninfo) = &spec.pageserver_connection_info {
+        // Stripe size GUC should be defined prior to connection string
+        if let Some(stripe_size) = conninfo.stripe_size {
+            writeln!(
+                file,
+                "# from compute spec's pageserver_connection_info.stripe_size field"
+            )?;
+            writeln!(file, "neon.stripe_size={stripe_size}")?;
+        }
+
+        let mut libpq_urls: Option<Vec<String>> = Some(Vec::new());
+        let num_shards = if conninfo.shard_count.0 == 0 {
+            1 // unsharded, treat it as a single shard
+        } else {
+            conninfo.shard_count.0
+        };
+
+        for shard_number in 0..num_shards {
+            let shard_index = ShardIndex {
+                shard_number: ShardNumber(shard_number),
+                shard_count: conninfo.shard_count,
+            };
+            let info = conninfo.shards.get(&shard_index).ok_or_else(|| {
+                anyhow::anyhow!(
+                    "shard {shard_index} missing from pageserver_connection_info shard map"
+                )
+            })?;
+
+            let first_pageserver = info
+                .pageservers
+                .first()
+                .expect("must have at least one pageserver");
+
+            // Add the libpq URL to the array, or if the URL is missing, reset the array
+            // forgetting any previous entries. All servers must have a libpq URL, or none
+            // at all.
+            if let Some(url) = &first_pageserver.libpq_url {
+                if let Some(ref mut urls) = libpq_urls {
+                    urls.push(url.clone());
+                }
+            } else {
+                libpq_urls = None
+            }
+        }
+        if let Some(libpq_urls) = libpq_urls {
+            writeln!(
+                file,
+                "# derived from compute spec's pageserver_connection_info field"
+            )?;
+            writeln!(
+                file,
+                "neon.pageserver_connstring={}",
+                escape_conf_value(&libpq_urls.join(","))
+            )?;
+        } else {
+            writeln!(file, "# no neon.pageserver_connstring")?;
+        }
+    } else {
+        // Stripe size GUC should be defined prior to connection string
+        if let Some(stripe_size) = spec.shard_stripe_size {
+            writeln!(file, "# from compute spec's shard_stripe_size field")?;
+            writeln!(file, "neon.stripe_size={stripe_size}")?;
+        }
+        if let Some(s) = &spec.pageserver_connstring {
+            writeln!(file, "# from compute spec's pageserver_connstring field")?;
+            writeln!(file, "neon.pageserver_connstring={}", escape_conf_value(s))?;
+        }
     }
+
     if !spec.safekeeper_connstrings.is_empty() {
         let mut neon_safekeepers_value = String::new();
         tracing::info!(
@@ -161,6 +236,12 @@ pub fn write_postgres_conf(
         }
     }
 
+    writeln!(
+        file,
+        "neon.privileged_role_name={}",
+        escape_conf_value(params.privileged_role_name.as_str())
+    )?;
+
     // If there are any extra options in the 'settings' field, append those
     if spec.cluster.settings.is_some() {
         writeln!(file, "# Managed by compute_ctl: begin")?;
@@ -276,6 +357,24 @@ pub fn write_postgres_conf(
         writeln!(file, "log_destination='stderr,syslog'")?;
     }
 
+    if lakebase_mode {
+        // Explicitly set the port based on the connstr, overriding any previous port setting.
+        // Note: It is important that we don't specify a different port again after this.
+        let port = postgres_port.expect("port must be present in connstr");
+        writeln!(file, "port = {port}")?;
+
+        // This is databricks specific settings.
+        // This should be at the end of the file but before `compute_ctl_temp_override.conf` below
+        // so that it can override any settings above.
+        // `compute_ctl_temp_override.conf` is intended to override any settings above during specific operations.
+        // To prevent potential breakage in the future, we keep it above `compute_ctl_temp_override.conf`.
+        writeln!(file, "# Databricks settings start")?;
+        if let Some(settings) = databricks_settings {
+            writeln!(file, "{}", settings.as_pg_settings())?;
+        }
+        writeln!(file, "# Databricks settings end")?;
+    }
+
     // This is essential to keep this line at the end of the file,
     // because it is intended to override any settings above.
     writeln!(file, "include_if_exists = 'compute_ctl_temp_override.conf'")?;

compute_tools/src/configurator.rs

@@ -1,23 +1,40 @@
-use std::sync::Arc;
+use std::fs::File;
 use std::thread;
+use std::{path::Path, sync::Arc};
 
-use compute_api::responses::ComputeStatus;
+use anyhow::Result;
+use compute_api::responses::{ComputeConfig, ComputeStatus};
 use tracing::{error, info, instrument};
 
-use crate::compute::ComputeNode;
+use crate::compute::{ComputeNode, ParsedSpec};
+use crate::spec::get_config_from_control_plane;
 
 #[instrument(skip_all)]
 fn configurator_main_loop(compute: &Arc<ComputeNode>) {
     info!("waiting for reconfiguration requests");
     loop {
         let mut state = compute.state.lock().unwrap();
+        /* BEGIN_HADRON */
+        // RefreshConfiguration should only be used inside the loop
+        assert_ne!(state.status, ComputeStatus::RefreshConfiguration);
+        /* END_HADRON */
 
-        // We have to re-check the status after re-acquiring the lock because it could be that
-        // the status has changed while we were waiting for the lock, and we might not need to
-        // wait on the condition variable. Otherwise, we might end up in some soft-/deadlock, i.e.
-        // we are waiting for a condition variable that will never be signaled.
-        if state.status != ComputeStatus::ConfigurationPending {
-            state = compute.state_changed.wait(state).unwrap();
+        if compute.params.lakebase_mode {
+            while state.status != ComputeStatus::ConfigurationPending
+                && state.status != ComputeStatus::RefreshConfigurationPending
+                && state.status != ComputeStatus::Failed
+            {
+                info!("configurator: compute status: {:?}, sleeping", state.status);
+                state = compute.state_changed.wait(state).unwrap();
+            }
+        } else {
+            // We have to re-check the status after re-acquiring the lock because it could be that
+            // the status has changed while we were waiting for the lock, and we might not need to
+            // wait on the condition variable. Otherwise, we might end up in some soft-/deadlock, i.e.
+            // we are waiting for a condition variable that will never be signaled.
+            if state.status != ComputeStatus::ConfigurationPending {
+                state = compute.state_changed.wait(state).unwrap();
+            }
         }
 
         // Re-check the status after waking up
@@ -37,6 +54,136 @@ fn configurator_main_loop(compute: &Arc<ComputeNode>) {
             // XXX: used to test that API is blocking
             // std::thread::sleep(std::time::Duration::from_millis(10000));
 
+            compute.set_status(new_status);
+        } else if state.status == ComputeStatus::RefreshConfigurationPending {
+            info!(
+                "compute node suspects its configuration is out of date, now refreshing configuration"
+            );
+            state.set_status(ComputeStatus::RefreshConfiguration, &compute.state_changed);
+            // Drop the lock guard here to avoid holding the lock while downloading config from the control plane / HCC.
+            // This is the only thread that can move compute_ctl out of the `RefreshConfiguration` state, so it
+            // is safe to drop the lock like this.
+            drop(state);
+
+            let get_config_result: anyhow::Result<ComputeConfig> =
+                if let Some(config_path) = &compute.params.config_path_test_only {
+                    // This path is only to make testing easier. In production we always get the config from the HCC.
+                    info!(
+                        "reloading config.json from path: {}",
+                        config_path.to_string_lossy()
+                    );
+                    let path = Path::new(config_path);
+                    if let Ok(file) = File::open(path) {
+                        match serde_json::from_reader::<File, ComputeConfig>(file) {
+                            Ok(config) => Ok(config),
+                            Err(e) => {
+                                error!("could not parse config file: {}", e);
+                                Err(anyhow::anyhow!("could not parse config file: {}", e))
+                            }
+                        }
+                    } else {
+                        error!(
+                            "could not open config file at path: {:?}",
+                            config_path.to_string_lossy()
+                        );
+                        Err(anyhow::anyhow!(
+                            "could not open config file at path: {}",
+                            config_path.to_string_lossy()
+                        ))
+                    }
+                } else if let Some(control_plane_uri) = &compute.params.control_plane_uri {
+                    get_config_from_control_plane(control_plane_uri, &compute.params.compute_id)
+                } else {
+                    Err(anyhow::anyhow!("config_path_test_only is not set"))
+                };
+
+            // Parse any received ComputeSpec and transpose the result into a Result<Option<ParsedSpec>>.
+            let parsed_spec_result: Result<Option<ParsedSpec>> =
+                get_config_result.and_then(|config| {
+                    if let Some(spec) = config.spec {
+                        if let Ok(pspec) = ParsedSpec::try_from(spec) {
+                            Ok(Some(pspec))
+                        } else {
+                            Err(anyhow::anyhow!("could not parse spec"))
+                        }
+                    } else {
+                        Ok(None)
+                    }
+                });
+
+            let new_status: ComputeStatus;
+            match parsed_spec_result {
+                // Control plane (HCM) returned a spec and we were able to parse it.
+                Ok(Some(pspec)) => {
+                    {
+                        let mut state = compute.state.lock().unwrap();
+                        // Defensive programming to make sure this thread is indeed the only one that can move the compute
+                        // node out of the `RefreshConfiguration` state. Would be nice if we can encode this invariant
+                        // into the type system.
+                        assert_eq!(state.status, ComputeStatus::RefreshConfiguration);
+
+                        if state
+                            .pspec
+                            .as_ref()
+                            .map(|ps| ps.pageserver_conninfo.clone())
+                            == Some(pspec.pageserver_conninfo.clone())
+                        {
+                            info!(
+                                "Refresh configuration: Retrieved spec is the same as the current spec. Waiting for control plane to update the spec before attempting reconfiguration."
+                            );
+                            state.status = ComputeStatus::Running;
+                            compute.state_changed.notify_all();
+                            drop(state);
+                            std::thread::sleep(std::time::Duration::from_secs(5));
+                            continue;
+                        }
+                        // state.pspec is consumed by compute.reconfigure() below. Note that compute.reconfigure() will acquire
+                        // the compute.state lock again so we need to have the lock guard go out of scope here. We could add a
+                        // "locked" variant of compute.reconfigure() that takes the lock guard as an argument to make this cleaner,
+                        // but it's not worth forking the codebase too much for this minor point alone right now.
+                        state.pspec = Some(pspec);
+                    }
+                    match compute.reconfigure() {
+                        Ok(_) => {
+                            info!("Refresh configuration: compute node configured");
+                            new_status = ComputeStatus::Running;
+                        }
+                        Err(e) => {
+                            error!(
+                                "Refresh configuration: could not configure compute node: {}",
+                                e
+                            );
+                            // Set the compute node back to the `RefreshConfigurationPending` state if the configuration
+                            // was not successful. It should be okay to treat this situation the same as if the loop
+                            // hasn't executed yet as long as the detection side keeps notifying.
+                            new_status = ComputeStatus::RefreshConfigurationPending;
+                        }
+                    }
+                }
+                // Control plane (HCM)'s response does not contain a spec. This is the "Empty" attachment case.
+                Ok(None) => {
+                    info!(
+                        "Compute Manager signaled that this compute is no longer attached to any storage. Exiting."
+                    );
+                    // We just immediately terminate the whole compute_ctl in this case. It's not necessary to attempt a
+                    // clean shutdown as Postgres is probably not responding anyway (which is why we are in this refresh
+                    // configuration state).
+                    std::process::exit(1);
+                }
+                // Various error cases:
+                // - The request to the control plane (HCM) either failed or returned a malformed spec.
+                // - compute_ctl itself is configured incorrectly (e.g., compute_id is not set).
+                Err(e) => {
+                    error!(
+                        "Refresh configuration: error getting a parsed spec: {:?}",
+                        e
+                    );
+                    new_status = ComputeStatus::RefreshConfigurationPending;
+                    // We may be dealing with an overloaded HCM if we end up in this path. Backoff 5 seconds before
+                    // retrying to avoid hammering the HCM.
+                    std::thread::sleep(std::time::Duration::from_secs(5));
+                }
+            }
             compute.set_status(new_status);
         } else if state.status == ComputeStatus::Failed {
             info!("compute node is now in Failed state, exiting");

compute_tools/src/hadron_metrics.rs

@@ -0,0 +1,60 @@
+use metrics::{
+    IntCounter, IntGaugeVec, core::Collector, proto::MetricFamily, register_int_counter,
+    register_int_gauge_vec,
+};
+use once_cell::sync::Lazy;
+
+// Counter keeping track of the number of PageStream request errors reported by Postgres.
+// An error is registered every time Postgres calls compute_ctl's /refresh_configuration API.
+// Postgres will invoke this API if it detected trouble with PageStream requests (get_page@lsn,
+// get_base_backup, etc.) it sends to any pageserver. An increase in this counter value typically
+// indicates Postgres downtime, as PageStream requests are critical for Postgres to function.
+pub static POSTGRES_PAGESTREAM_REQUEST_ERRORS: Lazy<IntCounter> = Lazy::new(|| {
+    register_int_counter!(
+        "pg_cctl_pagestream_request_errors_total",
+        "Number of PageStream request errors reported by the postgres process"
+    )
+    .expect("failed to define a metric")
+});
+
+// Counter keeping track of the number of compute configuration errors due to Postgres statement
+// timeouts. An error is registered every time `ComputeNode::reconfigure()` fails due to Postgres
+// error code 57014 (query cancelled). This statement timeout typically occurs when postgres is
+// stuck in a problematic retry loop when the PS is reject its connection requests (usually due
+// to PG pointing at the wrong PS). We should investigate the root cause when this counter value
+// increases by checking PG and PS logs.
+pub static COMPUTE_CONFIGURE_STATEMENT_TIMEOUT_ERRORS: Lazy<IntCounter> = Lazy::new(|| {
+    register_int_counter!(
+        "pg_cctl_configure_statement_timeout_errors_total",
+        "Number of compute configuration errors due to Postgres statement timeouts."
+    )
+    .expect("failed to define a metric")
+});
+
+pub static COMPUTE_ATTACHED: Lazy<IntGaugeVec> = Lazy::new(|| {
+    register_int_gauge_vec!(
+        "pg_cctl_attached",
+        "Compute node attached status (1 if attached)",
+        &[
+            "pg_compute_id",
+            "pg_instance_id",
+            "tenant_id",
+            "timeline_id"
+        ]
+    )
+    .expect("failed to define a metric")
+});
+
+pub fn collect() -> Vec<MetricFamily> {
+    let mut metrics = Vec::new();
+    metrics.extend(POSTGRES_PAGESTREAM_REQUEST_ERRORS.collect());
+    metrics.extend(COMPUTE_CONFIGURE_STATEMENT_TIMEOUT_ERRORS.collect());
+    metrics.extend(COMPUTE_ATTACHED.collect());
+    metrics
+}
+
+pub fn initialize_metrics() {
+    Lazy::force(&POSTGRES_PAGESTREAM_REQUEST_ERRORS);
+    Lazy::force(&COMPUTE_CONFIGURE_STATEMENT_TIMEOUT_ERRORS);
+    Lazy::force(&COMPUTE_ATTACHED);
+}

compute_tools/src/http/middleware/authorize.rs

@@ -16,13 +16,29 @@ use crate::http::JsonResponse;
 #[derive(Clone, Debug)]
 pub(in crate::http) struct Authorize {
     compute_id: String,
+    // BEGIN HADRON
+    // Hadron instance ID. Only set if it's a Lakebase V1 a.k.a. Hadron instance.
+    instance_id: Option<String>,
+    // END HADRON
     jwks: JwkSet,
     validation: Validation,
 }
 
 impl Authorize {
-    pub fn new(compute_id: String, jwks: JwkSet) -> Self {
+    pub fn new(compute_id: String, instance_id: Option<String>, jwks: JwkSet) -> Self {
         let mut validation = Validation::new(Algorithm::EdDSA);
+
+        // BEGIN HADRON
+        let use_rsa = jwks.keys.iter().any(|jwk| {
+            jwk.common
+                .key_algorithm
+                .is_some_and(|alg| alg == jsonwebtoken::jwk::KeyAlgorithm::RS256)
+        });
+        if use_rsa {
+            validation = Validation::new(Algorithm::RS256);
+        }
+        // END HADRON
+
         validation.validate_exp = true;
         // Unused by the control plane
         validation.validate_nbf = false;
@@ -34,6 +50,7 @@ impl Authorize {
 
         Self {
             compute_id,
+            instance_id,
             jwks,
             validation,
         }
@@ -47,10 +64,20 @@ impl AsyncAuthorizeRequest<Body> for Authorize {
 
     fn authorize(&mut self, mut request: Request<Body>) -> Self::Future {
         let compute_id = self.compute_id.clone();
+        let is_hadron_instance = self.instance_id.is_some();
         let jwks = self.jwks.clone();
         let validation = self.validation.clone();
 
         Box::pin(async move {
+            // BEGIN HADRON
+            // In Hadron deployments the "external" HTTP endpoint on compute_ctl can only be
+            // accessed by trusted components (enforced by dblet network policy), so we can bypass
+            // all auth here.
+            if is_hadron_instance {
+                return Ok(request);
+            }
+            // END HADRON
+
             let TypedHeader(Authorization(bearer)) = request
                 .extract_parts::<TypedHeader<Authorization<Bearer>>>()
                 .await

compute_tools/src/http/openapi_spec.yaml

@@ -96,7 +96,7 @@ paths:
         content:
           application/json:
             schema:
-                $ref: "#/components/schemas/SafekeepersLsn"
+              $ref: "#/components/schemas/ComputeSchemaWithLsn"
       responses:
         200:
           description: Promote succeeded or wasn't started
@@ -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:
@@ -297,14 +306,7 @@ paths:
         content:
           application/json:
             schema:
-              type: object
-              required:
-                - spec
-              properties:
-                spec:
-                  # XXX: I don't want to explain current spec in the OpenAPI format,
-                  # as it could be changed really soon. Consider doing it later.
-                  type: object
+              $ref: "#/components/schemas/ComputeSchema"
       responses:
         200:
           description: Compute configuration finished.
@@ -591,33 +593,37 @@ components:
           type: string
           example: "1.0.0"
 
-    SafekeepersLsn:
+    ComputeSchema:
       type: object
       required:
-        - safekeepers
+        - spec
+      properties:
+        spec:
+          type: object
+    ComputeSchemaWithLsn:
+      type: object
+      required:
+        - spec
         - wal_flush_lsn
       properties:
-        safekeepers:
-          description: Primary replica safekeepers
-          type: string
+        spec:
+          $ref: "#/components/schemas/ComputeState"
         wal_flush_lsn:
-          description: Primary last WAL flush LSN
           type: string
+          description: "last WAL flush LSN"
+          example: "0/028F10D8"
 
     LfcPrewarmState:
       type: object
       required:
         - status
-        - total
-        - prewarmed
-        - skipped
       properties:
         status:
-          description: Lfc prewarm status
-          enum: [not_prewarmed, prewarming, completed, failed]
+          description: LFC prewarm status
+          enum: [not_prewarmed, prewarming, completed, failed, skipped]
           type: string
         error:
-          description: Lfc prewarm error, if any
+          description: LFC prewarm error, if any
           type: string
         total:
           description: Total pages processed
@@ -628,6 +634,15 @@ components:
         skipped:
           description: Pages processed but not prewarmed
           type: integer
+        state_download_time_ms:
+          description: Time it takes to download LFC state to compute
+          type: integer
+        uncompress_time_ms:
+          description: Time it takes to uncompress LFC state
+          type: integer
+        prewarm_time_ms:
+          description: Time it takes to prewarm LFC state in Postgres
+          type: integer
 
     LfcOffloadState:
       type: object
@@ -635,12 +650,22 @@ components:
         - status
       properties:
         status:
-          description: Lfc offload status
-          enum: [not_offloaded, offloading, completed, failed]
+          description: LFC offload status
+          enum: [not_offloaded, offloading, completed, skipped, failed]
           type: string
         error:
-          description: Lfc offload error, if any
+          description: LFC offload error, if any
           type: string
+        state_query_time_ms:
+          description: Time it takes to get LFC state from Postgres
+          type: integer
+        compress_time_ms:
+          description: Time it takes to compress LFC state
+          type: integer
+        state_upload_time_ms:
+          description: Time it takes to upload LFC state to endpoint storage
+          type: integer
+
 
     PromoteState:
       type: object
@@ -654,6 +679,15 @@ components:
         error:
           description: Promote error, if any
           type: string
+        lsn_wait_time_ms:
+          description: Time it takes for secondary to catch up with primary WAL flush LSN
+          type: integer
+        pg_promote_time_ms:
+          description: Time it takes to call pg_promote on secondary
+          type: integer
+        reconfigure_time_ms:
+          description: Time it takes to reconfigure promoted secondary
+          type: integer
 
     SetRoleGrantsRequest:
       type: object

compute_tools/src/http/routes/configure.rs

@@ -43,7 +43,12 @@ pub(in crate::http) async fn configure(
         // configure request for tracing purposes.
         state.startup_span = Some(tracing::Span::current());
 
-        state.pspec = Some(pspec);
+        if compute.params.lakebase_mode {
+            ComputeNode::set_spec(&compute.params, &mut state, pspec);
+        } else {
+            state.pspec = Some(pspec);
+        }
+
         state.set_status(ComputeStatus::ConfigurationPending, &compute.state_changed);
         drop(state);
     }

compute_tools/src/http/routes/failpoints.rs

@@ -1,9 +1,8 @@
 use axum::response::{IntoResponse, Response};
 use http::StatusCode;
-use neon_failpoint::{configure_failpoint, configure_failpoint_with_context};
 use serde::{Deserialize, Serialize};
-use std::collections::HashMap;
 use tracing::info;
+use utils::failpoint_support::apply_failpoint;
 
 pub type ConfigureFailpointsRequest = Vec<FailpointConfig>;
 
@@ -12,16 +11,10 @@ pub type ConfigureFailpointsRequest = Vec<FailpointConfig>;
 pub struct FailpointConfig {
     /// Name of the fail point
     pub name: String,
-    /// List of actions to take, using the format described in neon_failpoint
+    /// List of actions to take, using the format described in `fail::cfg`
     ///
-    /// We support actions: "pause", "sleep(N)", "return", "return(value)", "exit", "off", "panic(message)"
-    /// Plus probability-based actions: "N%return(value)", "N%M*return(value)", "N%action", "N%M*action"
+    /// We also support `actions = "exit"` to cause the fail point to immediately exit.
     pub actions: String,
-    /// Optional context matching rules for conditional failpoints
-    /// Each key-value pair specifies a context key and a regex pattern to match against
-    /// All context matchers must match for the failpoint to trigger
-    #[serde(default, skip_serializing_if = "Option::is_none")]
-    pub context_matchers: Option<HashMap<String, String>>,
 }
 
 use crate::http::JsonResponse;
@@ -31,7 +24,7 @@ use crate::http::extract::Json;
 pub(in crate::http) async fn configure_failpoints(
     failpoints: Json<ConfigureFailpointsRequest>,
 ) -> Response {
-    if !neon_failpoint::has_failpoints() {
+    if !fail::has_failpoints() {
         return JsonResponse::error(
             StatusCode::PRECONDITION_FAILED,
             "Cannot manage failpoints because neon was compiled without failpoints support",
@@ -39,21 +32,16 @@ pub(in crate::http) async fn configure_failpoints(
     }
 
     for fp in &*failpoints {
-        info!(
-            "cfg failpoint: {} {} (context: {:?})",
-            fp.name, fp.actions, fp.context_matchers
-        );
+        info!("cfg failpoint: {} {}", fp.name, fp.actions);
 
-        let cfg_result = if let Some(context_matchers) = fp.context_matchers.clone() {
-            configure_failpoint_with_context(&fp.name, &fp.actions, context_matchers)
-        } else {
-            configure_failpoint(&fp.name, &fp.actions)
-        };
+        // We recognize one extra "action" that's not natively recognized
+        // by the failpoints crate: exit, to immediately kill the process
+        let cfg_result = apply_failpoint(&fp.name, &fp.actions);
 
         if let Err(e) = cfg_result {
             return JsonResponse::error(
                 StatusCode::BAD_REQUEST,
-                format!("failed to configure failpoint '{}': {e}", fp.name),
+                format!("failed to configure failpoints: {e}"),
             );
         }
     }

compute_tools/src/http/routes/hadron_liveness_probe.rs

@@ -0,0 +1,34 @@
+use crate::pg_isready::pg_isready;
+use crate::{compute::ComputeNode, http::JsonResponse};
+use axum::{extract::State, http::StatusCode, response::Response};
+use std::sync::Arc;
+
+/// NOTE: NOT ENABLED YET
+/// Detect if the compute is alive.
+/// Called by the liveness probe of the compute container.
+pub(in crate::http) async fn hadron_liveness_probe(
+    State(compute): State<Arc<ComputeNode>>,
+) -> Response {
+    let port = match compute.params.connstr.port() {
+        Some(port) => port,
+        None => {
+            return JsonResponse::error(
+                StatusCode::INTERNAL_SERVER_ERROR,
+                "Failed to get the port from the connection string",
+            );
+        }
+    };
+    match pg_isready(&compute.params.pg_isready_bin, port) {
+        Ok(_) => {
+            // The connection is successful, so the compute is alive.
+            // Return a 200 OK response.
+            JsonResponse::success(StatusCode::OK, "ok")
+        }
+        Err(e) => {
+            tracing::error!("Hadron liveness probe failed: {}", e);
+            // The connection failed, so the compute is not alive.
+            // Return a 500 Internal Server Error response.
+            JsonResponse::error(StatusCode::INTERNAL_SERVER_ERROR, e)
+        }
+    }
+}

compute_tools/src/http/routes/lfc.rs

@@ -1,12 +1,11 @@
-use crate::compute_prewarm::LfcPrewarmStateWithProgress;
 use crate::http::JsonResponse;
 use axum::response::{IntoResponse, Response};
 use axum::{Json, http::StatusCode};
 use axum_extra::extract::OptionalQuery;
-use compute_api::responses::LfcOffloadState;
+use compute_api::responses::{LfcOffloadState, LfcPrewarmState};
 type Compute = axum::extract::State<std::sync::Arc<crate::compute::ComputeNode>>;
 
-pub(in crate::http) async fn prewarm_state(compute: Compute) -> Json<LfcPrewarmStateWithProgress> {
+pub(in crate::http) async fn prewarm_state(compute: Compute) -> Json<LfcPrewarmState> {
     Json(compute.lfc_prewarm_state().await)
 }
 
@@ -46,3 +45,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/routes/metrics.rs

@@ -1,10 +1,19 @@
+use std::path::Path;
+use std::sync::Arc;
+
+use anyhow::Context;
 use axum::body::Body;
+use axum::extract::State;
 use axum::response::Response;
-use http::StatusCode;
 use http::header::CONTENT_TYPE;
+use http_body_util::BodyExt;
+use hyper::{Request, StatusCode};
 use metrics::proto::MetricFamily;
 use metrics::{Encoder, TextEncoder};
 
+use crate::communicator_socket_client::connect_communicator_socket;
+use crate::compute::ComputeNode;
+use crate::hadron_metrics;
 use crate::http::JsonResponse;
 use crate::metrics::collect;
 
@@ -13,11 +22,18 @@ pub(in crate::http) async fn get_metrics() -> Response {
     // When we call TextEncoder::encode() below, it will immediately return an
     // error if a metric family has no metrics, so we need to preemptively
     // filter out metric families with no metrics.
-    let metrics = collect()
+    let mut metrics = collect()
         .into_iter()
         .filter(|m| !m.get_metric().is_empty())
         .collect::<Vec<MetricFamily>>();
 
+    // Add Hadron metrics.
+    let hadron_metrics: Vec<MetricFamily> = hadron_metrics::collect()
+        .into_iter()
+        .filter(|m| !m.get_metric().is_empty())
+        .collect();
+    metrics.extend(hadron_metrics);
+
     let encoder = TextEncoder::new();
     let mut buffer = vec![];
 
@@ -31,3 +47,42 @@ pub(in crate::http) async fn get_metrics() -> Response {
         .body(Body::from(buffer))
         .unwrap()
 }
+
+/// Fetch and forward metrics from the Postgres neon extension's metrics
+/// exporter that are used by autoscaling-agent.
+///
+/// The neon extension exposes these metrics over a Unix domain socket
+/// in the data directory. That's not accessible directly from the outside
+/// world, so we have this endpoint in compute_ctl to expose it
+pub(in crate::http) async fn get_autoscaling_metrics(
+    State(compute): State<Arc<ComputeNode>>,
+) -> Result<Response, Response> {
+    let pgdata = Path::new(&compute.params.pgdata);
+
+    // Connect to the communicator process's metrics socket
+    let mut metrics_client = connect_communicator_socket(pgdata)
+        .await
+        .map_err(|e| JsonResponse::error(StatusCode::INTERNAL_SERVER_ERROR, format!("{e:#}")))?;
+
+    // Make a request for /autoscaling_metrics
+    let request = Request::builder()
+        .method("GET")
+        .uri("/autoscaling_metrics")
+        .header("Host", "localhost") // hyper requires Host, even though the server won't care
+        .body(Body::from(""))
+        .unwrap();
+    let resp = metrics_client
+        .send_request(request)
+        .await
+        .context("fetching metrics from Postgres metrics service")
+        .map_err(|e| JsonResponse::error(StatusCode::INTERNAL_SERVER_ERROR, format!("{e:#}")))?;
+
+    // Build a response that just forwards the response we got.
+    let mut response = Response::builder();
+    response = response.status(resp.status());
+    if let Some(content_type) = resp.headers().get(CONTENT_TYPE) {
+        response = response.header(CONTENT_TYPE, content_type);
+    }
+    let body = tonic::service::AxumBody::from_stream(resp.into_body().into_data_stream());
+    Ok(response.body(body).unwrap())
+}

compute_tools/src/http/routes/mod.rs

@@ -10,11 +10,13 @@ pub(in crate::http) mod extension_server;
 pub(in crate::http) mod extensions;
 pub(in crate::http) mod failpoints;
 pub(in crate::http) mod grants;
+pub(in crate::http) mod hadron_liveness_probe;
 pub(in crate::http) mod insights;
 pub(in crate::http) mod lfc;
 pub(in crate::http) mod metrics;
 pub(in crate::http) mod metrics_json;
 pub(in crate::http) mod promote;
+pub(in crate::http) mod refresh_configuration;
 pub(in crate::http) mod status;
 pub(in crate::http) mod terminate;
 

compute_tools/src/http/routes/promote.rs

@@ -1,14 +1,25 @@
 use crate::http::JsonResponse;
-use axum::Form;
+use axum::extract::Json;
+use compute_api::responses::PromoteConfig;
 use http::StatusCode;
 
 pub(in crate::http) async fn promote(
     compute: axum::extract::State<std::sync::Arc<crate::compute::ComputeNode>>,
-    Form(safekeepers_lsn): Form<compute_api::responses::SafekeepersLsn>,
+    Json(cfg): Json<PromoteConfig>,
 ) -> axum::response::Response {
-    let state = compute.promote(safekeepers_lsn).await;
-    if let compute_api::responses::PromoteState::Failed { error } = state {
-        return JsonResponse::error(StatusCode::INTERNAL_SERVER_ERROR, error);
+    // Return early at the cost of extra parsing spec
+    let pspec = match crate::compute::ParsedSpec::try_from(cfg.spec) {
+        Ok(p) => p,
+        Err(e) => return JsonResponse::error(StatusCode::BAD_REQUEST, e),
+    };
+
+    let cfg = PromoteConfig {
+        spec: pspec.spec,
+        wal_flush_lsn: cfg.wal_flush_lsn,
+    };
+    let state = compute.promote(cfg).await;
+    if let compute_api::responses::PromoteState::Failed { error: _ } = state {
+        return JsonResponse::create_response(StatusCode::INTERNAL_SERVER_ERROR, state);
     }
     JsonResponse::success(StatusCode::OK, state)
 }

compute_tools/src/http/routes/refresh_configuration.rs

@@ -0,0 +1,29 @@
+// This file is added by Hadron
+
+use std::sync::Arc;
+
+use axum::{
+    extract::State,
+    response::{IntoResponse, Response},
+};
+use http::StatusCode;
+
+use crate::compute::ComputeNode;
+use crate::hadron_metrics::POSTGRES_PAGESTREAM_REQUEST_ERRORS;
+use crate::http::JsonResponse;
+
+/// The /refresh_configuration POST method is used to nudge compute_ctl to pull a new spec
+/// from the HCC and attempt to reconfigure Postgres with the new spec. The method does not wait
+/// for the reconfiguration to complete. Rather, it simply delivers a signal that will cause
+/// configuration to be reloaded in a best effort manner. Invocation of this method does not
+/// guarantee that a reconfiguration will occur. The caller should consider keep sending this
+/// request while it believes that the compute configuration is out of date.
+pub(in crate::http) async fn refresh_configuration(
+    State(compute): State<Arc<ComputeNode>>,
+) -> Response {
+    POSTGRES_PAGESTREAM_REQUEST_ERRORS.inc();
+    match compute.signal_refresh_configuration().await {
+        Ok(_) => StatusCode::OK.into_response(),
+        Err(e) => JsonResponse::error(StatusCode::INTERNAL_SERVER_ERROR, e),
+    }
+}

compute_tools/src/http/routes/terminate.rs

@@ -1,7 +1,7 @@
 use crate::compute::{ComputeNode, forward_termination_signal};
 use crate::http::JsonResponse;
 use axum::extract::State;
-use axum::response::Response;
+use axum::response::{IntoResponse, Response};
 use axum_extra::extract::OptionalQuery;
 use compute_api::responses::{ComputeStatus, TerminateMode, TerminateResponse};
 use http::StatusCode;
@@ -33,7 +33,29 @@ pub(in crate::http) async fn terminate(
         if !matches!(state.status, ComputeStatus::Empty | ComputeStatus::Running) {
             return JsonResponse::invalid_status(state.status);
         }
+
+        // If compute is Empty, there's no Postgres to terminate. The regular compute_ctl termination path
+        // assumes Postgres to be configured and running, so we just special-handle this case by exiting
+        // the process directly.
+        if compute.params.lakebase_mode && state.status == ComputeStatus::Empty {
+            drop(state);
+            info!("terminating empty compute - will exit process");
+
+            // Queue a task to exit the process after 5 seconds. The 5-second delay aims to
+            // give enough time for the HTTP response to be sent so that HCM doesn't get an abrupt
+            // connection termination.
+            tokio::spawn(async {
+                tokio::time::sleep(tokio::time::Duration::from_secs(5)).await;
+                info!("exiting process after terminating empty compute");
+                std::process::exit(0);
+            });
+
+            return StatusCode::OK.into_response();
+        }
+
+        // For Running status, proceed with normal termination
         state.set_status(mode.into(), &compute.state_changed);
+        drop(state);
     }
 
     forward_termination_signal(false);

compute_tools/src/http/server.rs

@@ -23,7 +23,8 @@ use super::{
     middleware::authorize::Authorize,
     routes::{
         check_writability, configure, database_schema, dbs_and_roles, extension_server, extensions,
-        grants, insights, lfc, metrics, metrics_json, promote, status, terminate,
+        grants, hadron_liveness_probe, insights, lfc, metrics, metrics_json, promote,
+        refresh_configuration, status, terminate,
     },
 };
 use crate::compute::ComputeNode;
@@ -43,6 +44,7 @@ pub enum Server {
         port: u16,
         config: ComputeCtlConfig,
         compute_id: String,
+        instance_id: Option<String>,
     },
 }
 
@@ -67,7 +69,12 @@ impl From<&Server> for Router<Arc<ComputeNode>> {
                         post(extension_server::download_extension),
                     )
                     .route("/extensions", post(extensions::install_extension))
-                    .route("/grants", post(grants::add_grant));
+                    .route("/grants", post(grants::add_grant))
+                    // Hadron: Compute-initiated configuration refresh
+                    .route(
+                        "/refresh_configuration",
+                        post(refresh_configuration::refresh_configuration),
+                    );
 
                 // Add in any testing support
                 if cfg!(feature = "testing") {
@@ -79,13 +86,25 @@ impl From<&Server> for Router<Arc<ComputeNode>> {
                 router
             }
             Server::External {
-                config, compute_id, ..
+                config,
+                compute_id,
+                instance_id,
+                ..
             } => {
-                let unauthenticated_router =
-                    Router::<Arc<ComputeNode>>::new().route("/metrics", get(metrics::get_metrics));
+                let unauthenticated_router = Router::<Arc<ComputeNode>>::new()
+                    .route("/metrics", get(metrics::get_metrics))
+                    .route(
+                        "/autoscaling_metrics",
+                        get(metrics::get_autoscaling_metrics),
+                    );
 
                 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))
@@ -96,8 +115,13 @@ impl From<&Server> for Router<Arc<ComputeNode>> {
                     .route("/metrics.json", get(metrics_json::get_metrics))
                     .route("/status", get(status::get_status))
                     .route("/terminate", post(terminate::terminate))
+                    .route(
+                        "/hadron_liveness_probe",
+                        get(hadron_liveness_probe::hadron_liveness_probe),
+                    )
                     .layer(AsyncRequireAuthorizationLayer::new(Authorize::new(
                         compute_id.clone(),
+                        instance_id.clone(),
                         config.jwks.clone(),
                     )));
 

compute_tools/src/installed_extensions.rs

@@ -2,6 +2,8 @@ use std::collections::HashMap;
 
 use anyhow::Result;
 use compute_api::responses::{InstalledExtension, InstalledExtensions};
+use once_cell::sync::Lazy;
+use tokio_postgres::error::Error as PostgresError;
 use tokio_postgres::{Client, Config, NoTls};
 
 use crate::metrics::INSTALLED_EXTENSIONS;
@@ -10,14 +12,14 @@ use crate::metrics::INSTALLED_EXTENSIONS;
 /// and to make database listing query here more explicit.
 ///
 /// Limit the number of databases to 500 to avoid excessive load.
-async fn list_dbs(client: &mut Client) -> Result<Vec<String>> {
+async fn list_dbs(client: &mut Client) -> Result<Vec<String>, PostgresError> {
     // `pg_database.datconnlimit = -2` means that the database is in the
     // invalid state
     let databases = client
         .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",
             &[],
         )
@@ -37,7 +39,9 @@ async fn list_dbs(client: &mut Client) -> Result<Vec<String>> {
 /// Same extension can be installed in multiple databases with different versions,
 /// so we report a separate metric (number of databases where it is installed)
 /// for each extension version.
-pub async fn get_installed_extensions(mut conf: Config) -> Result<InstalledExtensions> {
+pub async fn get_installed_extensions(
+    mut conf: Config,
+) -> Result<InstalledExtensions, PostgresError> {
     conf.application_name("compute_ctl:get_installed_extensions");
     let databases: Vec<String> = {
         let (mut client, connection) = conf.connect(NoTls).await?;
@@ -63,7 +67,7 @@ pub async fn get_installed_extensions(mut conf: Config) -> Result<InstalledExten
 
         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?
@@ -116,3 +120,7 @@ pub async fn get_installed_extensions(mut conf: Config) -> Result<InstalledExten
         extensions: extensions_map.into_values().collect(),
     })
 }
+
+pub fn initialize_metrics() {
+    Lazy::force(&INSTALLED_EXTENSIONS);
+}

compute_tools/src/lib.rs

@@ -4,6 +4,7 @@
 #![deny(clippy::undocumented_unsafe_blocks)]
 
 pub mod checker;
+pub mod communicator_socket_client;
 pub mod config;
 pub mod configurator;
 pub mod http;
@@ -15,6 +16,7 @@ pub mod compute_prewarm;
 pub mod compute_promote;
 pub mod disk_quota;
 pub mod extension_server;
+pub mod hadron_metrics;
 pub mod installed_extensions;
 pub mod local_proxy;
 pub mod lsn_lease;
@@ -23,6 +25,7 @@ mod migration;
 pub mod monitor;
 pub mod params;
 pub mod pg_helpers;
+pub mod pg_isready;
 pub mod pgbouncer;
 pub mod rsyslog;
 pub mod spec;

compute_tools/src/logger.rs

@@ -1,7 +1,10 @@
 use std::collections::HashMap;
+use std::sync::{LazyLock, RwLock};
+use tracing::Subscriber;
 use tracing::info;
-use tracing_subscriber::layer::SubscriberExt;
+use tracing_appender;
 use tracing_subscriber::prelude::*;
+use tracing_subscriber::{fmt, layer::SubscriberExt, registry::LookupSpan};
 
 /// Initialize logging to stderr, and OpenTelemetry tracing and exporter.
 ///
@@ -13,31 +16,63 @@ use tracing_subscriber::prelude::*;
 /// set `OTEL_EXPORTER_OTLP_ENDPOINT=http://jaeger:4318`. See
 /// `tracing-utils` package description.
 ///
-pub async fn init_tracing_and_logging(default_log_level: &str) -> anyhow::Result<()> {
+pub fn init_tracing_and_logging(
+    default_log_level: &str,
+    log_dir_opt: &Option<String>,
+) -> anyhow::Result<(
+    Option<tracing_utils::Provider>,
+    Option<tracing_appender::non_blocking::WorkerGuard>,
+)> {
     // Initialize Logging
     let env_filter = tracing_subscriber::EnvFilter::try_from_default_env()
         .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new(default_log_level));
 
+    // Standard output streams
     let fmt_layer = tracing_subscriber::fmt::layer()
         .with_ansi(false)
         .with_target(false)
         .with_writer(std::io::stderr);
 
+    // Logs with file rotation. Files in `$log_dir/pgcctl.yyyy-MM-dd`
+    let (json_to_file_layer, _file_logs_guard) = if let Some(log_dir) = log_dir_opt {
+        std::fs::create_dir_all(log_dir)?;
+        let file_logs_appender = tracing_appender::rolling::RollingFileAppender::builder()
+            .rotation(tracing_appender::rolling::Rotation::DAILY)
+            .filename_prefix("pgcctl")
+            // Lib appends to existing files, so we will keep files for up to 2 days even on restart loops.
+            // At minimum, log-daemon will have 1 day to detect and upload a file (if created right before midnight).
+            .max_log_files(2)
+            .build(log_dir)
+            .expect("Initializing rolling file appender should succeed");
+        let (file_logs_writer, _file_logs_guard) =
+            tracing_appender::non_blocking(file_logs_appender);
+        let json_to_file_layer = tracing_subscriber::fmt::layer()
+            .with_ansi(false)
+            .with_target(false)
+            .event_format(PgJsonLogShapeFormatter)
+            .with_writer(file_logs_writer);
+        (Some(json_to_file_layer), Some(_file_logs_guard))
+    } else {
+        (None, None)
+    };
+
     // Initialize OpenTelemetry
-    let otlp_layer =
-        tracing_utils::init_tracing("compute_ctl", tracing_utils::ExportConfig::default()).await;
+    let provider =
+        tracing_utils::init_tracing("compute_ctl", tracing_utils::ExportConfig::default());
+    let otlp_layer = provider.as_ref().map(tracing_utils::layer);
 
     // Put it all together
     tracing_subscriber::registry()
         .with(env_filter)
         .with(otlp_layer)
         .with(fmt_layer)
+        .with(json_to_file_layer)
         .init();
     tracing::info!("logging and tracing started");
 
     utils::logging::replace_panic_hook_with_tracing_panic_hook().forget();
 
-    Ok(())
+    Ok((provider, _file_logs_guard))
 }
 
 /// Replace all newline characters with a special character to make it
@@ -92,3 +127,157 @@ pub fn startup_context_from_env() -> Option<opentelemetry::Context> {
         None
     }
 }
+
+/// Track relevant id's
+const UNKNOWN_IDS: &str = r#""pg_instance_id": "", "pg_compute_id": """#;
+static IDS: LazyLock<RwLock<String>> = LazyLock::new(|| RwLock::new(UNKNOWN_IDS.to_string()));
+
+pub fn update_ids(instance_id: &Option<String>, compute_id: &Option<String>) -> anyhow::Result<()> {
+    let ids = format!(
+        r#""pg_instance_id": "{}", "pg_compute_id": "{}""#,
+        instance_id.as_ref().map(|s| s.as_str()).unwrap_or_default(),
+        compute_id.as_ref().map(|s| s.as_str()).unwrap_or_default()
+    );
+    let mut guard = IDS
+        .write()
+        .map_err(|e| anyhow::anyhow!("Log set id's rwlock poisoned: {}", e))?;
+    *guard = ids;
+    Ok(())
+}
+
+/// Massage compute_ctl logs into PG json log shape so we can use the same Lumberjack setup.
+struct PgJsonLogShapeFormatter;
+impl<S, N> fmt::format::FormatEvent<S, N> for PgJsonLogShapeFormatter
+where
+    S: Subscriber + for<'a> LookupSpan<'a>,
+    N: for<'a> fmt::format::FormatFields<'a> + 'static,
+{
+    fn format_event(
+        &self,
+        ctx: &fmt::FmtContext<'_, S, N>,
+        mut writer: fmt::format::Writer<'_>,
+        event: &tracing::Event<'_>,
+    ) -> std::fmt::Result {
+        // Format values from the event's metadata, and open message string
+        let metadata = event.metadata();
+        {
+            let ids_guard = IDS.read();
+            let ids = ids_guard
+                .as_ref()
+                .map(|guard| guard.as_str())
+                // Surpress so that we don't lose all uploaded/ file logs if something goes super wrong. We would notice the missing id's.
+                .unwrap_or(UNKNOWN_IDS);
+            write!(
+                &mut writer,
+                r#"{{"timestamp": "{}", "error_severity": "{}", "file_name": "{}", "backend_type": "compute_ctl_self", {}, "message": "#,
+                chrono::Utc::now().format("%Y-%m-%d %H:%M:%S%.3f GMT"),
+                metadata.level(),
+                metadata.target(),
+                ids
+            )?;
+        }
+
+        let mut message = String::new();
+        let message_writer = fmt::format::Writer::new(&mut message);
+
+        // Gather the message
+        ctx.field_format().format_fields(message_writer, event)?;
+
+        // TODO: any better options than to copy-paste this OSS span formatter?
+        // impl<S, N, T> FormatEvent<S, N> for Format<Full, T>
+        // https://docs.rs/tracing-subscriber/latest/tracing_subscriber/fmt/trait.FormatEvent.html#impl-FormatEvent%3CS,+N%3E-for-Format%3CFull,+T%3E
+
+        // write message, close bracket, and new line
+        writeln!(writer, "{}}}", serde_json::to_string(&message).unwrap())
+    }
+}
+
+#[cfg(feature = "testing")]
+#[cfg(test)]
+mod test {
+    use super::*;
+    use std::{cell::RefCell, io};
+
+    // Use thread_local! instead of Mutex for test isolation
+    thread_local! {
+        static WRITER_OUTPUT: RefCell<String> = const { RefCell::new(String::new()) };
+    }
+
+    #[derive(Clone, Default)]
+    struct StaticStringWriter;
+
+    impl io::Write for StaticStringWriter {
+        fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
+            let output = String::from_utf8(buf.to_vec()).expect("Invalid UTF-8 in test output");
+            WRITER_OUTPUT.with(|s| s.borrow_mut().push_str(&output));
+            Ok(buf.len())
+        }
+
+        fn flush(&mut self) -> io::Result<()> {
+            Ok(())
+        }
+    }
+
+    impl fmt::MakeWriter<'_> for StaticStringWriter {
+        type Writer = Self;
+
+        fn make_writer(&self) -> Self::Writer {
+            Self
+        }
+    }
+
+    #[test]
+    fn test_log_pg_json_shape_formatter() {
+        // Use a scoped subscriber to prevent global state pollution
+        let subscriber = tracing_subscriber::registry().with(
+            tracing_subscriber::fmt::layer()
+                .with_ansi(false)
+                .with_target(false)
+                .event_format(PgJsonLogShapeFormatter)
+                .with_writer(StaticStringWriter),
+        );
+
+        let _ = update_ids(&Some("000".to_string()), &Some("111".to_string()));
+
+        // Clear any previous test state
+        WRITER_OUTPUT.with(|s| s.borrow_mut().clear());
+
+        let messages = [
+            "test message",
+            r#"json escape check:  name="BatchSpanProcessor.Flush.ExportError" reason="Other(reqwest::Error { kind: Request, url: \"http://localhost:4318/v1/traces\", source: hyper_
+            util::client::legacy::Error(Connect, ConnectError(\"tcp connect error\", Os { code: 111, kind: ConnectionRefused, message: \"Connection refused\" })) })" Failed during the export process"#,
+        ];
+
+        tracing::subscriber::with_default(subscriber, || {
+            for message in messages {
+                tracing::info!(message);
+            }
+        });
+        tracing::info!("not test message");
+
+        // Get captured output
+        let output = WRITER_OUTPUT.with(|s| s.borrow().clone());
+
+        let json_strings: Vec<&str> = output.lines().collect();
+        assert_eq!(
+            json_strings.len(),
+            messages.len(),
+            "Log didn't have the expected number of json strings."
+        );
+
+        let json_string_shape_regex = regex::Regex::new(
+            r#"\{"timestamp": "\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2}\.\d{3} GMT", "error_severity": "INFO", "file_name": ".+", "backend_type": "compute_ctl_self", "pg_instance_id": "000", "pg_compute_id": "111", "message": ".+"\}"#
+        ).unwrap();
+
+        for (i, expected_message) in messages.iter().enumerate() {
+            let json_string = json_strings[i];
+            assert!(
+                json_string_shape_regex.is_match(json_string),
+                "Json log didn't match expected pattern:\n{json_string}",
+            );
+            let parsed_json: serde_json::Value = serde_json::from_str(json_string).unwrap();
+            let actual_message = parsed_json["message"].as_str().unwrap();
+            assert_eq!(*expected_message, actual_message);
+        }
+    }
+}

compute_tools/src/lsn_lease.rs

@@ -4,14 +4,13 @@ use std::thread;
 use std::time::{Duration, SystemTime};
 
 use anyhow::{Result, bail};
-use compute_api::spec::{ComputeMode, PageserverProtocol};
-use itertools::Itertools as _;
+use compute_api::spec::{ComputeMode, PageserverConnectionInfo, PageserverProtocol};
 use pageserver_page_api as page_api;
 use postgres::{NoTls, SimpleQueryMessage};
 use tracing::{info, warn};
 use utils::id::{TenantId, TimelineId};
 use utils::lsn::Lsn;
-use utils::shard::{ShardCount, ShardNumber, TenantShardId};
+use utils::shard::TenantShardId;
 
 use crate::compute::ComputeNode;
 
@@ -78,17 +77,16 @@ fn acquire_lsn_lease_with_retry(
 
     loop {
         // Note: List of pageservers is dynamic, need to re-read configs before each attempt.
-        let (connstrings, auth) = {
+        let (conninfo, auth) = {
             let state = compute.state.lock().unwrap();
             let spec = state.pspec.as_ref().expect("spec must be set");
             (
-                spec.pageserver_connstr.clone(),
+                spec.pageserver_conninfo.clone(),
                 spec.storage_auth_token.clone(),
             )
         };
 
-        let result =
-            try_acquire_lsn_lease(&connstrings, auth.as_deref(), tenant_id, timeline_id, lsn);
+        let result = try_acquire_lsn_lease(conninfo, auth.as_deref(), tenant_id, timeline_id, lsn);
         match result {
             Ok(Some(res)) => {
                 return Ok(res);
@@ -112,35 +110,44 @@ fn acquire_lsn_lease_with_retry(
 
 /// Tries to acquire LSN leases on all Pageserver shards.
 fn try_acquire_lsn_lease(
-    connstrings: &str,
+    conninfo: PageserverConnectionInfo,
     auth: Option<&str>,
     tenant_id: TenantId,
     timeline_id: TimelineId,
     lsn: Lsn,
 ) -> Result<Option<SystemTime>> {
-    let connstrings = connstrings.split(',').collect_vec();
-    let shard_count = connstrings.len();
     let mut leases = Vec::new();
 
-    for (shard_number, &connstring) in connstrings.iter().enumerate() {
-        let tenant_shard_id = match shard_count {
-            0 | 1 => TenantShardId::unsharded(tenant_id),
-            shard_count => TenantShardId {
-                tenant_id,
-                shard_number: ShardNumber(shard_number as u8),
-                shard_count: ShardCount::new(shard_count as u8),
-            },
+    for (shard_index, shard) in conninfo.shards.into_iter() {
+        let tenant_shard_id = TenantShardId {
+            tenant_id,
+            shard_number: shard_index.shard_number,
+            shard_count: shard_index.shard_count,
         };
 
-        let lease = match PageserverProtocol::from_connstring(connstring)? {
-            PageserverProtocol::Libpq => {
-                acquire_lsn_lease_libpq(connstring, auth, tenant_shard_id, timeline_id, lsn)?
-            }
-            PageserverProtocol::Grpc => {
-                acquire_lsn_lease_grpc(connstring, auth, tenant_shard_id, timeline_id, lsn)?
-            }
-        };
-        leases.push(lease);
+        // XXX: If there are more than pageserver for the one shard, do we need to get a
+        // leas on all of them? Currently, that's what we assume, but this is hypothetical
+        // as of this writing, as we never pass the info for more than one pageserver per
+        // shard.
+        for pageserver in shard.pageservers {
+            let lease = match conninfo.prefer_protocol {
+                PageserverProtocol::Grpc => acquire_lsn_lease_grpc(
+                    &pageserver.grpc_url.unwrap(),
+                    auth,
+                    tenant_shard_id,
+                    timeline_id,
+                    lsn,
+                )?,
+                PageserverProtocol::Libpq => acquire_lsn_lease_libpq(
+                    &pageserver.libpq_url.unwrap(),
+                    auth,
+                    tenant_shard_id,
+                    timeline_id,
+                    lsn,
+                )?,
+            };
+            leases.push(lease);
+        }
     }
 
     Ok(leases.into_iter().min().flatten())

compute_tools/src/migration.rs

@@ -1,5 +1,5 @@
 use anyhow::{Context, Result};
-use neon_failpoint::fail_point;
+use fail::fail_point;
 use tokio_postgres::{Client, Transaction};
 use tracing::{error, info};
 
@@ -9,15 +9,20 @@ use crate::metrics::DB_MIGRATION_FAILED;
 pub(crate) struct MigrationRunner<'m> {
     client: &'m mut Client,
     migrations: &'m [&'m str],
+    lakebase_mode: bool,
 }
 
 impl<'m> MigrationRunner<'m> {
     /// Create a new migration runner
-    pub fn new(client: &'m mut Client, migrations: &'m [&'m str]) -> Self {
+    pub fn new(client: &'m mut Client, migrations: &'m [&'m str], lakebase_mode: bool) -> Self {
         // The neon_migration.migration_id::id column is a bigint, which is equivalent to an i64
         assert!(migrations.len() + 1 < i64::MAX as usize);
 
-        Self { client, migrations }
+        Self {
+            client,
+            migrations,
+            lakebase_mode,
+        }
     }
 
     /// Get the current value neon_migration.migration_id
@@ -40,14 +45,13 @@ impl<'m> MigrationRunner<'m> {
         // middle of applying a series of migrations fails in an expected
         // manner
         if cfg!(feature = "testing") {
-            let fail = async {
-                fail_point!("compute-migration", |fail_migration_id: Option<String>| {
+            let fail = (|| {
+                fail_point!("compute-migration", |fail_migration_id| {
                     migration_id == fail_migration_id.unwrap().parse::<i64>().unwrap()
                 });
 
                 false
-            }
-            .await;
+            })();
 
             if fail {
                 return Err(anyhow::anyhow!(format!(
@@ -72,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",
@@ -131,8 +135,13 @@ impl<'m> MigrationRunner<'m> {
             // ID is also the next index
             let migration_id = (current_migration + 1) as i64;
             let migration = self.migrations[current_migration];
+            let migration = if self.lakebase_mode {
+                migration.replace("neon_superuser", "databricks_superuser")
+            } else {
+                migration.to_string()
+            };
 
-            match Self::run_migration(self.client, migration_id, migration).await {
+            match Self::run_migration(self.client, migration_id, &migration).await {
                 Ok(_) => {
                     info!("Finished migration id={}", migration_id);
                 }

compute_tools/src/migrations/0001-add_bypass_rls_to_privileged_role.sql

@@ -0,0 +1 @@
+ALTER ROLE {privileged_role_name} BYPASSRLS;

compute_tools/src/migrations/0001-neon_superuser_bypass_rls.sql

@@ -1 +0,0 @@
-ALTER ROLE neon_superuser BYPASSRLS;

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, 'neon_superuser', '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, 'neon_superuser', '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_neon_superuser.sql

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

compute_tools/src/migrations/0003-grant_pg_create_subscription_to_privileged_role.sql

@@ -0,0 +1,6 @@
+DO $$
+BEGIN
+    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/0004-grant_pg_monitor_to_neon_superuser.sql

@@ -1 +0,0 @@
-GRANT pg_monitor TO neon_superuser WITH ADMIN OPTION;

compute_tools/src/migrations/0004-grant_pg_monitor_to_privileged_role.sql

@@ -0,0 +1 @@
+GRANT pg_monitor TO {privileged_role_name} WITH ADMIN OPTION;

compute_tools/src/migrations/0005-grant_all_on_tables_to_privileged_role.sql

@@ -1,4 +1,4 @@
 -- SKIP: Deemed insufficient for allowing relations created by extensions to be
---       interacted with by neon_superuser without permission issues.
+--       interacted with by {privileged_role_name} without permission issues.
 
-ALTER DEFAULT PRIVILEGES IN SCHEMA public GRANT ALL ON TABLES TO neon_superuser;
+ALTER DEFAULT PRIVILEGES IN SCHEMA public GRANT ALL ON TABLES TO {privileged_role_name};

compute_tools/src/migrations/0006-grant_all_on_sequences_to_privileged_role.sql

@@ -1,4 +1,4 @@
 -- SKIP: Deemed insufficient for allowing relations created by extensions to be
---       interacted with by neon_superuser without permission issues.
+--       interacted with by {privileged_role_name} without permission issues.
 
-ALTER DEFAULT PRIVILEGES IN SCHEMA public GRANT ALL ON SEQUENCES TO neon_superuser;
+ALTER DEFAULT PRIVILEGES IN SCHEMA public GRANT ALL ON SEQUENCES TO {privileged_role_name};

compute_tools/src/migrations/0007-grant_all_on_tables_with_grant_option_to_privileged_role.sql

@@ -1,3 +1,3 @@
 -- SKIP: Moved inline to the handle_grants() functions.
 
-ALTER DEFAULT PRIVILEGES IN SCHEMA public GRANT ALL ON TABLES TO neon_superuser WITH GRANT OPTION;
+ALTER DEFAULT PRIVILEGES IN SCHEMA public GRANT ALL ON TABLES TO {privileged_role_name} WITH GRANT OPTION;

compute_tools/src/migrations/0008-grant_all_on_sequences_with_grant_option_to_privileged_role.sql

@@ -1,3 +1,3 @@
 -- SKIP: Moved inline to the handle_grants() functions.
 
-ALTER DEFAULT PRIVILEGES IN SCHEMA public GRANT ALL ON SEQUENCES TO neon_superuser WITH GRANT OPTION;
+ALTER DEFAULT PRIVILEGES IN SCHEMA public GRANT ALL ON SEQUENCES TO {privileged_role_name} WITH GRANT OPTION;