fix customize config test

Signed-off-by: Alex Chi Z <chi@neon.tech>

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

.github/actionlint.yml

@@ -31,7 +31,7 @@ config-variables:
   - NEON_PROD_AWS_ACCOUNT_ID
   - PGREGRESS_PG16_PROJECT_ID
   - PGREGRESS_PG17_PROJECT_ID
-  - PREWARM_PGBENCH_SIZE
+  - 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/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=""

.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

@@ -418,7 +418,7 @@ jobs:
       statuses: write
       id-token: write # aws-actions/configure-aws-credentials
     env:
-      PGBENCH_SIZE: ${{ vars.PREWARM_PGBENCH_SIZE }}
+      PROJECT_ID: ${{ vars.PREWARM_PROJECT_ID }}
       POSTGRES_DISTRIB_DIR: /tmp/neon/pg_install
       DEFAULT_PG_VERSION: 17
       TEST_OUTPUT: /tmp/test_output

.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

@@ -632,6 +632,8 @@ 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
+          secrets: |
+            SUBZERO_ACCESS_TOKEN=${{ secrets.CI_ACCESS_TOKEN }}
           provenance: false
           push: true
           pull: true

.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,20 @@ jobs:
       options: --init --user root
     services:
       clickhouse:
-        image: clickhouse/clickhouse-server:24.6.3.64
+        image: clickhouse/clickhouse-server:24.8
+        env:
+          CLICKHOUSE_PASSWORD: ${{ needs.generate-ch-tmppw.outputs.tmp_val }}
         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 +95,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 +141,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

@@ -15,7 +15,6 @@ neon.iml
 /.neon
 /integration_tests/.neon
 compaction-suite-results.*
-pgxn/neon/communicator/communicator_bindings.h
 docker-compose/docker-compose-parallel.yml
 
 # Coverage
@@ -27,9 +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

Cargo.toml

@@ -35,7 +35,6 @@ members = [
     "libs/pq_proto",
     "libs/tenant_size_model",
     "libs/metrics",
-    "libs/neonart",
     "libs/postgres_connection",
     "libs/remote_storage",
     "libs/tracing-utils",
@@ -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",
@@ -93,7 +93,6 @@ clap = { version = "4.0", features = ["derive", "env"] }
 clashmap = { version = "1.0", features = ["raw-api"] }
 comfy-table = "7.1"
 const_format = "0.2"
-crossbeam-utils = "0.8.21"
 crc32c = "0.6"
 diatomic-waker = { version = "0.2.3" }
 either = "1.8"
@@ -137,6 +136,7 @@ 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.
@@ -144,27 +144,28 @@ 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"
 pbkdf2 = { version = "0.12.1", features = ["simple", "std"] }
 pem = "3.0.3"
-peekable = "0.3.0"
 pin-project-lite = "0.2"
 pprof = { version = "0.14", features = ["criterion", "flamegraph", "frame-pointer", "prost-codec"] }
 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"
@@ -190,7 +191,6 @@ smallvec = "1.11"
 smol_str = { version = "0.2.0", features = ["serde"] }
 socket2 = "0.5"
 spki = "0.7.3"
-spin = "0.9.8"
 strum = "0.26"
 strum_macros = "0.26"
 "subtle"  = "2.5.0"
@@ -202,6 +202,7 @@ thiserror = "1.0"
 tikv-jemallocator = { version = "0.6", features = ["profiling", "stats", "unprefixed_malloc_on_supported_platforms"] }
 tikv-jemalloc-ctl = { version = "0.6", features = ["stats"] }
 tokio = { version = "1.43.1", features = ["macros"] }
+tokio-epoll-uring = { git = "https://github.com/neondatabase/tokio-epoll-uring.git" , branch = "main" }
 tokio-io-timeout = "1.2.0"
 tokio-postgres-rustls = "0.12.0"
 tokio-rustls = { version = "0.26.0", default-features = false, features = ["tls12", "ring"]}
@@ -214,17 +215,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 }
@@ -243,9 +242,6 @@ x509-cert = { version = "0.2.5" }
 env_logger = "0.11"
 log = "0.4"
 
-tokio-epoll-uring = { git = "https://github.com/neondatabase/tokio-epoll-uring.git" , branch = "main" }
-uring-common = { git = "https://github.com/neondatabase/tokio-epoll-uring.git" , branch = "main" }
-
 ## Libraries from neondatabase/ git forks, ideally with changes to be upstreamed
 postgres = { git = "https://github.com/neondatabase/rust-postgres.git", branch = "neon" }
 postgres-protocol = { git = "https://github.com/neondatabase/rust-postgres.git", branch = "neon" }

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 \
+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 build \
+      --features $CARGO_FEATURES \
       --bin pg_sni_router  \
       --bin pageserver  \
       --bin pagectl  \

build-tools/Dockerfile

@@ -39,13 +39,13 @@ 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,22 +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 install -y nodejs \
+    && 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
@@ -215,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
@@ -254,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 && \
@@ -265,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
@@ -279,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
@@ -317,13 +314,13 @@ 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      --locked --version ${RUSTFILT_VERSION} && \
-    cargo install cargo-hakari  --locked --version ${CARGO_HAKARI_VERSION} && \
-    cargo install cargo-deny    --locked --version ${CARGO_DENY_VERSION} && \
-    cargo install cargo-hack    --locked --version ${CARGO_HACK_VERSION} && \
-    cargo install cargo-nextest --locked --version ${CARGO_NEXTEST_VERSION} && \
-    cargo install cargo-chef    --locked --version ${CARGO_CHEF_VERSION} && \
-    cargo install diesel_cli    --locked --version ${CARGO_DIESEL_CLI_VERSION} \
+    cargo install rustfilt      --locked --version "${RUSTFILT_VERSION}" && \
+    cargo install cargo-hakari  --locked --version "${CARGO_HAKARI_VERSION}" && \
+    cargo install cargo-deny    --locked --version "${CARGO_DENY_VERSION}" && \
+    cargo install cargo-hack    --locked --version "${CARGO_HACK_VERSION}" && \
+    cargo install cargo-nextest --locked --version "${CARGO_NEXTEST_VERSION}" && \
+    cargo install cargo-chef    --locked --version "${CARGO_CHEF_VERSION}" && \
+    cargo 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-lock.json

@@ -6,7 +6,7 @@
     "": {
       "name": "build-tools",
       "devDependencies": {
-        "@redocly/cli": "1.34.4",
+        "@redocly/cli": "1.34.5",
         "@sourcemeta/jsonschema": "10.0.0"
       }
     },
@@ -472,9 +472,9 @@
       }
     },
     "node_modules/@redocly/cli": {
-      "version": "1.34.4",
-      "resolved": "https://registry.npmjs.org/@redocly/cli/-/cli-1.34.4.tgz",
-      "integrity": "sha512-seH/GgrjSB1EeOsgJ/4Ct6Jk2N7sh12POn/7G8UQFARMyUMJpe1oHtBwT2ndfp4EFCpgBAbZ/82Iw6dwczNxEA==",
+      "version": "1.34.5",
+      "resolved": "https://registry.npmjs.org/@redocly/cli/-/cli-1.34.5.tgz",
+      "integrity": "sha512-5IEwxs7SGP5KEXjBKLU8Ffdz9by/KqNSeBk6YUVQaGxMXK//uYlTJIPntgUXbo1KAGG2d2q2XF8y4iFz6qNeiw==",
       "dev": true,
       "license": "MIT",
       "dependencies": {
@@ -484,14 +484,14 @@
         "@opentelemetry/sdk-trace-node": "1.26.0",
         "@opentelemetry/semantic-conventions": "1.27.0",
         "@redocly/config": "^0.22.0",
-        "@redocly/openapi-core": "1.34.4",
-        "@redocly/respect-core": "1.34.4",
+        "@redocly/openapi-core": "1.34.5",
+        "@redocly/respect-core": "1.34.5",
         "abort-controller": "^3.0.0",
         "chokidar": "^3.5.1",
         "colorette": "^1.2.0",
         "core-js": "^3.32.1",
         "dotenv": "16.4.7",
-        "form-data": "^4.0.0",
+        "form-data": "^4.0.4",
         "get-port-please": "^3.0.1",
         "glob": "^7.1.6",
         "handlebars": "^4.7.6",
@@ -522,9 +522,9 @@
       "license": "MIT"
     },
     "node_modules/@redocly/openapi-core": {
-      "version": "1.34.4",
-      "resolved": "https://registry.npmjs.org/@redocly/openapi-core/-/openapi-core-1.34.4.tgz",
-      "integrity": "sha512-hf53xEgpXIgWl3b275PgZU3OTpYh1RoD2LHdIfQ1JzBNTWsiNKczTEsI/4Tmh2N1oq9YcphhSMyk3lDh85oDjg==",
+      "version": "1.34.5",
+      "resolved": "https://registry.npmjs.org/@redocly/openapi-core/-/openapi-core-1.34.5.tgz",
+      "integrity": "sha512-0EbE8LRbkogtcCXU7liAyC00n9uNG9hJ+eMyHFdUsy9lB/WGqnEBgwjA9q2cyzAVcdTkQqTBBU1XePNnN3OijA==",
       "dev": true,
       "license": "MIT",
       "dependencies": {
@@ -544,21 +544,21 @@
       }
     },
     "node_modules/@redocly/respect-core": {
-      "version": "1.34.4",
-      "resolved": "https://registry.npmjs.org/@redocly/respect-core/-/respect-core-1.34.4.tgz",
-      "integrity": "sha512-MitKyKyQpsizA4qCVv+MjXL4WltfhFQAoiKiAzrVR1Kusro3VhYb6yJuzoXjiJhR0ukLP5QOP19Vcs7qmj9dZg==",
+      "version": "1.34.5",
+      "resolved": "https://registry.npmjs.org/@redocly/respect-core/-/respect-core-1.34.5.tgz",
+      "integrity": "sha512-GheC/g/QFztPe9UA9LamooSplQuy9pe0Yr8XGTqkz0ahivLDl7svoy/LSQNn1QH3XGtLKwFYMfTwFR2TAYyh5Q==",
       "dev": true,
       "license": "MIT",
       "dependencies": {
         "@faker-js/faker": "^7.6.0",
         "@redocly/ajv": "8.11.2",
-        "@redocly/openapi-core": "1.34.4",
+        "@redocly/openapi-core": "1.34.5",
         "better-ajv-errors": "^1.2.0",
         "colorette": "^2.0.20",
         "concat-stream": "^2.0.0",
         "cookie": "^0.7.2",
         "dotenv": "16.4.7",
-        "form-data": "4.0.0",
+        "form-data": "^4.0.4",
         "jest-diff": "^29.3.1",
         "jest-matcher-utils": "^29.3.1",
         "js-yaml": "4.1.0",
@@ -582,21 +582,6 @@
       "dev": true,
       "license": "MIT"
     },
-    "node_modules/@redocly/respect-core/node_modules/form-data": {
-      "version": "4.0.0",
-      "resolved": "https://registry.npmjs.org/form-data/-/form-data-4.0.0.tgz",
-      "integrity": "sha512-ETEklSGi5t0QMZuiXoA/Q6vcnxcLQP5vdugSpuAyi6SVGi2clPPp+xgEhuMaHC+zGgn31Kd235W35f7Hykkaww==",
-      "dev": true,
-      "license": "MIT",
-      "dependencies": {
-        "asynckit": "^0.4.0",
-        "combined-stream": "^1.0.8",
-        "mime-types": "^2.1.12"
-      },
-      "engines": {
-        "node": ">= 6"
-      }
-    },
     "node_modules/@sinclair/typebox": {
       "version": "0.27.8",
       "resolved": "https://registry.npmjs.org/@sinclair/typebox/-/typebox-0.27.8.tgz",
@@ -1345,9 +1330,9 @@
       "license": "MIT"
     },
     "node_modules/form-data": {
-      "version": "4.0.3",
-      "resolved": "https://registry.npmjs.org/form-data/-/form-data-4.0.3.tgz",
-      "integrity": "sha512-qsITQPfmvMOSAdeyZ+12I1c+CKSstAFAwu+97zrnWAbIr5u8wfsExUzCesVLC8NgHuRUqNN4Zy6UPWUTRGslcA==",
+      "version": "4.0.4",
+      "resolved": "https://registry.npmjs.org/form-data/-/form-data-4.0.4.tgz",
+      "integrity": "sha512-KrGhL9Q4zjj0kiUt5OO4Mr/A/jlI2jDYs5eHBpYHPcBEVSiipAvn2Ko2HnPe20rmcuuvMHNdZFp+4IlGTMF0Ow==",
       "dev": true,
       "license": "MIT",
       "dependencies": {

build-tools/package.json

@@ -2,7 +2,7 @@
   "name": "build-tools",
   "private": true,
   "devDependencies": {
-    "@redocly/cli": "1.34.4",
+    "@redocly/cli": "1.34.5",
     "@sourcemeta/jsonschema": "10.0.0"
   }
 }

compute/compute-node.Dockerfile

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

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

@@ -27,7 +27,10 @@ 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;
@@ -138,6 +139,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 {
@@ -188,11 +195,19 @@ 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)?;
 
+    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)?;
@@ -219,6 +234,13 @@ 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,
     )?;
@@ -227,11 +249,17 @@ fn main() -> Result<()> {
 
     scenario.teardown();
 
-    deinit_and_exit(exit_code);
+    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 || {
@@ -242,7 +270,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> {
@@ -267,25 +295,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");

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,8 +6,7 @@ use compute_api::responses::{
     LfcPrewarmState, PromoteState, TlsConfig,
 };
 use compute_api::spec::{
-    ComputeAudit, ComputeFeature, ComputeMode, ComputeSpec, ExtVersion, PageserverConnectionInfo,
-    PageserverProtocol, PageserverShardConnectionInfo, PageserverShardInfo, PgIdent,
+    ComputeAudit, ComputeFeature, ComputeMode, ComputeSpec, ExtVersion, PageserverProtocol, PgIdent,
 };
 use futures::StreamExt;
 use futures::future::join_all;
@@ -22,6 +21,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};
@@ -41,8 +41,9 @@ use utils::shard::{ShardCount, ShardIndex, ShardNumber};
 
 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;
@@ -114,6 +115,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<()>>>;
@@ -155,6 +167,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.
@@ -192,6 +205,7 @@ impl ComputeState {
     pub fn new() -> Self {
         Self {
             start_time: Utc::now(),
+            pg_start_time: None,
             status: ComputeStatus::Empty,
             last_active: None,
             error: None,
@@ -234,7 +248,7 @@ pub struct ParsedSpec {
     pub spec: ComputeSpec,
     pub tenant_id: TenantId,
     pub timeline_id: TimelineId,
-    pub pageserver_conninfo: PageserverConnectionInfo,
+    pub pageserver_connstr: String,
     pub safekeeper_connstrings: Vec<String>,
     pub storage_auth_token: Option<String>,
     /// k8s dns name and port
@@ -281,114 +295,26 @@ impl ParsedSpec {
     }
 }
 
-/// Extract PageserverConnectionInfo from a comma-separated list of libpq connection strings.
-///
-/// This is used for backwards-compatilibity, to parse the legacye `pageserver_connstr`
-/// field in the compute spec, or the 'neon.pageserver_connstring' GUC. Nowadays, the
-/// 'pageserver_connection_info' field should be used instead.
-fn extract_pageserver_conninfo_from_connstr(
-    connstr: &str,
-    stripe_size: Option<u32>,
-) -> Result<PageserverConnectionInfo, anyhow::Error> {
-    let shard_infos: Vec<_> = connstr
-        .split(',')
-        .map(|connstr| PageserverShardInfo {
-            pageservers: vec![PageserverShardConnectionInfo {
-                id: None,
-                libpq_url: Some(connstr.to_string()),
-                grpc_url: None,
-            }],
-        })
-        .collect();
-
-    match shard_infos.len() {
-        0 => anyhow::bail!("empty connection string"),
-        1 => {
-            // We assume that if there's only connection string, it means "unsharded",
-            // rather than a sharded system with just a single shard. The latter is
-            // possible in principle, but we never do it.
-            let shard_count = ShardCount::unsharded();
-            let only_shard = shard_infos.first().unwrap().clone();
-            let shards = vec![(ShardIndex::unsharded(), only_shard)];
-            Ok(PageserverConnectionInfo {
-                shard_count,
-                stripe_size: None,
-                shards: shards.into_iter().collect(),
-                prefer_protocol: PageserverProtocol::Libpq,
-            })
-        }
-        n => {
-            if stripe_size.is_none() {
-                anyhow::bail!("{n} shards but no stripe_size");
-            }
-            let shard_count = ShardCount(n.try_into()?);
-            let shards = shard_infos
-                .into_iter()
-                .enumerate()
-                .map(|(idx, shard_info)| {
-                    (
-                        ShardIndex {
-                            shard_count,
-                            shard_number: ShardNumber(
-                                idx.try_into().expect("shard number fits in u8"),
-                            ),
-                        },
-                        shard_info,
-                    )
-                })
-                .collect();
-            Ok(PageserverConnectionInfo {
-                shard_count,
-                stripe_size,
-                shards,
-                prefer_protocol: PageserverProtocol::Libpq,
-            })
-        }
-    }
-}
-
 impl TryFrom<ComputeSpec> for ParsedSpec {
-    type Error = anyhow::Error;
-    fn try_from(spec: ComputeSpec) -> Result<Self, anyhow::Error> {
+    type Error = String;
+    fn try_from(spec: ComputeSpec) -> Result<Self, String> {
         // Extract the options from the spec file that are needed to connect to
         // the storage system.
         //
-        // 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(extract_pageserver_conninfo_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(u32::from_str(&guc)?)
-                } else {
-                    None
-                };
-                pageserver_conninfo =
-                    Some(extract_pageserver_conninfo_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
+        // 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")?;
         let safekeeper_connstrings = if spec.safekeeper_connstrings.is_empty() {
             if matches!(spec.mode, ComputeMode::Primary) {
                 spec.cluster
                     .settings
                     .find("neon.safekeepers")
-                    .ok_or(anyhow::anyhow!("safekeeper connstrings should be provided"))?
+                    .ok_or("safekeeper connstrings should be provided")?
                     .split(',')
                     .map(|str| str.to_string())
                     .collect()
@@ -403,22 +329,22 @@ impl TryFrom<ComputeSpec> for ParsedSpec {
         let tenant_id: TenantId = if let Some(tenant_id) = spec.tenant_id {
             tenant_id
         } else {
-            let guc = spec
-                .cluster
+            spec.cluster
                 .settings
                 .find("neon.tenant_id")
-                .ok_or(anyhow::anyhow!("tenant id should be provided"))?;
-            TenantId::from_str(&guc).context("invalid tenant id")?
+                .ok_or("tenant id should be provided")
+                .map(|s| TenantId::from_str(&s))?
+                .or(Err("invalid tenant id"))?
         };
         let timeline_id: TimelineId = if let Some(timeline_id) = spec.timeline_id {
             timeline_id
         } else {
-            let guc = spec
-                .cluster
+            spec.cluster
                 .settings
                 .find("neon.timeline_id")
-                .ok_or(anyhow::anyhow!("timeline id should be provided"))?;
-            TimelineId::from_str(&guc).context(anyhow::anyhow!("invalid timeline id"))?
+                .ok_or("timeline id should be provided")
+                .map(|s| TimelineId::from_str(&s))?
+                .or(Err("invalid timeline id"))?
         };
 
         let endpoint_storage_addr: Option<String> = spec
@@ -432,7 +358,7 @@ impl TryFrom<ComputeSpec> for ParsedSpec {
 
         let res = ParsedSpec {
             spec,
-            pageserver_conninfo,
+            pageserver_connstr,
             safekeeper_connstrings,
             storage_auth_token,
             tenant_id,
@@ -442,7 +368,7 @@ impl TryFrom<ComputeSpec> for ParsedSpec {
         };
 
         // Now check validity of the parsed specification
-        res.validate().map_err(anyhow::Error::msg)?;
+        res.validate()?;
         Ok(res)
     }
 }
@@ -487,6 +413,52 @@ struct StartVmMonitorResult {
     vm_monitor: Option<JoinHandle<Result<()>>>,
 }
 
+/// 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,
+}
+
+impl DatabricksEnvVars {
+    pub fn new(compute_spec: &ComputeSpec, compute_id: Option<&String>) -> Self {
+        // 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.
+        let endpoint_id = 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,
+        }
+    }
+
+    /// 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)> {
+        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();
@@ -568,6 +540,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);
 
@@ -737,6 +710,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()
         }
@@ -1129,10 +1105,12 @@ 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 spec.pageserver_conninfo.prefer_protocol {
-            PageserverProtocol::Grpc => self.try_get_basebackup_grpc(spec, lsn)?,
+
+        let (connected, size) = match PageserverProtocol::from_connstring(shard0_connstr)? {
             PageserverProtocol::Libpq => self.try_get_basebackup_libpq(spec, lsn)?,
+            PageserverProtocol::Grpc => self.try_get_basebackup_grpc(spec, lsn)?,
         };
 
         self.fix_zenith_signal_neon_signal()?;
@@ -1170,32 +1148,23 @@ impl ComputeNode {
     /// 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_index = ShardIndex {
-            shard_number: ShardNumber(0),
-            shard_count: spec.pageserver_conninfo.shard_count,
+        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 = spec
-            .pageserver_conninfo
-            .shards
-            .get(&shard0_index)
-            .ok_or_else(|| {
-                anyhow::anyhow!("shard connection info missing for shard {}", shard0_index)
-            })?;
-        let pageserver = shard0
-            .pageservers
-            .first()
-            .expect("must have at least one pageserver");
-        let shard0_url = pageserver
-            .grpc_url
-            .clone()
-            .expect("no grpc_url for shard 0");
 
         let (reader, connected) = tokio::runtime::Handle::current().block_on(async move {
             let mut client = page_api::Client::connect(
-                shard0_url,
+                shard0_connstr,
                 spec.tenant_id,
                 spec.timeline_id,
-                shard0_index,
+                shard_index,
                 spec.storage_auth_token.clone(),
                 None, // NB: base backups use payload compression
             )
@@ -1227,26 +1196,8 @@ 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_index = ShardIndex {
-            shard_number: ShardNumber(0),
-            shard_count: spec.pageserver_conninfo.shard_count,
-        };
-        let shard0 = spec
-            .pageserver_conninfo
-            .shards
-            .get(&shard0_index)
-            .ok_or_else(|| {
-                anyhow::anyhow!("shard connection info missing for shard {}", shard0_index)
-            })?;
-        let pageserver = shard0
-            .pageservers
-            .first()
-            .expect("must have at least one pageserver");
-        let shard0_connstr = pageserver
-            .libpq_url
-            .clone()
-            .expect("no libpq_url for shard 0");
-        let mut config = postgres::Config::from_str(&shard0_connstr)?;
+        let shard0_connstr = spec.pageserver_connstr.split(',').next().unwrap();
+        let mut config = postgres::Config::from_str(shard0_connstr)?;
 
         // Use the storage auth token from the config file, if given.
         // Note: this overrides any password set in the connection string.
@@ -1332,7 +1283,10 @@ impl ComputeNode {
                     return result;
                 }
                 Err(ref e) if attempts < max_attempts => {
-                    warn!("Failed to get basebackup: {e:?} (attempt {attempts}/{max_attempts})");
+                    warn!(
+                        "Failed to get basebackup: {} (attempt {}/{})",
+                        e, attempts, max_attempts
+                    );
                     std::thread::sleep(std::time::Duration::from_millis(retry_period_ms as u64));
                     retry_period_ms *= 1.5;
                 }
@@ -1503,6 +1457,8 @@ 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()?;
@@ -1510,8 +1466,11 @@ impl ComputeNode {
             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
@@ -1540,11 +1499,39 @@ impl ComputeNode {
             }
         };
 
-        self.get_basebackup(compute_state, lsn)
-            .with_context(|| format!("failed to get basebackup@{lsn}"))?;
+        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
+            )
+        })?;
 
-        // 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
@@ -1657,14 +1644,31 @@ 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))
+            };
+
+            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");
@@ -1849,6 +1853,7 @@ 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();
@@ -1861,7 +1866,7 @@ impl ComputeNode {
                             eprintln!("connection error: {e}");
                         }
                     });
-                    if let Err(e) = handle_migrations(params, &mut client).await {
+                    if let Err(e) = handle_migrations(params, &mut client, lakebase_mode).await {
                         error!("Failed to run migrations: {}", e);
                     }
                 }
@@ -1877,6 +1882,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)]
@@ -1938,12 +1971,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()?;
@@ -2049,6 +2086,8 @@ impl ComputeNode {
                             // wait
                             ComputeStatus::Init
                             | ComputeStatus::Configuration
+                            | ComputeStatus::RefreshConfiguration
+                            | ComputeStatus::RefreshConfigurationPending
                             | ComputeStatus::Empty => {
                                 state = self.state_changed.wait(state).unwrap();
                             }
@@ -2484,22 +2523,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_conninfo = state
+        let old_pageserver_connstr = state
             .pspec
             .as_ref()
             .expect("spec must be set")
-            .pageserver_conninfo
+            .pageserver_connstr
             .clone();
         let mut unchanged = true;
         let _ = self
             .state_changed
             .wait_timeout_while(state, duration, |s| {
-                let pageserver_conninfo = &s
+                let pageserver_connstr = &s
                     .pspec
                     .as_ref()
                     .expect("spec must be set")
-                    .pageserver_conninfo;
-                unchanged = pageserver_conninfo == &old_pageserver_conninfo;
+                    .pageserver_connstr;
+                unchanged = pageserver_connstr == &old_pageserver_connstr;
                 unchanged
             })
             .unwrap();
@@ -2605,6 +2644,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<()> {
@@ -2729,10 +2796,7 @@ mod tests {
 
         match ParsedSpec::try_from(spec.clone()) {
             Ok(_p) => panic!("Failed to detect duplicate entry"),
-            Err(e) => assert!(
-                e.to_string()
-                    .starts_with("duplicate entry in safekeeper_connstrings:")
-            ),
+            Err(e) => assert!(e.starts_with("duplicate entry in safekeeper_connstrings:")),
         };
     }
 }

compute_tools/src/compute_prewarm.rs

@@ -90,6 +90,7 @@ impl ComputeNode {
     }
 
     /// If there is a prewarm request ongoing, return `false`, `true` otherwise.
+    /// Has a failpoint "compute-prewarm"
     pub fn prewarm_lfc(self: &Arc<Self>, from_endpoint: Option<String>) -> bool {
         {
             let state = &mut self.state.lock().unwrap().lfc_prewarm_state;
@@ -112,9 +113,8 @@ impl ComputeNode {
                 Err(err) => {
                     crate::metrics::LFC_PREWARM_ERRORS.inc();
                     error!(%err, "could not prewarm LFC");
-
                     LfcPrewarmState::Failed {
-                        error: err.to_string(),
+                        error: format!("{err:#}"),
                     }
                 }
             };
@@ -135,16 +135,20 @@ impl ComputeNode {
     async fn prewarm_impl(&self, from_endpoint: Option<String>) -> Result<bool> {
         let EndpointStoragePair { url, token } = self.endpoint_storage_pair(from_endpoint)?;
 
+        #[cfg(feature = "testing")]
+        fail::fail_point!("compute-prewarm", |_| {
+            bail!("prewarm configured to fail because of a failpoint")
+        });
+
         info!(%url, "requesting LFC state from endpoint storage");
         let request = Client::new().get(&url).bearer_auth(token);
         let res = request.send().await.context("querying endpoint storage")?;
-        let status = res.status();
-        match status {
+        match res.status() {
             StatusCode::OK => (),
             StatusCode::NOT_FOUND => {
                 return Ok(false);
             }
-            _ => bail!("{status} querying endpoint storage"),
+            status => bail!("{status} querying endpoint storage"),
         }
 
         let mut uncompressed = Vec::new();
@@ -205,7 +209,7 @@ impl ComputeNode {
         crate::metrics::LFC_OFFLOAD_ERRORS.inc();
         error!(%err, "could not offload LFC state to endpoint storage");
         self.state.lock().unwrap().lfc_offload_state = LfcOffloadState::Failed {
-            error: err.to_string(),
+            error: format!("{err:#}"),
         };
     }
 
@@ -213,16 +217,22 @@ impl ComputeNode {
         let EndpointStoragePair { url, token } = self.endpoint_storage_pair(None)?;
         info!(%url, "requesting LFC state from Postgres");
 
-        let mut compressed = Vec::new();
-        ComputeNode::get_maintenance_client(&self.tokio_conn_conf)
+        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(());
+        };
+
+        let mut compressed = Vec::new();
+        ZstdEncoder::new(state)
             .read_to_end(&mut compressed)
             .await
             .context("compressing LFC state")?;

compute_tools/src/compute_promote.rs

@@ -1,11 +1,12 @@
 use crate::compute::ComputeNode;
 use anyhow::{Context, Result, bail};
-use compute_api::{
-    responses::{LfcPrewarmState, PromoteState, SafekeepersLsn},
-    spec::ComputeMode,
-};
+use compute_api::responses::{LfcPrewarmState, PromoteConfig, PromoteState};
+use compute_api::spec::ComputeMode;
+use itertools::Itertools;
+use std::collections::HashMap;
 use std::{sync::Arc, time::Duration};
 use tokio::time::sleep;
+use tracing::info;
 use utils::lsn::Lsn;
 
 impl ComputeNode {
@@ -13,21 +14,22 @@ impl ComputeNode {
     /// and http client 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 {
+    /// Has a failpoint "compute-promotion"
+    pub async fn promote(self: &Arc<Self>, cfg: PromoteConfig) -> PromoteState {
         let cloned = self.clone();
+        let promote_fn = async move || {
+            let Err(err) = cloned.promote_impl(cfg).await else {
+                return PromoteState::Completed;
+            };
+            tracing::error!(%err, "promoting");
+            PromoteState::Failed {
+                error: format!("{err:#}"),
+            }
+        };
+
         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
         };
 
@@ -47,9 +49,7 @@ impl ComputeNode {
         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, mut cfg: PromoteConfig) -> Result<()> {
         {
             let state = self.state.lock().unwrap();
             let mode = &state.pspec.as_ref().unwrap().spec.mode;
@@ -73,7 +73,7 @@ impl ComputeNode {
             .await
             .context("connecting to postgres")?;
 
-        let primary_lsn = safekeepers_lsn.wal_flush_lsn;
+        let primary_lsn = cfg.wal_flush_lsn;
         let mut last_wal_replay_lsn: Lsn = Lsn::INVALID;
         const RETRIES: i32 = 20;
         for i in 0..=RETRIES {
@@ -86,7 +86,7 @@ impl ComputeNode {
             if last_wal_replay_lsn >= primary_lsn {
                 break;
             }
-            tracing::info!("Try {i}, replica lsn {last_wal_replay_lsn}, primary lsn {primary_lsn}");
+            info!("Try {i}, replica lsn {last_wal_replay_lsn}, primary lsn {primary_lsn}");
             sleep(Duration::from_secs(1)).await;
         }
         if last_wal_replay_lsn < primary_lsn {
@@ -96,7 +96,7 @@ impl ComputeNode {
         // 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, &[])
@@ -106,6 +106,12 @@ impl ComputeNode {
             .query("SELECT pg_reload_conf()", &[])
             .await
             .context("reloading postgres config")?;
+
+        #[cfg(feature = "testing")]
+        fail::fail_point!("compute-promotion", |_| {
+            bail!("promotion configured to fail because of a failpoint")
+        });
+
         let row = client
             .query_one("SELECT * FROM pg_promote()", &[])
             .await
@@ -125,8 +131,36 @@ 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(())
+        {
+            let mut state = self.state.lock().unwrap();
+            let spec = &mut state.pspec.as_mut().unwrap().spec;
+            spec.mode = ComputeMode::Primary;
+            let new_conf = cfg.spec.cluster.postgresql_conf.as_mut().unwrap();
+            let existing_conf = spec.cluster.postgresql_conf.as_ref().unwrap();
+            Self::merge_spec(new_conf, existing_conf);
+        }
+        info!("applied new spec, reconfiguring as primary");
+        self.reconfigure()
+    }
+
+    /// Merge old and new Postgres conf specs to apply on secondary.
+    /// Change new spec's port and safekeepers since they are supplied
+    /// differenly
+    fn merge_spec(new_conf: &mut String, existing_conf: &str) {
+        let mut new_conf_set: HashMap<&str, &str> = new_conf
+            .split_terminator('\n')
+            .map(|e| e.split_once("=").expect("invalid item"))
+            .collect();
+        new_conf_set.remove("neon.safekeepers");
+
+        let existing_conf_set: HashMap<&str, &str> = existing_conf
+            .split_terminator('\n')
+            .map(|e| e.split_once("=").expect("invalid item"))
+            .collect();
+        new_conf_set.insert("port", existing_conf_set["port"]);
+        *new_conf = new_conf_set
+            .iter()
+            .map(|(k, v)| format!("{k}={v}"))
+            .join("\n");
     }
 }

compute_tools/src/config.rs

@@ -7,16 +7,17 @@ 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> {
@@ -42,12 +43,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.
@@ -58,101 +63,15 @@ pub fn write_postgres_conf(
         writeln!(file, "{conf}")?;
     }
 
+    // Stripe size GUC should be defined prior to connection string
+    if let Some(stripe_size) = spec.shard_stripe_size {
+        writeln!(file, "neon.stripe_size={stripe_size}")?;
+    }
     // Add options for connecting to storage
     writeln!(file, "# Neon storage settings")?;
-    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_conninfo.stripe_size field"
-            )?;
-            writeln!(file, "neon.stripe_size={stripe_size}")?;
-        }
-
-        let mut libpq_urls: Option<Vec<String>> = Some(Vec::new());
-        let mut grpc_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
-            }
-            // Similarly for gRPC URLs
-            if let Some(url) = &first_pageserver.grpc_url {
-                if let Some(ref mut urls) = grpc_urls {
-                    urls.push(url.clone());
-                }
-            } else {
-                grpc_urls = None
-            }
-        }
-        if let Some(libpq_urls) = libpq_urls {
-            writeln!(
-                file,
-                "# derived from compute spec's pageserver_conninfo field"
-            )?;
-            writeln!(
-                file,
-                "neon.pageserver_connstring={}",
-                escape_conf_value(&libpq_urls.join(","))
-            )?;
-        } else {
-            writeln!(file, "# no neon.pageserver_connstring")?;
-        }
-        if let Some(grpc_urls) = grpc_urls {
-            writeln!(
-                file,
-                "# derived from compute spec's pageserver_conninfo field"
-            )?;
-            writeln!(
-                file,
-                "neon.pageserver_grpc_urls={}",
-                escape_conf_value(&grpc_urls.join(","))
-            )?;
-        } else {
-            writeln!(file, "# no neon.pageserver_grpc_urls")?;
-        }
-    } 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 let Some(s) = &spec.pageserver_connstring {
+        writeln!(file, "neon.pageserver_connstring={}", escape_conf_value(s))?;
     }
-
     if !spec.safekeeper_connstrings.is_empty() {
         let mut neon_safekeepers_value = String::new();
         tracing::info!(
@@ -373,6 +292,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,133 @@ 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_connstr.clone())
+                            == Some(pspec.pageserver_connstr.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
@@ -297,14 +297,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,18 +584,25 @@ 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

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/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/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,14 @@
 use crate::http::JsonResponse;
-use axum::Form;
+use axum::extract::Json;
 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<compute_api::responses::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);
+    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,10 +86,17 @@ 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))
@@ -96,8 +110,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,7 @@ 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};
 
@@ -119,3 +120,7 @@ pub async fn get_installed_extensions(
         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,13 +4,14 @@ use std::thread;
 use std::time::{Duration, SystemTime};
 
 use anyhow::{Result, bail};
-use compute_api::spec::{ComputeMode, PageserverConnectionInfo, PageserverProtocol};
+use compute_api::spec::{ComputeMode, PageserverProtocol};
+use itertools::Itertools as _;
 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::TenantShardId;
+use utils::shard::{ShardCount, ShardNumber, TenantShardId};
 
 use crate::compute::ComputeNode;
 
@@ -77,16 +78,17 @@ fn acquire_lsn_lease_with_retry(
 
     loop {
         // Note: List of pageservers is dynamic, need to re-read configs before each attempt.
-        let (conninfo, auth) = {
+        let (connstrings, auth) = {
             let state = compute.state.lock().unwrap();
             let spec = state.pspec.as_ref().expect("spec must be set");
             (
-                spec.pageserver_conninfo.clone(),
+                spec.pageserver_connstr.clone(),
                 spec.storage_auth_token.clone(),
             )
         };
 
-        let result = try_acquire_lsn_lease(conninfo, auth.as_deref(), tenant_id, timeline_id, lsn);
+        let result =
+            try_acquire_lsn_lease(&connstrings, auth.as_deref(), tenant_id, timeline_id, lsn);
         match result {
             Ok(Some(res)) => {
                 return Ok(res);
@@ -110,44 +112,35 @@ fn acquire_lsn_lease_with_retry(
 
 /// Tries to acquire LSN leases on all Pageserver shards.
 fn try_acquire_lsn_lease(
-    conninfo: PageserverConnectionInfo,
+    connstrings: &str,
     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_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,
+    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),
+            },
         };
 
-        // 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);
-        }
+        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);
     }
 
     Ok(leases.into_iter().min().flatten())

compute_tools/src/migration.rs

@@ -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
@@ -130,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/monitor.rs

@@ -11,6 +11,7 @@ use tracing::{Level, error, info, instrument, span};
 use crate::compute::ComputeNode;
 use crate::metrics::{PG_CURR_DOWNTIME_MS, PG_TOTAL_DOWNTIME_MS};
 
+const PG_DEFAULT_INIT_TIMEOUIT: Duration = Duration::from_secs(60);
 const MONITOR_CHECK_INTERVAL: Duration = Duration::from_millis(500);
 
 /// Struct to store runtime state of the compute monitor thread.
@@ -352,13 +353,47 @@ impl ComputeMonitor {
 // Hang on condition variable waiting until the compute status is `Running`.
 fn wait_for_postgres_start(compute: &ComputeNode) {
     let mut state = compute.state.lock().unwrap();
+    let pg_init_timeout = compute
+        .params
+        .pg_init_timeout
+        .unwrap_or(PG_DEFAULT_INIT_TIMEOUIT);
+
     while state.status != ComputeStatus::Running {
         info!("compute is not running, waiting before monitoring activity");
-        state = compute.state_changed.wait(state).unwrap();
+        if !compute.params.lakebase_mode {
+            state = compute.state_changed.wait(state).unwrap();
 
-        if state.status == ComputeStatus::Running {
-            break;
+            if state.status == ComputeStatus::Running {
+                break;
+            }
+            continue;
         }
+
+        if state.pg_start_time.is_some()
+            && Utc::now()
+                .signed_duration_since(state.pg_start_time.unwrap())
+                .to_std()
+                .unwrap_or_default()
+                > pg_init_timeout
+        {
+            // If Postgres isn't up and running with working PS/SK connections within POSTGRES_STARTUP_TIMEOUT, it is
+            // possible that we started Postgres with a wrong spec (so it is talking to the wrong PS/SK nodes). To prevent
+            // deadends we simply exit (panic) the compute node so it can restart with the latest spec.
+            //
+            // NB: We skip this check if we have not attempted to start PG yet (indicated by state.pg_start_up == None).
+            // This is to make sure the more appropriate errors are surfaced if we encounter issues before we even attempt
+            // to start PG (e.g., if we can't pull the spec, can't sync safekeepers, or can't get the basebackup).
+            error!(
+                "compute did not enter Running state in {} seconds, exiting",
+                pg_init_timeout.as_secs()
+            );
+            std::process::exit(1);
+        }
+        state = compute
+            .state_changed
+            .wait_timeout(state, Duration::from_secs(5))
+            .unwrap()
+            .0;
     }
 }
 

compute_tools/src/pg_helpers.rs

@@ -11,7 +11,9 @@ use std::time::{Duration, Instant};
 
 use anyhow::{Result, bail};
 use compute_api::responses::TlsConfig;
-use compute_api::spec::{Database, GenericOption, GenericOptions, PgIdent, Role};
+use compute_api::spec::{
+    Database, DatabricksSettings, GenericOption, GenericOptions, PgIdent, Role,
+};
 use futures::StreamExt;
 use indexmap::IndexMap;
 use ini::Ini;
@@ -184,6 +186,42 @@ impl DatabaseExt for Database {
     }
 }
 
+pub trait DatabricksSettingsExt {
+    fn as_pg_settings(&self) -> String;
+}
+
+impl DatabricksSettingsExt for DatabricksSettings {
+    fn as_pg_settings(&self) -> String {
+        // Postgres GUCs rendered from DatabricksSettings
+        vec![
+            // ssl_ca_file
+            Some(format!(
+                "ssl_ca_file = '{}'",
+                self.pg_compute_tls_settings.ca_file
+            )),
+            // [Optional] databricks.workspace_url
+            Some(format!(
+                "databricks.workspace_url = '{}'",
+                &self.databricks_workspace_host
+            )),
+            // todo(vikas.jain): these are not required anymore as they are moved to static
+            // conf but keeping these to avoid image mismatch between hcc and pg.
+            // Once hcc and pg are in sync, we can remove these.
+            //
+            // databricks.enable_databricks_identity_login
+            Some("databricks.enable_databricks_identity_login = true".to_string()),
+            // databricks.enable_sql_restrictions
+            Some("databricks.enable_sql_restrictions = true".to_string()),
+        ]
+        .into_iter()
+        // Removes `None`s
+        .flatten()
+        .collect::<Vec<String>>()
+        .join("\n")
+            + "\n"
+    }
+}
+
 /// Generic trait used to provide quoting / encoding for strings used in the
 /// Postgres SQL queries and DATABASE_URL.
 pub trait Escaping {

compute_tools/src/pg_isready.rs

@@ -0,0 +1,30 @@
+use anyhow::{Context, anyhow};
+
+// Run `/usr/local/bin/pg_isready -p {port}`
+// Check the connectivity of PG
+// Success means PG is listening on the port and accepting connections
+// Note that PG does not need to authenticate the connection, nor reserve a connection quota for it.
+// See https://www.postgresql.org/docs/current/app-pg-isready.html
+pub fn pg_isready(bin: &str, port: u16) -> anyhow::Result<()> {
+    let child_result = std::process::Command::new(bin)
+        .arg("-p")
+        .arg(port.to_string())
+        .spawn();
+
+    child_result
+        .context("spawn() failed")
+        .and_then(|mut child| child.wait().context("wait() failed"))
+        .and_then(|status| match status.success() {
+            true => Ok(()),
+            false => Err(anyhow!("process exited with {status}")),
+        })
+        // wrap any prior error with the overall context that we couldn't run the command
+        .with_context(|| format!("could not run `{bin} --port {port}`"))
+}
+
+// It's safe to assume pg_isready is under the same directory with postgres,
+// because it is a PG util bin installed along with postgres
+pub fn get_pg_isready_bin(pgbin: &str) -> String {
+    let split = pgbin.split("/").collect::<Vec<&str>>();
+    split[0..split.len() - 1].join("/") + "/pg_isready"
+}

compute_tools/src/spec.rs

@@ -1,4 +1,6 @@
 use std::fs::File;
+use std::fs::{self, Permissions};
+use std::os::unix::fs::PermissionsExt;
 use std::path::Path;
 
 use anyhow::{Result, anyhow, bail};
@@ -133,10 +135,25 @@ pub fn get_config_from_control_plane(base_uri: &str, compute_id: &str) -> Result
 }
 
 /// Check `pg_hba.conf` and update if needed to allow external connections.
-pub fn update_pg_hba(pgdata_path: &Path) -> Result<()> {
+pub fn update_pg_hba(pgdata_path: &Path, databricks_pg_hba: Option<&String>) -> Result<()> {
     // XXX: consider making it a part of config.json
     let pghba_path = pgdata_path.join("pg_hba.conf");
 
+    // Update pg_hba to contains databricks specfic settings before adding neon settings
+    // PG uses the first record that matches to perform authentication, so we need to have
+    // our rules before the default ones from neon.
+    // See https://www.postgresql.org/docs/current/auth-pg-hba-conf.html
+    if let Some(databricks_pg_hba) = databricks_pg_hba {
+        if config::line_in_file(
+            &pghba_path,
+            &format!("include_if_exists {}\n", *databricks_pg_hba),
+        )? {
+            info!("updated pg_hba.conf to include databricks_pg_hba.conf");
+        } else {
+            info!("pg_hba.conf already included databricks_pg_hba.conf");
+        }
+    }
+
     if config::line_in_file(&pghba_path, PG_HBA_ALL_MD5)? {
         info!("updated pg_hba.conf to allow external connections");
     } else {
@@ -146,6 +163,59 @@ pub fn update_pg_hba(pgdata_path: &Path) -> Result<()> {
     Ok(())
 }
 
+/// Check `pg_ident.conf` and update if needed to allow databricks config.
+pub fn update_pg_ident(pgdata_path: &Path, databricks_pg_ident: Option<&String>) -> Result<()> {
+    info!("checking pg_ident.conf");
+    let pghba_path = pgdata_path.join("pg_ident.conf");
+
+    // Update pg_ident to contains databricks specfic settings
+    if let Some(databricks_pg_ident) = databricks_pg_ident {
+        if config::line_in_file(
+            &pghba_path,
+            &format!("include_if_exists {}\n", *databricks_pg_ident),
+        )? {
+            info!("updated pg_ident.conf to include databricks_pg_ident.conf");
+        } else {
+            info!("pg_ident.conf already included databricks_pg_ident.conf");
+        }
+    }
+
+    Ok(())
+}
+
+/// Copy tls key_file and cert_file from k8s secret mount directory
+/// to pgdata and set private key file permissions as expected by Postgres.
+/// See this doc for expected permission <https://www.postgresql.org/docs/current/ssl-tcp.html>
+/// K8s secrets mount on dblet does not honor permission and ownership
+/// specified in the Volume or VolumeMount. So we need to explicitly copy the file and set the permissions.
+pub fn copy_tls_certificates(
+    key_file: &String,
+    cert_file: &String,
+    pgdata_path: &Path,
+) -> Result<()> {
+    let files = [cert_file, key_file];
+    for file in files.iter() {
+        let source = Path::new(file);
+        let dest = pgdata_path.join(source.file_name().unwrap());
+        if !dest.exists() {
+            std::fs::copy(source, &dest)?;
+            info!(
+                "Copying tls file: {} to {}",
+                &source.display(),
+                &dest.display()
+            );
+        }
+        if *file == key_file {
+            // Postgres requires private key to be readable only by the owner by having
+            // chmod 600 permissions.
+            let permissions = Permissions::from_mode(0o600);
+            fs::set_permissions(&dest, permissions)?;
+            info!("Setting permission on {}.", &dest.display());
+        }
+    }
+    Ok(())
+}
+
 /// Create a standby.signal file
 pub fn add_standby_signal(pgdata_path: &Path) -> Result<()> {
     // XXX: consider making it a part of config.json
@@ -170,7 +240,11 @@ pub async fn handle_neon_extension_upgrade(client: &mut Client) -> Result<()> {
 }
 
 #[instrument(skip_all)]
-pub async fn handle_migrations(params: ComputeNodeParams, client: &mut Client) -> Result<()> {
+pub async fn handle_migrations(
+    params: ComputeNodeParams,
+    client: &mut Client,
+    lakebase_mode: bool,
+) -> Result<()> {
     info!("handle migrations");
 
     // !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
@@ -234,7 +308,7 @@ pub async fn handle_migrations(params: ComputeNodeParams, client: &mut Client) -
         ),
     ];
 
-    MigrationRunner::new(client, &migrations)
+    MigrationRunner::new(client, &migrations, lakebase_mode)
         .run_migrations()
         .await?;
 

compute_tools/src/spec_apply.rs

@@ -411,7 +411,8 @@ impl ComputeNode {
             .map(|limit| match limit {
                 0..10 => limit,
                 10..30 => 10,
-                30.. => limit / 3,
+                30..300 => limit / 3,
+                300.. => 100,
             })
             // If we didn't find max_connections, default to 10 concurrent connections.
             .unwrap_or(10)

control_plane/src/bin/neon_local.rs

@@ -16,14 +16,9 @@ use std::time::Duration;
 use anyhow::{Context, Result, anyhow, bail};
 use clap::Parser;
 use compute_api::requests::ComputeClaimsScope;
-use compute_api::spec::{
-    ComputeMode, PageserverConnectionInfo, PageserverProtocol, PageserverShardInfo,
-};
+use compute_api::spec::{ComputeMode, PageserverProtocol};
 use control_plane::broker::StorageBroker;
 use control_plane::endpoint::{ComputeControlPlane, EndpointTerminateMode};
-use control_plane::endpoint::{
-    pageserver_conf_to_shard_conn_info, tenant_locate_response_to_conn_info,
-};
 use control_plane::endpoint_storage::{ENDPOINT_STORAGE_DEFAULT_ADDR, EndpointStorage};
 use control_plane::local_env;
 use control_plane::local_env::{
@@ -49,6 +44,7 @@ use pageserver_api::models::{
 };
 use pageserver_api::shard::{DEFAULT_STRIPE_SIZE, ShardCount, ShardStripeSize, TenantShardId};
 use postgres_backend::AuthType;
+use postgres_connection::parse_host_port;
 use safekeeper_api::membership::{SafekeeperGeneration, SafekeeperId};
 use safekeeper_api::{
     DEFAULT_HTTP_LISTEN_PORT as DEFAULT_SAFEKEEPER_HTTP_PORT,
@@ -56,11 +52,11 @@ use safekeeper_api::{
 };
 use storage_broker::DEFAULT_LISTEN_ADDR as DEFAULT_BROKER_ADDR;
 use tokio::task::JoinSet;
+use url::Host;
 use utils::auth::{Claims, Scope};
 use utils::id::{NodeId, TenantId, TenantTimelineId, TimelineId};
 use utils::lsn::Lsn;
 use utils::project_git_version;
-use utils::shard::ShardIndex;
 
 // Default id of a safekeeper node, if not specified on the command line.
 const DEFAULT_SAFEKEEPER_ID: NodeId = NodeId(1);
@@ -411,6 +407,12 @@ struct StorageControllerStartCmdArgs {
         help = "Base port for the storage controller instance idenfified by instance-id (defaults to pageserver cplane api)"
     )]
     base_port: Option<u16>,
+
+    #[clap(
+        long,
+        help = "Whether the storage controller should handle pageserver-reported local disk loss events."
+    )]
+    handle_ps_local_disk_loss: Option<bool>,
 }
 
 #[derive(clap::Args)]
@@ -558,7 +560,9 @@ enum EndpointCmd {
     Create(EndpointCreateCmdArgs),
     Start(EndpointStartCmdArgs),
     Reconfigure(EndpointReconfigureCmdArgs),
+    RefreshConfiguration(EndpointRefreshConfigurationArgs),
     Stop(EndpointStopCmdArgs),
+    UpdatePageservers(EndpointUpdatePageserversCmdArgs),
     GenerateJwt(EndpointGenerateJwtCmdArgs),
 }
 
@@ -719,6 +723,13 @@ struct EndpointReconfigureCmdArgs {
     safekeepers: Option<String>,
 }
 
+#[derive(clap::Args)]
+#[clap(about = "Refresh the endpoint's configuration by forcing it reload it's spec")]
+struct EndpointRefreshConfigurationArgs {
+    #[clap(help = "Postgres endpoint id")]
+    endpoint_id: String,
+}
+
 #[derive(clap::Args)]
 #[clap(about = "Stop an endpoint")]
 struct EndpointStopCmdArgs {
@@ -736,6 +747,16 @@ struct EndpointStopCmdArgs {
     mode: EndpointTerminateMode,
 }
 
+#[derive(clap::Args)]
+#[clap(about = "Update the pageservers in the spec file of the compute endpoint")]
+struct EndpointUpdatePageserversCmdArgs {
+    #[clap(help = "Postgres endpoint id")]
+    endpoint_id: String,
+
+    #[clap(short = 'p', long, help = "Specified pageserver id")]
+    pageserver_id: Option<NodeId>,
+}
+
 #[derive(clap::Args)]
 #[clap(about = "Generate a JWT for an endpoint")]
 struct EndpointGenerateJwtCmdArgs {
@@ -1515,7 +1536,7 @@ async fn handle_endpoint(subcmd: &EndpointCmd, env: &local_env::LocalEnv) -> Res
             let endpoint = cplane
                 .endpoints
                 .get(endpoint_id.as_str())
-                .ok_or_else(|| anyhow::anyhow!("endpoint {endpoint_id} not found"))?;
+                .ok_or_else(|| anyhow!("endpoint {endpoint_id} not found"))?;
 
             if !args.allow_multiple {
                 cplane.check_conflicting_endpoints(
@@ -1525,56 +1546,62 @@ async fn handle_endpoint(subcmd: &EndpointCmd, env: &local_env::LocalEnv) -> Res
                 )?;
             }
 
-            let prefer_protocol = if endpoint.grpc {
-                PageserverProtocol::Grpc
-            } else {
-                PageserverProtocol::Libpq
-            };
-
-            let mut pageserver_conninfo = if let Some(ps_id) = pageserver_id {
-                let conf = env.get_pageserver_conf(ps_id).unwrap();
-                let ps_conninfo = pageserver_conf_to_shard_conn_info(conf)?;
-
-                let shard_info = PageserverShardInfo {
-                    pageservers: vec![ps_conninfo],
+            let (pageservers, stripe_size) = if let Some(pageserver_id) = pageserver_id {
+                let conf = env.get_pageserver_conf(pageserver_id).unwrap();
+                // Use gRPC if requested.
+                let pageserver = if endpoint.grpc {
+                    let grpc_addr = conf.listen_grpc_addr.as_ref().expect("bad config");
+                    let (host, port) = parse_host_port(grpc_addr)?;
+                    let port = port.unwrap_or(DEFAULT_PAGESERVER_GRPC_PORT);
+                    (PageserverProtocol::Grpc, host, port)
+                } else {
+                    let (host, port) = parse_host_port(&conf.listen_pg_addr)?;
+                    let port = port.unwrap_or(5432);
+                    (PageserverProtocol::Libpq, host, port)
                 };
                 // If caller is telling us what pageserver to use, this is not a tenant which is
                 // fully managed by storage controller, therefore not sharded.
-                let shards: HashMap<_, _> = vec![(ShardIndex::unsharded(), shard_info)]
-                    .into_iter()
-                    .collect();
-                PageserverConnectionInfo {
-                    shard_count: ShardCount(0),
-                    stripe_size: None,
-                    shards,
-                    prefer_protocol,
-                }
+                (vec![pageserver], DEFAULT_STRIPE_SIZE)
             } else {
                 // Look up the currently attached location of the tenant, and its striping metadata,
                 // to pass these on to postgres.
                 let storage_controller = StorageController::from_env(env);
                 let locate_result = storage_controller.tenant_locate(endpoint.tenant_id).await?;
-                assert!(!locate_result.shards.is_empty());
-
-                // Initialize LSN leases for static computes.
-                if let ComputeMode::Static(lsn) = endpoint.mode {
-                    futures::future::try_join_all(locate_result.shards.iter().map(
-                        |shard| async move {
+                let pageservers = futures::future::try_join_all(
+                    locate_result.shards.into_iter().map(|shard| async move {
+                        if let ComputeMode::Static(lsn) = endpoint.mode {
+                            // Initialize LSN leases for static computes.
                             let conf = env.get_pageserver_conf(shard.node_id).unwrap();
                             let pageserver = PageServerNode::from_env(env, conf);
 
                             pageserver
                                 .http_client
                                 .timeline_init_lsn_lease(shard.shard_id, endpoint.timeline_id, lsn)
-                                .await
-                        },
-                    ))
-                    .await?;
-                }
+                                .await?;
+                        }
 
-                tenant_locate_response_to_conn_info(&locate_result)?
+                        let pageserver = if endpoint.grpc {
+                            (
+                                PageserverProtocol::Grpc,
+                                Host::parse(&shard.listen_grpc_addr.expect("no gRPC address"))?,
+                                shard.listen_grpc_port.expect("no gRPC port"),
+                            )
+                        } else {
+                            (
+                                PageserverProtocol::Libpq,
+                                Host::parse(&shard.listen_pg_addr)?,
+                                shard.listen_pg_port,
+                            )
+                        };
+                        anyhow::Ok(pageserver)
+                    }),
+                )
+                .await?;
+                let stripe_size = locate_result.shard_params.stripe_size;
+
+                (pageservers, stripe_size)
             };
-            pageserver_conninfo.prefer_protocol = prefer_protocol;
+            assert!(!pageservers.is_empty());
 
             let ps_conf = env.get_pageserver_conf(DEFAULT_PAGESERVER_ID)?;
             let auth_token = if matches!(ps_conf.pg_auth_type, AuthType::NeonJWT) {
@@ -1604,8 +1631,9 @@ async fn handle_endpoint(subcmd: &EndpointCmd, env: &local_env::LocalEnv) -> Res
                 endpoint_storage_addr,
                 safekeepers_generation,
                 safekeepers,
-                pageserver_conninfo,
+                pageservers,
                 remote_ext_base_url: remote_ext_base_url.clone(),
+                shard_stripe_size: stripe_size.0 as usize,
                 create_test_user: args.create_test_user,
                 start_timeout: args.start_timeout,
                 autoprewarm: args.autoprewarm,
@@ -1616,53 +1644,105 @@ async fn handle_endpoint(subcmd: &EndpointCmd, env: &local_env::LocalEnv) -> Res
             println!("Starting existing endpoint {endpoint_id}...");
             endpoint.start(args).await?;
         }
+        EndpointCmd::UpdatePageservers(args) => {
+            let endpoint_id = &args.endpoint_id;
+            let endpoint = cplane
+                .endpoints
+                .get(endpoint_id.as_str())
+                .with_context(|| format!("postgres endpoint {endpoint_id} is not found"))?;
+            let pageservers = match args.pageserver_id {
+                Some(pageserver_id) => {
+                    let pageserver =
+                        PageServerNode::from_env(env, env.get_pageserver_conf(pageserver_id)?);
+
+                    vec![(
+                        PageserverProtocol::Libpq,
+                        pageserver.pg_connection_config.host().clone(),
+                        pageserver.pg_connection_config.port(),
+                    )]
+                }
+                None => {
+                    let storage_controller = StorageController::from_env(env);
+                    storage_controller
+                        .tenant_locate(endpoint.tenant_id)
+                        .await?
+                        .shards
+                        .into_iter()
+                        .map(|shard| {
+                            (
+                                PageserverProtocol::Libpq,
+                                Host::parse(&shard.listen_pg_addr)
+                                    .expect("Storage controller reported malformed host"),
+                                shard.listen_pg_port,
+                            )
+                        })
+                        .collect::<Vec<_>>()
+                }
+            };
+
+            endpoint.update_pageservers_in_config(pageservers).await?;
+        }
         EndpointCmd::Reconfigure(args) => {
             let endpoint_id = &args.endpoint_id;
             let endpoint = cplane
                 .endpoints
                 .get(endpoint_id.as_str())
                 .with_context(|| format!("postgres endpoint {endpoint_id} is not found"))?;
-
-            let prefer_protocol = if endpoint.grpc {
-                PageserverProtocol::Grpc
-            } else {
-                PageserverProtocol::Libpq
-            };
-            let mut pageserver_conninfo = if let Some(ps_id) = args.endpoint_pageserver_id {
+            let pageservers = if let Some(ps_id) = args.endpoint_pageserver_id {
                 let conf = env.get_pageserver_conf(ps_id)?;
-                let ps_conninfo = pageserver_conf_to_shard_conn_info(conf)?;
-                let shard_info = PageserverShardInfo {
-                    pageservers: vec![ps_conninfo],
+                // Use gRPC if requested.
+                let pageserver = if endpoint.grpc {
+                    let grpc_addr = conf.listen_grpc_addr.as_ref().expect("bad config");
+                    let (host, port) = parse_host_port(grpc_addr)?;
+                    let port = port.unwrap_or(DEFAULT_PAGESERVER_GRPC_PORT);
+                    (PageserverProtocol::Grpc, host, port)
+                } else {
+                    let (host, port) = parse_host_port(&conf.listen_pg_addr)?;
+                    let port = port.unwrap_or(5432);
+                    (PageserverProtocol::Libpq, host, port)
                 };
-
-                // If caller is telling us what pageserver to use, this is not a tenant which is
-                // fully managed by storage controller, therefore not sharded.
-                let shards: HashMap<_, _> = vec![(ShardIndex::unsharded(), shard_info)]
-                    .into_iter()
-                    .collect();
-                PageserverConnectionInfo {
-                    shard_count: ShardCount::unsharded(),
-                    stripe_size: None,
-                    shards,
-                    prefer_protocol,
-                }
+                vec![pageserver]
             } else {
-                // Look up the currently attached location of the tenant, and its striping metadata,
-                // to pass these on to postgres.
                 let storage_controller = StorageController::from_env(env);
-                let locate_result = storage_controller.tenant_locate(endpoint.tenant_id).await?;
-
-                tenant_locate_response_to_conn_info(&locate_result)?
+                storage_controller
+                    .tenant_locate(endpoint.tenant_id)
+                    .await?
+                    .shards
+                    .into_iter()
+                    .map(|shard| {
+                        // Use gRPC if requested.
+                        if endpoint.grpc {
+                            (
+                                PageserverProtocol::Grpc,
+                                Host::parse(&shard.listen_grpc_addr.expect("no gRPC address"))
+                                    .expect("bad hostname"),
+                                shard.listen_grpc_port.expect("no gRPC port"),
+                            )
+                        } else {
+                            (
+                                PageserverProtocol::Libpq,
+                                Host::parse(&shard.listen_pg_addr).expect("bad hostname"),
+                                shard.listen_pg_port,
+                            )
+                        }
+                    })
+                    .collect::<Vec<_>>()
             };
-            pageserver_conninfo.prefer_protocol = prefer_protocol;
-
             // If --safekeepers argument is given, use only the listed
             // safekeeper nodes; otherwise all from the env.
             let safekeepers = parse_safekeepers(&args.safekeepers)?;
             endpoint
-                .reconfigure(Some(&pageserver_conninfo), safekeepers, None)
+                .reconfigure(Some(pageservers), None, safekeepers, None)
                 .await?;
         }
+        EndpointCmd::RefreshConfiguration(args) => {
+            let endpoint_id = &args.endpoint_id;
+            let endpoint = cplane
+                .endpoints
+                .get(endpoint_id.as_str())
+                .with_context(|| format!("postgres endpoint {endpoint_id} is not found"))?;
+            endpoint.refresh_configuration().await?;
+        }
         EndpointCmd::Stop(args) => {
             let endpoint_id = &args.endpoint_id;
             let endpoint = cplane
@@ -1800,6 +1880,7 @@ async fn handle_storage_controller(
                 instance_id: args.instance_id,
                 base_port: args.base_port,
                 start_timeout: args.start_timeout,
+                handle_ps_local_disk_loss: args.handle_ps_local_disk_loss,
             };
 
             if let Err(e) = svc.start(start_args).await {

control_plane/src/endpoint.rs

@@ -37,7 +37,7 @@
 //!         <other PostgreSQL files>
 //! ```
 //!
-use std::collections::{BTreeMap, HashMap};
+use std::collections::BTreeMap;
 use std::fmt::Display;
 use std::net::{IpAddr, Ipv4Addr, SocketAddr, TcpStream};
 use std::path::PathBuf;
@@ -58,12 +58,8 @@ use compute_api::responses::{
 };
 use compute_api::spec::{
     Cluster, ComputeAudit, ComputeFeature, ComputeMode, ComputeSpec, Database, PageserverProtocol,
-    PageserverShardInfo, PgIdent, RemoteExtSpec, Role,
+    PgIdent, RemoteExtSpec, Role,
 };
-
-// re-export these, because they're used in the reconfigure() function
-pub use compute_api::spec::{PageserverConnectionInfo, PageserverShardConnectionInfo};
-
 use jsonwebtoken::jwk::{
     AlgorithmParameters, CommonParameters, EllipticCurve, Jwk, JwkSet, KeyAlgorithm, KeyOperations,
     OctetKeyPairParameters, OctetKeyPairType, PublicKeyUse,
@@ -78,11 +74,9 @@ use sha2::{Digest, Sha256};
 use spki::der::Decode;
 use spki::{SubjectPublicKeyInfo, SubjectPublicKeyInfoRef};
 use tracing::debug;
+use url::Host;
 use utils::id::{NodeId, TenantId, TimelineId};
-use utils::shard::{ShardIndex, ShardNumber};
-
-use pageserver_api::config::DEFAULT_GRPC_LISTEN_PORT as DEFAULT_PAGESERVER_GRPC_PORT;
-use postgres_connection::parse_host_port;
+use utils::shard::ShardStripeSize;
 
 use crate::local_env::LocalEnv;
 use crate::postgresql_conf::PostgresConf;
@@ -393,8 +387,9 @@ pub struct EndpointStartArgs {
     pub endpoint_storage_addr: String,
     pub safekeepers_generation: Option<SafekeeperGeneration>,
     pub safekeepers: Vec<NodeId>,
-    pub pageserver_conninfo: PageserverConnectionInfo,
+    pub pageservers: Vec<(PageserverProtocol, Host, u16)>,
     pub remote_ext_base_url: Option<String>,
+    pub shard_stripe_size: usize,
     pub create_test_user: bool,
     pub start_timeout: Duration,
     pub autoprewarm: bool,
@@ -667,6 +662,14 @@ impl Endpoint {
         }
     }
 
+    fn build_pageserver_connstr(pageservers: &[(PageserverProtocol, Host, u16)]) -> String {
+        pageservers
+            .iter()
+            .map(|(scheme, host, port)| format!("{scheme}://no_user@{host}:{port}"))
+            .collect::<Vec<_>>()
+            .join(",")
+    }
+
     /// Map safekeepers ids to the actual connection strings.
     fn build_safekeepers_connstrs(&self, sk_ids: Vec<NodeId>) -> Result<Vec<String>> {
         let mut safekeeper_connstrings = Vec::new();
@@ -712,6 +715,9 @@ impl Endpoint {
             std::fs::remove_dir_all(self.pgdata())?;
         }
 
+        let pageserver_connstring = Self::build_pageserver_connstr(&args.pageservers);
+        assert!(!pageserver_connstring.is_empty());
+
         let safekeeper_connstrings = self.build_safekeepers_connstrs(args.safekeepers)?;
 
         // check for file remote_extensions_spec.json
@@ -726,46 +732,6 @@ impl Endpoint {
             remote_extensions = None;
         };
 
-        // For the sake of backwards-compatibility, also fill in 'pageserver_connstring'
-        //
-        // XXX: I believe this is not really needed, except to make
-        // test_forward_compatibility happy.
-        //
-        // Use a closure so that we can conviniently return None in the middle of the
-        // loop.
-        let pageserver_connstring = (|| {
-            let num_shards = if args.pageserver_conninfo.shard_count.is_unsharded() {
-                1
-            } else {
-                args.pageserver_conninfo.shard_count.0
-            };
-            let mut connstrings = Vec::new();
-            for shard_no in 0..num_shards {
-                let shard_index = ShardIndex {
-                    shard_count: args.pageserver_conninfo.shard_count,
-                    shard_number: ShardNumber(shard_no),
-                };
-                let shard = args
-                    .pageserver_conninfo
-                    .shards
-                    .get(&shard_index)
-                    .expect(&format!(
-                        "shard {} not found in pageserver_connection_info",
-                        shard_index
-                    ));
-                let pageserver = shard
-                    .pageservers
-                    .first()
-                    .expect("must have at least one pageserver");
-                if let Some(libpq_url) = &pageserver.libpq_url {
-                    connstrings.push(libpq_url.clone());
-                } else {
-                    return None;
-                }
-            }
-            Some(connstrings.join(","))
-        })();
-
         // Create config file
         let config = {
             let mut spec = ComputeSpec {
@@ -810,14 +776,13 @@ impl Endpoint {
                 branch_id: None,
                 endpoint_id: Some(self.endpoint_id.clone()),
                 mode: self.mode,
-                pageserver_connection_info: Some(args.pageserver_conninfo.clone()),
-                pageserver_connstring,
+                pageserver_connstring: Some(pageserver_connstring),
                 safekeepers_generation: args.safekeepers_generation.map(|g| g.into_inner()),
                 safekeeper_connstrings,
                 storage_auth_token: args.auth_token.clone(),
                 remote_extensions,
                 pgbouncer_settings: None,
-                shard_stripe_size: args.pageserver_conninfo.stripe_size, // redundant with pageserver_connection_info.stripe_size
+                shard_stripe_size: Some(args.shard_stripe_size),
                 local_proxy_config: None,
                 reconfigure_concurrency: self.reconfigure_concurrency,
                 drop_subscriptions_before_start: self.drop_subscriptions_before_start,
@@ -828,6 +793,7 @@ impl Endpoint {
                 autoprewarm: args.autoprewarm,
                 offload_lfc_interval_seconds: args.offload_lfc_interval_seconds,
                 suspend_timeout_seconds: -1, // Only used in neon_local.
+                databricks_settings: None,
             };
 
             // this strange code is needed to support respec() in tests
@@ -972,7 +938,9 @@ impl Endpoint {
                         | ComputeStatus::Configuration
                         | ComputeStatus::TerminationPendingFast
                         | ComputeStatus::TerminationPendingImmediate
-                        | ComputeStatus::Terminated => {
+                        | ComputeStatus::Terminated
+                        | ComputeStatus::RefreshConfigurationPending
+                        | ComputeStatus::RefreshConfiguration => {
                             bail!("unexpected compute status: {:?}", state.status)
                         }
                     }
@@ -995,6 +963,29 @@ impl Endpoint {
         Ok(())
     }
 
+    // Update the pageservers in the spec file of the endpoint. This is useful to test the spec refresh scenario.
+    pub async fn update_pageservers_in_config(
+        &self,
+        pageservers: Vec<(PageserverProtocol, Host, u16)>,
+    ) -> Result<()> {
+        let config_path = self.endpoint_path().join("config.json");
+        let mut config: ComputeConfig = {
+            let file = std::fs::File::open(&config_path)?;
+            serde_json::from_reader(file)?
+        };
+
+        let pageserver_connstring = Self::build_pageserver_connstr(&pageservers);
+        assert!(!pageserver_connstring.is_empty());
+        let mut spec = config.spec.unwrap();
+        spec.pageserver_connstring = Some(pageserver_connstring);
+        config.spec = Some(spec);
+
+        let file = std::fs::File::create(&config_path)?;
+        serde_json::to_writer_pretty(file, &config)?;
+
+        Ok(())
+    }
+
     // Call the /status HTTP API
     pub async fn get_status(&self) -> Result<ComputeStatusResponse> {
         let client = reqwest::Client::new();
@@ -1029,7 +1020,8 @@ impl Endpoint {
 
     pub async fn reconfigure(
         &self,
-        pageserver_conninfo: Option<&PageserverConnectionInfo>,
+        pageservers: Option<Vec<(PageserverProtocol, Host, u16)>>,
+        stripe_size: Option<ShardStripeSize>,
         safekeepers: Option<Vec<NodeId>>,
         safekeeper_generation: Option<SafekeeperGeneration>,
     ) -> Result<()> {
@@ -1044,15 +1036,15 @@ impl Endpoint {
         let postgresql_conf = self.read_postgresql_conf()?;
         spec.cluster.postgresql_conf = Some(postgresql_conf);
 
-        if let Some(pageserver_conninfo) = pageserver_conninfo {
-            // If pageservers are provided, we need to ensure that they are not empty.
-            // This is a requirement for the compute_ctl configuration.
-            anyhow::ensure!(
-                !pageserver_conninfo.shards.is_empty(),
-                "no pageservers provided"
-            );
-            spec.pageserver_connection_info = Some(pageserver_conninfo.clone());
-            spec.shard_stripe_size = pageserver_conninfo.stripe_size;
+        // If pageservers are not specified, don't change them.
+        if let Some(pageservers) = pageservers {
+            anyhow::ensure!(!pageservers.is_empty(), "no pageservers provided");
+
+            let pageserver_connstr = Self::build_pageserver_connstr(&pageservers);
+            spec.pageserver_connstring = Some(pageserver_connstr);
+            if stripe_size.is_some() {
+                spec.shard_stripe_size = stripe_size.map(|s| s.0 as usize);
+            }
         }
 
         // If safekeepers are not specified, don't change them.
@@ -1101,9 +1093,11 @@ impl Endpoint {
 
     pub async fn reconfigure_pageservers(
         &self,
-        pageservers: &PageserverConnectionInfo,
+        pageservers: Vec<(PageserverProtocol, Host, u16)>,
+        stripe_size: Option<ShardStripeSize>,
     ) -> Result<()> {
-        self.reconfigure(Some(pageservers), None, None).await
+        self.reconfigure(Some(pageservers), stripe_size, None, None)
+            .await
     }
 
     pub async fn reconfigure_safekeepers(
@@ -1111,7 +1105,7 @@ impl Endpoint {
         safekeepers: Vec<NodeId>,
         generation: SafekeeperGeneration,
     ) -> Result<()> {
-        self.reconfigure(None, Some(safekeepers), Some(generation))
+        self.reconfigure(None, None, Some(safekeepers), Some(generation))
             .await
     }
 
@@ -1157,6 +1151,33 @@ impl Endpoint {
         Ok(response)
     }
 
+    pub async fn refresh_configuration(&self) -> Result<()> {
+        let client = reqwest::Client::builder()
+            .timeout(Duration::from_secs(30))
+            .build()
+            .unwrap();
+        let response = client
+            .post(format!(
+                "http://{}:{}/refresh_configuration",
+                self.internal_http_address.ip(),
+                self.internal_http_address.port()
+            ))
+            .send()
+            .await?;
+
+        let status = response.status();
+        if !(status.is_client_error() || status.is_server_error()) {
+            Ok(())
+        } else {
+            let url = response.url().to_owned();
+            let msg = match response.text().await {
+                Ok(err_body) => format!("Error: {err_body}"),
+                Err(_) => format!("Http error ({}) at {}.", status.as_u16(), url),
+            };
+            Err(anyhow::anyhow!(msg))
+        }
+    }
+
     pub fn connstr(&self, user: &str, db_name: &str) -> String {
         format!(
             "postgresql://{}@{}:{}/{}",
@@ -1167,68 +1188,3 @@ impl Endpoint {
         )
     }
 }
-
-pub fn pageserver_conf_to_shard_conn_info(
-    conf: &crate::local_env::PageServerConf,
-) -> Result<PageserverShardConnectionInfo> {
-    let libpq_url = {
-        let (host, port) = parse_host_port(&conf.listen_pg_addr)?;
-        let port = port.unwrap_or(5432);
-        Some(format!("postgres://no_user@{host}:{port}"))
-    };
-    let grpc_url = if let Some(grpc_addr) = &conf.listen_grpc_addr {
-        let (host, port) = parse_host_port(grpc_addr)?;
-        let port = port.unwrap_or(DEFAULT_PAGESERVER_GRPC_PORT);
-        Some(format!("grpc://no_user@{host}:{port}"))
-    } else {
-        None
-    };
-    Ok(PageserverShardConnectionInfo {
-        id: Some(conf.id.to_string()),
-        libpq_url,
-        grpc_url,
-    })
-}
-
-pub fn tenant_locate_response_to_conn_info(
-    response: &pageserver_api::controller_api::TenantLocateResponse,
-) -> Result<PageserverConnectionInfo> {
-    let mut shards = HashMap::new();
-    for shard in response.shards.iter() {
-        tracing::info!("parsing {}", shard.listen_pg_addr);
-        let libpq_url = {
-            let host = &shard.listen_pg_addr;
-            let port = shard.listen_pg_port;
-            Some(format!("postgres://no_user@{host}:{port}"))
-        };
-        let grpc_url = if let Some(grpc_addr) = &shard.listen_grpc_addr {
-            let host = grpc_addr;
-            let port = shard.listen_grpc_port.expect("no gRPC port");
-            Some(format!("grpc://no_user@{host}:{port}"))
-        } else {
-            None
-        };
-
-        let shard_info = PageserverShardInfo {
-            pageservers: vec![PageserverShardConnectionInfo {
-                id: Some(shard.node_id.to_string()),
-                libpq_url,
-                grpc_url,
-            }],
-        };
-
-        shards.insert(shard.shard_id.to_index(), shard_info);
-    }
-
-    let stripe_size = if response.shard_params.count.is_unsharded() {
-        None
-    } else {
-        Some(response.shard_params.stripe_size.0)
-    };
-    Ok(PageserverConnectionInfo {
-        shard_count: response.shard_params.count,
-        stripe_size,
-        shards,
-        prefer_protocol: PageserverProtocol::default(),
-    })
-}

control_plane/src/pageserver.rs

@@ -571,6 +571,11 @@ impl PageServerNode {
                 .map(|x| x.parse::<bool>())
                 .transpose()
                 .context("Failed to parse 'basebackup_cache_enabled' as bool")?,
+            rel_size_v1_access_disabled: settings
+                .remove("rel_size_v1_access_disabled")
+                .map(|x| x.parse::<bool>())
+                .transpose()
+                .context("Failed to parse 'rel_size_v1_access_disabled' as bool")?,
         };
         if !settings.is_empty() {
             bail!("Unrecognized tenant settings: {settings:?}")

control_plane/src/storage_controller.rs

@@ -56,6 +56,7 @@ pub struct NeonStorageControllerStartArgs {
     pub instance_id: u8,
     pub base_port: Option<u16>,
     pub start_timeout: humantime::Duration,
+    pub handle_ps_local_disk_loss: Option<bool>,
 }
 
 impl NeonStorageControllerStartArgs {
@@ -64,6 +65,7 @@ impl NeonStorageControllerStartArgs {
             instance_id: 1,
             base_port: None,
             start_timeout,
+            handle_ps_local_disk_loss: None,
         }
     }
 }
@@ -669,6 +671,10 @@ impl StorageController {
 
         println!("Starting storage controller at {scheme}://{host}:{listen_port}");
 
+        if start_args.handle_ps_local_disk_loss.unwrap_or_default() {
+            args.push("--handle-ps-local-disk-loss".to_string());
+        }
+
         background_process::start_process(
             COMMAND,
             &instance_dir,

deny.toml

@@ -35,6 +35,7 @@ reason = "The paste crate is a build-only dependency with no runtime components.
 # More documentation for the licenses section can be found here:
 # https://embarkstudios.github.io/cargo-deny/checks/licenses/cfg.html
 [licenses]
+version = 2
 allow = [
     "0BSD",
     "Apache-2.0",

endpoint_storage/src/app.rs

@@ -233,7 +233,7 @@ mod tests {
                 .unwrap()
                 .as_millis();
             use rand::Rng;
-            let random = rand::thread_rng().r#gen::<u32>();
+            let random = rand::rng().random::<u32>();
 
             let s3_config = remote_storage::S3Config {
                 bucket_name: var(REAL_S3_BUCKET).unwrap(),

libs/compute_api/src/responses.rs

@@ -108,11 +108,10 @@ pub enum PromoteState {
     Failed { error: String },
 }
 
-#[derive(Deserialize, Serialize, Default, Debug, Clone)]
+#[derive(Deserialize, Default, Debug)]
 #[serde(rename_all = "snake_case")]
-/// Result of /safekeepers_lsn
-pub struct SafekeepersLsn {
-    pub safekeepers: String,
+pub struct PromoteConfig {
+    pub spec: ComputeSpec,
     pub wal_flush_lsn: utils::lsn::Lsn,
 }
 
@@ -173,6 +172,11 @@ pub enum ComputeStatus {
     TerminationPendingImmediate,
     // Terminated Postgres
     Terminated,
+    // A spec refresh is being requested
+    RefreshConfigurationPending,
+    // A spec refresh is being applied. We cannot refresh configuration again until the current
+    // refresh is done, i.e., signal_refresh_configuration() will return 500 error.
+    RefreshConfiguration,
 }
 
 #[derive(Deserialize, Serialize)]
@@ -185,6 +189,10 @@ impl Display for ComputeStatus {
         match self {
             ComputeStatus::Empty => f.write_str("empty"),
             ComputeStatus::ConfigurationPending => f.write_str("configuration-pending"),
+            ComputeStatus::RefreshConfiguration => f.write_str("refresh-configuration"),
+            ComputeStatus::RefreshConfigurationPending => {
+                f.write_str("refresh-configuration-pending")
+            }
             ComputeStatus::Init => f.write_str("init"),
             ComputeStatus::Running => f.write_str("running"),
             ComputeStatus::Configuration => f.write_str("configuration"),

libs/compute_api/src/spec.rs

@@ -14,7 +14,6 @@ use serde::{Deserialize, Serialize};
 use url::Url;
 use utils::id::{TenantId, TimelineId};
 use utils::lsn::Lsn;
-use utils::shard::{ShardCount, ShardIndex};
 
 use crate::responses::TlsConfig;
 
@@ -106,17 +105,6 @@ pub struct ComputeSpec {
     // updated to fill these fields, we can make these non optional.
     pub tenant_id: Option<TenantId>,
     pub timeline_id: Option<TimelineId>,
-
-    /// Pageserver information can be passed in three different ways:
-    /// 1. Here in `pageserver_connection_info`
-    /// 2. In the `pageserver_connstring` field.
-    /// 3. in `cluster.settings`.
-    ///
-    /// The goal is to use method 1. everywhere. But for backwards-compatibility with old
-    /// versions of the control plane, `compute_ctl` will check 2. and 3. if the
-    /// `pageserver_connection_info` field is missing.
-    pub pageserver_connection_info: Option<PageserverConnectionInfo>,
-
     pub pageserver_connstring: Option<String>,
 
     // More neon ids that we expose to the compute_ctl
@@ -153,7 +141,7 @@ pub struct ComputeSpec {
 
     // Stripe size for pageserver sharding, in pages
     #[serde(default)]
-    pub shard_stripe_size: Option<u32>,
+    pub shard_stripe_size: Option<usize>,
 
     /// Local Proxy configuration used for JWT authentication
     #[serde(default)]
@@ -205,6 +193,9 @@ pub struct ComputeSpec {
     ///
     /// We use this value to derive other values, such as the installed extensions metric.
     pub suspend_timeout_seconds: i64,
+
+    // Databricks specific options for compute instance.
+    pub databricks_settings: Option<DatabricksSettings>,
 }
 
 /// Feature flag to signal `compute_ctl` to enable certain experimental functionality.
@@ -226,32 +217,6 @@ pub enum ComputeFeature {
     UnknownFeature,
 }
 
-#[derive(Clone, Debug, Deserialize, Serialize, Eq, PartialEq)]
-pub struct PageserverConnectionInfo {
-    /// NB: 0 for unsharded tenants, 1 for sharded tenants with 1 shard, following storage
-    pub shard_count: ShardCount,
-
-    /// INVARIANT: null if shard_count is 0, otherwise non-null and immutable
-    pub stripe_size: Option<u32>,
-
-    pub shards: HashMap<ShardIndex, PageserverShardInfo>,
-
-    #[serde(default)]
-    pub prefer_protocol: PageserverProtocol,
-}
-
-#[derive(Clone, Debug, Deserialize, Serialize, Eq, PartialEq)]
-pub struct PageserverShardInfo {
-    pub pageservers: Vec<PageserverShardConnectionInfo>,
-}
-
-#[derive(Clone, Debug, Deserialize, Serialize, Eq, PartialEq)]
-pub struct PageserverShardConnectionInfo {
-    pub id: Option<String>,
-    pub libpq_url: Option<String>,
-    pub grpc_url: Option<String>,
-}
-
 #[derive(Clone, Debug, Default, Deserialize, Serialize)]
 pub struct RemoteExtSpec {
     pub public_extensions: Option<Vec<String>>,
@@ -369,12 +334,6 @@ impl ComputeMode {
     }
 }
 
-impl Display for ComputeMode {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        f.write_str(self.to_type_str())
-    }
-}
-
 /// Log level for audit logging
 #[derive(Clone, Debug, Default, Eq, PartialEq, Deserialize, Serialize)]
 pub enum ComputeAudit {
@@ -460,6 +419,32 @@ pub struct GenericOption {
     pub vartype: String,
 }
 
+/// Postgres compute TLS settings.
+#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
+pub struct PgComputeTlsSettings {
+    // Absolute path to the certificate file for server-side TLS.
+    pub cert_file: String,
+    // Absolute path to the private key file for server-side TLS.
+    pub key_file: String,
+    // Absolute path to the certificate authority file for verifying client certificates.
+    pub ca_file: String,
+}
+
+/// Databricks specific options for compute instance.
+/// This is used to store any other settings that needs to be propagate to Compute
+/// but should not be persisted to ComputeSpec in the database.
+#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
+pub struct DatabricksSettings {
+    pub pg_compute_tls_settings: PgComputeTlsSettings,
+    // Absolute file path to databricks_pg_hba.conf file.
+    pub databricks_pg_hba: String,
+    // Absolute file path to databricks_pg_ident.conf file.
+    pub databricks_pg_ident: String,
+    // Hostname portion of the Databricks workspace URL of the endpoint, or empty string if not known.
+    // A valid hostname is required for the compute instance to support PAT logins.
+    pub databricks_workspace_host: String,
+}
+
 /// Optional collection of `GenericOption`'s. Type alias allows us to
 /// declare a `trait` on it.
 pub type GenericOptions = Option<Vec<GenericOption>>;
@@ -485,15 +470,13 @@ pub struct JwksSettings {
     pub jwt_audience: Option<String>,
 }
 
-/// Protocol used to connect to a Pageserver.
-#[derive(Clone, Copy, Debug, Default, Deserialize, Serialize, PartialEq, Eq)]
+/// Protocol used to connect to a Pageserver. Parsed from the connstring scheme.
+#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
 pub enum PageserverProtocol {
     /// The original protocol based on libpq and COPY. Uses postgresql:// or postgres:// scheme.
     #[default]
-    #[serde(rename = "libpq")]
     Libpq,
     /// A newer, gRPC-based protocol. Uses grpc:// scheme.
-    #[serde(rename = "grpc")]
     Grpc,
 }
 

libs/consumption_metrics/src/lib.rs

@@ -90,7 +90,7 @@ impl<'a> IdempotencyKey<'a> {
         IdempotencyKey {
             now: Utc::now(),
             node_id,
-            nonce: rand::thread_rng().gen_range(0..=9999),
+            nonce: rand::rng().random_range(0..=9999),
         }
     }
 

libs/desim/src/node_os.rs

@@ -41,7 +41,7 @@ impl NodeOs {
 
     /// Generate a random number in range [0, max).
     pub fn random(&self, max: u64) -> u64 {
-        self.internal.rng.lock().gen_range(0..max)
+        self.internal.rng.lock().random_range(0..max)
     }
 
     /// Append a new event to the world event log.

libs/desim/src/options.rs

@@ -32,10 +32,10 @@ impl Delay {
     /// Generate a random delay in range [min, max]. Return None if the
     /// message should be dropped.
     pub fn delay(&self, rng: &mut StdRng) -> Option<u64> {
-        if rng.gen_bool(self.fail_prob) {
+        if rng.random_bool(self.fail_prob) {
             return None;
         }
-        Some(rng.gen_range(self.min..=self.max))
+        Some(rng.random_range(self.min..=self.max))
     }
 }
 

libs/desim/src/world.rs

@@ -69,7 +69,7 @@ impl World {
     /// Create a new random number generator.
     pub fn new_rng(&self) -> StdRng {
         let mut rng = self.rng.lock();
-        StdRng::from_rng(rng.deref_mut()).unwrap()
+        StdRng::from_rng(rng.deref_mut())
     }
 
     /// Create a new node.

libs/metrics/Cargo.toml

@@ -17,5 +17,5 @@ procfs.workspace = true
 measured-process.workspace = true
 
 [dev-dependencies]
-rand = "0.8"
-rand_distr = "0.4.3"
+rand.workspace = true
+rand_distr = "0.5"

libs/metrics/src/hll.rs

@@ -260,7 +260,7 @@ mod tests {
 
     #[test]
     fn test_cardinality_small() {
-        let (actual, estimate) = test_cardinality(100, Zipf::new(100, 1.2f64).unwrap());
+        let (actual, estimate) = test_cardinality(100, Zipf::new(100.0, 1.2f64).unwrap());
 
         assert_eq!(actual, [46, 30, 32]);
         assert!(51.3 < estimate[0] && estimate[0] < 51.4);
@@ -270,7 +270,7 @@ mod tests {
 
     #[test]
     fn test_cardinality_medium() {
-        let (actual, estimate) = test_cardinality(10000, Zipf::new(10000, 1.2f64).unwrap());
+        let (actual, estimate) = test_cardinality(10000, Zipf::new(10000.0, 1.2f64).unwrap());
 
         assert_eq!(actual, [2529, 1618, 1629]);
         assert!(2309.1 < estimate[0] && estimate[0] < 2309.2);
@@ -280,7 +280,8 @@ mod tests {
 
     #[test]
     fn test_cardinality_large() {
-        let (actual, estimate) = test_cardinality(1_000_000, Zipf::new(1_000_000, 1.2f64).unwrap());
+        let (actual, estimate) =
+            test_cardinality(1_000_000, Zipf::new(1_000_000.0, 1.2f64).unwrap());
 
         assert_eq!(actual, [129077, 79579, 79630]);
         assert!(126067.2 < estimate[0] && estimate[0] < 126067.3);
@@ -290,7 +291,7 @@ mod tests {
 
     #[test]
     fn test_cardinality_small2() {
-        let (actual, estimate) = test_cardinality(100, Zipf::new(200, 0.8f64).unwrap());
+        let (actual, estimate) = test_cardinality(100, Zipf::new(200.0, 0.8f64).unwrap());
 
         assert_eq!(actual, [92, 58, 60]);
         assert!(116.1 < estimate[0] && estimate[0] < 116.2);
@@ -300,7 +301,7 @@ mod tests {
 
     #[test]
     fn test_cardinality_medium2() {
-        let (actual, estimate) = test_cardinality(10000, Zipf::new(20000, 0.8f64).unwrap());
+        let (actual, estimate) = test_cardinality(10000, Zipf::new(20000.0, 0.8f64).unwrap());
 
         assert_eq!(actual, [8201, 5131, 5051]);
         assert!(6846.4 < estimate[0] && estimate[0] < 6846.5);
@@ -310,7 +311,8 @@ mod tests {
 
     #[test]
     fn test_cardinality_large2() {
-        let (actual, estimate) = test_cardinality(1_000_000, Zipf::new(2_000_000, 0.8f64).unwrap());
+        let (actual, estimate) =
+            test_cardinality(1_000_000, Zipf::new(2_000_000.0, 0.8f64).unwrap());
 
         assert_eq!(actual, [777847, 482069, 482246]);
         assert!(699437.4 < estimate[0] && estimate[0] < 699437.5);

libs/metrics/src/lib.rs

@@ -129,6 +129,12 @@ impl<L: LabelGroup> InfoMetric<L> {
     }
 }
 
+impl<L: LabelGroup + Default> Default for InfoMetric<L, GaugeState> {
+    fn default() -> Self {
+        InfoMetric::new(L::default())
+    }
+}
+
 impl<L: LabelGroup, M: MetricType<Metadata = ()>> InfoMetric<L, M> {
     pub fn with_metric(label: L, metric: M) -> Self {
         Self {

libs/neon-shmem/Cargo.toml

@@ -16,5 +16,5 @@ rustc-hash.workspace = true
 tempfile = "3.14.0"
 
 [dev-dependencies]
-rand = "0.9"
-rand_distr = "0.5.1"
+rand.workspace = true
+rand_distr = "0.5.1"

libs/neon-shmem/benches/hmap_resize.rs

@@ -1,330 +0,0 @@
-use criterion::{BatchSize, BenchmarkId, Criterion, criterion_group, criterion_main};
-use neon_shmem::hash::HashMapAccess;
-use neon_shmem::hash::HashMapInit;
-use neon_shmem::hash::entry::Entry;
-use rand::distr::{Distribution, StandardUniform};
-use rand::prelude::*;
-use std::default::Default;
-use std::hash::BuildHasher;
-
-// Taken from bindings to C code
-
-#[derive(Clone, Debug, Hash, Eq, PartialEq)]
-#[repr(C)]
-pub struct FileCacheKey {
-    pub _spc_id: u32,
-    pub _db_id: u32,
-    pub _rel_number: u32,
-    pub _fork_num: u32,
-    pub _block_num: u32,
-}
-
-impl Distribution<FileCacheKey> for StandardUniform {
-    // questionable, but doesn't need to be good randomness
-    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> FileCacheKey {
-        FileCacheKey {
-            _spc_id: rng.random(),
-            _db_id: rng.random(),
-            _rel_number: rng.random(),
-            _fork_num: rng.random(),
-            _block_num: rng.random(),
-        }
-    }
-}
-
-#[derive(Clone, Debug)]
-#[repr(C)]
-pub struct FileCacheEntry {
-    pub _offset: u32,
-    pub _access_count: u32,
-    pub _prev: *mut FileCacheEntry,
-    pub _next: *mut FileCacheEntry,
-    pub _state: [u32; 8],
-}
-
-impl FileCacheEntry {
-    fn dummy() -> Self {
-        Self {
-            _offset: 0,
-            _access_count: 0,
-            _prev: std::ptr::null_mut(),
-            _next: std::ptr::null_mut(),
-            _state: [0; 8],
-        }
-    }
-}
-
-// Utilities for applying operations.
-
-#[derive(Clone, Debug)]
-struct TestOp<K, V>(K, Option<V>);
-
-fn apply_op<K: Clone + std::hash::Hash + Eq, V, S: std::hash::BuildHasher>(
-    op: TestOp<K, V>,
-    map: &mut HashMapAccess<K, V, S>,
-) {
-    let entry = map.entry(op.0);
-
-    match op.1 {
-        Some(new) => match entry {
-            Entry::Occupied(mut e) => Some(e.insert(new)),
-            Entry::Vacant(e) => {
-                _ = e.insert(new).unwrap();
-                None
-            }
-        },
-        None => match entry {
-            Entry::Occupied(e) => Some(e.remove()),
-            Entry::Vacant(_) => None,
-        },
-    };
-}
-
-// Hash utilities
-
-struct SeaRandomState {
-    k1: u64,
-    k2: u64,
-    k3: u64,
-    k4: u64,
-}
-
-impl std::hash::BuildHasher for SeaRandomState {
-    type Hasher = seahash::SeaHasher;
-
-    fn build_hasher(&self) -> Self::Hasher {
-        seahash::SeaHasher::with_seeds(self.k1, self.k2, self.k3, self.k4)
-    }
-}
-
-impl SeaRandomState {
-    fn new() -> Self {
-        let mut rng = rand::rng();
-        Self {
-            k1: rng.random(),
-            k2: rng.random(),
-            k3: rng.random(),
-            k4: rng.random(),
-        }
-    }
-}
-
-fn small_benchs(c: &mut Criterion) {
-    let mut group = c.benchmark_group("Small maps");
-    group.sample_size(10);
-
-    group.bench_function("small_rehash", |b| {
-        let ideal_filled = 4_000_000;
-        let size = 5_000_000;
-        let mut writer = HashMapInit::new_resizeable(size, size * 2).attach_writer();
-        let mut rng = rand::rng();
-        while writer.get_num_buckets_in_use() < ideal_filled as usize {
-            let key: FileCacheKey = rng.random();
-            let val = FileCacheEntry::dummy();
-            apply_op(TestOp(key, Some(val)), &mut writer);
-        }
-        b.iter(|| writer.shuffle());
-    });
-
-    group.bench_function("small_rehash_xxhash", |b| {
-        let ideal_filled = 4_000_000;
-        let size = 5_000_000;
-        let mut writer = HashMapInit::new_resizeable(size, size * 2)
-            .with_hasher(twox_hash::xxhash64::RandomState::default())
-            .attach_writer();
-        let mut rng = rand::rng();
-        while writer.get_num_buckets_in_use() < ideal_filled as usize {
-            let key: FileCacheKey = rng.random();
-            let val = FileCacheEntry::dummy();
-            apply_op(TestOp(key, Some(val)), &mut writer);
-        }
-        b.iter(|| writer.shuffle());
-    });
-
-    group.bench_function("small_rehash_ahash", |b| {
-        let ideal_filled = 4_000_000;
-        let size = 5_000_000;
-        let mut writer = HashMapInit::new_resizeable(size, size * 2)
-            .with_hasher(ahash::RandomState::default())
-            .attach_writer();
-        let mut rng = rand::rng();
-        while writer.get_num_buckets_in_use() < ideal_filled as usize {
-            let key: FileCacheKey = rng.random();
-            let val = FileCacheEntry::dummy();
-            apply_op(TestOp(key, Some(val)), &mut writer);
-        }
-        b.iter(|| writer.shuffle());
-    });
-
-    group.bench_function("small_rehash_seahash", |b| {
-        let ideal_filled = 4_000_000;
-        let size = 5_000_000;
-        let mut writer = HashMapInit::new_resizeable(size, size * 2)
-            .with_hasher(SeaRandomState::new())
-            .attach_writer();
-        let mut rng = rand::rng();
-        while writer.get_num_buckets_in_use() < ideal_filled as usize {
-            let key: FileCacheKey = rng.random();
-            let val = FileCacheEntry::dummy();
-            apply_op(TestOp(key, Some(val)), &mut writer);
-        }
-        b.iter(|| writer.shuffle());
-    });
-
-    group.finish();
-}
-
-fn real_benchs(c: &mut Criterion) {
-    let mut group = c.benchmark_group("Realistic workloads");
-    group.sample_size(10);
-    group.bench_function("real_bulk_insert", |b| {
-        let size = 125_000_000;
-        let ideal_filled = 100_000_000;
-        let mut rng = rand::rng();
-        b.iter_batched(
-            || HashMapInit::new_resizeable(size, size * 2).attach_writer(),
-            |writer| {
-                for _ in 0..ideal_filled {
-                    let key: FileCacheKey = rng.random();
-                    let val = FileCacheEntry::dummy();
-                    let entry = writer.entry(key);
-                    std::hint::black_box(match entry {
-                        Entry::Occupied(mut e) => {
-                            e.insert(val);
-                        }
-                        Entry::Vacant(e) => {
-                            _ = e.insert(val).unwrap();
-                        }
-                    })
-                }
-            },
-            BatchSize::SmallInput,
-        )
-    });
-
-    group.bench_function("real_rehash", |b| {
-        let size = 125_000_000;
-        let ideal_filled = 100_000_000;
-        let mut writer = HashMapInit::new_resizeable(size, size).attach_writer();
-        let mut rng = rand::rng();
-        while writer.get_num_buckets_in_use() < ideal_filled {
-            let key: FileCacheKey = rng.random();
-            let val = FileCacheEntry::dummy();
-            apply_op(TestOp(key, Some(val)), &mut writer);
-        }
-        b.iter(|| writer.shuffle());
-    });
-
-    group.bench_function("real_rehash_hashbrown", |b| {
-        let size = 125_000_000;
-        let ideal_filled = 100_000_000;
-        let mut writer = hashbrown::raw::RawTable::new();
-        let mut rng = rand::rng();
-        let hasher = rustc_hash::FxBuildHasher::default();
-        unsafe {
-            writer
-                .resize(
-                    size,
-                    |(k, _)| hasher.hash_one(&k),
-                    hashbrown::raw::Fallibility::Infallible,
-                )
-                .unwrap();
-        }
-        while writer.len() < ideal_filled as usize {
-            let key: FileCacheKey = rng.random();
-            let val = FileCacheEntry::dummy();
-            writer.insert(hasher.hash_one(&key), (key, val), |(k, _)| {
-                hasher.hash_one(&k)
-            });
-        }
-        b.iter(|| unsafe {
-            writer.table.rehash_in_place(
-                &|table, index| {
-                    hasher.hash_one(
-                        &table
-                            .bucket::<(FileCacheKey, FileCacheEntry)>(index)
-                            .as_ref()
-                            .0,
-                    )
-                },
-                std::mem::size_of::<(FileCacheKey, FileCacheEntry)>(),
-                if std::mem::needs_drop::<(FileCacheKey, FileCacheEntry)>() {
-                    Some(|ptr| std::ptr::drop_in_place(ptr as *mut (FileCacheKey, FileCacheEntry)))
-                } else {
-                    None
-                },
-            )
-        });
-    });
-
-    for elems in [2, 4, 8, 16, 32, 64, 96, 112] {
-        group.bench_with_input(
-            BenchmarkId::new("real_rehash_varied", elems),
-            &elems,
-            |b, &size| {
-                let ideal_filled = size * 1_000_000;
-                let size = 125_000_000;
-                let mut writer = HashMapInit::new_resizeable(size, size).attach_writer();
-                let mut rng = rand::rng();
-                while writer.get_num_buckets_in_use() < ideal_filled as usize {
-                    let key: FileCacheKey = rng.random();
-                    let val = FileCacheEntry::dummy();
-                    apply_op(TestOp(key, Some(val)), &mut writer);
-                }
-                b.iter(|| writer.shuffle());
-            },
-        );
-        group.bench_with_input(
-            BenchmarkId::new("real_rehash_varied_hashbrown", elems),
-            &elems,
-            |b, &size| {
-                let ideal_filled = size * 1_000_000;
-                let size = 125_000_000;
-                let mut writer = hashbrown::raw::RawTable::new();
-                let mut rng = rand::rng();
-                let hasher = rustc_hash::FxBuildHasher::default();
-                unsafe {
-                    writer
-                        .resize(
-                            size,
-                            |(k, _)| hasher.hash_one(&k),
-                            hashbrown::raw::Fallibility::Infallible,
-                        )
-                        .unwrap();
-                }
-                while writer.len() < ideal_filled as usize {
-                    let key: FileCacheKey = rng.random();
-                    let val = FileCacheEntry::dummy();
-                    writer.insert(hasher.hash_one(&key), (key, val), |(k, _)| {
-                        hasher.hash_one(&k)
-                    });
-                }
-                b.iter(|| unsafe {
-                    writer.table.rehash_in_place(
-                        &|table, index| {
-                            hasher.hash_one(
-                                &table
-                                    .bucket::<(FileCacheKey, FileCacheEntry)>(index)
-                                    .as_ref()
-                                    .0,
-                            )
-                        },
-                        std::mem::size_of::<(FileCacheKey, FileCacheEntry)>(),
-                        if std::mem::needs_drop::<(FileCacheKey, FileCacheEntry)>() {
-                            Some(|ptr| {
-                                std::ptr::drop_in_place(ptr as *mut (FileCacheKey, FileCacheEntry))
-                            })
-                        } else {
-                            None
-                        },
-                    )
-                });
-            },
-        );
-    }
-
-    group.finish();
-}
-
-criterion_group!(benches, small_benchs, real_benchs);
-criterion_main!(benches);

libs/neon-shmem/src/hash.rs

@@ -16,7 +16,6 @@
 //!
 //! Concurrency is managed very simply: the entire map is guarded by one shared-memory RwLock.
 
-use std::fmt::Debug;
 use std::hash::{BuildHasher, Hash};
 use std::mem::MaybeUninit;
 
@@ -57,22 +56,6 @@ pub struct HashMapInit<'a, K, V, S = rustc_hash::FxBuildHasher> {
     num_buckets: u32,
 }
 
-impl<'a, K, V, S> Debug for HashMapInit<'a, K, V, S>
-where
-    K: Debug,
-    V: Debug,
-{
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        f.debug_struct("HashMapInit")
-            .field("shmem_handle", &self.shmem_handle)
-            .field("shared_ptr", &self.shared_ptr)
-            .field("shared_size", &self.shared_size)
-            // .field("hasher", &self.hasher)
-            .field("num_buckets", &self.num_buckets)
-            .finish()
-    }
-}
-
 /// This is a per-process handle to a hash table that (possibly) lives in shared memory.
 /// If a child process is launched with fork(), the child process should
 /// get its own HashMapAccess by calling HashMapInit::attach_writer/reader().
@@ -88,20 +71,6 @@ pub struct HashMapAccess<'a, K, V, S = rustc_hash::FxBuildHasher> {
 unsafe impl<K: Sync, V: Sync, S> Sync for HashMapAccess<'_, K, V, S> {}
 unsafe impl<K: Send, V: Send, S> Send for HashMapAccess<'_, K, V, S> {}
 
-impl<'a, K, V, S> Debug for HashMapAccess<'a, K, V, S>
-where
-    K: Debug,
-    V: Debug,
-{
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        f.debug_struct("HashMapAccess")
-            .field("shmem_handle", &self.shmem_handle)
-            .field("shared_ptr", &self.shared_ptr)
-            // .field("hasher", &self.hasher)
-            .finish()
-    }
-}
-
 impl<'a, K: Clone + Hash + Eq, V, S> HashMapInit<'a, K, V, S> {
     /// Change the 'hasher' used by the hash table.
     ///
@@ -329,7 +298,7 @@ where
 
     /// Get a reference to the entry containing a key.
     ///
-    /// NB: This takes a write lock as there's no way to distinguish whether the intention
+    /// NB: THis takes a write lock as there's no way to distinguish whether the intention
     /// is to use the entry for reading or for writing in advance.
     pub fn entry(&self, key: K) -> Entry<'a, '_, K, V> {
         let hash = self.get_hash_value(&key);
@@ -541,7 +510,7 @@ where
     /// # Panics
     /// Panics if called on a map initialized with [`HashMapInit::with_fixed`] or if `num_buckets` is
     /// greater than the number of buckets in the map.
-    pub fn begin_shrink(&self, num_buckets: u32) {
+    pub fn begin_shrink(&mut self, num_buckets: u32) {
         let mut map = unsafe { self.shared_ptr.as_mut() }.unwrap().write();
         assert!(
             num_buckets <= map.get_num_buckets() as u32,

libs/neon-shmem/src/hash/core.rs

@@ -1,6 +1,5 @@
 //! Simple hash table with chaining.
 
-use std::fmt::Debug;
 use std::hash::Hash;
 use std::mem::MaybeUninit;
 
@@ -18,19 +17,6 @@ pub(crate) struct Bucket<K, V> {
     pub(crate) inner: Option<(K, V)>,
 }
 
-impl<K, V> Debug for Bucket<K, V>
-where
-    K: Debug,
-    V: Debug,
-{
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        f.debug_struct("Bucket")
-            .field("next", &self.next)
-            .field("inner", &self.inner)
-            .finish()
-    }
-}
-
 /// Core hash table implementation.
 pub(crate) struct CoreHashMap<'a, K, V> {
     /// Dictionary used to map hashes to bucket indices.
@@ -45,22 +31,6 @@ pub(crate) struct CoreHashMap<'a, K, V> {
     pub(crate) buckets_in_use: u32,
 }
 
-impl<'a, K, V> Debug for CoreHashMap<'a, K, V>
-where
-    K: Debug,
-    V: Debug,
-{
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        f.debug_struct("CoreHashMap")
-            .field("dictionary", &self.dictionary)
-            .field("buckets", &self.buckets)
-            .field("free_head", &self.free_head)
-            .field("alloc_limit", &self.alloc_limit)
-            .field("buckets_in_use", &self.buckets_in_use)
-            .finish()
-    }
-}
-
 /// Error for when there are no empty buckets left but one is needed.
 #[derive(Debug, PartialEq)]
 pub struct FullError;

libs/neon-shmem/src/shmem.rs

@@ -21,7 +21,6 @@ use nix::unistd::ftruncate as nix_ftruncate;
 /// the underlying file is resized. Do not access the area beyond the current size. Currently, that
 /// will cause the file to be expanded, but we might use `mprotect()` etc. to enforce that in the
 /// future.
-#[derive(Debug)]
 pub struct ShmemHandle {
     /// memfd file descriptor
     fd: OwnedFd,
@@ -36,7 +35,6 @@ pub struct ShmemHandle {
 }
 
 /// This is stored at the beginning in the shared memory area.
-#[derive(Debug)]
 struct SharedStruct {
     max_size: usize,
 

libs/neonart/Cargo.toml

@@ -1,14 +0,0 @@
-[package]
-name = "neonart"
-version = "0.1.0"
-edition.workspace = true
-license.workspace = true
-
-[dependencies]
-crossbeam-utils.workspace = true
-spin.workspace = true
-tracing.workspace = true
-
-[dev-dependencies]
-rand = "0.9.1"
-rand_distr = "0.5.1"

libs/neonart/src/algorithm.rs

@@ -1,599 +0,0 @@
-mod lock_and_version;
-pub(crate) mod node_ptr;
-mod node_ref;
-
-use std::vec::Vec;
-
-use crate::algorithm::lock_and_version::ConcurrentUpdateError;
-use crate::algorithm::node_ptr::MAX_PREFIX_LEN;
-use crate::algorithm::node_ref::{NewNodeRef, NodeRef, ReadLockedNodeRef, WriteLockedNodeRef};
-use crate::allocator::OutOfMemoryError;
-
-use crate::TreeWriteGuard;
-use crate::UpdateAction;
-use crate::allocator::ArtAllocator;
-use crate::epoch::EpochPin;
-use crate::{Key, Value};
-
-pub(crate) type RootPtr<V> = node_ptr::NodePtr<V>;
-
-#[derive(Debug)]
-pub enum ArtError {
-    ConcurrentUpdate, // need to retry
-    OutOfMemory,
-}
-
-impl From<ConcurrentUpdateError> for ArtError {
-    fn from(_: ConcurrentUpdateError) -> ArtError {
-        ArtError::ConcurrentUpdate
-    }
-}
-
-impl From<OutOfMemoryError> for ArtError {
-    fn from(_: OutOfMemoryError) -> ArtError {
-        ArtError::OutOfMemory
-    }
-}
-
-pub fn new_root<V: Value>(
-    allocator: &impl ArtAllocator<V>,
-) -> Result<RootPtr<V>, OutOfMemoryError> {
-    node_ptr::new_root(allocator)
-}
-
-pub(crate) fn search<'e, K: Key, V: Value>(
-    key: &K,
-    root: RootPtr<V>,
-    epoch_pin: &'e EpochPin,
-) -> Option<&'e V> {
-    loop {
-        let root_ref = NodeRef::from_root_ptr(root);
-        if let Ok(result) = lookup_recurse(key.as_bytes(), root_ref, None, epoch_pin) {
-            break result;
-        }
-        // retry
-    }
-}
-
-pub(crate) fn iter_next<'e, V: Value>(
-    key: &[u8],
-    root: RootPtr<V>,
-    epoch_pin: &'e EpochPin,
-) -> Option<(Vec<u8>, &'e V)> {
-    loop {
-        let mut path = Vec::new();
-        let root_ref = NodeRef::from_root_ptr(root);
-
-        match next_recurse(key, &mut path, root_ref, epoch_pin) {
-            Ok(Some(v)) => {
-                assert_eq!(path.len(), key.len());
-                break Some((path, v));
-            }
-            Ok(None) => break None,
-            Err(ConcurrentUpdateError()) => {
-                // retry
-                continue;
-            }
-        }
-    }
-}
-
-pub(crate) fn update_fn<'e, 'g, K: Key, V: Value, A: ArtAllocator<V>, F>(
-    key: &K,
-    value_fn: F,
-    root: RootPtr<V>,
-    guard: &'g mut TreeWriteGuard<'e, K, V, A>,
-) -> Result<(), OutOfMemoryError>
-where
-    F: FnOnce(Option<&V>) -> UpdateAction<V>,
-{
-    let value_fn_cell = std::cell::Cell::new(Some(value_fn));
-    loop {
-        let root_ref = NodeRef::from_root_ptr(root);
-        let this_value_fn = |arg: Option<&V>| value_fn_cell.take().unwrap()(arg);
-        let key_bytes = key.as_bytes();
-
-        match update_recurse(
-            key_bytes,
-            this_value_fn,
-            root_ref,
-            None,
-            None,
-            guard,
-            0,
-            key_bytes,
-        ) {
-            Ok(()) => break Ok(()),
-            Err(ArtError::ConcurrentUpdate) => {
-                continue; // retry
-            }
-            Err(ArtError::OutOfMemory) => break Err(OutOfMemoryError()),
-        }
-    }
-}
-
-// Error means you must retry.
-//
-// This corresponds to the 'lookupOpt' function in the paper
-#[allow(clippy::only_used_in_recursion)]
-fn lookup_recurse<'e, V: Value>(
-    key: &[u8],
-    node: NodeRef<'e, V>,
-    parent: Option<ReadLockedNodeRef<V>>,
-    epoch_pin: &'e EpochPin,
-) -> Result<Option<&'e V>, ConcurrentUpdateError> {
-    let rnode = node.read_lock_or_restart()?;
-    if let Some(parent) = parent {
-        parent.read_unlock_or_restart()?;
-    }
-
-    // check if the prefix matches, may increment level
-    let prefix_len = if let Some(prefix_len) = rnode.prefix_matches(key) {
-        prefix_len
-    } else {
-        rnode.read_unlock_or_restart()?;
-        return Ok(None);
-    };
-
-    if rnode.is_leaf() {
-        assert_eq!(key.len(), prefix_len);
-        let vptr = rnode.get_leaf_value_ptr()?;
-        // safety: It's OK to return a ref of the pointer because we checked the version
-        // and the lifetime of 'epoch_pin' enforces that the reference is only accessible
-        // as long as the epoch is pinned.
-        let v = unsafe { vptr.as_ref().unwrap() };
-        return Ok(Some(v));
-    }
-
-    let key = &key[prefix_len..];
-
-    // find child (or leaf value)
-    let next_node = rnode.find_child_or_restart(key[0])?;
-
-    match next_node {
-        None => Ok(None), // key not found
-        Some(child) => lookup_recurse(&key[1..], child, Some(rnode), epoch_pin),
-    }
-}
-
-#[allow(clippy::only_used_in_recursion)]
-fn next_recurse<'e, V: Value>(
-    min_key: &[u8],
-    path: &mut Vec<u8>,
-    node: NodeRef<'e, V>,
-    epoch_pin: &'e EpochPin,
-) -> Result<Option<&'e V>, ConcurrentUpdateError> {
-    let rnode = node.read_lock_or_restart()?;
-    let prefix = rnode.get_prefix();
-    if !prefix.is_empty() {
-        path.extend_from_slice(prefix);
-    }
-
-    use std::cmp::Ordering;
-    let comparison = path.as_slice().cmp(&min_key[0..path.len()]);
-    if comparison == Ordering::Less {
-        rnode.read_unlock_or_restart()?;
-        return Ok(None);
-    }
-
-    if rnode.is_leaf() {
-        assert_eq!(path.len(), min_key.len());
-        let vptr = rnode.get_leaf_value_ptr()?;
-        // safety: It's OK to return a ref of the pointer because we checked the version
-        // and the lifetime of 'epoch_pin' enforces that the reference is only accessible
-        // as long as the epoch is pinned.
-        let v = unsafe { vptr.as_ref().unwrap() };
-        return Ok(Some(v));
-    }
-
-    let mut min_key_byte = match comparison {
-        Ordering::Less => unreachable!(), // checked this above already
-        Ordering::Equal => min_key[path.len()],
-        Ordering::Greater => 0,
-    };
-
-    loop {
-        match rnode.find_next_child_or_restart(min_key_byte)? {
-            None => {
-                return Ok(None);
-            }
-            Some((key_byte, child_ref)) => {
-                let path_len = path.len();
-                path.push(key_byte);
-                let result = next_recurse(min_key, path, child_ref, epoch_pin)?;
-                if result.is_some() {
-                    return Ok(result);
-                }
-                if key_byte == u8::MAX {
-                    return Ok(None);
-                }
-                path.truncate(path_len);
-                min_key_byte = key_byte + 1;
-            }
-        }
-    }
-}
-
-// This corresponds to the 'insertOpt' function in the paper
-#[allow(clippy::only_used_in_recursion)]
-#[allow(clippy::too_many_arguments)]
-pub(crate) fn update_recurse<'e, K: Key, V: Value, A: ArtAllocator<V>, F>(
-    key: &[u8],
-    value_fn: F,
-    node: NodeRef<'e, V>,
-    rparent: Option<(ReadLockedNodeRef<V>, u8)>,
-    rgrandparent: Option<(ReadLockedNodeRef<V>, u8)>,
-    guard: &'_ mut TreeWriteGuard<'e, K, V, A>,
-    level: usize,
-    orig_key: &[u8],
-) -> Result<(), ArtError>
-where
-    F: FnOnce(Option<&V>) -> UpdateAction<V>,
-{
-    let rnode = node.read_lock_or_restart()?;
-
-    let prefix_match_len = rnode.prefix_matches(key);
-    if prefix_match_len.is_none() {
-        let (rparent, parent_key) = rparent.expect("direct children of the root have no prefix");
-        let mut wparent = rparent.upgrade_to_write_lock_or_restart()?;
-        let mut wnode = rnode.upgrade_to_write_lock_or_restart()?;
-
-        match value_fn(None) {
-            UpdateAction::Nothing => {}
-            UpdateAction::Insert(new_value) => {
-                insert_split_prefix(key, new_value, &mut wnode, &mut wparent, parent_key, guard)?;
-            }
-            UpdateAction::Remove => {
-                panic!("unexpected Remove action on insertion");
-            }
-        }
-        wnode.write_unlock();
-        wparent.write_unlock();
-        return Ok(());
-    }
-    let prefix_match_len = prefix_match_len.unwrap();
-    let key = &key[prefix_match_len..];
-    let level = level + prefix_match_len;
-
-    if rnode.is_leaf() {
-        assert_eq!(key.len(), 0);
-        let (rparent, parent_key) = rparent.expect("root cannot be leaf");
-        let mut wparent = rparent.upgrade_to_write_lock_or_restart()?;
-        let mut wnode = rnode.upgrade_to_write_lock_or_restart()?;
-
-        // safety: Now that we have acquired the write lock, we have exclusive access to the
-        // value. XXX: There might be concurrent reads though?
-        let value_mut = wnode.get_leaf_value_mut();
-
-        match value_fn(Some(value_mut)) {
-            UpdateAction::Nothing => {
-                wparent.write_unlock();
-                wnode.write_unlock();
-            }
-            UpdateAction::Insert(_) => panic!("cannot insert over existing value"),
-            UpdateAction::Remove => {
-                guard.remember_obsolete_node(wnode.as_ptr());
-                wparent.delete_child(parent_key);
-                wnode.write_unlock_obsolete();
-
-                if let Some(rgrandparent) = rgrandparent {
-                    // FIXME: Ignore concurrency error. It doesn't lead to
-                    // corruption, but it means we might leak something. Until
-                    // another update cleans it up.
-                    let _ = cleanup_parent(wparent, rgrandparent, guard);
-                }
-            }
-        }
-
-        return Ok(());
-    }
-
-    let next_node = rnode.find_child_or_restart(key[0])?;
-
-    if next_node.is_none() {
-        if rnode.is_full() {
-            let (rparent, parent_key) = rparent.expect("root node cannot become full");
-            let mut wparent = rparent.upgrade_to_write_lock_or_restart()?;
-            let wnode = rnode.upgrade_to_write_lock_or_restart()?;
-
-            match value_fn(None) {
-                UpdateAction::Nothing => {
-                    wnode.write_unlock();
-                    wparent.write_unlock();
-                }
-                UpdateAction::Insert(new_value) => {
-                    insert_and_grow(key, new_value, wnode, &mut wparent, parent_key, guard)?;
-                    wparent.write_unlock();
-                }
-                UpdateAction::Remove => {
-                    panic!("unexpected Remove action on insertion");
-                }
-            };
-        } else {
-            let mut wnode = rnode.upgrade_to_write_lock_or_restart()?;
-            if let Some((rparent, _)) = rparent {
-                rparent.read_unlock_or_restart()?;
-            }
-            match value_fn(None) {
-                UpdateAction::Nothing => {}
-                UpdateAction::Insert(new_value) => {
-                    insert_to_node(&mut wnode, key, new_value, guard)?;
-                }
-                UpdateAction::Remove => {
-                    panic!("unexpected Remove action on insertion");
-                }
-            };
-            wnode.write_unlock();
-        }
-        Ok(())
-    } else {
-        let next_child = next_node.unwrap(); // checked above it's not None
-        if let Some((ref rparent, _)) = rparent {
-            rparent.check_or_restart()?;
-        }
-
-        // recurse to next level
-        update_recurse(
-            &key[1..],
-            value_fn,
-            next_child,
-            Some((rnode, key[0])),
-            rparent,
-            guard,
-            level + 1,
-            orig_key,
-        )
-    }
-}
-
-#[derive(Clone)]
-enum PathElement {
-    Prefix(Vec<u8>),
-    KeyByte(u8),
-}
-
-impl std::fmt::Debug for PathElement {
-    fn fmt(&self, fmt: &mut std::fmt::Formatter<'_>) -> Result<(), std::fmt::Error> {
-        match self {
-            PathElement::Prefix(prefix) => write!(fmt, "{prefix:?}"),
-            PathElement::KeyByte(key_byte) => write!(fmt, "{key_byte}"),
-        }
-    }
-}
-
-pub(crate) fn dump_tree<V: Value + std::fmt::Debug>(
-    root: RootPtr<V>,
-    epoch_pin: &'_ EpochPin,
-    dst: &mut dyn std::io::Write,
-) {
-    let root_ref = NodeRef::from_root_ptr(root);
-
-    let _ = dump_recurse(&[], root_ref, epoch_pin, 0, dst);
-}
-
-// TODO: return an Err if writeln!() returns error, instead of unwrapping
-#[allow(clippy::only_used_in_recursion)]
-fn dump_recurse<'e, V: Value + std::fmt::Debug>(
-    path: &[PathElement],
-    node: NodeRef<'e, V>,
-    epoch_pin: &'e EpochPin,
-    level: usize,
-    dst: &mut dyn std::io::Write,
-) -> Result<(), ConcurrentUpdateError> {
-    let indent = str::repeat(" ", level);
-
-    let rnode = node.read_lock_or_restart()?;
-    let mut path = Vec::from(path);
-    let prefix = rnode.get_prefix();
-    if !prefix.is_empty() {
-        path.push(PathElement::Prefix(Vec::from(prefix)));
-    }
-
-    if rnode.is_leaf() {
-        let vptr = rnode.get_leaf_value_ptr()?;
-        // safety: It's OK to return a ref of the pointer because we checked the version
-        // and the lifetime of 'epoch_pin' enforces that the reference is only accessible
-        // as long as the epoch is pinned.
-        let val = unsafe { vptr.as_ref().unwrap() };
-        writeln!(dst, "{indent} {path:?}: {val:?}").unwrap();
-        return Ok(());
-    }
-
-    for key_byte in 0..=u8::MAX {
-        match rnode.find_child_or_restart(key_byte)? {
-            None => continue,
-            Some(child_ref) => {
-                let rchild = child_ref.read_lock_or_restart()?;
-                writeln!(
-                    dst,
-                    "{} {:?}, {}: prefix {:?}",
-                    indent,
-                    &path,
-                    key_byte,
-                    rchild.get_prefix()
-                )
-                .unwrap();
-
-                let mut child_path = path.clone();
-                child_path.push(PathElement::KeyByte(key_byte));
-
-                dump_recurse(&child_path, child_ref, epoch_pin, level + 1, dst)?;
-            }
-        }
-    }
-
-    Ok(())
-}
-
-///```text
-///        [fooba]r -> value
-///
-/// [foo]b -> [a]r  -> value
-///      e -> [ls]e -> value
-///```
-fn insert_split_prefix<K: Key, V: Value, A: ArtAllocator<V>>(
-    key: &[u8],
-    value: V,
-    node: &mut WriteLockedNodeRef<V>,
-    parent: &mut WriteLockedNodeRef<V>,
-    parent_key: u8,
-    guard: &'_ TreeWriteGuard<K, V, A>,
-) -> Result<(), OutOfMemoryError> {
-    let old_node = node;
-    let old_prefix = old_node.get_prefix();
-    let common_prefix_len = common_prefix(key, old_prefix);
-
-    // Allocate a node for the new value.
-    let new_value_node = allocate_node_for_value(
-        &key[common_prefix_len + 1..],
-        value,
-        guard.tree_writer.allocator,
-    )?;
-
-    // Allocate a new internal node with the common prefix
-    // FIXME: deallocate 'new_value_node' on OOM
-    let mut prefix_node =
-        node_ref::new_internal(&key[..common_prefix_len], guard.tree_writer.allocator)?;
-
-    // Add the old node and the new nodes to the new internal node
-    prefix_node.insert_old_child(old_prefix[common_prefix_len], old_node);
-    prefix_node.insert_new_child(key[common_prefix_len], new_value_node);
-
-    // Modify the prefix of the old child in place
-    old_node.truncate_prefix(old_prefix.len() - common_prefix_len - 1);
-
-    // replace the pointer in the parent
-    parent.replace_child(parent_key, prefix_node.into_ptr());
-
-    Ok(())
-}
-
-fn insert_to_node<K: Key, V: Value, A: ArtAllocator<V>>(
-    wnode: &mut WriteLockedNodeRef<V>,
-    key: &[u8],
-    value: V,
-    guard: &'_ TreeWriteGuard<K, V, A>,
-) -> Result<(), OutOfMemoryError> {
-    let value_child = allocate_node_for_value(&key[1..], value, guard.tree_writer.allocator)?;
-    wnode.insert_child(key[0], value_child.into_ptr());
-    Ok(())
-}
-
-// On entry: 'parent' and 'node' are locked
-fn insert_and_grow<'e, 'g, K: Key, V: Value, A: ArtAllocator<V>>(
-    key: &[u8],
-    value: V,
-    wnode: WriteLockedNodeRef<V>,
-    parent: &mut WriteLockedNodeRef<V>,
-    parent_key_byte: u8,
-    guard: &'g mut TreeWriteGuard<'e, K, V, A>,
-) -> Result<(), ArtError> {
-    let mut bigger_node = wnode.grow(guard.tree_writer.allocator)?;
-
-    // FIXME: deallocate 'bigger_node' on OOM
-    let value_child = allocate_node_for_value(&key[1..], value, guard.tree_writer.allocator)?;
-    bigger_node.insert_new_child(key[0], value_child);
-
-    // Replace the pointer in the parent
-    parent.replace_child(parent_key_byte, bigger_node.into_ptr());
-
-    guard.remember_obsolete_node(wnode.as_ptr());
-    wnode.write_unlock_obsolete();
-
-    Ok(())
-}
-
-fn cleanup_parent<'e, 'g, K: Key, V: Value, A: ArtAllocator<V>>(
-    wparent: WriteLockedNodeRef<V>,
-    rgrandparent: (ReadLockedNodeRef<V>, u8),
-    guard: &'g mut TreeWriteGuard<'e, K, V, A>,
-) -> Result<(), ArtError> {
-    let (rgrandparent, grandparent_key_byte) = rgrandparent;
-
-    // If the parent becomes completely empty after the deletion, remove the parent from the
-    // grandparent. (This case is possible because we reserve only 8 bytes for the prefix.)
-    // TODO: not implemented.
-
-    // If the parent has only one child, replace the parent with the remaining child. (This is not
-    // possible if the child's prefix field cannot absorb the parent's)
-    if wparent.num_children() == 1 {
-        // Try to lock the remaining child. This can fail if the child is updated
-        // concurrently.
-        let (key_byte, remaining_child) = wparent.find_remaining_child();
-
-        let mut wremaining_child = remaining_child.write_lock_or_restart()?;
-
-        if 1 + wremaining_child.get_prefix().len() + wparent.get_prefix().len() <= MAX_PREFIX_LEN {
-            let mut wgrandparent = rgrandparent.upgrade_to_write_lock_or_restart()?;
-
-            // Ok, we have locked the leaf, the parent, the grandparent, and the parent's only
-            // remaining leaf. Proceed with the updates.
-
-            // Update the prefix on the remaining leaf
-            wremaining_child.prepend_prefix(wparent.get_prefix(), key_byte);
-
-            // Replace the pointer in the grandparent to point directly to the remaining leaf
-            wgrandparent.replace_child(grandparent_key_byte, wremaining_child.as_ptr());
-
-            // Mark the parent as deleted.
-            guard.remember_obsolete_node(wparent.as_ptr());
-            wparent.write_unlock_obsolete();
-            return Ok(());
-        }
-    }
-
-    // If the parent's children would fit on a smaller node type after the deletion, replace it with
-    // a smaller node.
-    if wparent.can_shrink() {
-        let mut wgrandparent = rgrandparent.upgrade_to_write_lock_or_restart()?;
-        let smaller_node = wparent.shrink(guard.tree_writer.allocator)?;
-
-        // Replace the pointer in the grandparent
-        wgrandparent.replace_child(grandparent_key_byte, smaller_node.into_ptr());
-
-        guard.remember_obsolete_node(wparent.as_ptr());
-        wparent.write_unlock_obsolete();
-        return Ok(());
-    }
-
-    // nothing to do
-    wparent.write_unlock();
-    Ok(())
-}
-
-// Allocate a new leaf node to hold 'value'. If the key is long, we
-// may need to allocate new internal nodes to hold it too
-fn allocate_node_for_value<'a, V: Value, A: ArtAllocator<V>>(
-    key: &[u8],
-    value: V,
-    allocator: &'a A,
-) -> Result<NewNodeRef<'a, V, A>, OutOfMemoryError> {
-    let mut prefix_off = key.len().saturating_sub(MAX_PREFIX_LEN);
-
-    let leaf_node = node_ref::new_leaf(&key[prefix_off..key.len()], value, allocator)?;
-
-    let mut node = leaf_node;
-    while prefix_off > 0 {
-        // Need another internal node
-        let remain_prefix = &key[0..prefix_off];
-
-        prefix_off = remain_prefix.len().saturating_sub(MAX_PREFIX_LEN + 1);
-        let mut internal_node = node_ref::new_internal(
-            &remain_prefix[prefix_off..remain_prefix.len() - 1],
-            allocator,
-        )?;
-        internal_node.insert_new_child(*remain_prefix.last().unwrap(), node);
-        node = internal_node;
-    }
-
-    Ok(node)
-}
-
-fn common_prefix(a: &[u8], b: &[u8]) -> usize {
-    for i in 0..MAX_PREFIX_LEN {
-        if a[i] != b[i] {
-            return i;
-        }
-    }
-    panic!("prefixes are equal");
-}

libs/neonart/src/algorithm/lock_and_version.rs

@@ -1,117 +0,0 @@
-//! Each node in the tree has contains one atomic word that stores three things:
-//!
-//! Bit 0: set if the node is "obsolete". An obsolete node has been removed from the tree,
-//!        but might still be accessed by concurrent readers until the epoch expires.
-//! Bit 1: set if the node is currently write-locked. Used as a spinlock.
-//! Bits 2-63: Version number, incremented every time the node is modified.
-//!
-//! AtomicLockAndVersion represents that.
-
-use std::sync::atomic::{AtomicU64, Ordering};
-
-pub(crate) struct ConcurrentUpdateError();
-
-pub(crate) struct AtomicLockAndVersion {
-    inner: AtomicU64,
-}
-
-impl AtomicLockAndVersion {
-    pub(crate) fn new() -> AtomicLockAndVersion {
-        AtomicLockAndVersion {
-            inner: AtomicU64::new(0),
-        }
-    }
-}
-
-impl AtomicLockAndVersion {
-    pub(crate) fn read_lock_or_restart(&self) -> Result<u64, ConcurrentUpdateError> {
-        let version = self.await_node_unlocked();
-        if is_obsolete(version) {
-            return Err(ConcurrentUpdateError());
-        }
-        Ok(version)
-    }
-
-    pub(crate) fn check_or_restart(&self, version: u64) -> Result<(), ConcurrentUpdateError> {
-        self.read_unlock_or_restart(version)
-    }
-
-    pub(crate) fn read_unlock_or_restart(&self, version: u64) -> Result<(), ConcurrentUpdateError> {
-        if self.inner.load(Ordering::Acquire) != version {
-            return Err(ConcurrentUpdateError());
-        }
-        Ok(())
-    }
-
-    pub(crate) fn upgrade_to_write_lock_or_restart(
-        &self,
-        version: u64,
-    ) -> Result<(), ConcurrentUpdateError> {
-        if self
-            .inner
-            .compare_exchange(
-                version,
-                set_locked_bit(version),
-                Ordering::Acquire,
-                Ordering::Relaxed,
-            )
-            .is_err()
-        {
-            return Err(ConcurrentUpdateError());
-        }
-        Ok(())
-    }
-
-    pub(crate) fn write_lock_or_restart(&self) -> Result<(), ConcurrentUpdateError> {
-        let old = self.inner.load(Ordering::Relaxed);
-        if is_obsolete(old) || is_locked(old) {
-            return Err(ConcurrentUpdateError());
-        }
-        if self
-            .inner
-            .compare_exchange(
-                old,
-                set_locked_bit(old),
-                Ordering::Acquire,
-                Ordering::Relaxed,
-            )
-            .is_err()
-        {
-            return Err(ConcurrentUpdateError());
-        }
-        Ok(())
-    }
-
-    pub(crate) fn write_unlock(&self) {
-        // reset locked bit and overflow into version
-        self.inner.fetch_add(2, Ordering::Release);
-    }
-
-    pub(crate) fn write_unlock_obsolete(&self) {
-        // set obsolete, reset locked, overflow into version
-        self.inner.fetch_add(3, Ordering::Release);
-    }
-
-    // Helper functions
-    fn await_node_unlocked(&self) -> u64 {
-        let mut version = self.inner.load(Ordering::Acquire);
-        while is_locked(version) {
-            // spinlock
-            std::thread::yield_now();
-            version = self.inner.load(Ordering::Acquire)
-        }
-        version
-    }
-}
-
-fn set_locked_bit(version: u64) -> u64 {
-    version + 2
-}
-
-fn is_obsolete(version: u64) -> bool {
-    (version & 1) == 1
-}
-
-fn is_locked(version: u64) -> bool {
-    (version & 2) == 2
-}

libs/neonart/src/algorithm/node_ref.rs

@@ -1,349 +0,0 @@
-use std::fmt::Debug;
-use std::marker::PhantomData;
-
-use super::node_ptr;
-use super::node_ptr::NodePtr;
-use crate::EpochPin;
-use crate::Value;
-use crate::algorithm::lock_and_version::AtomicLockAndVersion;
-use crate::algorithm::lock_and_version::ConcurrentUpdateError;
-use crate::allocator::ArtAllocator;
-use crate::allocator::OutOfMemoryError;
-
-pub struct NodeRef<'e, V> {
-    ptr: NodePtr<V>,
-
-    phantom: PhantomData<&'e EpochPin<'e>>,
-}
-
-impl<'e, V> Debug for NodeRef<'e, V> {
-    fn fmt(&self, fmt: &mut std::fmt::Formatter<'_>) -> Result<(), std::fmt::Error> {
-        write!(fmt, "{:?}", self.ptr)
-    }
-}
-
-impl<'e, V: Value> NodeRef<'e, V> {
-    pub(crate) fn from_root_ptr(root_ptr: NodePtr<V>) -> NodeRef<'e, V> {
-        NodeRef {
-            ptr: root_ptr,
-            phantom: PhantomData,
-        }
-    }
-
-    pub(crate) fn read_lock_or_restart(
-        &self,
-    ) -> Result<ReadLockedNodeRef<'e, V>, ConcurrentUpdateError> {
-        let version = self.lockword().read_lock_or_restart()?;
-        Ok(ReadLockedNodeRef {
-            ptr: self.ptr,
-            version,
-            phantom: self.phantom,
-        })
-    }
-
-    pub(crate) fn write_lock_or_restart(
-        &self,
-    ) -> Result<WriteLockedNodeRef<'e, V>, ConcurrentUpdateError> {
-        self.lockword().write_lock_or_restart()?;
-        Ok(WriteLockedNodeRef {
-            ptr: self.ptr,
-            phantom: self.phantom,
-        })
-    }
-
-    fn lockword(&self) -> &AtomicLockAndVersion {
-        self.ptr.lockword()
-    }
-}
-
-/// A reference to a node that has been optimistically read-locked. The functions re-check
-/// the version after each read.
-pub struct ReadLockedNodeRef<'e, V> {
-    ptr: NodePtr<V>,
-    version: u64,
-
-    phantom: PhantomData<&'e EpochPin<'e>>,
-}
-
-impl<'e, V: Value> ReadLockedNodeRef<'e, V> {
-    pub(crate) fn is_leaf(&self) -> bool {
-        self.ptr.is_leaf()
-    }
-
-    pub(crate) fn is_full(&self) -> bool {
-        self.ptr.is_full()
-    }
-
-    pub(crate) fn get_prefix(&self) -> &[u8] {
-        self.ptr.get_prefix()
-    }
-
-    /// Note: because we're only holding a read lock, the prefix can change concurrently.
-    /// You must be prepared to restart, if read_unlock() returns error later.
-    ///
-    /// Returns the length of the prefix, or None if it's not a match
-    pub(crate) fn prefix_matches(&self, key: &[u8]) -> Option<usize> {
-        self.ptr.prefix_matches(key)
-    }
-
-    pub(crate) fn find_child_or_restart(
-        &self,
-        key_byte: u8,
-    ) -> Result<Option<NodeRef<'e, V>>, ConcurrentUpdateError> {
-        let child_or_value = self.ptr.find_child(key_byte);
-        self.ptr.lockword().check_or_restart(self.version)?;
-
-        match child_or_value {
-            None => Ok(None),
-            Some(child_ptr) => Ok(Some(NodeRef {
-                ptr: child_ptr,
-                phantom: self.phantom,
-            })),
-        }
-    }
-
-    pub(crate) fn find_next_child_or_restart(
-        &self,
-        min_key_byte: u8,
-    ) -> Result<Option<(u8, NodeRef<'e, V>)>, ConcurrentUpdateError> {
-        let child_or_value = self.ptr.find_next_child(min_key_byte);
-        self.ptr.lockword().check_or_restart(self.version)?;
-
-        match child_or_value {
-            None => Ok(None),
-            Some((k, child_ptr)) => Ok(Some((
-                k,
-                NodeRef {
-                    ptr: child_ptr,
-                    phantom: self.phantom,
-                },
-            ))),
-        }
-    }
-
-    pub(crate) fn get_leaf_value_ptr(&self) -> Result<*const V, ConcurrentUpdateError> {
-        let result = self.ptr.get_leaf_value();
-        self.ptr.lockword().check_or_restart(self.version)?;
-
-        // Extend the lifetime.
-        let result = std::ptr::from_ref(result);
-
-        Ok(result)
-    }
-
-    pub(crate) fn upgrade_to_write_lock_or_restart(
-        self,
-    ) -> Result<WriteLockedNodeRef<'e, V>, ConcurrentUpdateError> {
-        self.ptr
-            .lockword()
-            .upgrade_to_write_lock_or_restart(self.version)?;
-
-        Ok(WriteLockedNodeRef {
-            ptr: self.ptr,
-            phantom: self.phantom,
-        })
-    }
-
-    pub(crate) fn read_unlock_or_restart(self) -> Result<(), ConcurrentUpdateError> {
-        self.ptr.lockword().check_or_restart(self.version)?;
-        Ok(())
-    }
-
-    pub(crate) fn check_or_restart(&self) -> Result<(), ConcurrentUpdateError> {
-        self.ptr.lockword().check_or_restart(self.version)?;
-        Ok(())
-    }
-}
-
-/// A reference to a node that has been optimistically read-locked. The functions re-check
-/// the version after each read.
-pub struct WriteLockedNodeRef<'e, V> {
-    ptr: NodePtr<V>,
-    phantom: PhantomData<&'e EpochPin<'e>>,
-}
-
-impl<'e, V: Value> WriteLockedNodeRef<'e, V> {
-    pub(crate) fn can_shrink(&self) -> bool {
-        self.ptr.can_shrink()
-    }
-
-    pub(crate) fn num_children(&self) -> usize {
-        self.ptr.num_children()
-    }
-
-    pub(crate) fn write_unlock(mut self) {
-        self.ptr.lockword().write_unlock();
-        self.ptr = NodePtr::null();
-    }
-
-    pub(crate) fn write_unlock_obsolete(mut self) {
-        self.ptr.lockword().write_unlock_obsolete();
-        self.ptr = NodePtr::null();
-    }
-
-    pub(crate) fn get_prefix(&self) -> &[u8] {
-        self.ptr.get_prefix()
-    }
-
-    pub(crate) fn truncate_prefix(&mut self, new_prefix_len: usize) {
-        self.ptr.truncate_prefix(new_prefix_len)
-    }
-
-    pub(crate) fn prepend_prefix(&mut self, prefix: &[u8], prefix_byte: u8) {
-        self.ptr.prepend_prefix(prefix, prefix_byte)
-    }
-
-    pub(crate) fn insert_child(&mut self, key_byte: u8, child: NodePtr<V>) {
-        self.ptr.insert_child(key_byte, child)
-    }
-
-    pub(crate) fn get_leaf_value_mut(&mut self) -> &mut V {
-        self.ptr.get_leaf_value_mut()
-    }
-
-    pub(crate) fn grow<'a, A>(
-        &self,
-        allocator: &'a A,
-    ) -> Result<NewNodeRef<'a, V, A>, OutOfMemoryError>
-    where
-        A: ArtAllocator<V>,
-    {
-        let new_node = self.ptr.grow(allocator)?;
-        Ok(NewNodeRef {
-            ptr: new_node,
-            allocator,
-            extra_nodes: Vec::new(),
-        })
-    }
-
-    pub(crate) fn shrink<'a, A>(
-        &self,
-        allocator: &'a A,
-    ) -> Result<NewNodeRef<'a, V, A>, OutOfMemoryError>
-    where
-        A: ArtAllocator<V>,
-    {
-        let new_node = self.ptr.shrink(allocator)?;
-        Ok(NewNodeRef {
-            ptr: new_node,
-            allocator,
-            extra_nodes: Vec::new(),
-        })
-    }
-
-    pub(crate) fn as_ptr(&self) -> NodePtr<V> {
-        self.ptr
-    }
-
-    pub(crate) fn replace_child(&mut self, key_byte: u8, replacement: NodePtr<V>) {
-        self.ptr.replace_child(key_byte, replacement);
-    }
-
-    pub(crate) fn delete_child(&mut self, key_byte: u8) {
-        self.ptr.delete_child(key_byte);
-    }
-
-    pub(crate) fn find_remaining_child(&self) -> (u8, NodeRef<'e, V>) {
-        assert_eq!(self.num_children(), 1);
-        let child_or_value = self.ptr.find_next_child(0);
-
-        match child_or_value {
-            None => panic!("could not find only child in node"),
-            Some((k, child_ptr)) => (
-                k,
-                NodeRef {
-                    ptr: child_ptr,
-                    phantom: self.phantom,
-                },
-            ),
-        }
-    }
-}
-
-impl<'e, V> Drop for WriteLockedNodeRef<'e, V> {
-    fn drop(&mut self) {
-        if !self.ptr.is_null() {
-            self.ptr.lockword().write_unlock();
-        }
-    }
-}
-
-pub(crate) struct NewNodeRef<'a, V, A>
-where
-    V: Value,
-    A: ArtAllocator<V>,
-{
-    ptr: NodePtr<V>,
-    allocator: &'a A,
-
-    extra_nodes: Vec<NodePtr<V>>,
-}
-
-impl<'a, V, A> NewNodeRef<'a, V, A>
-where
-    V: Value,
-    A: ArtAllocator<V>,
-{
-    pub(crate) fn insert_old_child(&mut self, key_byte: u8, child: &WriteLockedNodeRef<V>) {
-        self.ptr.insert_child(key_byte, child.as_ptr())
-    }
-
-    pub(crate) fn into_ptr(mut self) -> NodePtr<V> {
-        let ptr = self.ptr;
-        self.ptr = NodePtr::null();
-        ptr
-    }
-
-    pub(crate) fn insert_new_child(&mut self, key_byte: u8, child: NewNodeRef<'a, V, A>) {
-        let child_ptr = child.into_ptr();
-        self.ptr.insert_child(key_byte, child_ptr);
-        self.extra_nodes.push(child_ptr);
-    }
-}
-
-impl<'a, V, A> Drop for NewNodeRef<'a, V, A>
-where
-    V: Value,
-    A: ArtAllocator<V>,
-{
-    /// This drop implementation deallocates the newly allocated node, if into_ptr() was not called.
-    fn drop(&mut self) {
-        if !self.ptr.is_null() {
-            self.ptr.deallocate(self.allocator);
-            for p in self.extra_nodes.iter() {
-                p.deallocate(self.allocator);
-            }
-        }
-    }
-}
-
-pub(crate) fn new_internal<'a, V, A>(
-    prefix: &[u8],
-    allocator: &'a A,
-) -> Result<NewNodeRef<'a, V, A>, OutOfMemoryError>
-where
-    V: Value,
-    A: ArtAllocator<V>,
-{
-    Ok(NewNodeRef {
-        ptr: node_ptr::new_internal(prefix, allocator)?,
-        allocator,
-        extra_nodes: Vec::new(),
-    })
-}
-
-pub(crate) fn new_leaf<'a, V, A>(
-    prefix: &[u8],
-    value: V,
-    allocator: &'a A,
-) -> Result<NewNodeRef<'a, V, A>, OutOfMemoryError>
-where
-    V: Value,
-    A: ArtAllocator<V>,
-{
-    Ok(NewNodeRef {
-        ptr: node_ptr::new_leaf(prefix, value, allocator)?,
-        allocator,
-        extra_nodes: Vec::new(),
-    })
-}

libs/neonart/src/allocator.rs

@@ -1,156 +0,0 @@
-pub mod block;
-mod multislab;
-mod slab;
-pub mod r#static;
-
-use std::alloc::Layout;
-use std::marker::PhantomData;
-use std::mem::MaybeUninit;
-use std::sync::atomic::Ordering;
-
-use crate::allocator::multislab::MultiSlabAllocator;
-use crate::allocator::r#static::alloc_from_slice;
-
-use spin;
-
-use crate::Tree;
-pub use crate::algorithm::node_ptr::{
-    NodeInternal4, NodeInternal16, NodeInternal48, NodeInternal256, NodeLeaf,
-};
-
-#[derive(Debug)]
-pub struct OutOfMemoryError();
-
-pub trait ArtAllocator<V: crate::Value> {
-    fn alloc_tree(&self) -> *mut Tree<V>;
-
-    fn alloc_node_internal4(&self) -> *mut NodeInternal4<V>;
-    fn alloc_node_internal16(&self) -> *mut NodeInternal16<V>;
-    fn alloc_node_internal48(&self) -> *mut NodeInternal48<V>;
-    fn alloc_node_internal256(&self) -> *mut NodeInternal256<V>;
-    fn alloc_node_leaf(&self) -> *mut NodeLeaf<V>;
-
-    fn dealloc_node_internal4(&self, ptr: *mut NodeInternal4<V>);
-    fn dealloc_node_internal16(&self, ptr: *mut NodeInternal16<V>);
-    fn dealloc_node_internal48(&self, ptr: *mut NodeInternal48<V>);
-    fn dealloc_node_internal256(&self, ptr: *mut NodeInternal256<V>);
-    fn dealloc_node_leaf(&self, ptr: *mut NodeLeaf<V>);
-}
-
-pub struct ArtMultiSlabAllocator<'t, V>
-where
-    V: crate::Value,
-{
-    tree_area: spin::Mutex<Option<&'t mut MaybeUninit<Tree<V>>>>,
-
-    pub(crate) inner: MultiSlabAllocator<'t, 5>,
-
-    phantom_val: PhantomData<V>,
-}
-
-impl<'t, V: crate::Value> ArtMultiSlabAllocator<'t, V> {
-    const LAYOUTS: [Layout; 5] = [
-        Layout::new::<NodeInternal4<V>>(),
-        Layout::new::<NodeInternal16<V>>(),
-        Layout::new::<NodeInternal48<V>>(),
-        Layout::new::<NodeInternal256<V>>(),
-        Layout::new::<NodeLeaf<V>>(),
-    ];
-
-    pub fn new(area: &'t mut [MaybeUninit<u8>]) -> &'t mut ArtMultiSlabAllocator<'t, V> {
-        let (allocator_area, remain) = alloc_from_slice::<ArtMultiSlabAllocator<V>>(area);
-        let (tree_area, remain) = alloc_from_slice::<Tree<V>>(remain);
-
-        allocator_area.write(ArtMultiSlabAllocator {
-            tree_area: spin::Mutex::new(Some(tree_area)),
-            inner: MultiSlabAllocator::new(remain, &Self::LAYOUTS),
-            phantom_val: PhantomData,
-        })
-    }
-}
-
-impl<'t, V: crate::Value> ArtAllocator<V> for ArtMultiSlabAllocator<'t, V> {
-    fn alloc_tree(&self) -> *mut Tree<V> {
-        let mut t = self.tree_area.lock();
-        if let Some(tree_area) = t.take() {
-            return tree_area.as_mut_ptr().cast();
-        }
-        panic!("cannot allocate more than one tree");
-    }
-
-    fn alloc_node_internal4(&self) -> *mut NodeInternal4<V> {
-        self.inner.alloc_slab(0).cast()
-    }
-    fn alloc_node_internal16(&self) -> *mut NodeInternal16<V> {
-        self.inner.alloc_slab(1).cast()
-    }
-    fn alloc_node_internal48(&self) -> *mut NodeInternal48<V> {
-        self.inner.alloc_slab(2).cast()
-    }
-    fn alloc_node_internal256(&self) -> *mut NodeInternal256<V> {
-        self.inner.alloc_slab(3).cast()
-    }
-    fn alloc_node_leaf(&self) -> *mut NodeLeaf<V> {
-        self.inner.alloc_slab(4).cast()
-    }
-
-    fn dealloc_node_internal4(&self, ptr: *mut NodeInternal4<V>) {
-        self.inner.dealloc_slab(0, ptr.cast())
-    }
-
-    fn dealloc_node_internal16(&self, ptr: *mut NodeInternal16<V>) {
-        self.inner.dealloc_slab(1, ptr.cast())
-    }
-    fn dealloc_node_internal48(&self, ptr: *mut NodeInternal48<V>) {
-        self.inner.dealloc_slab(2, ptr.cast())
-    }
-    fn dealloc_node_internal256(&self, ptr: *mut NodeInternal256<V>) {
-        self.inner.dealloc_slab(3, ptr.cast())
-    }
-    fn dealloc_node_leaf(&self, ptr: *mut NodeLeaf<V>) {
-        self.inner.dealloc_slab(4, ptr.cast())
-    }
-}
-
-impl<'t, V: crate::Value> ArtMultiSlabAllocator<'t, V> {
-    pub(crate) fn get_statistics(&self) -> ArtMultiSlabStats {
-        ArtMultiSlabStats {
-            num_internal4: self.inner.slab_descs[0]
-                .num_allocated
-                .load(Ordering::Relaxed),
-            num_internal16: self.inner.slab_descs[1]
-                .num_allocated
-                .load(Ordering::Relaxed),
-            num_internal48: self.inner.slab_descs[2]
-                .num_allocated
-                .load(Ordering::Relaxed),
-            num_internal256: self.inner.slab_descs[3]
-                .num_allocated
-                .load(Ordering::Relaxed),
-            num_leaf: self.inner.slab_descs[4]
-                .num_allocated
-                .load(Ordering::Relaxed),
-
-            num_blocks_internal4: self.inner.slab_descs[0].num_blocks.load(Ordering::Relaxed),
-            num_blocks_internal16: self.inner.slab_descs[1].num_blocks.load(Ordering::Relaxed),
-            num_blocks_internal48: self.inner.slab_descs[2].num_blocks.load(Ordering::Relaxed),
-            num_blocks_internal256: self.inner.slab_descs[3].num_blocks.load(Ordering::Relaxed),
-            num_blocks_leaf: self.inner.slab_descs[4].num_blocks.load(Ordering::Relaxed),
-        }
-    }
-}
-
-#[derive(Clone, Debug)]
-pub struct ArtMultiSlabStats {
-    pub num_internal4: u64,
-    pub num_internal16: u64,
-    pub num_internal48: u64,
-    pub num_internal256: u64,
-    pub num_leaf: u64,
-
-    pub num_blocks_internal4: u64,
-    pub num_blocks_internal16: u64,
-    pub num_blocks_internal48: u64,
-    pub num_blocks_internal256: u64,
-    pub num_blocks_leaf: u64,
-}

libs/neonart/src/allocator/block.rs

@@ -1,191 +0,0 @@
-//! Simple allocator of fixed-size blocks
-
-use std::mem::MaybeUninit;
-use std::sync::atomic::{AtomicU64, Ordering};
-
-use spin;
-
-pub const BLOCK_SIZE: usize = 16 * 1024;
-
-const INVALID_BLOCK: u64 = u64::MAX;
-
-pub(crate) struct BlockAllocator<'t> {
-    blocks_ptr: &'t [MaybeUninit<u8>],
-    num_blocks: u64,
-    num_initialized: AtomicU64,
-
-    freelist_head: spin::Mutex<u64>,
-}
-
-struct FreeListBlock {
-    inner: spin::Mutex<FreeListBlockInner>,
-}
-
-struct FreeListBlockInner {
-    next: u64,
-
-    num_free_blocks: u64,
-    free_blocks: [u64; 100], // FIXME: fill the rest of the block
-}
-
-impl<'t> BlockAllocator<'t> {
-    pub(crate) fn new(area: &'t mut [MaybeUninit<u8>]) -> Self {
-        // Use all the space for the blocks
-        let padding = area.as_ptr().align_offset(BLOCK_SIZE);
-        let remain = &mut area[padding..];
-
-        let num_blocks = (remain.len() / BLOCK_SIZE) as u64;
-
-        BlockAllocator {
-            blocks_ptr: remain,
-            num_blocks,
-            num_initialized: AtomicU64::new(0),
-            freelist_head: spin::Mutex::new(INVALID_BLOCK),
-        }
-    }
-
-    /// safety: you must hold a lock on the pointer to this block, otherwise it might get
-    /// reused for another kind of block
-    fn read_freelist_block(&self, blkno: u64) -> &FreeListBlock {
-        let ptr: *const FreeListBlock = self.get_block_ptr(blkno).cast();
-        unsafe { ptr.as_ref().unwrap() }
-    }
-
-    fn get_block_ptr(&self, blkno: u64) -> *mut u8 {
-        assert!(blkno < self.num_blocks);
-        unsafe {
-            self.blocks_ptr
-                .as_ptr()
-                .byte_offset(blkno as isize * BLOCK_SIZE as isize)
-        }
-        .cast_mut()
-        .cast()
-    }
-
-    #[allow(clippy::mut_from_ref)]
-    pub(crate) fn alloc_block(&self) -> &mut [MaybeUninit<u8>] {
-        // FIXME: handle OOM
-        let blkno = self.alloc_block_internal();
-        if blkno == INVALID_BLOCK {
-            panic!("out of memory");
-        }
-
-        let ptr: *mut MaybeUninit<u8> = self.get_block_ptr(blkno).cast();
-        unsafe { std::slice::from_raw_parts_mut(ptr, BLOCK_SIZE) }
-    }
-
-    fn alloc_block_internal(&self) -> u64 {
-        //  check the free list.
-        {
-            let mut freelist_head = self.freelist_head.lock();
-            if *freelist_head != INVALID_BLOCK {
-                let freelist_block = self.read_freelist_block(*freelist_head);
-
-                // acquire lock on the freelist block before releasing the lock on the parent (i.e. lock coupling)
-                let mut g = freelist_block.inner.lock();
-
-                if g.num_free_blocks > 0 {
-                    g.num_free_blocks -= 1;
-                    let result = g.free_blocks[g.num_free_blocks as usize];
-                    return result;
-                } else {
-                    // consume the freelist block itself
-                    let result = *freelist_head;
-                    *freelist_head = g.next;
-                    // This freelist block is now unlinked and can be repurposed
-                    drop(g);
-                    return result;
-                }
-            }
-        }
-
-        // If there are some blocks left that we've never used, pick next such block
-        let mut next_uninitialized = self.num_initialized.load(Ordering::Relaxed);
-        while next_uninitialized < self.num_blocks {
-            match self.num_initialized.compare_exchange(
-                next_uninitialized,
-                next_uninitialized + 1,
-                Ordering::Relaxed,
-                Ordering::Relaxed,
-            ) {
-                Ok(_) => {
-                    return next_uninitialized;
-                }
-                Err(old) => {
-                    next_uninitialized = old;
-                    continue;
-                }
-            }
-        }
-
-        // out of blocks
-        INVALID_BLOCK
-    }
-
-    // TODO: this is currently unused. The slab allocator never releases blocks
-    #[allow(dead_code)]
-    pub(crate) fn release_block(&self, block_ptr: *mut u8) {
-        let blockno = unsafe { block_ptr.byte_offset_from(self.blocks_ptr) / BLOCK_SIZE as isize };
-        self.release_block_internal(blockno as u64);
-    }
-
-    fn release_block_internal(&self, blockno: u64) {
-        let mut freelist_head = self.freelist_head.lock();
-        if *freelist_head != INVALID_BLOCK {
-            let freelist_block = self.read_freelist_block(*freelist_head);
-
-            // acquire lock on the freelist block before releasing the lock on the parent (i.e. lock coupling)
-            let mut g = freelist_block.inner.lock();
-
-            let num_free_blocks = g.num_free_blocks;
-            if num_free_blocks < g.free_blocks.len() as u64 {
-                g.free_blocks[num_free_blocks as usize] = blockno;
-                g.num_free_blocks += 1;
-                return;
-            }
-        }
-
-        // Convert the block into a new freelist block
-        let block_ptr: *mut FreeListBlock = self.get_block_ptr(blockno).cast();
-        let init = FreeListBlock {
-            inner: spin::Mutex::new(FreeListBlockInner {
-                next: *freelist_head,
-                num_free_blocks: 0,
-                free_blocks: [INVALID_BLOCK; 100],
-            }),
-        };
-        unsafe { (*block_ptr) = init };
-        *freelist_head = blockno;
-    }
-
-    // for debugging
-    pub(crate) fn get_statistics(&self) -> BlockAllocatorStats {
-        let mut num_free_blocks = 0;
-
-        let mut _prev_lock = None;
-        let head_lock = self.freelist_head.lock();
-        let mut next_blk = *head_lock;
-        let mut _head_lock = Some(head_lock);
-        while next_blk != INVALID_BLOCK {
-            let freelist_block = self.read_freelist_block(next_blk);
-            let lock = freelist_block.inner.lock();
-            num_free_blocks += lock.num_free_blocks;
-            next_blk = lock.next;
-            _prev_lock = Some(lock); // hold the lock until we've read the next block
-            _head_lock = None;
-        }
-
-        BlockAllocatorStats {
-            num_blocks: self.num_blocks,
-            num_initialized: self.num_initialized.load(Ordering::Relaxed),
-            num_free_blocks,
-        }
-    }
-}
-
-#[derive(Clone, Debug)]
-pub struct BlockAllocatorStats {
-    pub num_blocks: u64,
-    pub num_initialized: u64,
-    pub num_free_blocks: u64,
-}

libs/neonart/src/allocator/multislab.rs

@@ -1,33 +0,0 @@
-use std::alloc::Layout;
-use std::mem::MaybeUninit;
-
-use crate::allocator::block::BlockAllocator;
-use crate::allocator::slab::SlabDesc;
-
-pub struct MultiSlabAllocator<'t, const N: usize> {
-    pub(crate) block_allocator: BlockAllocator<'t>,
-
-    pub(crate) slab_descs: [SlabDesc; N],
-}
-
-impl<'t, const N: usize> MultiSlabAllocator<'t, N> {
-    pub(crate) fn new(
-        area: &'t mut [MaybeUninit<u8>],
-        layouts: &[Layout; N],
-    ) -> MultiSlabAllocator<'t, N> {
-        let block_allocator = BlockAllocator::new(area);
-        MultiSlabAllocator {
-            block_allocator,
-
-            slab_descs: std::array::from_fn(|i| SlabDesc::new(&layouts[i])),
-        }
-    }
-
-    pub(crate) fn alloc_slab(&self, slab_idx: usize) -> *mut u8 {
-        self.slab_descs[slab_idx].alloc_chunk(&self.block_allocator)
-    }
-
-    pub(crate) fn dealloc_slab(&self, slab_idx: usize, ptr: *mut u8) {
-        self.slab_descs[slab_idx].dealloc_chunk(ptr, &self.block_allocator)
-    }
-}

libs/neonart/src/allocator/slab.rs

@@ -1,433 +0,0 @@
-//! A slab allocator that carves out fixed-size chunks from larger blocks.
-//!
-//!
-
-use std::alloc::Layout;
-use std::mem::MaybeUninit;
-use std::ops::Deref;
-use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
-
-use spin;
-
-use super::alloc_from_slice;
-use super::block::BlockAllocator;
-
-use crate::allocator::block::BLOCK_SIZE;
-
-pub(crate) struct SlabDesc {
-    pub(crate) layout: Layout,
-
-    block_lists: spin::RwLock<BlockLists>,
-
-    pub(crate) num_blocks: AtomicU64,
-    pub(crate) num_allocated: AtomicU64,
-}
-
-// FIXME: Not sure if SlabDesc is really Sync or Send. It probably is when it's empty, but
-// 'block_lists' contains pointers when it's not empty. In the current use as part of the
-// the art tree, SlabDescs are only moved during initialization.
-unsafe impl Sync for SlabDesc {}
-unsafe impl Send for SlabDesc {}
-
-#[derive(Default, Debug)]
-struct BlockLists {
-    full_blocks: BlockList,
-    nonfull_blocks: BlockList,
-}
-
-impl BlockLists {
-    // Unlink a node. It must be in either one of the two lists.
-    unsafe fn unlink(&mut self, elem: *mut SlabBlockHeader) {
-        let list = unsafe {
-            if (*elem).next.is_null() {
-                if self.full_blocks.tail == elem {
-                    Some(&mut self.full_blocks)
-                } else {
-                    Some(&mut self.nonfull_blocks)
-                }
-            } else if (*elem).prev.is_null() {
-                if self.full_blocks.head == elem {
-                    Some(&mut self.full_blocks)
-                } else {
-                    Some(&mut self.nonfull_blocks)
-                }
-            } else {
-                None
-            }
-        };
-        unsafe { unlink_slab_block(list, elem) };
-    }
-}
-
-unsafe fn unlink_slab_block(mut list: Option<&mut BlockList>, elem: *mut SlabBlockHeader) {
-    unsafe {
-        if (*elem).next.is_null() {
-            assert_eq!(list.as_ref().unwrap().tail, elem);
-            list.as_mut().unwrap().tail = (*elem).prev;
-        } else {
-            assert_eq!((*(*elem).next).prev, elem);
-            (*(*elem).next).prev = (*elem).prev;
-        }
-        if (*elem).prev.is_null() {
-            assert_eq!(list.as_ref().unwrap().head, elem);
-            list.as_mut().unwrap().head = (*elem).next;
-        } else {
-            assert_eq!((*(*elem).prev).next, elem);
-            (*(*elem).prev).next = (*elem).next;
-        }
-    }
-}
-
-#[derive(Debug)]
-struct BlockList {
-    head: *mut SlabBlockHeader,
-    tail: *mut SlabBlockHeader,
-}
-
-impl Default for BlockList {
-    fn default() -> Self {
-        BlockList {
-            head: std::ptr::null_mut(),
-            tail: std::ptr::null_mut(),
-        }
-    }
-}
-
-impl BlockList {
-    unsafe fn push_head(&mut self, elem: *mut SlabBlockHeader) {
-        unsafe {
-            if self.is_empty() {
-                self.tail = elem;
-                (*elem).next = std::ptr::null_mut();
-            } else {
-                (*elem).next = self.head;
-                (*self.head).prev = elem;
-            }
-            (*elem).prev = std::ptr::null_mut();
-            self.head = elem;
-        }
-    }
-
-    fn is_empty(&self) -> bool {
-        self.head.is_null()
-    }
-
-    unsafe fn unlink(&mut self, elem: *mut SlabBlockHeader) {
-        unsafe { unlink_slab_block(Some(self), elem) }
-    }
-
-    #[cfg(test)]
-    fn dump(&self) {
-        let mut next = self.head;
-
-        while !next.is_null() {
-            let n = unsafe { next.as_ref() }.unwrap();
-            eprintln!(
-                "  blk {:?} (free {}/{})",
-                next,
-                n.num_free_chunks.load(Ordering::Relaxed),
-                n.num_chunks
-            );
-            next = n.next;
-        }
-    }
-}
-
-impl SlabDesc {
-    pub(crate) fn new(layout: &Layout) -> SlabDesc {
-        SlabDesc {
-            layout: *layout,
-            block_lists: spin::RwLock::new(BlockLists::default()),
-            num_allocated: AtomicU64::new(0),
-            num_blocks: AtomicU64::new(0),
-        }
-    }
-}
-
-#[derive(Debug)]
-struct SlabBlockHeader {
-    free_chunks_head: spin::Mutex<*mut FreeChunk>,
-    num_free_chunks: AtomicU32,
-    num_chunks: u32, // this is really a constant for a given Layout
-
-    // these fields are protected by the lock on the BlockLists
-    prev: *mut SlabBlockHeader,
-    next: *mut SlabBlockHeader,
-}
-
-struct FreeChunk {
-    next: *mut FreeChunk,
-}
-
-enum ReadOrWriteGuard<'a, T> {
-    Read(spin::RwLockReadGuard<'a, T>),
-    Write(spin::RwLockWriteGuard<'a, T>),
-}
-
-impl<'a, T> Deref for ReadOrWriteGuard<'a, T> {
-    type Target = T;
-
-    fn deref(&self) -> &<Self as Deref>::Target {
-        match self {
-            ReadOrWriteGuard::Read(g) => g.deref(),
-            ReadOrWriteGuard::Write(g) => g.deref(),
-        }
-    }
-}
-
-impl SlabDesc {
-    pub fn alloc_chunk(&self, block_allocator: &BlockAllocator) -> *mut u8 {
-        // Are there any free chunks?
-        let mut acquire_write = false;
-        'outer: loop {
-            let mut block_lists_guard = if acquire_write {
-                ReadOrWriteGuard::Write(self.block_lists.write())
-            } else {
-                ReadOrWriteGuard::Read(self.block_lists.read())
-            };
-            'inner: loop {
-                let block_ptr = block_lists_guard.nonfull_blocks.head;
-                if block_ptr.is_null() {
-                    break 'outer;
-                }
-                unsafe {
-                    let mut free_chunks_head = (*block_ptr).free_chunks_head.lock();
-                    if !(*free_chunks_head).is_null() {
-                        let result = *free_chunks_head;
-                        (*free_chunks_head) = (*result).next;
-                        let _old = (*block_ptr).num_free_chunks.fetch_sub(1, Ordering::Relaxed);
-
-                        self.num_allocated.fetch_add(1, Ordering::Relaxed);
-                        return result.cast();
-                    }
-                }
-
-                // The block at the head of the list was full. Grab write lock and retry
-                match block_lists_guard {
-                    ReadOrWriteGuard::Read(_) => {
-                        acquire_write = true;
-                        continue 'outer;
-                    }
-                    ReadOrWriteGuard::Write(ref mut g) => {
-                        // move the node to the list of full blocks
-                        unsafe {
-                            g.nonfull_blocks.unlink(block_ptr);
-                            g.full_blocks.push_head(block_ptr);
-                        };
-                        continue 'inner;
-                    }
-                }
-            }
-        }
-
-        // no free chunks. Allocate a new block (and the chunk from that)
-        let (new_block, new_chunk) = self.alloc_block_and_chunk(block_allocator);
-        self.num_blocks.fetch_add(1, Ordering::Relaxed);
-
-        // Add the block to the list in the SlabDesc
-        unsafe {
-            let mut block_lists_guard = self.block_lists.write();
-            block_lists_guard.nonfull_blocks.push_head(new_block);
-        }
-        self.num_allocated.fetch_add(1, Ordering::Relaxed);
-        new_chunk
-    }
-
-    pub fn dealloc_chunk(&self, chunk_ptr: *mut u8, _block_allocator: &BlockAllocator) {
-        // Find the block it belongs to. You can find the block from the address. (And knowing the
-        // layout, you could calculate the chunk number too.)
-        let block_ptr: *mut SlabBlockHeader = {
-            let block_addr = (chunk_ptr.addr() / BLOCK_SIZE) * BLOCK_SIZE;
-            chunk_ptr.with_addr(block_addr).cast()
-        };
-        let chunk_ptr: *mut FreeChunk = chunk_ptr.cast();
-
-        // Mark the chunk as free in 'freechunks' list
-        let num_chunks;
-        let num_free_chunks;
-        unsafe {
-            let mut free_chunks_head = (*block_ptr).free_chunks_head.lock();
-            (*chunk_ptr).next = *free_chunks_head;
-            *free_chunks_head = chunk_ptr;
-
-            num_free_chunks = (*block_ptr).num_free_chunks.fetch_add(1, Ordering::Relaxed) + 1;
-            num_chunks = (*block_ptr).num_chunks;
-        }
-
-        if num_free_chunks == 1 {
-            // If the block was full previously, add it to the nonfull blocks list. Note that
-            // we're not holding the lock anymore, so it can immediately become full again.
-            // That's harmless, it will be moved back to the full list again when a call
-            // to alloc_chunk() sees it.
-            let mut block_lists = self.block_lists.write();
-            unsafe {
-                block_lists.unlink(block_ptr);
-                block_lists.nonfull_blocks.push_head(block_ptr);
-            };
-        } else if num_free_chunks == num_chunks {
-            // If the block became completely empty, move it to the free list
-            // TODO
-            // FIXME: we're still holding the spinlock. It's not exactly safe to return it to
-            // the free blocks list, is it? Defer it as garbage to wait out concurrent updates?
-            //block_allocator.release_block()
-        }
-
-        // update stats
-        self.num_allocated.fetch_sub(1, Ordering::Relaxed);
-    }
-
-    fn alloc_block_and_chunk(
-        &self,
-        block_allocator: &BlockAllocator,
-    ) -> (*mut SlabBlockHeader, *mut u8) {
-        // fixme: handle OOM
-        let block_slice: &mut [MaybeUninit<u8>] = block_allocator.alloc_block();
-        let (block_header, remain) = alloc_from_slice::<SlabBlockHeader>(block_slice);
-
-        let padding = remain.as_ptr().align_offset(self.layout.align());
-
-        let num_chunks = (remain.len() - padding) / self.layout.size();
-
-        let first_chunk_ptr: *mut FreeChunk = remain[padding..].as_mut_ptr().cast();
-
-        unsafe {
-            let mut chunk_ptr = first_chunk_ptr;
-            for _ in 0..num_chunks - 1 {
-                let next_chunk_ptr = chunk_ptr.byte_add(self.layout.size());
-                (*chunk_ptr).next = next_chunk_ptr;
-                chunk_ptr = next_chunk_ptr;
-            }
-            (*chunk_ptr).next = std::ptr::null_mut();
-
-            let result_chunk = first_chunk_ptr;
-
-            let block_header = block_header.write(SlabBlockHeader {
-                free_chunks_head: spin::Mutex::new((*first_chunk_ptr).next),
-                prev: std::ptr::null_mut(),
-                next: std::ptr::null_mut(),
-                num_chunks: num_chunks as u32,
-                num_free_chunks: AtomicU32::new(num_chunks as u32 - 1),
-            });
-
-            (block_header, result_chunk.cast())
-        }
-    }
-
-    #[cfg(test)]
-    fn dump(&self) {
-        eprintln!(
-            "slab dump ({} blocks, {} allocated chunks)",
-            self.num_blocks.load(Ordering::Relaxed),
-            self.num_allocated.load(Ordering::Relaxed)
-        );
-        let lists = self.block_lists.read();
-
-        eprintln!("nonfull blocks:");
-        lists.nonfull_blocks.dump();
-        eprintln!("full blocks:");
-        lists.full_blocks.dump();
-    }
-}
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-
-    use rand::Rng;
-    use rand_distr::Zipf;
-
-    struct TestObject {
-        val: usize,
-        _dummy: [u8; BLOCK_SIZE / 4],
-    }
-
-    struct TestObjectSlab<'a>(SlabDesc, BlockAllocator<'a>);
-    impl<'a> TestObjectSlab<'a> {
-        fn new(block_allocator: BlockAllocator) -> TestObjectSlab {
-            TestObjectSlab(SlabDesc::new(&Layout::new::<TestObject>()), block_allocator)
-        }
-
-        fn alloc(&self, val: usize) -> *mut TestObject {
-            let obj: *mut TestObject = self.0.alloc_chunk(&self.1).cast();
-            unsafe { (*obj).val = val };
-            obj
-        }
-
-        fn dealloc(&self, obj: *mut TestObject) {
-            self.0.dealloc_chunk(obj.cast(), &self.1)
-        }
-    }
-
-    #[test]
-    fn test_slab_alloc() {
-        const MEM_SIZE: usize = 100000000;
-        let mut area = Box::new_uninit_slice(MEM_SIZE);
-        let block_allocator = BlockAllocator::new(&mut area);
-
-        let slab = TestObjectSlab::new(block_allocator);
-
-        let mut all: Vec<*mut TestObject> = Vec::new();
-        for i in 0..11 {
-            all.push(slab.alloc(i));
-        }
-        #[allow(clippy::needless_range_loop)]
-        for i in 0..11 {
-            assert!(unsafe { (*all[i]).val == i });
-        }
-
-        let distribution = Zipf::new(10.0, 1.1).unwrap();
-        let mut rng = rand::rng();
-        for _ in 0..100000 {
-            slab.0.dump();
-            let idx = rng.sample(distribution) as usize;
-            let ptr: *mut TestObject = all[idx];
-            if !ptr.is_null() {
-                assert_eq!(unsafe { (*ptr).val }, idx);
-                slab.dealloc(ptr);
-                all[idx] = std::ptr::null_mut();
-            } else {
-                all[idx] = slab.alloc(idx);
-            }
-        }
-    }
-
-    fn new_test_blk(i: u32) -> *mut SlabBlockHeader {
-        Box::into_raw(Box::new(SlabBlockHeader {
-            free_chunks_head: spin::Mutex::new(std::ptr::null_mut()),
-            num_free_chunks: AtomicU32::new(0),
-            num_chunks: i,
-            prev: std::ptr::null_mut(),
-            next: std::ptr::null_mut(),
-        }))
-    }
-
-    #[test]
-    fn test_block_linked_list() {
-        // note: these are leaked, but that's OK for tests
-        let a = new_test_blk(0);
-        let b = new_test_blk(1);
-
-        let mut list = BlockList::default();
-        assert!(list.is_empty());
-
-        unsafe {
-            list.push_head(a);
-            assert!(!list.is_empty());
-            list.unlink(a);
-        }
-        assert!(list.is_empty());
-
-        unsafe {
-            list.push_head(b);
-            list.push_head(a);
-            assert_eq!(list.head, a);
-            assert_eq!((*a).next, b);
-            assert_eq!((*b).prev, a);
-            assert_eq!(list.tail, b);
-
-            list.unlink(a);
-            list.unlink(b);
-            assert!(list.is_empty());
-        }
-    }
-}

libs/neonart/src/allocator/static.rs

@@ -1,44 +0,0 @@
-use std::mem::MaybeUninit;
-
-pub fn alloc_from_slice<T>(
-    area: &mut [MaybeUninit<u8>],
-) -> (&mut MaybeUninit<T>, &mut [MaybeUninit<u8>]) {
-    let layout = std::alloc::Layout::new::<T>();
-
-    let area_start = area.as_mut_ptr();
-
-    // pad to satisfy alignment requirements
-    let padding = area_start.align_offset(layout.align());
-    if padding + layout.size() > area.len() {
-        panic!("out of memory");
-    }
-    let area = &mut area[padding..];
-    let (result_area, remain) = area.split_at_mut(layout.size());
-
-    let result_ptr: *mut MaybeUninit<T> = result_area.as_mut_ptr().cast();
-    let result = unsafe { result_ptr.as_mut().unwrap() };
-
-    (result, remain)
-}
-
-pub fn alloc_array_from_slice<T>(
-    area: &mut [MaybeUninit<u8>],
-    len: usize,
-) -> (&mut [MaybeUninit<T>], &mut [MaybeUninit<u8>]) {
-    let layout = std::alloc::Layout::new::<T>();
-
-    let area_start = area.as_mut_ptr();
-
-    // pad to satisfy alignment requirements
-    let padding = area_start.align_offset(layout.align());
-    if padding + layout.size() * len > area.len() {
-        panic!("out of memory");
-    }
-    let area = &mut area[padding..];
-    let (result_area, remain) = area.split_at_mut(layout.size() * len);
-
-    let result_ptr: *mut MaybeUninit<T> = result_area.as_mut_ptr().cast();
-    let result = unsafe { std::slice::from_raw_parts_mut(result_ptr.as_mut().unwrap(), len) };
-
-    (result, remain)
-}

libs/neonart/src/epoch.rs

@@ -1,142 +0,0 @@
-//! This is similar to crossbeam_epoch crate, but works in shared memory
-
-use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
-
-use crossbeam_utils::CachePadded;
-
-const NUM_SLOTS: usize = 1000;
-
-/// This is the struct that is stored in shmem
-///
-/// bit 0: is it pinned or not?
-/// rest of the bits are the epoch counter.
-pub struct EpochShared {
-    global_epoch: AtomicU64,
-    participants: [CachePadded<AtomicU64>; NUM_SLOTS],
-
-    broadcast_lock: spin::Mutex<()>,
-}
-
-impl EpochShared {
-    pub fn new() -> EpochShared {
-        EpochShared {
-            global_epoch: AtomicU64::new(2),
-            participants: [const { CachePadded::new(AtomicU64::new(2)) }; NUM_SLOTS],
-            broadcast_lock: spin::Mutex::new(()),
-        }
-    }
-
-    pub fn register(&self) -> LocalHandle {
-        LocalHandle {
-            global: self,
-            last_slot: AtomicUsize::new(0), // todo: choose more intelligently
-        }
-    }
-
-    fn release_pin(&self, slot: usize, _epoch: u64) {
-        let global_epoch = self.global_epoch.load(Ordering::Relaxed);
-        self.participants[slot].store(global_epoch, Ordering::Relaxed);
-    }
-
-    fn pin_internal(&self, slot_hint: usize) -> (usize, u64) {
-        // pick a slot
-        let mut slot = slot_hint;
-        let epoch = loop {
-            let old = self.participants[slot].fetch_or(1, Ordering::Relaxed);
-            if old & 1 == 0 {
-                // Got this slot
-                break old;
-            }
-
-            // the slot was busy by another thread / process. try a different slot
-            slot += 1;
-            if slot == NUM_SLOTS {
-                slot = 0;
-            }
-            continue;
-        };
-        (slot, epoch)
-    }
-
-    pub(crate) fn advance(&self) -> u64 {
-        // Advance the global epoch
-        let old_epoch = self.global_epoch.fetch_add(2, Ordering::Relaxed);
-        // Anyone that release their pin after this will update their slot.
-        old_epoch + 2
-    }
-
-    pub(crate) fn broadcast(&self) {
-        let Some(_guard) = self.broadcast_lock.try_lock() else {
-            return;
-        };
-
-        let epoch = self.global_epoch.load(Ordering::Relaxed);
-        let old_epoch = epoch.wrapping_sub(2);
-
-        // Update all free slots.
-        for i in 0..NUM_SLOTS {
-            // TODO: check result, as a sanity check. It should either be the old epoch, or pinned
-            let _ = self.participants[i].compare_exchange(
-                old_epoch,
-                epoch,
-                Ordering::Relaxed,
-                Ordering::Relaxed,
-            );
-        }
-
-        // FIXME: memory fence here, since we used Relaxed?
-    }
-
-    pub(crate) fn get_oldest(&self) -> u64 {
-        // Read all slots.
-        let now = self.global_epoch.load(Ordering::Relaxed);
-        let mut oldest = now;
-        for i in 0..NUM_SLOTS {
-            let this_epoch = self.participants[i].load(Ordering::Relaxed);
-            let delta = now.wrapping_sub(this_epoch);
-            if delta > u64::MAX / 2 {
-                // this is very recent
-            } else if delta > now.wrapping_sub(oldest) {
-                oldest = this_epoch;
-            }
-        }
-        oldest
-    }
-
-    pub(crate) fn get_current(&self) -> u64 {
-        self.global_epoch.load(Ordering::Relaxed)
-    }
-}
-
-pub(crate) struct EpochPin<'e> {
-    slot: usize,
-    pub(crate) epoch: u64,
-
-    handle: &'e LocalHandle<'e>,
-}
-
-impl<'e> Drop for EpochPin<'e> {
-    fn drop(&mut self) {
-        self.handle.global.release_pin(self.slot, self.epoch);
-    }
-}
-
-pub struct LocalHandle<'g> {
-    global: &'g EpochShared,
-
-    last_slot: AtomicUsize,
-}
-
-impl<'g> LocalHandle<'g> {
-    pub fn pin(&self) -> EpochPin {
-        let (slot, epoch) = self
-            .global
-            .pin_internal(self.last_slot.load(Ordering::Relaxed));
-        self.last_slot.store(slot, Ordering::Relaxed);
-        EpochPin {
-            handle: self,
-            epoch,
-            slot,
-        }
-    }
-}

libs/neonart/src/lib.rs

@@ -1,583 +0,0 @@
-//! Adaptive Radix Tree (ART) implementation, with Optimistic Lock Coupling.
-//!
-//! The data structure is described in these two papers:
-//!
-//! [1] Leis, V. & Kemper, Alfons & Neumann, Thomas. (2013).
-//!     The adaptive radix tree: ARTful indexing for main-memory databases.
-//!     Proceedings - International Conference on Data Engineering. 38-49. 10.1109/ICDE.2013.6544812.
-//!     https://db.in.tum.de/~leis/papers/ART.pdf
-//!
-//! [2] Leis, Viktor & Scheibner, Florian & Kemper, Alfons & Neumann, Thomas. (2016).
-//!     The ART of practical synchronization.
-//!     1-8. 10.1145/2933349.2933352.
-//!     https://db.in.tum.de/~leis/papers/artsync.pdf
-//!
-//! [1] describes the base data structure, and [2] describes the Optimistic Lock Coupling that we
-//! use.
-//!
-//! The papers mention a few different variants. We have made the following choices in this
-//! implementation:
-//!
-//! - All keys have the same length
-//!
-//! - Single-value leaves.
-//!
-//! - For collapsing inner nodes, we use the Pessimistic approach, where each inner node stores a
-//!   variable length "prefix", which stores the keys of all the one-way nodes which have been
-//!   removed. However, similar to the "hybrid" approach described in the paper, each node only has
-//!   space for a constant-size prefix of 8 bytes. If a node would have a longer prefix, then we
-//!   create create one-way nodes to store them. (There was no particular reason for this choice,
-//!   the "hybrid" approach described in the paper might be better.)
-//!
-//! - For concurrency, we use Optimistic Lock Coupling. The paper [2] also describes another method,
-//!   ROWEX, which generally performs better when there is contention, but that is not important
-//!   for use and Optimisic Lock Coupling is simpler to implement.
-//!
-//! ## Requirements
-//!
-//! This data structure is currently used for the integrated LFC, relsize and last-written LSN cache
-//! in the compute communicator, part of the 'neon' Postgres extension. We have some unique
-//! requirements, which is why we had to write our own. Namely:
-//!
-//! - The data structure has to live in fixed-sized shared memory segment. That rules out any
-//!   built-in Rust collections and most crates. (Except possibly with the 'allocator_api' rust
-//!   feature, which still nightly-only experimental as of this writing).
-//!
-//! - The data structure is accessed from multiple processes. Only one process updates the data
-//!   structure, but other processes perform reads. That rules out using built-in Rust locking
-//!   primitives like Mutex and RwLock, and most crates too.
-//!
-//! - Within the one process with write-access, multiple threads can perform updates concurrently.
-//!   That rules out using PostgreSQL LWLocks for the locking.
-//!
-//! The implementation is generic, and doesn't depend on any PostgreSQL specifics, but it has been
-//! written with that usage and the above constraints in mind. Some noteworthy assumptions:
-//!
-//! - Contention is assumed to be rare. In the integrated cache in PostgreSQL, there's higher level
-//!   locking in the PostgreSQL buffer manager, which ensures that two backends should not try to
-//!   read / write the same page at the same time. (Prefetching can conflict with actual reads,
-//!   however.)
-//!
-//!  - The keys in the integrated cache are 17 bytes long.
-//!
-//! ## Usage
-//!
-//! Because this is designed to be used as a Postgres shared memory data structure, initialization
-//! happens in three stages:
-//!
-//! 0. A fixed area of shared memory is allocated at postmaster startup.
-//!
-//! 1. TreeInitStruct::new() is called to initialize it, still in Postmaster process, before any
-//!    other process or thread is running. It returns a TreeInitStruct, which is inherited by all
-//!    the processes through fork().
-//!
-//! 2. One process may have write-access to the struct, by calling
-//!    [TreeInitStruct::attach_writer]. (That process is the communicator process.)
-//!
-//! 3. Other processes get read-access to the struct, by calling [TreeInitStruct::attach_reader]
-//!
-//! "Write access" means that you can insert / update / delete values in the tree.
-//!
-//! NOTE: The Values stored in the tree are sometimes moved, when a leaf node fills up and a new
-//! larger node needs to be allocated. The versioning and epoch-based allocator ensure that the data
-//! structure stays consistent, but if the Value has interior mutability, like atomic fields,
-//! updates to such fields might be lost if the leaf node is concurrently moved! If that becomes a
-//! problem, the version check could be passed up to the caller, so that the caller could detect the
-//! lost updates and retry the operation.
-//!
-//! ## Implementation
-//!
-//! node_ptr: Provides low-level implementations of the four different node types (eight actually,
-//! since there is an Internal and Leaf variant of each)
-//!
-//! lock_and_version.rs: Provides an abstraction for the combined lock and version counter on each
-//! node.
-//!
-//! node_ref.rs: The code in node_ptr.rs deals with raw pointers. node_ref.rs provides more type-safe
-//!   abstractions on top.
-//!
-//! algorithm.rs: Contains the functions to implement lookups and updates in the tree
-//!
-//! allocator.rs: Provides a facility to allocate memory for the tree nodes. (We must provide our
-//!   own abstraction for that because we need the data structure to live in a pre-allocated shared
-//!   memory segment).
-//!
-//! epoch.rs: The data structure requires that when a node is removed from the tree, it is not
-//!   immediately deallocated, but stays around for as long as concurrent readers might still have
-//!   pointers to them. This is enforced by an epoch system. This is similar to
-//!   e.g. crossbeam_epoch, but we couldn't use that either because it has to work across processes
-//!   communicating over the shared memory segment.
-//!
-//! ## See also
-//!
-//! There are some existing Rust ART implementations out there, but none of them filled all
-//! the requirements:
-//!
-//! - https://github.com/XiangpengHao/congee
-//! - https://github.com/declanvk/blart
-//!
-//! ## TODO
-//!
-//! - Removing values has not been implemented
-
-mod algorithm;
-pub mod allocator;
-mod epoch;
-
-use algorithm::RootPtr;
-use algorithm::node_ptr::NodePtr;
-
-use std::collections::VecDeque;
-use std::fmt::Debug;
-use std::marker::PhantomData;
-use std::ptr::NonNull;
-use std::sync::atomic::{AtomicBool, Ordering};
-
-use crate::epoch::EpochPin;
-
-#[cfg(test)]
-mod tests;
-
-use allocator::ArtAllocator;
-pub use allocator::ArtMultiSlabAllocator;
-pub use allocator::OutOfMemoryError;
-
-/// Fixed-length key type.
-///
-pub trait Key: Debug {
-    const KEY_LEN: usize;
-
-    fn as_bytes(&self) -> &[u8];
-}
-
-/// Values stored in the tree
-///
-/// Values need to be Cloneable, because when a node "grows", the value is copied to a new node and
-/// the old sticks around until all readers that might see the old value are gone.
-// fixme obsolete, no longer needs Clone
-pub trait Value {}
-
-const MAX_GARBAGE: usize = 1024;
-
-/// The root of the tree, plus other tree-wide data. This is stored in the shared memory.
-pub struct Tree<V: Value> {
-    /// For simplicity, so that we never need to grow or shrink the root, the root node is always an
-    /// Internal256 node. Also, it never has a prefix (that's actually a bit wasteful, incurring one
-    /// indirection to every lookup)
-    root: RootPtr<V>,
-
-    writer_attached: AtomicBool,
-
-    epoch: epoch::EpochShared,
-}
-
-unsafe impl<V: Value + Sync> Sync for Tree<V> {}
-unsafe impl<V: Value + Send> Send for Tree<V> {}
-
-struct GarbageQueue<V>(VecDeque<(NodePtr<V>, u64)>);
-
-unsafe impl<V: Value + Sync> Sync for GarbageQueue<V> {}
-unsafe impl<V: Value + Send> Send for GarbageQueue<V> {}
-
-impl<V> GarbageQueue<V> {
-    fn new() -> GarbageQueue<V> {
-        GarbageQueue(VecDeque::with_capacity(MAX_GARBAGE))
-    }
-
-    fn remember_obsolete_node(&mut self, ptr: NodePtr<V>, epoch: u64) {
-        self.0.push_front((ptr, epoch));
-    }
-
-    fn next_obsolete(&mut self, cutoff_epoch: u64) -> Option<NodePtr<V>> {
-        if let Some(back) = self.0.back() {
-            if back.1 < cutoff_epoch {
-                return Some(self.0.pop_back().unwrap().0);
-            }
-        }
-        None
-    }
-}
-
-/// Struct created at postmaster startup
-pub struct TreeInitStruct<'t, K: Key, V: Value, A: ArtAllocator<V>> {
-    tree: &'t Tree<V>,
-
-    allocator: &'t A,
-
-    phantom_key: PhantomData<K>,
-}
-
-/// The worker process has a reference to this. The write operations are only safe
-/// from the worker process
-pub struct TreeWriteAccess<'t, K: Key, V: Value, A: ArtAllocator<V>>
-where
-    K: Key,
-    V: Value,
-{
-    tree: &'t Tree<V>,
-
-    pub allocator: &'t A,
-
-    epoch_handle: epoch::LocalHandle<'t>,
-
-    phantom_key: PhantomData<K>,
-
-    /// Obsolete nodes that cannot be recycled until their epoch expires.
-    garbage: spin::Mutex<GarbageQueue<V>>,
-}
-
-/// The backends have a reference to this. It cannot be used to modify the tree
-pub struct TreeReadAccess<'t, K: Key, V: Value>
-where
-    K: Key,
-    V: Value,
-{
-    tree: &'t Tree<V>,
-
-    epoch_handle: epoch::LocalHandle<'t>,
-
-    phantom_key: PhantomData<K>,
-}
-
-impl<'t, K: Key, V: Value, A: ArtAllocator<V>> TreeInitStruct<'t, K, V, A> {
-    pub fn new(allocator: &'t A) -> TreeInitStruct<'t, K, V, A> {
-        let tree_ptr = allocator.alloc_tree();
-        let tree_ptr = NonNull::new(tree_ptr).expect("out of memory");
-        let init = Tree {
-            root: algorithm::new_root(allocator).expect("out of memory"),
-            writer_attached: AtomicBool::new(false),
-            epoch: epoch::EpochShared::new(),
-        };
-        unsafe { tree_ptr.write(init) };
-
-        TreeInitStruct {
-            tree: unsafe { tree_ptr.as_ref() },
-            allocator,
-            phantom_key: PhantomData,
-        }
-    }
-
-    pub fn attach_writer(self) -> TreeWriteAccess<'t, K, V, A> {
-        let previously_attached = self.tree.writer_attached.swap(true, Ordering::Relaxed);
-        if previously_attached {
-            panic!("writer already attached");
-        }
-        TreeWriteAccess {
-            tree: self.tree,
-            allocator: self.allocator,
-            phantom_key: PhantomData,
-            epoch_handle: self.tree.epoch.register(),
-            garbage: spin::Mutex::new(GarbageQueue::new()),
-        }
-    }
-
-    pub fn attach_reader(self) -> TreeReadAccess<'t, K, V> {
-        TreeReadAccess {
-            tree: self.tree,
-            phantom_key: PhantomData,
-            epoch_handle: self.tree.epoch.register(),
-        }
-    }
-}
-
-impl<'t, K: Key, V: Value, A: ArtAllocator<V>> TreeWriteAccess<'t, K, V, A> {
-    pub fn start_write<'g>(&'t self) -> TreeWriteGuard<'g, K, V, A>
-    where
-        't: 'g,
-    {
-        TreeWriteGuard {
-            tree_writer: self,
-            epoch_pin: self.epoch_handle.pin(),
-            phantom_key: PhantomData,
-            created_garbage: false,
-        }
-    }
-
-    pub fn start_read(&'t self) -> TreeReadGuard<'t, K, V> {
-        TreeReadGuard {
-            tree: self.tree,
-            epoch_pin: self.epoch_handle.pin(),
-            phantom_key: PhantomData,
-        }
-    }
-}
-
-impl<'t, K: Key, V: Value> TreeReadAccess<'t, K, V> {
-    pub fn start_read(&'t self) -> TreeReadGuard<'t, K, V> {
-        TreeReadGuard {
-            tree: self.tree,
-            epoch_pin: self.epoch_handle.pin(),
-            phantom_key: PhantomData,
-        }
-    }
-}
-
-pub struct TreeReadGuard<'e, K, V>
-where
-    K: Key,
-    V: Value,
-{
-    tree: &'e Tree<V>,
-
-    epoch_pin: EpochPin<'e>,
-    phantom_key: PhantomData<K>,
-}
-
-impl<'e, K: Key, V: Value> TreeReadGuard<'e, K, V> {
-    pub fn get(&'e self, key: &K) -> Option<&'e V> {
-        algorithm::search(key, self.tree.root, &self.epoch_pin)
-    }
-}
-
-pub struct TreeWriteGuard<'e, K, V, A>
-where
-    K: Key,
-    V: Value,
-    A: ArtAllocator<V>,
-{
-    tree_writer: &'e TreeWriteAccess<'e, K, V, A>,
-
-    epoch_pin: EpochPin<'e>,
-    phantom_key: PhantomData<K>,
-
-    created_garbage: bool,
-}
-
-pub enum UpdateAction<V> {
-    Nothing,
-    Insert(V),
-    Remove,
-}
-
-impl<'e, K: Key, V: Value, A: ArtAllocator<V>> TreeWriteGuard<'e, K, V, A> {
-    /// Get a value
-    pub fn get(&'e mut self, key: &K) -> Option<&'e V> {
-        algorithm::search(key, self.tree_writer.tree.root, &self.epoch_pin)
-    }
-
-    /// Insert a value
-    pub fn insert(self, key: &K, value: V) -> Result<bool, OutOfMemoryError> {
-        let mut success = None;
-
-        self.update_with_fn(key, |existing| {
-            if existing.is_some() {
-                success = Some(false);
-                UpdateAction::Nothing
-            } else {
-                success = Some(true);
-                UpdateAction::Insert(value)
-            }
-        })?;
-        Ok(success.expect("value_fn not called"))
-    }
-
-    /// Remove value. Returns true if it existed
-    pub fn remove(self, key: &K) -> bool {
-        let mut result = false;
-        // FIXME: It's not clear if OOM is expected while removing. It seems
-        // not nice, but shrinking a node can OOM. Then again, we could opt
-        // to not shrink a node if we cannot allocate, to live a little longer.
-        self.update_with_fn(key, |existing| match existing {
-            Some(_) => {
-                result = true;
-                UpdateAction::Remove
-            }
-            None => UpdateAction::Nothing,
-        })
-        .expect("out of memory while removing");
-        result
-    }
-
-    /// Try to remove value and return the old value.
-    pub fn remove_and_return(self, key: &K) -> Option<V>
-    where
-        V: Clone,
-    {
-        let mut old = None;
-        self.update_with_fn(key, |existing| {
-            old = existing.cloned();
-            UpdateAction::Remove
-        })
-        .expect("out of memory while removing");
-        old
-    }
-
-    /// Update key using the given function. All the other modifying operations are based on this.
-    ///
-    /// The function is passed a reference to the existing value, if any. If the function
-    /// returns None, the value is removed from the tree (or if there was no existing value,
-    /// does nothing). If the function returns Some, the existing value is replaced, of if there
-    /// was no existing value, it is inserted. FIXME: update comment
-    pub fn update_with_fn<F>(mut self, key: &K, value_fn: F) -> Result<(), OutOfMemoryError>
-    where
-        F: FnOnce(Option<&V>) -> UpdateAction<V>,
-    {
-        algorithm::update_fn(key, value_fn, self.tree_writer.tree.root, &mut self)?;
-
-        if self.created_garbage {
-            let _ = self.collect_garbage();
-        }
-        Ok(())
-    }
-
-    fn remember_obsolete_node(&mut self, ptr: NodePtr<V>) {
-        self.tree_writer
-            .garbage
-            .lock()
-            .remember_obsolete_node(ptr, self.epoch_pin.epoch);
-        self.created_garbage = true;
-    }
-
-    // returns number of nodes recycled
-    fn collect_garbage(&self) -> usize {
-        self.tree_writer.tree.epoch.advance();
-        self.tree_writer.tree.epoch.broadcast();
-
-        let cutoff_epoch = self.tree_writer.tree.epoch.get_oldest();
-
-        let mut result = 0;
-        let mut garbage_queue = self.tree_writer.garbage.lock();
-        while let Some(ptr) = garbage_queue.next_obsolete(cutoff_epoch) {
-            ptr.deallocate(self.tree_writer.allocator);
-            result += 1;
-        }
-        result
-    }
-}
-
-pub struct TreeIterator<K>
-where
-    K: Key + for<'a> From<&'a [u8]>,
-{
-    done: bool,
-    pub next_key: Vec<u8>,
-    max_key: Option<Vec<u8>>,
-
-    phantom_key: PhantomData<K>,
-}
-
-impl<K> TreeIterator<K>
-where
-    K: Key + for<'a> From<&'a [u8]>,
-{
-    pub fn new_wrapping() -> TreeIterator<K> {
-        TreeIterator {
-            done: false,
-            next_key: vec![0; K::KEY_LEN],
-            max_key: None,
-            phantom_key: PhantomData,
-        }
-    }
-
-    pub fn new(range: &std::ops::Range<K>) -> TreeIterator<K> {
-        let result = TreeIterator {
-            done: false,
-            next_key: Vec::from(range.start.as_bytes()),
-            max_key: Some(Vec::from(range.end.as_bytes())),
-            phantom_key: PhantomData,
-        };
-        assert_eq!(result.next_key.len(), K::KEY_LEN);
-        assert_eq!(result.max_key.as_ref().unwrap().len(), K::KEY_LEN);
-
-        result
-    }
-
-    pub fn next<'g, V>(&mut self, read_guard: &'g TreeReadGuard<'g, K, V>) -> Option<(K, &'g V)>
-    where
-        V: Value,
-    {
-        if self.done {
-            return None;
-        }
-
-        let mut wrapped_around = false;
-        loop {
-            assert_eq!(self.next_key.len(), K::KEY_LEN);
-            if let Some((k, v)) =
-                algorithm::iter_next(&self.next_key, read_guard.tree.root, &read_guard.epoch_pin)
-            {
-                assert_eq!(k.len(), K::KEY_LEN);
-                assert_eq!(self.next_key.len(), K::KEY_LEN);
-
-                // Check if we reached the end of the range
-                if let Some(max_key) = &self.max_key {
-                    if k.as_slice() >= max_key.as_slice() {
-                        self.done = true;
-                        break None;
-                    }
-                }
-
-                // increment the key
-                self.next_key = k.clone();
-                increment_key(self.next_key.as_mut_slice());
-                let k = k.as_slice().into();
-
-                break Some((k, v));
-            } else {
-                if self.max_key.is_some() {
-                    self.done = true;
-                } else {
-                    // Start from beginning
-                    if !wrapped_around {
-                        for i in 0..K::KEY_LEN {
-                            self.next_key[i] = 0;
-                        }
-                        wrapped_around = true;
-                        continue;
-                    } else {
-                        // The tree is completely empty
-                        // FIXME: perhaps we should remember the starting point instead.
-                        // Currently this will scan some ranges twice.
-                        break None;
-                    }
-                }
-                break None;
-            }
-        }
-    }
-}
-
-fn increment_key(key: &mut [u8]) -> bool {
-    for i in (0..key.len()).rev() {
-        let (byte, overflow) = key[i].overflowing_add(1);
-        key[i] = byte;
-        if !overflow {
-            return false;
-        }
-    }
-    true
-}
-
-// Debugging functions
-impl<'e, K: Key, V: Value + Debug, A: ArtAllocator<V>> TreeWriteGuard<'e, K, V, A> {
-    pub fn dump(&mut self, dst: &mut dyn std::io::Write) {
-        algorithm::dump_tree(self.tree_writer.tree.root, &self.epoch_pin, dst)
-    }
-}
-impl<'e, K: Key, V: Value + Debug> TreeReadGuard<'e, K, V> {
-    pub fn dump(&mut self, dst: &mut dyn std::io::Write) {
-        algorithm::dump_tree(self.tree.root, &self.epoch_pin, dst)
-    }
-}
-impl<'e, K: Key, V: Value> TreeWriteAccess<'e, K, V, ArtMultiSlabAllocator<'e, V>> {
-    pub fn get_statistics(&self) -> ArtTreeStatistics {
-        self.allocator.get_statistics();
-        ArtTreeStatistics {
-            blocks: self.allocator.inner.block_allocator.get_statistics(),
-            slabs: self.allocator.get_statistics(),
-            epoch: self.tree.epoch.get_current(),
-            oldest_epoch: self.tree.epoch.get_oldest(),
-            num_garbage: self.garbage.lock().0.len() as u64,
-        }
-    }
-}
-
-#[derive(Clone, Debug)]
-pub struct ArtTreeStatistics {
-    pub blocks: allocator::block::BlockAllocatorStats,
-    pub slabs: allocator::ArtMultiSlabStats,
-
-    pub epoch: u64,
-    pub oldest_epoch: u64,
-    pub num_garbage: u64,
-}

libs/neonart/src/tests.rs

@@ -1,236 +0,0 @@
-use std::collections::BTreeMap;
-use std::collections::HashSet;
-use std::fmt::{Debug, Formatter};
-use std::sync::atomic::{AtomicUsize, Ordering};
-
-use crate::ArtAllocator;
-use crate::ArtMultiSlabAllocator;
-use crate::TreeInitStruct;
-use crate::TreeIterator;
-use crate::TreeWriteAccess;
-use crate::UpdateAction;
-
-use crate::{Key, Value};
-
-use rand::Rng;
-use rand::seq::SliceRandom;
-use rand_distr::Zipf;
-
-const TEST_KEY_LEN: usize = 16;
-
-#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
-struct TestKey([u8; TEST_KEY_LEN]);
-
-impl TestKey {
-    const MIN: TestKey = TestKey([0; TEST_KEY_LEN]);
-    const MAX: TestKey = TestKey([u8::MAX; TEST_KEY_LEN]);
-}
-
-impl Key for TestKey {
-    const KEY_LEN: usize = TEST_KEY_LEN;
-    fn as_bytes(&self) -> &[u8] {
-        &self.0
-    }
-}
-
-impl From<&TestKey> for u128 {
-    fn from(val: &TestKey) -> u128 {
-        u128::from_be_bytes(val.0)
-    }
-}
-
-impl From<u128> for TestKey {
-    fn from(val: u128) -> TestKey {
-        TestKey(val.to_be_bytes())
-    }
-}
-
-impl<'a> From<&'a [u8]> for TestKey {
-    fn from(bytes: &'a [u8]) -> TestKey {
-        TestKey(bytes.try_into().unwrap())
-    }
-}
-
-impl Value for usize {}
-
-fn test_inserts<K: Into<TestKey> + Copy>(keys: &[K]) {
-    const MEM_SIZE: usize = 10000000;
-    let mut area = Box::new_uninit_slice(MEM_SIZE);
-
-    let allocator = ArtMultiSlabAllocator::new(&mut area);
-
-    let init_struct = TreeInitStruct::<TestKey, usize, _>::new(allocator);
-    let tree_writer = init_struct.attach_writer();
-
-    for (idx, k) in keys.iter().enumerate() {
-        let w = tree_writer.start_write();
-        let res = w.insert(&(*k).into(), idx);
-        assert!(res.is_ok());
-    }
-
-    for (idx, k) in keys.iter().enumerate() {
-        let r = tree_writer.start_read();
-        let value = r.get(&(*k).into());
-        assert_eq!(value, Some(idx).as_ref());
-    }
-
-    eprintln!("stats: {:?}", tree_writer.get_statistics());
-}
-
-#[test]
-fn dense() {
-    // This exercises splitting a node with prefix
-    let keys: &[u128] = &[0, 1, 2, 3, 256];
-    test_inserts(keys);
-
-    // Dense keys
-    let mut keys: Vec<u128> = (0..10000).collect();
-    test_inserts(&keys);
-
-    // Do the same in random orders
-    for _ in 1..10 {
-        keys.shuffle(&mut rand::rng());
-        test_inserts(&keys);
-    }
-}
-
-#[test]
-fn sparse() {
-    // sparse keys
-    let mut keys: Vec<TestKey> = Vec::new();
-    let mut used_keys = HashSet::new();
-    for _ in 0..10000 {
-        loop {
-            let key = rand::random::<u128>();
-            if used_keys.contains(&key) {
-                continue;
-            }
-            used_keys.insert(key);
-            keys.push(key.into());
-            break;
-        }
-    }
-    test_inserts(&keys);
-}
-
-struct TestValue(AtomicUsize);
-
-impl TestValue {
-    fn new(val: usize) -> TestValue {
-        TestValue(AtomicUsize::new(val))
-    }
-
-    fn load(&self) -> usize {
-        self.0.load(Ordering::Relaxed)
-    }
-}
-
-impl Value for TestValue {}
-
-impl Clone for TestValue {
-    fn clone(&self) -> TestValue {
-        TestValue::new(self.load())
-    }
-}
-
-impl Debug for TestValue {
-    fn fmt(&self, fmt: &mut Formatter<'_>) -> Result<(), std::fmt::Error> {
-        write!(fmt, "{:?}", self.load())
-    }
-}
-
-#[derive(Clone, Debug)]
-struct TestOp(TestKey, Option<usize>);
-
-fn apply_op<A: ArtAllocator<TestValue>>(
-    op: &TestOp,
-    tree: &TreeWriteAccess<TestKey, TestValue, A>,
-    shadow: &mut BTreeMap<TestKey, usize>,
-) {
-    eprintln!("applying op: {op:?}");
-
-    // apply the change to the shadow tree first
-    let shadow_existing = if let Some(v) = op.1 {
-        shadow.insert(op.0, v)
-    } else {
-        shadow.remove(&op.0)
-    };
-
-    // apply to Art tree
-    let w = tree.start_write();
-    w.update_with_fn(&op.0, |existing| {
-        assert_eq!(existing.map(TestValue::load), shadow_existing);
-
-        match (existing, op.1) {
-            (None, None) => UpdateAction::Nothing,
-            (None, Some(new_val)) => UpdateAction::Insert(TestValue::new(new_val)),
-            (Some(_old_val), None) => UpdateAction::Remove,
-            (Some(old_val), Some(new_val)) => {
-                old_val.0.store(new_val, Ordering::Relaxed);
-                UpdateAction::Nothing
-            }
-        }
-    })
-    .expect("out of memory");
-}
-
-fn test_iter<A: ArtAllocator<TestValue>>(
-    tree: &TreeWriteAccess<TestKey, TestValue, A>,
-    shadow: &BTreeMap<TestKey, usize>,
-) {
-    let mut shadow_iter = shadow.iter();
-    let mut iter = TreeIterator::new(&(TestKey::MIN..TestKey::MAX));
-
-    loop {
-        let shadow_item = shadow_iter.next().map(|(k, v)| (*k, *v));
-        let r = tree.start_read();
-        let item = iter.next(&r);
-
-        if shadow_item != item.map(|(k, v)| (k, v.load())) {
-            eprintln!("FAIL: iterator returned {item:?}, expected {shadow_item:?}");
-            tree.start_read().dump(&mut std::io::stderr());
-
-            eprintln!("SHADOW:");
-            for si in shadow {
-                eprintln!("key: {:?}, val: {}", si.0, si.1);
-            }
-            panic!("FAIL: iterator returned {item:?}, expected {shadow_item:?}");
-        }
-        if item.is_none() {
-            break;
-        }
-    }
-}
-
-#[test]
-fn random_ops() {
-    const MEM_SIZE: usize = 10000000;
-    let mut area = Box::new_uninit_slice(MEM_SIZE);
-
-    let allocator = ArtMultiSlabAllocator::new(&mut area);
-
-    let init_struct = TreeInitStruct::<TestKey, TestValue, _>::new(allocator);
-    let tree_writer = init_struct.attach_writer();
-
-    let mut shadow: std::collections::BTreeMap<TestKey, usize> = BTreeMap::new();
-
-    let distribution = Zipf::new(u128::MAX as f64, 1.1).unwrap();
-    let mut rng = rand::rng();
-    for i in 0..100000 {
-        let mut key: TestKey = (rng.sample(distribution) as u128).into();
-
-        if rng.random_bool(0.10) {
-            key = TestKey::from(u128::from(&key) | 0xffffffff);
-        }
-
-        let op = TestOp(key, if rng.random_bool(0.75) { Some(i) } else { None });
-
-        apply_op(&op, &tree_writer, &mut shadow);
-
-        if i % 1000 == 0 {
-            eprintln!("{i} ops processed");
-            eprintln!("stats: {:?}", tree_writer.get_statistics());
-            test_iter(&tree_writer, &shadow);
-        }
-    }
-}

libs/pageserver_api/src/config.rs

@@ -394,7 +394,7 @@ impl From<&OtelExporterConfig> for tracing_utils::ExportConfig {
         tracing_utils::ExportConfig {
             endpoint: Some(val.endpoint.clone()),
             protocol: val.protocol.into(),
-            timeout: val.timeout,
+            timeout: Some(val.timeout),
         }
     }
 }
@@ -651,6 +651,9 @@ pub struct TenantConfigToml {
     // FIXME: Remove skip_serializing_if when the feature is stable.
     #[serde(skip_serializing_if = "std::ops::Not::not")]
     pub basebackup_cache_enabled: bool,
+
+    #[serde(skip_serializing_if = "std::ops::Not::not")]
+    pub rel_size_v1_access_disabled: bool,
 }
 
 pub mod defaults {
@@ -959,6 +962,7 @@ impl Default for TenantConfigToml {
             sampling_ratio: None,
             relsize_snapshot_cache_capacity: DEFAULT_RELSIZE_SNAPSHOT_CACHE_CAPACITY,
             basebackup_cache_enabled: false,
+            rel_size_v1_access_disabled: false,
         }
     }
 }

libs/pageserver_api/src/controller_api.rs

@@ -596,6 +596,7 @@ pub struct TimelineImportRequest {
     pub timeline_id: TimelineId,
     pub start_lsn: Lsn,
     pub sk_set: Vec<NodeId>,
+    pub force_upsert: bool,
 }
 
 #[derive(serde::Serialize, serde::Deserialize, Clone)]

libs/pageserver_api/src/key.rs

@@ -519,6 +519,15 @@ pub fn rel_dir_to_key(spcnode: Oid, dbnode: Oid) -> Key {
     }
 }
 
+pub const REL_DIR_MIGRATION_KEY: Key = Key {
+    field1: REL_DIR_KEY_PREFIX,
+    field2: 0,
+    field3: 0,
+    field4: 0,
+    field5: 0,
+    field6: 0,
+};
+
 #[inline(always)]
 pub fn rel_tag_sparse_key(spcnode: Oid, dbnode: Oid, relnode: Oid, forknum: u8) -> Key {
     Key {
@@ -981,12 +990,12 @@ mod tests {
         let mut rng = rand::rngs::StdRng::seed_from_u64(42);
 
         let key = Key {
-            field1: rng.r#gen(),
-            field2: rng.r#gen(),
-            field3: rng.r#gen(),
-            field4: rng.r#gen(),
-            field5: rng.r#gen(),
-            field6: rng.r#gen(),
+            field1: rng.random(),
+            field2: rng.random(),
+            field3: rng.random(),
+            field4: rng.random(),
+            field5: rng.random(),
+            field6: rng.random(),
         };
 
         assert_eq!(key, Key::from_str(&format!("{key}")).unwrap());

libs/pageserver_api/src/models.rs

@@ -443,9 +443,9 @@ pub struct ImportPgdataIdempotencyKey(pub String);
 impl ImportPgdataIdempotencyKey {
     pub fn random() -> Self {
         use rand::Rng;
-        use rand::distributions::Alphanumeric;
+        use rand::distr::Alphanumeric;
         Self(
-            rand::thread_rng()
+            rand::rng()
                 .sample_iter(&Alphanumeric)
                 .take(20)
                 .map(char::from)
@@ -646,6 +646,8 @@ pub struct TenantConfigPatch {
     pub relsize_snapshot_cache_capacity: FieldPatch<usize>,
     #[serde(skip_serializing_if = "FieldPatch::is_noop")]
     pub basebackup_cache_enabled: FieldPatch<bool>,
+    #[serde(skip_serializing_if = "FieldPatch::is_noop")]
+    pub rel_size_v1_access_disabled: FieldPatch<bool>,
 }
 
 /// Like [`crate::config::TenantConfigToml`], but preserves the information
@@ -783,6 +785,9 @@ pub struct TenantConfig {
 
     #[serde(skip_serializing_if = "Option::is_none")]
     pub basebackup_cache_enabled: Option<bool>,
+
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub rel_size_v1_access_disabled: Option<bool>,
 }
 
 impl TenantConfig {
@@ -830,6 +835,7 @@ impl TenantConfig {
             mut sampling_ratio,
             mut relsize_snapshot_cache_capacity,
             mut basebackup_cache_enabled,
+            mut rel_size_v1_access_disabled,
         } = self;
 
         patch.checkpoint_distance.apply(&mut checkpoint_distance);
@@ -939,6 +945,9 @@ impl TenantConfig {
         patch
             .basebackup_cache_enabled
             .apply(&mut basebackup_cache_enabled);
+        patch
+            .rel_size_v1_access_disabled
+            .apply(&mut rel_size_v1_access_disabled);
 
         Ok(Self {
             checkpoint_distance,
@@ -980,6 +989,7 @@ impl TenantConfig {
             sampling_ratio,
             relsize_snapshot_cache_capacity,
             basebackup_cache_enabled,
+            rel_size_v1_access_disabled,
         })
     }
 
@@ -1094,6 +1104,9 @@ impl TenantConfig {
             basebackup_cache_enabled: self
                 .basebackup_cache_enabled
                 .unwrap_or(global_conf.basebackup_cache_enabled),
+            rel_size_v1_access_disabled: self
+                .rel_size_v1_access_disabled
+                .unwrap_or(global_conf.rel_size_v1_access_disabled),
         }
     }
 }
@@ -1500,6 +1513,7 @@ pub struct TimelineArchivalConfigRequest {
 #[derive(Serialize, Deserialize, PartialEq, Eq, Clone)]
 pub struct TimelinePatchIndexPartRequest {
     pub rel_size_migration: Option<RelSizeMigration>,
+    pub rel_size_migrated_at: Option<Lsn>,
     pub gc_compaction_last_completed_lsn: Option<Lsn>,
     pub applied_gc_cutoff_lsn: Option<Lsn>,
     #[serde(default)]
@@ -1525,18 +1539,21 @@ pub struct OffloadedTimelineInfo {
     pub archived_at: chrono::DateTime<chrono::Utc>,
 }
 
-#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
+/// The state of the rel size migration. This is persisted in the DbDir key and index part. Do not change without considering
+/// compatibility.
+#[derive(Default, Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
 #[serde(rename_all = "camelCase")]
 pub enum RelSizeMigration {
     /// The tenant is using the old rel_size format.
     /// Note that this enum is persisted as `Option<RelSizeMigration>` in the index part, so
     /// `None` is the same as `Some(RelSizeMigration::Legacy)`.
+    #[default]
     Legacy,
     /// The tenant is migrating to the new rel_size format. Both old and new rel_size format are
-    /// persisted in the index part. The read path will read both formats and merge them.
+    /// persisted in the storage. The read path will read both formats and validate them.
     Migrating,
     /// The tenant has migrated to the new rel_size format. Only the new rel_size format is persisted
-    /// in the index part, and the read path will not read the old format.
+    /// in the storage, and the read path will not read the old format.
     Migrated,
 }
 
@@ -1619,6 +1636,7 @@ pub struct TimelineInfo {
 
     /// The status of the rel_size migration.
     pub rel_size_migration: Option<RelSizeMigration>,
+    pub rel_size_migrated_at: Option<Lsn>,
 
     /// Whether the timeline is invisible in synthetic size calculations.
     pub is_invisible: Option<bool>,

libs/pageserver_api/src/shard.rs

@@ -69,22 +69,6 @@ impl Hash for ShardIdentity {
     }
 }
 
-/// Stripe size in number of pages
-#[derive(Clone, Copy, Serialize, Deserialize, Eq, PartialEq, Debug)]
-pub struct ShardStripeSize(pub u32);
-
-impl Default for ShardStripeSize {
-    fn default() -> Self {
-        DEFAULT_STRIPE_SIZE
-    }
-}
-
-impl std::fmt::Display for ShardStripeSize {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        self.0.fmt(f)
-    }
-}
-
 /// Layout version: for future upgrades where we might change how the key->shard mapping works
 #[derive(Clone, Copy, Serialize, Deserialize, Eq, PartialEq, Hash, Debug)]
 pub struct ShardLayout(u8);

libs/pageserver_api/src/upcall_api.rs

@@ -21,6 +21,14 @@ pub struct ReAttachRequest {
     /// if the node already has a node_id set.
     #[serde(skip_serializing_if = "Option::is_none", default)]
     pub register: Option<NodeRegisterRequest>,
+
+    /// Hadron: Optional flag to indicate whether the node is starting with an empty local disk.
+    /// Will be set to true if the node couldn't find any local tenant data on startup, could be
+    /// due to the node starting for the first time or due to a local SSD failure/disk wipe event.
+    /// The flag may be used by the storage controller to update its observed state of the world
+    /// to make sure that it sends explicit location_config calls to the node following the
+    /// re-attach request.
+    pub empty_local_disk: Option<bool>,
 }
 
 #[derive(Serialize, Deserialize, Debug)]

libs/pq_proto/src/lib.rs

@@ -203,12 +203,12 @@ impl fmt::Display for CancelKeyData {
     }
 }
 
-use rand::distributions::{Distribution, Standard};
-impl Distribution<CancelKeyData> for Standard {
+use rand::distr::{Distribution, StandardUniform};
+impl Distribution<CancelKeyData> for StandardUniform {
     fn sample<R: rand::Rng + ?Sized>(&self, rng: &mut R) -> CancelKeyData {
         CancelKeyData {
-            backend_pid: rng.r#gen(),
-            cancel_key: rng.r#gen(),
+            backend_pid: rng.random(),
+            cancel_key: rng.random(),
         }
     }
 }

libs/proxy/postgres-protocol2/src/authentication/sasl.rs

@@ -155,10 +155,10 @@ pub struct ScramSha256 {
 
 fn nonce() -> String {
     // rand 0.5's ThreadRng is cryptographically secure
-    let mut rng = rand::thread_rng();
+    let mut rng = rand::rng();
     (0..NONCE_LENGTH)
         .map(|_| {
-            let mut v = rng.gen_range(0x21u8..0x7e);
+            let mut v = rng.random_range(0x21u8..0x7e);
             if v == 0x2c {
                 v = 0x7e
             }

libs/proxy/postgres-protocol2/src/message/backend.rs

@@ -74,7 +74,6 @@ impl Header {
 }
 
 /// An enum representing Postgres backend messages.
-#[non_exhaustive]
 pub enum Message {
     AuthenticationCleartextPassword,
     AuthenticationGss,
@@ -145,16 +144,7 @@ impl Message {
             PARSE_COMPLETE_TAG => Message::ParseComplete,
             BIND_COMPLETE_TAG => Message::BindComplete,
             CLOSE_COMPLETE_TAG => Message::CloseComplete,
-            NOTIFICATION_RESPONSE_TAG => {
-                let process_id = buf.read_i32::<BigEndian>()?;
-                let channel = buf.read_cstr()?;
-                let message = buf.read_cstr()?;
-                Message::NotificationResponse(NotificationResponseBody {
-                    process_id,
-                    channel,
-                    message,
-                })
-            }
+            NOTIFICATION_RESPONSE_TAG => Message::NotificationResponse(NotificationResponseBody {}),
             COPY_DONE_TAG => Message::CopyDone,
             COMMAND_COMPLETE_TAG => {
                 let tag = buf.read_cstr()?;
@@ -543,28 +533,7 @@ impl NoticeResponseBody {
     }
 }
 
-pub struct NotificationResponseBody {
-    process_id: i32,
-    channel: Bytes,
-    message: Bytes,
-}
-
-impl NotificationResponseBody {
-    #[inline]
-    pub fn process_id(&self) -> i32 {
-        self.process_id
-    }
-
-    #[inline]
-    pub fn channel(&self) -> io::Result<&str> {
-        get_str(&self.channel)
-    }
-
-    #[inline]
-    pub fn message(&self) -> io::Result<&str> {
-        get_str(&self.message)
-    }
-}
+pub struct NotificationResponseBody {}
 
 pub struct ParameterDescriptionBody {
     storage: Bytes,

libs/proxy/postgres-protocol2/src/password/mod.rs

@@ -28,7 +28,7 @@ const SCRAM_DEFAULT_SALT_LEN: usize = 16;
 /// special characters that would require escaping in an SQL command.
 pub async fn scram_sha_256(password: &[u8]) -> String {
     let mut salt: [u8; SCRAM_DEFAULT_SALT_LEN] = [0; SCRAM_DEFAULT_SALT_LEN];
-    let mut rng = rand::thread_rng();
+    let mut rng = rand::rng();
     rng.fill_bytes(&mut salt);
     scram_sha_256_salt(password, salt).await
 }

libs/proxy/subzero_core/.gitignore

@@ -0,0 +1,2 @@
+target
+Cargo.lock

libs/proxy/subzero_core/Cargo.toml

@@ -0,0 +1,12 @@
+# This is a stub for the subzero-core crate.
+[package]
+name = "subzero-core"
+version = "3.0.1"
+edition = "2024"
+publish = false # "private"!
+
+[features]
+default = []
+postgresql = []
+
+[dependencies]

libs/proxy/subzero_core/src/lib.rs

@@ -0,0 +1 @@
+// This is a stub for the subzero-core crate.

libs/proxy/tokio-postgres2/src/client.rs

@@ -13,8 +13,9 @@ use serde::{Deserialize, Serialize};
 use tokio::sync::mpsc;
 
 use crate::cancel_token::RawCancelToken;
-use crate::codec::{BackendMessages, FrontendMessage};
+use crate::codec::{BackendMessages, FrontendMessage, RecordNotices};
 use crate::config::{Host, SslMode};
+use crate::connection::gc_bytesmut;
 use crate::query::RowStream;
 use crate::simple_query::SimpleQueryStream;
 use crate::types::{Oid, Type};
@@ -95,20 +96,13 @@ impl InnerClient {
         Ok(PartialQuery(Some(self)))
     }
 
-    // pub fn send_with_sync<F>(&mut self, f: F) -> Result<&mut Responses, Error>
-    // where
-    //     F: FnOnce(&mut BytesMut) -> Result<(), Error>,
-    // {
-    //     self.start()?.send_with_sync(f)
-    // }
-
     pub fn send_simple_query(&mut self, query: &str) -> Result<&mut Responses, Error> {
         self.responses.waiting += 1;
 
         self.buffer.clear();
         // simple queries do not need sync.
         frontend::query(query, &mut self.buffer).map_err(Error::encode)?;
-        let buf = self.buffer.split().freeze();
+        let buf = self.buffer.split();
         self.send_message(FrontendMessage::Raw(buf))
     }
 
@@ -125,7 +119,7 @@ impl Drop for PartialQuery<'_> {
         if let Some(client) = self.0.take() {
             client.buffer.clear();
             frontend::sync(&mut client.buffer);
-            let buf = client.buffer.split().freeze();
+            let buf = client.buffer.split();
             let _ = client.send_message(FrontendMessage::Raw(buf));
         }
     }
@@ -141,7 +135,7 @@ impl<'a> PartialQuery<'a> {
         client.buffer.clear();
         f(&mut client.buffer)?;
         frontend::flush(&mut client.buffer);
-        let buf = client.buffer.split().freeze();
+        let buf = client.buffer.split();
         client.send_message(FrontendMessage::Raw(buf))
     }
 
@@ -154,7 +148,7 @@ impl<'a> PartialQuery<'a> {
         client.buffer.clear();
         f(&mut client.buffer)?;
         frontend::sync(&mut client.buffer);
-        let buf = client.buffer.split().freeze();
+        let buf = client.buffer.split();
         let _ = client.send_message(FrontendMessage::Raw(buf));
 
         Ok(&mut self.0.take().unwrap().responses)
@@ -191,6 +185,7 @@ impl Client {
         ssl_mode: SslMode,
         process_id: i32,
         secret_key: i32,
+        write_buf: BytesMut,
     ) -> Client {
         Client {
             inner: InnerClient {
@@ -201,7 +196,7 @@ impl Client {
                     waiting: 0,
                     received: 0,
                 },
-                buffer: Default::default(),
+                buffer: write_buf,
             },
             cached_typeinfo: Default::default(),
 
@@ -221,6 +216,18 @@ impl Client {
         &mut self.inner
     }
 
+    pub fn record_notices(&mut self, limit: usize) -> mpsc::UnboundedReceiver<Box<str>> {
+        let (tx, rx) = mpsc::unbounded_channel();
+
+        let notices = RecordNotices { sender: tx, limit };
+        self.inner
+            .sender
+            .send(FrontendMessage::RecordNotices(notices))
+            .ok();
+
+        rx
+    }
+
     /// Pass text directly to the Postgres backend to allow it to sort out typing itself and
     /// to save a roundtrip
     pub async fn query_raw_txt<S, I>(
@@ -280,8 +287,35 @@ impl Client {
         simple_query::batch_execute(self.inner_mut(), query).await
     }
 
-    pub async fn discard_all(&mut self) -> Result<ReadyForQueryStatus, Error> {
-        self.batch_execute("discard all").await
+    /// Similar to `discard_all`, but it does not clear any query plans
+    ///
+    /// This runs in the background, so it can be executed without `await`ing.
+    pub fn reset_session_background(&mut self) -> Result<(), Error> {
+        // "CLOSE ALL": closes any cursors
+        // "SET SESSION AUTHORIZATION DEFAULT": resets the current_user back to the session_user
+        // "RESET ALL": resets any GUCs back to their session defaults.
+        // "DEALLOCATE ALL": deallocates any prepared statements
+        // "UNLISTEN *": stops listening on all channels
+        // "SELECT pg_advisory_unlock_all();": unlocks all advisory locks
+        // "DISCARD TEMP;": drops all temporary tables
+        // "DISCARD SEQUENCES;": deallocates all cached sequence state
+
+        let _responses = self.inner_mut().send_simple_query(
+            "ROLLBACK;
+            CLOSE ALL;
+            SET SESSION AUTHORIZATION DEFAULT;
+            RESET ALL;
+            DEALLOCATE ALL;
+            UNLISTEN *;
+            SELECT pg_advisory_unlock_all();
+            DISCARD TEMP;
+            DISCARD SEQUENCES;",
+        )?;
+
+        // Clean up memory usage.
+        gc_bytesmut(&mut self.inner_mut().buffer);
+
+        Ok(())
     }
 
     /// Begins a new database transaction.

libs/proxy/tokio-postgres2/src/codec.rs

@@ -1,16 +1,26 @@
 use std::io;
 
-use bytes::{Bytes, BytesMut};
+use bytes::BytesMut;
 use fallible_iterator::FallibleIterator;
 use postgres_protocol2::message::backend;
+use tokio::sync::mpsc::UnboundedSender;
 use tokio_util::codec::{Decoder, Encoder};
 
 pub enum FrontendMessage {
-    Raw(Bytes),
+    Raw(BytesMut),
+    RecordNotices(RecordNotices),
+}
+
+pub struct RecordNotices {
+    pub sender: UnboundedSender<Box<str>>,
+    pub limit: usize,
 }
 
 pub enum BackendMessage {
-    Normal { messages: BackendMessages },
+    Normal {
+        messages: BackendMessages,
+        ready: bool,
+    },
     Async(backend::Message),
 }
 
@@ -33,14 +43,11 @@ impl FallibleIterator for BackendMessages {
 
 pub struct PostgresCodec;
 
-impl Encoder<FrontendMessage> for PostgresCodec {
+impl Encoder<BytesMut> for PostgresCodec {
     type Error = io::Error;
 
-    fn encode(&mut self, item: FrontendMessage, dst: &mut BytesMut) -> io::Result<()> {
-        match item {
-            FrontendMessage::Raw(buf) => dst.extend_from_slice(&buf),
-        }
-
+    fn encode(&mut self, item: BytesMut, dst: &mut BytesMut) -> io::Result<()> {
+        dst.unsplit(item);
         Ok(())
     }
 }
@@ -52,6 +59,7 @@ impl Decoder for PostgresCodec {
     fn decode(&mut self, src: &mut BytesMut) -> Result<Option<BackendMessage>, io::Error> {
         let mut idx = 0;
 
+        let mut ready = false;
         while let Some(header) = backend::Header::parse(&src[idx..])? {
             let len = header.len() as usize + 1;
             if src[idx..].len() < len {
@@ -75,6 +83,7 @@ impl Decoder for PostgresCodec {
             idx += len;
 
             if header.tag() == backend::READY_FOR_QUERY_TAG {
+                ready = true;
                 break;
             }
         }
@@ -84,6 +93,7 @@ impl Decoder for PostgresCodec {
         } else {
             Ok(Some(BackendMessage::Normal {
                 messages: BackendMessages(src.split_to(idx)),
+                ready,
             }))
         }
     }