make use of esc string opt

We also don't need to return the 'rest' slice, we can just insert it into the vec

Lastly, we can simplify the hotloop by making write_char_escape cold and moving the vec write inside this fn

.config/hakari.toml

@@ -21,14 +21,13 @@ 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/actions/prepare-for-subzero/action.yml

@@ -1,28 +0,0 @@
-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,10 +86,6 @@ 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
@@ -120,7 +116,7 @@ jobs:
           ARCH: ${{ inputs.arch }}
           SANITIZERS: ${{ inputs.sanitizers }}
         run: |
-          CARGO_FLAGS="--locked --features testing,rest_broker"
+          CARGO_FLAGS="--locked --features testing"
           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,10 +46,6 @@ 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

@@ -1,384 +0,0 @@
-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/build-macos.yml

@@ -54,10 +54,6 @@ 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,8 +632,6 @@ 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,7 +72,6 @@ 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/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 test-proxy-bench branch
+      - test-proxy-bench # Runs on pushes to branches starting with test-proxy-bench
   # 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,58 +55,30 @@ 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: 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: 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
+      run: ${PROXY_BENCH_PATH}/run.sh
 
-    - name: Ingest Bench Results
+    - name: Ingest Bench Results # neon repo script
       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
-
-    - 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
+      run: echo "Proxy bench job failed" && exit 1

.gitignore

@@ -26,14 +26,9 @@ 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

@@ -49,7 +49,6 @@ members = [
     "libs/proxy/tokio-postgres2",
     "endpoint_storage",
     "pgxn/neon/communicator",
-    "proxy/subzero_core",
 ]
 
 [workspace.package]
@@ -143,10 +142,10 @@ notify = "6.0.0"
 num_cpus = "1.15"
 num-traits = "0.2.19"
 once_cell = "1.13"
-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"
+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"
 parking_lot = "0.12"
 parquet = { version = "53", default-features = false, features = ["zstd"] }
 parquet_derive = "53"
@@ -158,13 +157,11 @@ 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.9"
-# Remove after p256 is updated to 0.14.
-rand_core = "=0.6"
+rand = "0.8"
 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_30"] }
+reqwest-tracing = { version = "0.5", features = ["opentelemetry_0_27"] }
 reqwest-middleware = "0.4"
 reqwest-retry = "0.7"
 routerify = "3"
@@ -214,15 +211,17 @@ 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"] }
-tower-otel = { version = "0.6", features = ["axum"] }
+
+# This revision uses opentelemetry 0.27. There's no tag for it.
+tower-otel = { git = "https://github.com/mattiapenati/tower-otel", rev = "56a7321053bcb72443888257b622ba0d43a11fcd" }
+
 tower-service = "0.3.3"
 tracing = "0.1"
 tracing-error = "0.2"
 tracing-log = "0.2"
-tracing-opentelemetry = "0.31"
+tracing-opentelemetry = "0.28"
 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 }

Dockerfile

@@ -63,14 +63,7 @@ WORKDIR /home/nonroot
 
 COPY --chown=nonroot . .
 
-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
+RUN cargo chef prepare --recipe-path recipe.json
 
 # Main build image
 FROM $REPOSITORY/$IMAGE:$TAG AS build
@@ -78,33 +71,20 @@ 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 --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 \
+RUN set -e \
     && 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  --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 \
+RUN set -e \
     && 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/package-lock.json

@@ -6,7 +6,7 @@
     "": {
       "name": "build-tools",
       "devDependencies": {
-        "@redocly/cli": "1.34.5",
+        "@redocly/cli": "1.34.4",
         "@sourcemeta/jsonschema": "10.0.0"
       }
     },
@@ -472,9 +472,9 @@
       }
     },
     "node_modules/@redocly/cli": {
-      "version": "1.34.5",
-      "resolved": "https://registry.npmjs.org/@redocly/cli/-/cli-1.34.5.tgz",
-      "integrity": "sha512-5IEwxs7SGP5KEXjBKLU8Ffdz9by/KqNSeBk6YUVQaGxMXK//uYlTJIPntgUXbo1KAGG2d2q2XF8y4iFz6qNeiw==",
+      "version": "1.34.4",
+      "resolved": "https://registry.npmjs.org/@redocly/cli/-/cli-1.34.4.tgz",
+      "integrity": "sha512-seH/GgrjSB1EeOsgJ/4Ct6Jk2N7sh12POn/7G8UQFARMyUMJpe1oHtBwT2ndfp4EFCpgBAbZ/82Iw6dwczNxEA==",
       "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.5",
-        "@redocly/respect-core": "1.34.5",
+        "@redocly/openapi-core": "1.34.4",
+        "@redocly/respect-core": "1.34.4",
         "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.4",
+        "form-data": "^4.0.0",
         "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.5",
-      "resolved": "https://registry.npmjs.org/@redocly/openapi-core/-/openapi-core-1.34.5.tgz",
-      "integrity": "sha512-0EbE8LRbkogtcCXU7liAyC00n9uNG9hJ+eMyHFdUsy9lB/WGqnEBgwjA9q2cyzAVcdTkQqTBBU1XePNnN3OijA==",
+      "version": "1.34.4",
+      "resolved": "https://registry.npmjs.org/@redocly/openapi-core/-/openapi-core-1.34.4.tgz",
+      "integrity": "sha512-hf53xEgpXIgWl3b275PgZU3OTpYh1RoD2LHdIfQ1JzBNTWsiNKczTEsI/4Tmh2N1oq9YcphhSMyk3lDh85oDjg==",
       "dev": true,
       "license": "MIT",
       "dependencies": {
@@ -544,21 +544,21 @@
       }
     },
     "node_modules/@redocly/respect-core": {
-      "version": "1.34.5",
-      "resolved": "https://registry.npmjs.org/@redocly/respect-core/-/respect-core-1.34.5.tgz",
-      "integrity": "sha512-GheC/g/QFztPe9UA9LamooSplQuy9pe0Yr8XGTqkz0ahivLDl7svoy/LSQNn1QH3XGtLKwFYMfTwFR2TAYyh5Q==",
+      "version": "1.34.4",
+      "resolved": "https://registry.npmjs.org/@redocly/respect-core/-/respect-core-1.34.4.tgz",
+      "integrity": "sha512-MitKyKyQpsizA4qCVv+MjXL4WltfhFQAoiKiAzrVR1Kusro3VhYb6yJuzoXjiJhR0ukLP5QOP19Vcs7qmj9dZg==",
       "dev": true,
       "license": "MIT",
       "dependencies": {
         "@faker-js/faker": "^7.6.0",
         "@redocly/ajv": "8.11.2",
-        "@redocly/openapi-core": "1.34.5",
+        "@redocly/openapi-core": "1.34.4",
         "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.4",
+        "form-data": "4.0.0",
         "jest-diff": "^29.3.1",
         "jest-matcher-utils": "^29.3.1",
         "js-yaml": "4.1.0",
@@ -582,6 +582,21 @@
       "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",
@@ -1330,9 +1345,9 @@
       "license": "MIT"
     },
     "node_modules/form-data": {
-      "version": "4.0.4",
-      "resolved": "https://registry.npmjs.org/form-data/-/form-data-4.0.4.tgz",
-      "integrity": "sha512-KrGhL9Q4zjj0kiUt5OO4Mr/A/jlI2jDYs5eHBpYHPcBEVSiipAvn2Ko2HnPe20rmcuuvMHNdZFp+4IlGTMF0Ow==",
+      "version": "4.0.3",
+      "resolved": "https://registry.npmjs.org/form-data/-/form-data-4.0.3.tgz",
+      "integrity": "sha512-qsITQPfmvMOSAdeyZ+12I1c+CKSstAFAwu+97zrnWAbIr5u8wfsExUzCesVLC8NgHuRUqNN4Zy6UPWUTRGslcA==",
       "dev": true,
       "license": "MIT",
       "dependencies": {

build-tools/package.json

@@ -2,7 +2,7 @@
   "name": "build-tools",
   "private": true,
   "devDependencies": {
-    "@redocly/cli": "1.34.5",
+    "@redocly/cli": "1.34.4",
     "@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://archive.debian.org/debian bullseye-backports main" > /etc/apt/sources.list.d/bullseye-backports.list; \
+        echo "deb http://deb.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,13 +26,7 @@ commands:
   - name: postgres-exporter
     user: nobody
     sysvInitAction: respawn
-    # 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'
+    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'
   - name: pgbouncer-exporter
     user: postgres
     sysvInitAction: respawn

compute/vm-image-spec-bullseye.yaml

@@ -26,13 +26,7 @@ commands:
   - name: postgres-exporter
     user: nobody
     sysvInitAction: respawn
-    # 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'
+    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'
   - name: pgbouncer-exporter
     user: postgres
     sysvInitAction: respawn

compute_tools/Cargo.toml

@@ -27,10 +27,7 @@ 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
@@ -47,7 +44,6 @@ 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
@@ -62,7 +58,6 @@ 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,14 +52,8 @@ 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,10 +49,9 @@ 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;
@@ -139,12 +138,6 @@ 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 {
@@ -195,19 +188,11 @@ fn main() -> Result<()> {
         .build()?;
     let _rt_guard = runtime.enter();
 
-    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)?;
+    runtime.block_on(init(cli.dev))?;
 
     // 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)?;
@@ -234,13 +219,6 @@ 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,
     )?;
@@ -249,17 +227,11 @@ fn main() -> Result<()> {
 
     scenario.teardown();
 
-    deinit_and_exit(tracing_provider, exit_code);
+    deinit_and_exit(exit_code);
 }
 
-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)?;
+async fn init(dev_mode: bool) -> Result<()> {
+    init_tracing_and_logging(DEFAULT_LOG_LEVEL).await?;
 
     let mut signals = Signals::new([SIGINT, SIGTERM, SIGQUIT])?;
     thread::spawn(move || {
@@ -270,7 +242,7 @@ fn init(
 
     info!("compute build_tag: {}", &BUILD_TAG.to_string());
 
-    Ok((provider, file_logs_guard))
+    Ok(())
 }
 
 fn get_config(cli: &Cli) -> Result<ComputeConfig> {
@@ -295,27 +267,25 @@ fn get_config(cli: &Cli) -> Result<ComputeConfig> {
     }
 }
 
-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");
-        }
+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");
     }
 
     info!("shutting down");

compute_tools/src/communicator_socket_client.rs

@@ -1,98 +0,0 @@
-//! 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

@@ -21,7 +21,6 @@ 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,9 +40,8 @@ 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::{self, startup_context_from_env};
+use crate::logger::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;
@@ -115,17 +113,6 @@ 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<()>>>;
@@ -167,7 +154,6 @@ 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.
@@ -205,7 +191,6 @@ impl ComputeState {
     pub fn new() -> Self {
         Self {
             start_time: Utc::now(),
-            pg_start_time: None,
             status: ComputeStatus::Empty,
             last_active: None,
             error: None,
@@ -494,7 +479,6 @@ 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);
 
@@ -664,9 +648,6 @@ 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()
         }
@@ -1460,7 +1441,7 @@ impl ComputeNode {
         })?;
 
         // Update pg_hba.conf received with basebackup.
-        update_pg_hba(pgdata_path, None)?;
+        update_pg_hba(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
@@ -1765,7 +1746,6 @@ 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();
@@ -1778,7 +1758,7 @@ impl ComputeNode {
                             eprintln!("connection error: {e}");
                         }
                     });
-                    if let Err(e) = handle_migrations(params, &mut client, lakebase_mode).await {
+                    if let Err(e) = handle_migrations(params, &mut client).await {
                         error!("Failed to run migrations: {}", e);
                     }
                 }
@@ -1794,34 +1774,6 @@ 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)]
@@ -1994,8 +1946,6 @@ impl ComputeNode {
                             // wait
                             ComputeStatus::Init
                             | ComputeStatus::Configuration
-                            | ComputeStatus::RefreshConfiguration
-                            | ComputeStatus::RefreshConfigurationPending
                             | ComputeStatus::Empty => {
                                 state = self.state_changed.wait(state).unwrap();
                             }
@@ -2552,34 +2502,6 @@ 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<()> {

compute_tools/src/compute_prewarm.rs

@@ -90,7 +90,6 @@ 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;
@@ -113,8 +112,9 @@ impl ComputeNode {
                 Err(err) => {
                     crate::metrics::LFC_PREWARM_ERRORS.inc();
                     error!(%err, "could not prewarm LFC");
+
                     LfcPrewarmState::Failed {
-                        error: format!("{err:#}"),
+                        error: err.to_string(),
                     }
                 }
             };
@@ -135,20 +135,16 @@ 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")?;
-        match res.status() {
+        let status = res.status();
+        match status {
             StatusCode::OK => (),
             StatusCode::NOT_FOUND => {
                 return Ok(false);
             }
-            status => bail!("{status} querying endpoint storage"),
+            _ => bail!("{status} querying endpoint storage"),
         }
 
         let mut uncompressed = Vec::new();
@@ -209,7 +205,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: format!("{err:#}"),
+            error: err.to_string(),
         };
     }
 
@@ -217,22 +213,16 @@ impl ComputeNode {
         let EndpointStoragePair { url, token } = self.endpoint_storage_pair(None)?;
         info!(%url, "requesting LFC state from Postgres");
 
-        let row = ComputeNode::get_maintenance_client(&self.tokio_conn_conf)
+        let mut compressed = Vec::new();
+        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")?;
-        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)
+            .context("querying LFC state")?
+            .try_get::<usize, &[u8]>(0)
+            .context("deserializing LFC state")
+            .map(ZstdEncoder::new)?
             .read_to_end(&mut compressed)
             .await
             .context("compressing LFC state")?;

compute_tools/src/compute_promote.rs

@@ -1,12 +1,11 @@
 use crate::compute::ComputeNode;
 use anyhow::{Context, Result, bail};
-use compute_api::responses::{LfcPrewarmState, PromoteConfig, PromoteState};
-use compute_api::spec::ComputeMode;
-use itertools::Itertools;
-use std::collections::HashMap;
+use compute_api::{
+    responses::{LfcPrewarmState, PromoteState, SafekeepersLsn},
+    spec::ComputeMode,
+};
 use std::{sync::Arc, time::Duration};
 use tokio::time::sleep;
-use tracing::info;
 use utils::lsn::Lsn;
 
 impl ComputeNode {
@@ -14,22 +13,21 @@ 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
-    /// Has a failpoint "compute-promotion"
-    pub async fn promote(self: &Arc<Self>, cfg: PromoteConfig) -> PromoteState {
+    pub async fn promote(self: &Arc<Self>, safekeepers_lsn: SafekeepersLsn) -> 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(promote_fn().await) });
+            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(),
+                        }
+                    }
+                })
+            });
             rx
         };
 
@@ -49,7 +47,9 @@ impl ComputeNode {
         task.borrow().clone()
     }
 
-    async fn promote_impl(&self, mut cfg: PromoteConfig) -> Result<()> {
+    // 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<()> {
         {
             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 = cfg.wal_flush_lsn;
+        let primary_lsn = safekeepers_lsn.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;
             }
-            info!("Try {i}, replica lsn {last_wal_replay_lsn}, primary lsn {primary_lsn}");
+            tracing::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='{}'",
-            cfg.spec.safekeeper_connstrings.join(",")
+            safekeepers_lsn.safekeepers
         );
         client
             .query(&safekeepers_sql, &[])
@@ -106,12 +106,6 @@ 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
@@ -131,36 +125,8 @@ impl ComputeNode {
             bail!("replica in read only mode after promotion");
         }
 
-        {
-            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");
+        let mut state = self.state.lock().unwrap();
+        state.pspec.as_mut().unwrap().spec.mode = ComputeMode::Primary;
+        Ok(())
     }
 }

compute_tools/src/configurator.rs

@@ -1,40 +1,23 @@
-use std::fs::File;
+use std::sync::Arc;
 use std::thread;
-use std::{path::Path, sync::Arc};
 
-use anyhow::Result;
-use compute_api::responses::{ComputeConfig, ComputeStatus};
+use compute_api::responses::ComputeStatus;
 use tracing::{error, info, instrument};
 
-use crate::compute::{ComputeNode, ParsedSpec};
-use crate::spec::get_config_from_control_plane;
+use crate::compute::ComputeNode;
 
 #[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 */
 
-        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();
-            }
+        // 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
@@ -54,133 +37,6 @@ 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

@@ -1,60 +0,0 @@
-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,29 +16,13 @@ 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, instance_id: Option<String>, jwks: JwkSet) -> Self {
+    pub fn new(compute_id: 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;
@@ -50,7 +34,6 @@ impl Authorize {
 
         Self {
             compute_id,
-            instance_id,
             jwks,
             validation,
         }
@@ -64,20 +47,10 @@ 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/ComputeSchemaWithLsn"
+                $ref: "#/components/schemas/SafekeepersLsn"
       responses:
         200:
           description: Promote succeeded or wasn't started
@@ -297,7 +297,14 @@ paths:
         content:
           application/json:
             schema:
-              $ref: "#/components/schemas/ComputeSchema"
+              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
       responses:
         200:
           description: Compute configuration finished.
@@ -584,25 +591,18 @@ components:
           type: string
           example: "1.0.0"
 
-    ComputeSchema:
+    SafekeepersLsn:
       type: object
       required:
-        - spec
-      properties:
-        spec:
-          type: object
-    ComputeSchemaWithLsn:
-      type: object
-      required:
-        - spec
+        - safekeepers
         - wal_flush_lsn
       properties:
-        spec:
-          $ref: "#/components/schemas/ComputeState"
-        wal_flush_lsn:
+        safekeepers:
+          description: Primary replica safekeepers
+          type: string
+        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,12 +43,7 @@ pub(in crate::http) async fn configure(
         // configure request for tracing purposes.
         state.startup_span = Some(tracing::Span::current());
 
-        if compute.params.lakebase_mode {
-            ComputeNode::set_spec(&compute.params, &mut state, pspec);
-        } else {
-            state.pspec = Some(pspec);
-        }
-
+        state.pspec = Some(pspec);
         state.set_status(ComputeStatus::ConfigurationPending, &compute.state_changed);
         drop(state);
     }

compute_tools/src/http/routes/hadron_liveness_probe.rs

@@ -1,34 +0,0 @@
-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,19 +1,10 @@
-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;
 
@@ -22,18 +13,11 @@ 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 mut metrics = collect()
+    let 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![];
 
@@ -47,42 +31,3 @@ 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,13 +10,11 @@ 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::extract::Json;
+use axum::Form;
 use http::StatusCode;
 
 pub(in crate::http) async fn promote(
     compute: axum::extract::State<std::sync::Arc<crate::compute::ComputeNode>>,
-    Json(cfg): Json<compute_api::responses::PromoteConfig>,
+    Form(safekeepers_lsn): Form<compute_api::responses::SafekeepersLsn>,
 ) -> axum::response::Response {
-    let state = compute.promote(cfg).await;
-    if let compute_api::responses::PromoteState::Failed { error: _ } = state {
-        return JsonResponse::create_response(StatusCode::INTERNAL_SERVER_ERROR, state);
+    let state = compute.promote(safekeepers_lsn).await;
+    if let compute_api::responses::PromoteState::Failed { error } = state {
+        return JsonResponse::error(StatusCode::INTERNAL_SERVER_ERROR, error);
     }
     JsonResponse::success(StatusCode::OK, state)
 }

compute_tools/src/http/routes/refresh_configuration.rs

@@ -1,29 +0,0 @@
-// 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::{IntoResponse, Response};
+use axum::response::Response;
 use axum_extra::extract::OptionalQuery;
 use compute_api::responses::{ComputeStatus, TerminateMode, TerminateResponse};
 use http::StatusCode;
@@ -33,29 +33,7 @@ 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,8 +23,7 @@ use super::{
     middleware::authorize::Authorize,
     routes::{
         check_writability, configure, database_schema, dbs_and_roles, extension_server, extensions,
-        grants, hadron_liveness_probe, insights, lfc, metrics, metrics_json, promote,
-        refresh_configuration, status, terminate,
+        grants, insights, lfc, metrics, metrics_json, promote, status, terminate,
     },
 };
 use crate::compute::ComputeNode;
@@ -44,7 +43,6 @@ pub enum Server {
         port: u16,
         config: ComputeCtlConfig,
         compute_id: String,
-        instance_id: Option<String>,
     },
 }
 
@@ -69,12 +67,7 @@ impl From<&Server> for Router<Arc<ComputeNode>> {
                         post(extension_server::download_extension),
                     )
                     .route("/extensions", post(extensions::install_extension))
-                    .route("/grants", post(grants::add_grant))
-                    // Hadron: Compute-initiated configuration refresh
-                    .route(
-                        "/refresh_configuration",
-                        post(refresh_configuration::refresh_configuration),
-                    );
+                    .route("/grants", post(grants::add_grant));
 
                 // Add in any testing support
                 if cfg!(feature = "testing") {
@@ -86,17 +79,10 @@ impl From<&Server> for Router<Arc<ComputeNode>> {
                 router
             }
             Server::External {
-                config,
-                compute_id,
-                instance_id,
-                ..
+                config, compute_id, ..
             } => {
-                let unauthenticated_router = Router::<Arc<ComputeNode>>::new()
-                    .route("/metrics", get(metrics::get_metrics))
-                    .route(
-                        "/autoscaling_metrics",
-                        get(metrics::get_autoscaling_metrics),
-                    );
+                let unauthenticated_router =
+                    Router::<Arc<ComputeNode>>::new().route("/metrics", get(metrics::get_metrics));
 
                 let authenticated_router = Router::<Arc<ComputeNode>>::new()
                     .route("/lfc/prewarm", get(lfc::prewarm_state).post(lfc::prewarm))
@@ -110,13 +96,8 @@ 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,7 +2,6 @@ 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};
 
@@ -120,7 +119,3 @@ 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,7 +4,6 @@
 #![deny(clippy::undocumented_unsafe_blocks)]
 
 pub mod checker;
-pub mod communicator_socket_client;
 pub mod config;
 pub mod configurator;
 pub mod http;
@@ -16,7 +15,6 @@ 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;
@@ -25,7 +23,6 @@ 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,10 +1,7 @@
 use std::collections::HashMap;
-use std::sync::{LazyLock, RwLock};
-use tracing::Subscriber;
 use tracing::info;
-use tracing_appender;
+use tracing_subscriber::layer::SubscriberExt;
 use tracing_subscriber::prelude::*;
-use tracing_subscriber::{fmt, layer::SubscriberExt, registry::LookupSpan};
 
 /// Initialize logging to stderr, and OpenTelemetry tracing and exporter.
 ///
@@ -16,63 +13,31 @@ use tracing_subscriber::{fmt, layer::SubscriberExt, registry::LookupSpan};
 /// set `OTEL_EXPORTER_OTLP_ENDPOINT=http://jaeger:4318`. See
 /// `tracing-utils` package description.
 ///
-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>,
-)> {
+pub async fn init_tracing_and_logging(default_log_level: &str) -> anyhow::Result<()> {
     // 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 provider =
-        tracing_utils::init_tracing("compute_ctl", tracing_utils::ExportConfig::default());
-    let otlp_layer = provider.as_ref().map(tracing_utils::layer);
+    let otlp_layer =
+        tracing_utils::init_tracing("compute_ctl", tracing_utils::ExportConfig::default()).await;
 
     // 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((provider, _file_logs_guard))
+    Ok(())
 }
 
 /// Replace all newline characters with a special character to make it
@@ -127,157 +92,3 @@ 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/migration.rs

@@ -9,20 +9,15 @@ 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], lakebase_mode: bool) -> Self {
+    pub fn new(client: &'m mut Client, migrations: &'m [&'m str]) -> 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,
-            lakebase_mode,
-        }
+        Self { client, migrations }
     }
 
     /// Get the current value neon_migration.migration_id
@@ -135,13 +130,8 @@ 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,7 +11,6 @@ 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.
@@ -353,47 +352,13 @@ 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");
-        if !compute.params.lakebase_mode {
-            state = compute.state_changed.wait(state).unwrap();
+        state = compute.state_changed.wait(state).unwrap();
 
-            if state.status == ComputeStatus::Running {
-                break;
-            }
-            continue;
+        if state.status == ComputeStatus::Running {
+            break;
         }
-
-        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,9 +11,7 @@ use std::time::{Duration, Instant};
 
 use anyhow::{Result, bail};
 use compute_api::responses::TlsConfig;
-use compute_api::spec::{
-    Database, DatabricksSettings, GenericOption, GenericOptions, PgIdent, Role,
-};
+use compute_api::spec::{Database, GenericOption, GenericOptions, PgIdent, Role};
 use futures::StreamExt;
 use indexmap::IndexMap;
 use ini::Ini;
@@ -186,42 +184,6 @@ 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

@@ -1,30 +0,0 @@
-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,6 +1,4 @@
 use std::fs::File;
-use std::fs::{self, Permissions};
-use std::os::unix::fs::PermissionsExt;
 use std::path::Path;
 
 use anyhow::{Result, anyhow, bail};
@@ -135,25 +133,10 @@ 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, databricks_pg_hba: Option<&String>) -> Result<()> {
+pub fn update_pg_hba(pgdata_path: &Path) -> 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/16/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 {
@@ -163,59 +146,6 @@ pub fn update_pg_hba(pgdata_path: &Path, databricks_pg_hba: Option<&String>) ->
     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
@@ -240,11 +170,7 @@ pub async fn handle_neon_extension_upgrade(client: &mut Client) -> Result<()> {
 }
 
 #[instrument(skip_all)]
-pub async fn handle_migrations(
-    params: ComputeNodeParams,
-    client: &mut Client,
-    lakebase_mode: bool,
-) -> Result<()> {
+pub async fn handle_migrations(params: ComputeNodeParams, client: &mut Client) -> Result<()> {
     info!("handle migrations");
 
     // !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
@@ -308,7 +234,7 @@ pub async fn handle_migrations(
         ),
     ];
 
-    MigrationRunner::new(client, &migrations, lakebase_mode)
+    MigrationRunner::new(client, &migrations)
         .run_migrations()
         .await?;
 

compute_tools/src/spec_apply.rs

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

control_plane/README.md

@@ -8,10 +8,10 @@ code changes locally, but not suitable for running production systems.
 
 ## Example: Start with Postgres 16
 
-To create and start a local development environment with Postgres 16, you will need to provide `--pg-version` flag to 2 of the start-up commands.
+To create and start a local development environment with Postgres 16, you will need to provide `--pg-version` flag to 3 of the start-up commands.
 
 ```shell
-cargo neon init
+cargo neon init --pg-version 16
 cargo neon start
 cargo neon tenant create --set-default --pg-version 16
 cargo neon endpoint create main --pg-version 16

control_plane/src/bin/neon_local.rs

@@ -407,12 +407,6 @@ 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)]
@@ -560,9 +554,7 @@ enum EndpointCmd {
     Create(EndpointCreateCmdArgs),
     Start(EndpointStartCmdArgs),
     Reconfigure(EndpointReconfigureCmdArgs),
-    RefreshConfiguration(EndpointRefreshConfigurationArgs),
     Stop(EndpointStopCmdArgs),
-    UpdatePageservers(EndpointUpdatePageserversCmdArgs),
     GenerateJwt(EndpointGenerateJwtCmdArgs),
 }
 
@@ -723,13 +715,6 @@ 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 {
@@ -747,16 +732,6 @@ 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 {
@@ -1536,7 +1511,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!("endpoint {endpoint_id} not found"))?;
+                .ok_or_else(|| anyhow::anyhow!("endpoint {endpoint_id} not found"))?;
 
             if !args.allow_multiple {
                 cplane.check_conflicting_endpoints(
@@ -1644,44 +1619,6 @@ 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
@@ -1735,14 +1672,6 @@ async fn handle_endpoint(subcmd: &EndpointCmd, env: &local_env::LocalEnv) -> Res
                 .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
@@ -1880,7 +1809,6 @@ 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

@@ -65,6 +65,7 @@ use jsonwebtoken::jwk::{
     OctetKeyPairParameters, OctetKeyPairType, PublicKeyUse,
 };
 use nix::sys::signal::{Signal, kill};
+use pageserver_api::shard::ShardStripeSize;
 use pem::Pem;
 use reqwest::header::CONTENT_TYPE;
 use safekeeper_api::PgMajorVersion;
@@ -76,7 +77,6 @@ use spki::{SubjectPublicKeyInfo, SubjectPublicKeyInfoRef};
 use tracing::debug;
 use url::Host;
 use utils::id::{NodeId, TenantId, TimelineId};
-use utils::shard::ShardStripeSize;
 
 use crate::local_env::LocalEnv;
 use crate::postgresql_conf::PostgresConf;
@@ -937,9 +937,7 @@ impl Endpoint {
                         | ComputeStatus::Configuration
                         | ComputeStatus::TerminationPendingFast
                         | ComputeStatus::TerminationPendingImmediate
-                        | ComputeStatus::Terminated
-                        | ComputeStatus::RefreshConfigurationPending
-                        | ComputeStatus::RefreshConfiguration => {
+                        | ComputeStatus::Terminated => {
                             bail!("unexpected compute status: {:?}", state.status)
                         }
                     }
@@ -962,29 +960,6 @@ 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();
@@ -1150,33 +1125,6 @@ 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://{}@{}:{}/{}",

control_plane/src/storage_controller.rs

@@ -56,7 +56,6 @@ 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 {
@@ -65,7 +64,6 @@ impl NeonStorageControllerStartArgs {
             instance_id: 1,
             base_port: None,
             start_timeout,
-            handle_ps_local_disk_loss: None,
         }
     }
 }
@@ -671,10 +669,6 @@ 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,7 +35,6 @@ 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::rng().random::<u32>();
+            let random = rand::thread_rng().r#gen::<u32>();
 
             let s3_config = remote_storage::S3Config {
                 bucket_name: var(REAL_S3_BUCKET).unwrap(),

libs/compute_api/src/responses.rs

@@ -108,10 +108,11 @@ pub enum PromoteState {
     Failed { error: String },
 }
 
-#[derive(Deserialize, Default, Debug)]
+#[derive(Deserialize, Serialize, Default, Debug, Clone)]
 #[serde(rename_all = "snake_case")]
-pub struct PromoteConfig {
-    pub spec: ComputeSpec,
+/// Result of /safekeepers_lsn
+pub struct SafekeepersLsn {
+    pub safekeepers: String,
     pub wal_flush_lsn: utils::lsn::Lsn,
 }
 
@@ -172,11 +173,6 @@ 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)]
@@ -189,10 +185,6 @@ 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

@@ -416,32 +416,6 @@ 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>>;

libs/consumption_metrics/src/lib.rs

@@ -90,7 +90,7 @@ impl<'a> IdempotencyKey<'a> {
         IdempotencyKey {
             now: Utc::now(),
             node_id,
-            nonce: rand::rng().random_range(0..=9999),
+            nonce: rand::thread_rng().gen_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().random_range(0..max)
+        self.internal.rng.lock().gen_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.random_bool(self.fail_prob) {
+        if rng.gen_bool(self.fail_prob) {
             return None;
         }
-        Some(rng.random_range(self.min..=self.max))
+        Some(rng.gen_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())
+        StdRng::from_rng(rng.deref_mut()).unwrap()
     }
 
     /// Create a new node.

libs/metrics/Cargo.toml

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

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

libs/neon-shmem/Cargo.toml

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

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: Some(val.timeout),
+            timeout: val.timeout,
         }
     }
 }

libs/pageserver_api/src/controller_api.rs

@@ -596,7 +596,6 @@ 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

@@ -981,12 +981,12 @@ mod tests {
         let mut rng = rand::rngs::StdRng::seed_from_u64(42);
 
         let key = Key {
-            field1: rng.random(),
-            field2: rng.random(),
-            field3: rng.random(),
-            field4: rng.random(),
-            field5: rng.random(),
-            field6: rng.random(),
+            field1: rng.r#gen(),
+            field2: rng.r#gen(),
+            field3: rng.r#gen(),
+            field4: rng.r#gen(),
+            field5: rng.r#gen(),
+            field6: rng.r#gen(),
         };
 
         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::distr::Alphanumeric;
+        use rand::distributions::Alphanumeric;
         Self(
-            rand::rng()
+            rand::thread_rng()
                 .sample_iter(&Alphanumeric)
                 .take(20)
                 .map(char::from)
@@ -1500,7 +1500,6 @@ 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)]
@@ -1534,10 +1533,10 @@ pub enum RelSizeMigration {
     /// `None` is the same as `Some(RelSizeMigration::Legacy)`.
     Legacy,
     /// The tenant is migrating to the new rel_size format. Both old and new rel_size format are
-    /// persisted in the storage. The read path will read both formats and validate them.
+    /// persisted in the index part. The read path will read both formats and merge them.
     Migrating,
     /// The tenant has migrated to the new rel_size format. Only the new rel_size format is persisted
-    /// in the storage, and the read path will not read the old format.
+    /// in the index part, and the read path will not read the old format.
     Migrated,
 }
 
@@ -1620,7 +1619,6 @@ 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,6 +69,22 @@ 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,14 +21,6 @@ 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/postgres_ffi/Cargo.toml

@@ -9,7 +9,6 @@ regex.workspace = true
 bytes.workspace = true
 anyhow.workspace = true
 crc32c.workspace = true
-criterion.workspace = true
 once_cell.workspace = true
 log.workspace = true
 memoffset.workspace = true
@@ -22,6 +21,7 @@ tracing.workspace = true
 postgres_versioninfo.workspace = true
 
 [dev-dependencies]
+criterion.workspace = true
 env_logger.workspace = true
 postgres.workspace = true
 

libs/pq_proto/src/lib.rs

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

libs/proxy/json/Cargo.toml

@@ -10,3 +10,8 @@ itoa = "1"
 
 [dev-dependencies]
 futures = "0.3"
+criterion.workspace = true
+
+[[bench]]
+name = "escape"
+harness = false

libs/proxy/json/benches/escape.rs

@@ -0,0 +1,53 @@
+use std::hint::black_box;
+
+use criterion::{Bencher, Criterion, criterion_group, criterion_main};
+
+#[derive(Debug)]
+#[allow(dead_code)]
+struct Foo {
+    some_field: Bar,
+    some_other_field: String,
+}
+
+#[derive(Debug)]
+#[allow(dead_code)]
+struct Bar {
+    nested_fields: String,
+    some_other_value: i32,
+}
+
+pub fn escape(c: &mut Criterion) {
+    c.bench_function("small", |b| bench_json_encode_inner(b, "object_key"));
+    c.bench_function("small_static", |b| {
+        bench_json_encode_inner(b, json::esc!("object_key"));
+    });
+    c.bench_function("large_fmt", |b| {
+        let value = Foo {
+            some_field: Bar {
+                nested_fields: "could not connect to database, control plane error \"foo bar\""
+                    .to_string()
+                    .to_string(),
+                some_other_value: -1,
+            },
+            some_other_field: "error".to_string(),
+        };
+
+        bench_json_encode_inner(b, format_args!("{:?}", &value));
+    });
+}
+
+criterion_group!(benches, escape);
+criterion_main!(benches);
+
+fn bench_json_encode_inner(b: &mut Bencher<'_>, v: impl json::ValueEncoder + Copy) {
+    let mut output = Vec::new();
+
+    // write it once so we don't alloc during the benchmark.
+    json::ValueSer::new(&mut output).value(black_box(v));
+
+    b.iter(|| {
+        output.clear();
+        json::ValueSer::new(&mut output).value(black_box(v));
+        black_box(&output[..]);
+    });
+}

libs/proxy/json/src/lib.rs

@@ -81,7 +81,8 @@ mod macros;
 mod str;
 mod value;
 
-pub use value::{Null, ValueEncoder};
+pub use str::EscapedStr;
+pub use value::{KeyEncoder, Null, ValueEncoder};
 
 #[must_use]
 /// Serialize a single json value.
@@ -164,7 +165,9 @@ impl<'buf> ObjectSer<'buf> {
     /// Start a new object entry with the given string key, returning a [`ValueSer`] for the associated value.
     #[inline]
     pub fn key(&mut self, key: impl KeyEncoder) -> ValueSer<'_> {
-        key.write_key(self)
+        // we create a psuedo value to write the key into.
+        let start = self.start;
+        self.entry_inner(|buf| key.encode(ValueSer { buf, start }))
     }
 
     /// Write an entry (key-value pair) to the object.
@@ -211,10 +214,6 @@ impl<'buf> ObjectSer<'buf> {
     }
 }
 
-pub trait KeyEncoder {
-    fn write_key<'a>(self, obj: &'a mut ObjectSer) -> ValueSer<'a>;
-}
-
 #[must_use]
 /// Serialize a json object.
 pub struct ListSer<'buf> {
@@ -279,14 +278,14 @@ impl<'buf> ListSer<'buf> {
 
 #[cfg(test)]
 mod tests {
-    use crate::{Null, ValueSer};
+    use crate::{Null, ValueSer, esc};
 
     #[test]
     fn object() {
         let mut buf = vec![];
         let mut object = ValueSer::new(&mut buf).object();
-        object.entry("foo", "bar");
-        object.entry("baz", Null);
+        object.entry(esc!("foo"), "bar");
+        object.entry(esc!("baz"), Null);
         object.finish();
 
         assert_eq!(buf, br#"{"foo":"bar","baz":null}"#);
@@ -307,8 +306,8 @@ mod tests {
     fn object_macro() {
         let res = crate::value_to_string!(|obj| {
             crate::value_as_object!(|obj| {
-                obj.entry("foo", "bar");
-                obj.entry("baz", Null);
+                obj.entry(esc!("foo"), "bar");
+                obj.entry(esc!("baz"), Null);
             })
         });
 
@@ -364,7 +363,7 @@ mod tests {
                 let entry = obj.key("2");
                 let entry = {
                     let mut nested_obj = entry.object();
-                    nested_obj.entry("foo", "bar");
+                    nested_obj.entry(esc!("foo"), "bar");
                     nested_obj.rollback()
                 };
 

libs/proxy/json/src/macros.rs

@@ -84,3 +84,11 @@ macro_rules! value_as_list {
         res
     }};
 }
+
+/// A helper macro that ensures the provided string literal does not need escaping.
+#[macro_export]
+macro_rules! esc {
+    ($str:literal) => {
+        const { $crate::EscapedStr::from_static($str) }
+    };
+}

libs/proxy/json/src/str.rs

@@ -10,58 +10,98 @@
 
 use std::fmt::{self, Write};
 
-/// Represents a character escape code in a type-safe manner.
-pub enum CharEscape {
-    /// An escaped quote `"`
-    Quote,
-    /// An escaped reverse solidus `\`
-    ReverseSolidus,
-    // /// An escaped solidus `/`
-    // Solidus,
-    /// An escaped backspace character (usually escaped as `\b`)
-    Backspace,
-    /// An escaped form feed character (usually escaped as `\f`)
-    FormFeed,
-    /// An escaped line feed character (usually escaped as `\n`)
-    LineFeed,
-    /// An escaped carriage return character (usually escaped as `\r`)
-    CarriageReturn,
-    /// An escaped tab character (usually escaped as `\t`)
-    Tab,
-    /// An escaped ASCII plane control character (usually escaped as
-    /// `\u00XX` where `XX` are two hex characters)
-    AsciiControl(u8),
-}
+use crate::{KeyEncoder, ValueEncoder, ValueSer};
 
-impl CharEscape {
-    #[inline]
-    fn from_escape_table(escape: u8, byte: u8) -> CharEscape {
-        match escape {
-            self::BB => CharEscape::Backspace,
-            self::TT => CharEscape::Tab,
-            self::NN => CharEscape::LineFeed,
-            self::FF => CharEscape::FormFeed,
-            self::RR => CharEscape::CarriageReturn,
-            self::QU => CharEscape::Quote,
-            self::BS => CharEscape::ReverseSolidus,
-            self::UU => CharEscape::AsciiControl(byte),
-            _ => unreachable!(),
+#[repr(transparent)]
+pub struct EscapedStr([u8]);
+
+impl EscapedStr {
+    /// Assumes the string does not need escaping.
+    ///
+    /// # Panics
+    ///
+    /// This will panic if the string does need escaping.
+    #[inline(always)]
+    pub const fn from_static(s: &'static str) -> &'static Self {
+        let bytes = s.as_bytes();
+
+        let mut i = 0;
+        while i < bytes.len() {
+            let byte = bytes[i];
+
+            if byte < 0x20 || byte == b'"' || byte == b'\\' {
+                panic!("the string needs escaping");
+            }
+
+            i += 1;
         }
+
+        // safety: this EscapedStr is transparent over [u8].
+        unsafe { std::mem::transmute::<&[u8], &EscapedStr>(bytes) }
+    }
+
+    /// Escapes the string eagerly.
+    pub fn escape(s: &str) -> Box<Self> {
+        let mut writer = Vec::with_capacity(s.len());
+
+        Collect { buf: &mut writer }
+            .write_str(s)
+            .expect("formatting should not error");
+
+        let bytes = writer.into_boxed_slice();
+
+        // safety: this EscapedStr is transparent over [u8].
+        unsafe { std::mem::transmute::<Box<[u8]>, Box<EscapedStr>>(bytes) }
     }
 }
 
-pub(crate) fn format_escaped_str(writer: &mut Vec<u8>, value: &str) {
+impl KeyEncoder for &EscapedStr {}
+impl ValueEncoder for &EscapedStr {
+    fn encode(self, v: crate::ValueSer<'_>) {
+        let buf = &mut *v.buf;
+        buf.reserve(2 + self.0.len());
+
+        buf.push(b'"');
+        buf.extend_from_slice(&self.0);
+        buf.push(b'"');
+
+        v.finish();
+    }
+}
+
+impl KeyEncoder for &str {}
+impl ValueEncoder for &str {
+    #[inline]
+    fn encode(self, v: ValueSer<'_>) {
+        format_escaped_str(v.buf, self);
+        v.finish();
+    }
+}
+
+impl KeyEncoder for fmt::Arguments<'_> {}
+impl ValueEncoder for fmt::Arguments<'_> {
+    #[inline]
+    fn encode(self, v: ValueSer<'_>) {
+        if let Some(s) = self.as_str() {
+            format_escaped_str(v.buf, s);
+        } else {
+            format_escaped_fmt(v.buf, self);
+        }
+        v.finish();
+    }
+}
+
+fn format_escaped_str(writer: &mut Vec<u8>, value: &str) {
     writer.reserve(2 + value.len());
 
     writer.push(b'"');
 
-    let rest = format_escaped_str_contents(writer, value);
-    writer.extend_from_slice(rest);
+    format_escaped_str_contents(writer, value);
 
     writer.push(b'"');
 }
 
-pub(crate) fn format_escaped_fmt(writer: &mut Vec<u8>, args: fmt::Arguments) {
+fn format_escaped_fmt(writer: &mut Vec<u8>, args: fmt::Arguments) {
     writer.push(b'"');
 
     Collect { buf: writer }
@@ -77,33 +117,36 @@ struct Collect<'buf> {
 
 impl fmt::Write for Collect<'_> {
     fn write_str(&mut self, s: &str) -> fmt::Result {
-        let last = format_escaped_str_contents(self.buf, s);
-        self.buf.extend(last);
+        format_escaped_str_contents(self.buf, s);
         Ok(())
     }
 }
 
 // writes any escape sequences, and returns the suffix still needed to be written.
-fn format_escaped_str_contents<'a>(writer: &mut Vec<u8>, value: &'a str) -> &'a [u8] {
-    let bytes = value.as_bytes();
+fn format_escaped_str_contents(writer: &mut Vec<u8>, value: &str) {
+    let mut bytes = value.as_bytes();
 
-    let mut start = 0;
-
-    for (i, &byte) in bytes.iter().enumerate() {
+    let mut i = 0;
+    while i < bytes.len() {
+        let byte = bytes[i];
         let escape = ESCAPE[byte as usize];
+
+        i += 1;
         if escape == 0 {
             continue;
         }
 
-        writer.extend_from_slice(&bytes[start..i]);
+        // hitting an escape character is unlikely.
+        cold();
 
-        let char_escape = CharEscape::from_escape_table(escape, byte);
-        write_char_escape(writer, char_escape);
+        let string_run;
+        (string_run, bytes) = bytes.split_at(i);
+        i = 0;
 
-        start = i + 1;
+        write_char_escape(writer, string_run);
     }
 
-    &bytes[start..]
+    writer.extend_from_slice(bytes);
 }
 
 const BB: u8 = b'b'; // \x08
@@ -138,29 +181,38 @@ static ESCAPE: [u8; 256] = [
     __, __, __, __, __, __, __, __, __, __, __, __, __, __, __, __, // F
 ];
 
-fn write_char_escape(writer: &mut Vec<u8>, char_escape: CharEscape) {
-    let s = match char_escape {
-        CharEscape::Quote => b"\\\"",
-        CharEscape::ReverseSolidus => b"\\\\",
-        // CharEscape::Solidus => b"\\/",
-        CharEscape::Backspace => b"\\b",
-        CharEscape::FormFeed => b"\\f",
-        CharEscape::LineFeed => b"\\n",
-        CharEscape::CarriageReturn => b"\\r",
-        CharEscape::Tab => b"\\t",
-        CharEscape::AsciiControl(byte) => {
-            static HEX_DIGITS: [u8; 16] = *b"0123456789abcdef";
-            let bytes = &[
-                b'\\',
-                b'u',
-                b'0',
-                b'0',
-                HEX_DIGITS[(byte >> 4) as usize],
-                HEX_DIGITS[(byte & 0xF) as usize],
-            ];
-            return writer.extend_from_slice(bytes);
-        }
-    };
+#[cold]
+fn cold() {}
 
+fn write_char_escape(writer: &mut Vec<u8>, bytes: &[u8]) {
+    debug_assert!(
+        !bytes.is_empty(),
+        "caller guarantees that bytes is non empty"
+    );
+
+    let (&byte, string_run) = bytes.split_last().unwrap_or((&0, b""));
+
+    let escape = ESCAPE[byte as usize];
+    debug_assert_ne!(escape, 0, "caller guarantees that escape will be non-zero");
+
+    // the escape char from the escape table is the correct replacement
+    // character.
+    let mut bytes = [b'\\', escape, b'0', b'0', b'0', b'0'];
+    let mut s = &bytes[0..2];
+
+    // if the replacement character is 'u', then we need
+    // to write the unicode encoding
+    if escape == UU {
+        static HEX_DIGITS: [u8; 16] = *b"0123456789abcdef";
+
+        // we rarely encounter characters that must be escaped as unicode.
+        cold();
+
+        bytes[4] = HEX_DIGITS[(byte >> 4) as usize];
+        bytes[5] = HEX_DIGITS[(byte & 0xF) as usize];
+        s = &bytes;
+    }
+
+    writer.extend_from_slice(string_run);
     writer.extend_from_slice(s);
 }

libs/proxy/json/src/value.rs

@@ -1,11 +1,12 @@
-use core::fmt;
 use std::collections::{BTreeMap, HashMap};
 
-use crate::str::{format_escaped_fmt, format_escaped_str};
-use crate::{KeyEncoder, ObjectSer, ValueSer, value_as_list, value_as_object};
+use crate::{ValueSer, value_as_list, value_as_object};
+
+/// Marker trait for values that are valid keys
+pub trait KeyEncoder: ValueEncoder {}
 
 /// Write a value to the underlying json representation.
-pub trait ValueEncoder {
+pub trait ValueEncoder: Sized {
     fn encode(self, v: ValueSer<'_>);
 }
 
@@ -24,23 +25,11 @@ impl<T: Copy + ValueEncoder> ValueEncoder for &T {
     }
 }
 
-impl ValueEncoder for &str {
+impl KeyEncoder for String {}
+impl ValueEncoder for String {
     #[inline]
     fn encode(self, v: ValueSer<'_>) {
-        format_escaped_str(v.buf, self);
-        v.finish();
-    }
-}
-
-impl ValueEncoder for fmt::Arguments<'_> {
-    #[inline]
-    fn encode(self, v: ValueSer<'_>) {
-        if let Some(s) = self.as_str() {
-            format_escaped_str(v.buf, s);
-        } else {
-            format_escaped_fmt(v.buf, self);
-        }
-        v.finish();
+        self.as_str().encode(v);
     }
 }
 
@@ -94,26 +83,8 @@ impl<T: ValueEncoder> ValueEncoder for Option<T> {
     }
 }
 
-impl KeyEncoder for &str {
-    #[inline]
-    fn write_key<'a>(self, obj: &'a mut ObjectSer) -> ValueSer<'a> {
-        let obj = &mut *obj;
-        obj.entry_inner(|b| format_escaped_str(b, self))
-    }
-}
-
-impl KeyEncoder for fmt::Arguments<'_> {
-    #[inline]
-    fn write_key<'a>(self, obj: &'a mut ObjectSer) -> ValueSer<'a> {
-        if let Some(key) = self.as_str() {
-            obj.entry_inner(|b| format_escaped_str(b, key))
-        } else {
-            obj.entry_inner(|b| format_escaped_fmt(b, self))
-        }
-    }
-}
-
 /// Represents the JSON null value.
+#[derive(Clone, Copy)]
 pub struct Null;
 
 impl ValueEncoder for Null {

libs/proxy/json/tests/test.rs

@@ -0,0 +1,180 @@
+use json::Null;
+use json::ValueEncoder;
+use std::collections::BTreeMap;
+
+macro_rules! treemap {
+    () => {
+        BTreeMap::new()
+    };
+    ($($k:expr => $v:expr),+ $(,)?) => {
+        {
+            let mut m = BTreeMap::new();
+            $(
+                m.insert($k, $v);
+            )+
+            m
+        }
+    };
+}
+
+fn test_encode_ok<'a, T>(errors: impl IntoIterator<Item = (T, &'a str)>)
+where
+    T: ValueEncoder,
+{
+    for (value, out) in errors {
+        let s = json::value_to_string!(|v| value.encode(v));
+        assert_eq!(&*s, out);
+    }
+}
+
+#[test]
+fn test_write_null() {
+    let tests = [(Null, "null")];
+    test_encode_ok(tests);
+}
+
+#[test]
+fn test_write_u64() {
+    let tests = [(3u64, "3"), (u64::MAX, &u64::MAX.to_string())];
+    test_encode_ok(tests);
+}
+
+#[test]
+fn test_write_i64() {
+    let tests = [
+        (3i64, "3"),
+        (-2i64, "-2"),
+        (-1234i64, "-1234"),
+        (i64::MIN, &i64::MIN.to_string()),
+    ];
+    test_encode_ok(tests);
+}
+
+#[test]
+fn test_write_f64() {
+    let tests = [
+        (3.0, "3.0"),
+        (3.1, "3.1"),
+        (-1.5, "-1.5"),
+        (0.5, "0.5"),
+        (f64::MIN, "-1.7976931348623157e308"),
+        (f64::MAX, "1.7976931348623157e308"),
+        (f64::EPSILON, "2.220446049250313e-16"),
+    ];
+    test_encode_ok(tests);
+}
+
+#[test]
+fn test_write_str() {
+    let tests = [
+        // normal
+        ("", "\"\""),
+        ("foo", "\"foo\""),
+        // ascii escape chars.
+        ("\"", "\"\\\"\""),
+        ("\x08", "\"\\b\""),
+        ("\n", "\"\\n\""),
+        ("\r", "\"\\r\""),
+        ("\t", "\"\\t\""),
+        ("\x07", "\"\\u0007\""),
+        // unicode not escaped.
+        ("\u{12ab}", "\"\u{12ab}\""),
+        ("\u{AB12}", "\"\u{AB12}\""),
+        ("\u{1F395}", "\"\u{1F395}\""),
+    ];
+    test_encode_ok(tests);
+}
+
+#[test]
+fn test_write_bool() {
+    let tests = [(true, "true"), (false, "false")];
+    test_encode_ok(tests);
+}
+
+#[test]
+fn test_write_list() {
+    test_encode_ok([
+        (vec![], "[]"),
+        (vec![true], "[true]"),
+        (vec![true, false], "[true,false]"),
+    ]);
+
+    test_encode_ok([
+        (vec![vec![], vec![], vec![]], "[[],[],[]]"),
+        (vec![vec![1, 2, 3], vec![], vec![]], "[[1,2,3],[],[]]"),
+        (vec![vec![], vec![1, 2, 3], vec![]], "[[],[1,2,3],[]]"),
+        (vec![vec![], vec![], vec![1, 2, 3]], "[[],[],[1,2,3]]"),
+    ]);
+}
+
+#[test]
+fn test_write_object() {
+    test_encode_ok([
+        (treemap!(), "{}"),
+        (treemap!("a".to_owned() => true), "{\"a\":true}"),
+        (
+            treemap!(
+                "a".to_owned() => true,
+                "b".to_owned() => false,
+            ),
+            "{\"a\":true,\"b\":false}",
+        ),
+    ]);
+
+    test_encode_ok([
+        (
+            treemap![
+                "a".to_owned() => treemap![],
+                "b".to_owned() => treemap![],
+                "c".to_owned() => treemap![],
+            ],
+            "{\"a\":{},\"b\":{},\"c\":{}}",
+        ),
+        (
+            treemap![
+                "a".to_owned() => treemap![
+                    "a".to_owned() => treemap!["a" => vec![1,2,3]],
+                    "b".to_owned() => treemap![],
+                    "c".to_owned() => treemap![],
+                ],
+                "b".to_owned() => treemap![],
+                "c".to_owned() => treemap![],
+            ],
+            "{\"a\":{\"a\":{\"a\":[1,2,3]},\"b\":{},\"c\":{}},\"b\":{},\"c\":{}}",
+        ),
+        (
+            treemap![
+                "a".to_owned() => treemap![],
+                "b".to_owned() => treemap![
+                    "a".to_owned() => treemap!["a" => vec![1,2,3]],
+                    "b".to_owned() => treemap![],
+                    "c".to_owned() => treemap![],
+                ],
+                "c".to_owned() => treemap![],
+            ],
+            "{\"a\":{},\"b\":{\"a\":{\"a\":[1,2,3]},\"b\":{},\"c\":{}},\"c\":{}}",
+        ),
+        (
+            treemap![
+                "a".to_owned() => treemap![],
+                "b".to_owned() => treemap![],
+                "c".to_owned() => treemap![
+                    "a".to_owned() => treemap!["a" => vec![1,2,3]],
+                    "b".to_owned() => treemap![],
+                    "c".to_owned() => treemap![],
+                ],
+            ],
+            "{\"a\":{},\"b\":{},\"c\":{\"a\":{\"a\":[1,2,3]},\"b\":{},\"c\":{}}}",
+        ),
+    ]);
+}
+
+#[test]
+fn test_write_option() {
+    test_encode_ok([(None, "null"), (Some("jodhpurs"), "\"jodhpurs\"")]);
+
+    test_encode_ok([
+        (None, "null"),
+        (Some(vec!["foo", "bar"]), "[\"foo\",\"bar\"]"),
+    ]);
+}

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::rng();
+    let mut rng = rand::thread_rng();
     (0..NONCE_LENGTH)
         .map(|_| {
-            let mut v = rng.random_range(0x21u8..0x7e);
+            let mut v = rng.gen_range(0x21u8..0x7e);
             if v == 0x2c {
                 v = 0x7e
             }

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

@@ -74,6 +74,7 @@ impl Header {
 }
 
 /// An enum representing Postgres backend messages.
+#[non_exhaustive]
 pub enum Message {
     AuthenticationCleartextPassword,
     AuthenticationGss,
@@ -144,7 +145,16 @@ impl Message {
             PARSE_COMPLETE_TAG => Message::ParseComplete,
             BIND_COMPLETE_TAG => Message::BindComplete,
             CLOSE_COMPLETE_TAG => Message::CloseComplete,
-            NOTIFICATION_RESPONSE_TAG => Message::NotificationResponse(NotificationResponseBody {}),
+            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,
+                })
+            }
             COPY_DONE_TAG => Message::CopyDone,
             COMMAND_COMPLETE_TAG => {
                 let tag = buf.read_cstr()?;
@@ -533,7 +543,28 @@ impl NoticeResponseBody {
     }
 }
 
-pub struct NotificationResponseBody {}
+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 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::rng();
+    let mut rng = rand::thread_rng();
     rng.fill_bytes(&mut salt);
     scram_sha_256_salt(password, salt).await
 }

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

@@ -13,9 +13,8 @@ use serde::{Deserialize, Serialize};
 use tokio::sync::mpsc;
 
 use crate::cancel_token::RawCancelToken;
-use crate::codec::{BackendMessages, FrontendMessage, RecordNotices};
+use crate::codec::{BackendMessages, FrontendMessage};
 use crate::config::{Host, SslMode};
-use crate::connection::gc_bytesmut;
 use crate::query::RowStream;
 use crate::simple_query::SimpleQueryStream;
 use crate::types::{Oid, Type};
@@ -96,13 +95,20 @@ 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();
+        let buf = self.buffer.split().freeze();
         self.send_message(FrontendMessage::Raw(buf))
     }
 
@@ -119,7 +125,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();
+            let buf = client.buffer.split().freeze();
             let _ = client.send_message(FrontendMessage::Raw(buf));
         }
     }
@@ -135,7 +141,7 @@ impl<'a> PartialQuery<'a> {
         client.buffer.clear();
         f(&mut client.buffer)?;
         frontend::flush(&mut client.buffer);
-        let buf = client.buffer.split();
+        let buf = client.buffer.split().freeze();
         client.send_message(FrontendMessage::Raw(buf))
     }
 
@@ -148,7 +154,7 @@ impl<'a> PartialQuery<'a> {
         client.buffer.clear();
         f(&mut client.buffer)?;
         frontend::sync(&mut client.buffer);
-        let buf = client.buffer.split();
+        let buf = client.buffer.split().freeze();
         let _ = client.send_message(FrontendMessage::Raw(buf));
 
         Ok(&mut self.0.take().unwrap().responses)
@@ -215,18 +221,6 @@ 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>(
@@ -286,35 +280,8 @@ impl Client {
         simple_query::batch_execute(self.inner_mut(), query).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(())
+    pub async fn discard_all(&mut self) -> Result<ReadyForQueryStatus, Error> {
+        self.batch_execute("discard all").await
     }
 
     /// Begins a new database transaction.

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

@@ -1,26 +1,16 @@
 use std::io;
 
-use bytes::BytesMut;
+use bytes::{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(BytesMut),
-    RecordNotices(RecordNotices),
-}
-
-pub struct RecordNotices {
-    pub sender: UnboundedSender<Box<str>>,
-    pub limit: usize,
+    Raw(Bytes),
 }
 
 pub enum BackendMessage {
-    Normal {
-        messages: BackendMessages,
-        ready: bool,
-    },
+    Normal { messages: BackendMessages },
     Async(backend::Message),
 }
 
@@ -43,18 +33,14 @@ impl FallibleIterator for BackendMessages {
 
 pub struct PostgresCodec;
 
-impl Encoder<BytesMut> for PostgresCodec {
+impl Encoder<FrontendMessage> for PostgresCodec {
     type Error = io::Error;
 
-    fn encode(&mut self, item: BytesMut, dst: &mut BytesMut) -> io::Result<()> {
-        // When it comes to request/response workflows, we usually flush the entire write
-        // buffer in order to wait for the response before we send a new request.
-        // Therefore we can avoid the copy and just replace the buffer.
-        if dst.is_empty() {
-            *dst = item;
-        } else {
-            dst.extend_from_slice(&item);
+    fn encode(&mut self, item: FrontendMessage, dst: &mut BytesMut) -> io::Result<()> {
+        match item {
+            FrontendMessage::Raw(buf) => dst.extend_from_slice(&buf),
         }
+
         Ok(())
     }
 }
@@ -66,7 +52,6 @@ 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 {
@@ -90,7 +75,6 @@ impl Decoder for PostgresCodec {
             idx += len;
 
             if header.tag() == backend::READY_FOR_QUERY_TAG {
-                ready = true;
                 break;
             }
         }
@@ -100,7 +84,6 @@ impl Decoder for PostgresCodec {
         } else {
             Ok(Some(BackendMessage::Normal {
                 messages: BackendMessages(src.split_to(idx)),
-                ready,
             }))
         }
     }

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

@@ -11,8 +11,9 @@ use tokio::io::{AsyncRead, AsyncWrite};
 use tokio::net::TcpStream;
 
 use crate::connect::connect;
-use crate::connect_raw::{self, StartupStream};
+use crate::connect_raw::{RawConnection, connect_raw};
 use crate::connect_tls::connect_tls;
+use crate::maybe_tls_stream::MaybeTlsStream;
 use crate::tls::{MakeTlsConnect, TlsConnect, TlsStream};
 use crate::{Client, Connection, Error};
 
@@ -243,27 +244,24 @@ impl Config {
         &self,
         stream: S,
         tls: T,
-    ) -> Result<StartupStream<S, T::Stream>, Error>
+    ) -> Result<RawConnection<S, T::Stream>, Error>
     where
         S: AsyncRead + AsyncWrite + Unpin,
         T: TlsConnect<S>,
     {
         let stream = connect_tls(stream, self.ssl_mode, tls).await?;
-        let mut stream = StartupStream::new(stream);
-        connect_raw::authenticate(&mut stream, self).await?;
-
-        Ok(stream)
+        connect_raw(stream, self).await
     }
 
-    pub fn authenticate<S, T>(
+    pub async fn authenticate<S, T>(
         &self,
-        stream: &mut StartupStream<S, T>,
-    ) -> impl Future<Output = Result<(), Error>>
+        stream: MaybeTlsStream<S, T>,
+    ) -> Result<RawConnection<S, T>, Error>
     where
         S: AsyncRead + AsyncWrite + Unpin,
         T: TlsStream + Unpin,
     {
-        connect_raw::authenticate(stream, self)
+        connect_raw(stream, self).await
     }
 }
 

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

@@ -1,17 +1,17 @@
 use std::net::IpAddr;
 
-use futures_util::TryStreamExt;
 use postgres_protocol2::message::backend::Message;
-use tokio::io::{AsyncRead, AsyncWrite};
 use tokio::net::TcpStream;
 use tokio::sync::mpsc;
 
 use crate::client::SocketConfig;
-use crate::config::{Host, SslMode};
-use crate::connect_raw::StartupStream;
+use crate::codec::BackendMessage;
+use crate::config::Host;
+use crate::connect_raw::connect_raw;
 use crate::connect_socket::connect_socket;
+use crate::connect_tls::connect_tls;
 use crate::tls::{MakeTlsConnect, TlsConnect};
-use crate::{Client, Config, Connection, Error};
+use crate::{Client, Config, Connection, Error, RawConnection};
 
 pub async fn connect<T>(
     tls: &T,
@@ -45,36 +45,20 @@ where
     T: TlsConnect<TcpStream>,
 {
     let socket = connect_socket(host_addr, host, port, config.connect_timeout).await?;
-    let stream = config.tls_and_authenticate(socket, tls).await?;
-    managed(
+    let stream = connect_tls(socket, config.ssl_mode, tls).await?;
+    let RawConnection {
         stream,
-        host_addr,
-        host.clone(),
-        port,
-        config.ssl_mode,
-        config.connect_timeout,
-    )
-    .await
-}
-
-pub async fn managed<TlsStream>(
-    mut stream: StartupStream<TcpStream, TlsStream>,
-    host_addr: Option<IpAddr>,
-    host: Host,
-    port: u16,
-    ssl_mode: SslMode,
-    connect_timeout: Option<std::time::Duration>,
-) -> Result<(Client, Connection<TcpStream, TlsStream>), Error>
-where
-    TlsStream: AsyncRead + AsyncWrite + Unpin,
-{
-    let (process_id, secret_key) = wait_until_ready(&mut stream).await?;
+        parameters,
+        delayed_notice,
+        process_id,
+        secret_key,
+    } = connect_raw(stream, config).await?;
 
     let socket_config = SocketConfig {
         host_addr,
-        host,
+        host: host.clone(),
         port,
-        connect_timeout,
+        connect_timeout: config.connect_timeout,
     };
 
     let (client_tx, conn_rx) = mpsc::unbounded_channel();
@@ -83,37 +67,18 @@ where
         client_tx,
         client_rx,
         socket_config,
-        ssl_mode,
+        config.ssl_mode,
         process_id,
         secret_key,
     );
 
-    let stream = stream.into_framed();
-    let connection = Connection::new(stream, conn_tx, conn_rx);
+    // delayed notices are always sent as "Async" messages.
+    let delayed = delayed_notice
+        .into_iter()
+        .map(|m| BackendMessage::Async(Message::NoticeResponse(m)))
+        .collect();
+
+    let connection = Connection::new(stream, delayed, parameters, conn_tx, conn_rx);
 
     Ok((client, connection))
 }
-
-async fn wait_until_ready<S, T>(stream: &mut StartupStream<S, T>) -> Result<(i32, i32), Error>
-where
-    S: AsyncRead + AsyncWrite + Unpin,
-    T: AsyncRead + AsyncWrite + Unpin,
-{
-    let mut process_id = 0;
-    let mut secret_key = 0;
-
-    loop {
-        match stream.try_next().await.map_err(Error::io)? {
-            Some(Message::BackendKeyData(body)) => {
-                process_id = body.process_id();
-                secret_key = body.secret_key();
-            }
-            // These values are currently not used by `Client`/`Connection`. Ignore them.
-            Some(Message::ParameterStatus(_)) | Some(Message::NoticeResponse(_)) => {}
-            Some(Message::ReadyForQuery(_)) => return Ok((process_id, secret_key)),
-            Some(Message::ErrorResponse(body)) => return Err(Error::db(body)),
-            Some(_) => return Err(Error::unexpected_message()),
-            None => return Err(Error::closed()),
-        }
-    }
-}

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

@@ -1,27 +1,52 @@
+use std::collections::HashMap;
 use std::io;
 use std::pin::Pin;
-use std::task::{Context, Poll, ready};
+use std::task::{Context, Poll};
 
 use bytes::BytesMut;
 use fallible_iterator::FallibleIterator;
-use futures_util::{SinkExt, Stream, TryStreamExt};
+use futures_util::{Sink, SinkExt, Stream, TryStreamExt, ready};
 use postgres_protocol2::authentication::sasl;
 use postgres_protocol2::authentication::sasl::ScramSha256;
-use postgres_protocol2::message::backend::{AuthenticationSaslBody, Message};
+use postgres_protocol2::message::backend::{AuthenticationSaslBody, Message, NoticeResponseBody};
 use postgres_protocol2::message::frontend;
-use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
-use tokio_util::codec::{Framed, FramedParts};
+use tokio::io::{AsyncRead, AsyncWrite};
+use tokio_util::codec::Framed;
 
 use crate::Error;
-use crate::codec::PostgresCodec;
+use crate::codec::{BackendMessage, BackendMessages, FrontendMessage, PostgresCodec};
 use crate::config::{self, AuthKeys, Config};
-use crate::connection::{GC_THRESHOLD, INITIAL_CAPACITY};
 use crate::maybe_tls_stream::MaybeTlsStream;
 use crate::tls::TlsStream;
 
 pub struct StartupStream<S, T> {
     inner: Framed<MaybeTlsStream<S, T>, PostgresCodec>,
-    read_buf: BytesMut,
+    buf: BackendMessages,
+    delayed_notice: Vec<NoticeResponseBody>,
+}
+
+impl<S, T> Sink<FrontendMessage> for StartupStream<S, T>
+where
+    S: AsyncRead + AsyncWrite + Unpin,
+    T: AsyncRead + AsyncWrite + Unpin,
+{
+    type Error = io::Error;
+
+    fn poll_ready(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        Pin::new(&mut self.inner).poll_ready(cx)
+    }
+
+    fn start_send(mut self: Pin<&mut Self>, item: FrontendMessage) -> io::Result<()> {
+        Pin::new(&mut self.inner).start_send(item)
+    }
+
+    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        Pin::new(&mut self.inner).poll_flush(cx)
+    }
+
+    fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        Pin::new(&mut self.inner).poll_close(cx)
+    }
 }
 
 impl<S, T> Stream for StartupStream<S, T>
@@ -31,109 +56,81 @@ where
 {
     type Item = io::Result<Message>;
 
-    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
-        // We don't use `self.inner.poll_next()` as that might over-read into the read buffer.
-
-        // read 1 byte tag, 4 bytes length.
-        let header = ready!(self.as_mut().poll_fill_buf_exact(cx, 5)?);
-
-        let len = u32::from_be_bytes(header[1..5].try_into().unwrap());
-        if len < 4 {
-            return Poll::Ready(Some(Err(std::io::Error::other(
-                "postgres message too small",
-            ))));
-        }
-        if len >= 65536 {
-            return Poll::Ready(Some(Err(std::io::Error::other(
-                "postgres message too large",
-            ))));
-        }
-
-        // the tag is an additional byte.
-        let _message = ready!(self.as_mut().poll_fill_buf_exact(cx, len as usize + 1)?);
-
-        // Message::parse will remove the all the bytes from the buffer.
-        Poll::Ready(Message::parse(&mut self.read_buf).transpose())
-    }
-}
-
-impl<S, T> StartupStream<S, T>
-where
-    S: AsyncRead + AsyncWrite + Unpin,
-    T: AsyncRead + AsyncWrite + Unpin,
-{
-    /// Fill the buffer until it's the exact length provided. No additional data will be read from the socket.
-    ///
-    /// If the current buffer length is greater, nothing happens.
-    fn poll_fill_buf_exact(
-        self: Pin<&mut Self>,
+    fn poll_next(
+        mut self: Pin<&mut Self>,
         cx: &mut Context<'_>,
-        len: usize,
-    ) -> Poll<Result<&[u8], std::io::Error>> {
-        let this = self.get_mut();
-        let mut stream = Pin::new(this.inner.get_mut());
-
-        let mut n = this.read_buf.len();
-        while n < len {
-            this.read_buf.resize(len, 0);
-
-            let mut buf = ReadBuf::new(&mut this.read_buf[..]);
-            buf.set_filled(n);
-
-            if stream.as_mut().poll_read(cx, &mut buf)?.is_pending() {
-                this.read_buf.truncate(n);
-                return Poll::Pending;
+    ) -> Poll<Option<io::Result<Message>>> {
+        loop {
+            match self.buf.next() {
+                Ok(Some(message)) => return Poll::Ready(Some(Ok(message))),
+                Ok(None) => {}
+                Err(e) => return Poll::Ready(Some(Err(e))),
             }
 
-            if buf.filled().len() == n {
-                return Poll::Ready(Err(std::io::Error::new(
-                    std::io::ErrorKind::UnexpectedEof,
-                    "early eof",
-                )));
+            match ready!(Pin::new(&mut self.inner).poll_next(cx)) {
+                Some(Ok(BackendMessage::Normal { messages, .. })) => self.buf = messages,
+                Some(Ok(BackendMessage::Async(message))) => return Poll::Ready(Some(Ok(message))),
+                Some(Err(e)) => return Poll::Ready(Some(Err(e))),
+                None => return Poll::Ready(None),
             }
-            n = buf.filled().len();
-
-            this.read_buf.truncate(n);
-        }
-
-        Poll::Ready(Ok(&this.read_buf[..len]))
-    }
-
-    pub fn into_framed(mut self) -> Framed<MaybeTlsStream<S, T>, PostgresCodec> {
-        *self.inner.read_buffer_mut() = self.read_buf;
-        self.inner
-    }
-
-    pub fn new(io: MaybeTlsStream<S, T>) -> Self {
-        let mut parts = FramedParts::new(io, PostgresCodec);
-        parts.write_buf = BytesMut::with_capacity(INITIAL_CAPACITY);
-
-        let mut inner = Framed::from_parts(parts);
-
-        // This is the default already, but nice to be explicit.
-        // We divide by two because writes will overshoot the boundary.
-        // We don't want constant overshoots to cause us to constantly re-shrink the buffer.
-        inner.set_backpressure_boundary(GC_THRESHOLD / 2);
-
-        Self {
-            inner,
-            read_buf: BytesMut::with_capacity(INITIAL_CAPACITY),
         }
     }
 }
 
-pub(crate) async fn authenticate<S, T>(
-    stream: &mut StartupStream<S, T>,
+pub struct RawConnection<S, T> {
+    pub stream: Framed<MaybeTlsStream<S, T>, PostgresCodec>,
+    pub parameters: HashMap<String, String>,
+    pub delayed_notice: Vec<NoticeResponseBody>,
+    pub process_id: i32,
+    pub secret_key: i32,
+}
+
+pub async fn connect_raw<S, T>(
+    stream: MaybeTlsStream<S, T>,
     config: &Config,
-) -> Result<(), Error>
+) -> Result<RawConnection<S, T>, Error>
 where
     S: AsyncRead + AsyncWrite + Unpin,
     T: TlsStream + Unpin,
 {
-    frontend::startup_message(&config.server_params, stream.inner.write_buffer_mut())
-        .map_err(Error::encode)?;
+    let mut stream = StartupStream {
+        inner: Framed::new(stream, PostgresCodec),
+        buf: BackendMessages::empty(),
+        delayed_notice: Vec::new(),
+    };
 
-    stream.inner.flush().await.map_err(Error::io)?;
+    startup(&mut stream, config).await?;
+    authenticate(&mut stream, config).await?;
+    let (process_id, secret_key, parameters) = read_info(&mut stream).await?;
+
+    Ok(RawConnection {
+        stream: stream.inner,
+        parameters,
+        delayed_notice: stream.delayed_notice,
+        process_id,
+        secret_key,
+    })
+}
+
+async fn startup<S, T>(stream: &mut StartupStream<S, T>, config: &Config) -> Result<(), Error>
+where
+    S: AsyncRead + AsyncWrite + Unpin,
+    T: AsyncRead + AsyncWrite + Unpin,
+{
+    let mut buf = BytesMut::new();
+    frontend::startup_message(&config.server_params, &mut buf).map_err(Error::encode)?;
+
+    stream
+        .send(FrontendMessage::Raw(buf.freeze()))
+        .await
+        .map_err(Error::io)
+}
+
+async fn authenticate<S, T>(stream: &mut StartupStream<S, T>, config: &Config) -> Result<(), Error>
+where
+    S: AsyncRead + AsyncWrite + Unpin,
+    T: TlsStream + Unpin,
+{
     match stream.try_next().await.map_err(Error::io)? {
         Some(Message::AuthenticationOk) => {
             can_skip_channel_binding(config)?;
@@ -147,8 +144,7 @@ where
                 .as_ref()
                 .ok_or_else(|| Error::config("password missing".into()))?;
 
-            frontend::password_message(pass, stream.inner.write_buffer_mut())
-                .map_err(Error::encode)?;
+            authenticate_password(stream, pass).await?;
         }
         Some(Message::AuthenticationSasl(body)) => {
             authenticate_sasl(stream, body, config).await?;
@@ -167,7 +163,6 @@ where
         None => return Err(Error::closed()),
     }
 
-    stream.inner.flush().await.map_err(Error::io)?;
     match stream.try_next().await.map_err(Error::io)? {
         Some(Message::AuthenticationOk) => Ok(()),
         Some(Message::ErrorResponse(body)) => Err(Error::db(body)),
@@ -185,6 +180,23 @@ fn can_skip_channel_binding(config: &Config) -> Result<(), Error> {
     }
 }
 
+async fn authenticate_password<S, T>(
+    stream: &mut StartupStream<S, T>,
+    password: &[u8],
+) -> Result<(), Error>
+where
+    S: AsyncRead + AsyncWrite + Unpin,
+    T: AsyncRead + AsyncWrite + Unpin,
+{
+    let mut buf = BytesMut::new();
+    frontend::password_message(password, &mut buf).map_err(Error::encode)?;
+
+    stream
+        .send(FrontendMessage::Raw(buf.freeze()))
+        .await
+        .map_err(Error::io)
+}
+
 async fn authenticate_sasl<S, T>(
     stream: &mut StartupStream<S, T>,
     body: AuthenticationSaslBody,
@@ -239,10 +251,13 @@ where
         return Err(Error::config("password or auth keys missing".into()));
     };
 
-    frontend::sasl_initial_response(mechanism, scram.message(), stream.inner.write_buffer_mut())
-        .map_err(Error::encode)?;
+    let mut buf = BytesMut::new();
+    frontend::sasl_initial_response(mechanism, scram.message(), &mut buf).map_err(Error::encode)?;
+    stream
+        .send(FrontendMessage::Raw(buf.freeze()))
+        .await
+        .map_err(Error::io)?;
 
-    stream.inner.flush().await.map_err(Error::io)?;
     let body = match stream.try_next().await.map_err(Error::io)? {
         Some(Message::AuthenticationSaslContinue(body)) => body,
         Some(Message::ErrorResponse(body)) => return Err(Error::db(body)),
@@ -255,10 +270,13 @@ where
         .await
         .map_err(|e| Error::authentication(e.into()))?;
 
-    frontend::sasl_response(scram.message(), stream.inner.write_buffer_mut())
-        .map_err(Error::encode)?;
+    let mut buf = BytesMut::new();
+    frontend::sasl_response(scram.message(), &mut buf).map_err(Error::encode)?;
+    stream
+        .send(FrontendMessage::Raw(buf.freeze()))
+        .await
+        .map_err(Error::io)?;
 
-    stream.inner.flush().await.map_err(Error::io)?;
     let body = match stream.try_next().await.map_err(Error::io)? {
         Some(Message::AuthenticationSaslFinal(body)) => body,
         Some(Message::ErrorResponse(body)) => return Err(Error::db(body)),
@@ -272,3 +290,35 @@ where
 
     Ok(())
 }
+
+async fn read_info<S, T>(
+    stream: &mut StartupStream<S, T>,
+) -> Result<(i32, i32, HashMap<String, String>), Error>
+where
+    S: AsyncRead + AsyncWrite + Unpin,
+    T: AsyncRead + AsyncWrite + Unpin,
+{
+    let mut process_id = 0;
+    let mut secret_key = 0;
+    let mut parameters = HashMap::new();
+
+    loop {
+        match stream.try_next().await.map_err(Error::io)? {
+            Some(Message::BackendKeyData(body)) => {
+                process_id = body.process_id();
+                secret_key = body.secret_key();
+            }
+            Some(Message::ParameterStatus(body)) => {
+                parameters.insert(
+                    body.name().map_err(Error::parse)?.to_string(),
+                    body.value().map_err(Error::parse)?.to_string(),
+                );
+            }
+            Some(Message::NoticeResponse(body)) => stream.delayed_notice.push(body),
+            Some(Message::ReadyForQuery(_)) => return Ok((process_id, secret_key, parameters)),
+            Some(Message::ErrorResponse(body)) => return Err(Error::db(body)),
+            Some(_) => return Err(Error::unexpected_message()),
+            None => return Err(Error::closed()),
+        }
+    }
+}

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

@@ -1,23 +1,22 @@
+use std::collections::{HashMap, VecDeque};
 use std::future::Future;
 use std::pin::Pin;
 use std::task::{Context, Poll};
 
 use bytes::BytesMut;
-use fallible_iterator::FallibleIterator;
-use futures_util::{Sink, StreamExt, ready};
-use postgres_protocol2::message::backend::{Message, NoticeResponseBody};
+use futures_util::{Sink, Stream, ready};
+use postgres_protocol2::message::backend::Message;
 use postgres_protocol2::message::frontend;
 use tokio::io::{AsyncRead, AsyncWrite};
 use tokio::sync::mpsc;
 use tokio_util::codec::Framed;
 use tokio_util::sync::PollSender;
-use tracing::trace;
+use tracing::{info, trace};
 
-use crate::Error;
-use crate::codec::{
-    BackendMessage, BackendMessages, FrontendMessage, PostgresCodec, RecordNotices,
-};
+use crate::codec::{BackendMessage, BackendMessages, FrontendMessage, PostgresCodec};
+use crate::error::DbError;
 use crate::maybe_tls_stream::MaybeTlsStream;
+use crate::{AsyncMessage, Error, Notification};
 
 #[derive(PartialEq, Debug)]
 enum State {
@@ -34,39 +33,18 @@ enum State {
 /// occurred, or because its associated `Client` has dropped and all outstanding work has completed.
 #[must_use = "futures do nothing unless polled"]
 pub struct Connection<S, T> {
-    stream: Framed<MaybeTlsStream<S, T>, PostgresCodec>,
+    /// HACK: we need this in the Neon Proxy.
+    pub stream: Framed<MaybeTlsStream<S, T>, PostgresCodec>,
+    /// HACK: we need this in the Neon Proxy to forward params.
+    pub parameters: HashMap<String, String>,
 
     sender: PollSender<BackendMessages>,
     receiver: mpsc::UnboundedReceiver<FrontendMessage>,
-    notices: Option<RecordNotices>,
 
-    pending_response: Option<BackendMessages>,
+    pending_responses: VecDeque<BackendMessage>,
     state: State,
 }
 
-pub const INITIAL_CAPACITY: usize = 2 * 1024;
-pub const GC_THRESHOLD: usize = 16 * 1024;
-
-/// Gargabe collect the [`BytesMut`] if it has too much spare capacity.
-pub fn gc_bytesmut(buf: &mut BytesMut) {
-    // We use a different mode to shrink the buf when above the threshold.
-    // When above the threshold, we only re-allocate when the buf has 2x spare capacity.
-    let reclaim = GC_THRESHOLD.checked_sub(buf.len()).unwrap_or(buf.len());
-
-    // `try_reclaim` tries to get the capacity from any shared `BytesMut`s,
-    // before then comparing the length against the capacity.
-    if buf.try_reclaim(reclaim) {
-        let capacity = usize::max(buf.len(), INITIAL_CAPACITY);
-
-        // Allocate a new `BytesMut` so that we deallocate the old version.
-        let mut new = BytesMut::with_capacity(capacity);
-        new.extend_from_slice(buf);
-        *buf = new;
-    }
-}
-
-pub enum Never {}
-
 impl<S, T> Connection<S, T>
 where
     S: AsyncRead + AsyncWrite + Unpin,
@@ -74,51 +52,72 @@ where
 {
     pub(crate) fn new(
         stream: Framed<MaybeTlsStream<S, T>, PostgresCodec>,
+        pending_responses: VecDeque<BackendMessage>,
+        parameters: HashMap<String, String>,
         sender: mpsc::Sender<BackendMessages>,
         receiver: mpsc::UnboundedReceiver<FrontendMessage>,
     ) -> Connection<S, T> {
         Connection {
             stream,
+            parameters,
             sender: PollSender::new(sender),
             receiver,
-            notices: None,
-            pending_response: None,
+            pending_responses,
             state: State::Active,
         }
     }
 
+    fn poll_response(
+        &mut self,
+        cx: &mut Context<'_>,
+    ) -> Poll<Option<Result<BackendMessage, Error>>> {
+        if let Some(message) = self.pending_responses.pop_front() {
+            trace!("retrying pending response");
+            return Poll::Ready(Some(Ok(message)));
+        }
+
+        Pin::new(&mut self.stream)
+            .poll_next(cx)
+            .map(|o| o.map(|r| r.map_err(Error::io)))
+    }
+
     /// Read and process messages from the connection to postgres.
     /// client <- postgres
-    fn poll_read(&mut self, cx: &mut Context<'_>) -> Poll<Result<Never, Error>> {
+    fn poll_read(&mut self, cx: &mut Context<'_>) -> Poll<Result<AsyncMessage, Error>> {
         loop {
-            let messages = match self.pending_response.take() {
-                Some(messages) => messages,
-                None => {
-                    let message = match self.stream.poll_next_unpin(cx) {
-                        Poll::Pending => return Poll::Pending,
-                        Poll::Ready(None) => return Poll::Ready(Err(Error::closed())),
-                        Poll::Ready(Some(Err(e))) => return Poll::Ready(Err(Error::io(e))),
-                        Poll::Ready(Some(Ok(message))) => message,
-                    };
-
-                    match message {
-                        BackendMessage::Async(Message::NoticeResponse(body)) => {
-                            self.handle_notice(body)?;
-                            continue;
-                        }
-                        BackendMessage::Async(_) => continue,
-                        BackendMessage::Normal { messages, ready } => {
-                            // if we read a ReadyForQuery from postgres, let's try GC the read buffer.
-                            if ready {
-                                gc_bytesmut(self.stream.read_buffer_mut());
-                            }
-
-                            messages
-                        }
-                    }
+            let message = match self.poll_response(cx)? {
+                Poll::Ready(Some(message)) => message,
+                Poll::Ready(None) => return Poll::Ready(Err(Error::closed())),
+                Poll::Pending => {
+                    trace!("poll_read: waiting on response");
+                    return Poll::Pending;
                 }
             };
 
+            let messages = match message {
+                BackendMessage::Async(Message::NoticeResponse(body)) => {
+                    let error = DbError::parse(&mut body.fields()).map_err(Error::parse)?;
+                    return Poll::Ready(Ok(AsyncMessage::Notice(error)));
+                }
+                BackendMessage::Async(Message::NotificationResponse(body)) => {
+                    let notification = Notification {
+                        process_id: body.process_id(),
+                        channel: body.channel().map_err(Error::parse)?.to_string(),
+                        payload: body.message().map_err(Error::parse)?.to_string(),
+                    };
+                    return Poll::Ready(Ok(AsyncMessage::Notification(notification)));
+                }
+                BackendMessage::Async(Message::ParameterStatus(body)) => {
+                    self.parameters.insert(
+                        body.name().map_err(Error::parse)?.to_string(),
+                        body.value().map_err(Error::parse)?.to_string(),
+                    );
+                    continue;
+                }
+                BackendMessage::Async(_) => unreachable!(),
+                BackendMessage::Normal { messages } => messages,
+            };
+
             match self.sender.poll_reserve(cx) {
                 Poll::Ready(Ok(())) => {
                     let _ = self.sender.send_item(messages);
@@ -127,7 +126,8 @@ where
                     return Poll::Ready(Err(Error::closed()));
                 }
                 Poll::Pending => {
-                    self.pending_response = Some(messages);
+                    self.pending_responses
+                        .push_back(BackendMessage::Normal { messages });
                     trace!("poll_read: waiting on sender");
                     return Poll::Pending;
                 }
@@ -135,31 +135,6 @@ where
         }
     }
 
-    fn handle_notice(&mut self, body: NoticeResponseBody) -> Result<(), Error> {
-        let Some(notices) = &mut self.notices else {
-            return Ok(());
-        };
-
-        let mut fields = body.fields();
-        while let Some(field) = fields.next().map_err(Error::parse)? {
-            // loop until we find the message field
-            if field.type_() == b'M' {
-                // if the message field is within the limit, send it.
-                if let Some(new_limit) = notices.limit.checked_sub(field.value().len()) {
-                    match notices.sender.send(field.value().into()) {
-                        // set the new limit.
-                        Ok(()) => notices.limit = new_limit,
-                        // closed.
-                        Err(_) => self.notices = None,
-                    }
-                }
-                break;
-            }
-        }
-
-        Ok(())
-    }
-
     /// Fetch the next client request and enqueue the response sender.
     fn poll_request(&mut self, cx: &mut Context<'_>) -> Poll<Option<FrontendMessage>> {
         if self.receiver.is_closed() {
@@ -193,19 +168,22 @@ where
 
             match self.poll_request(cx) {
                 // send the message to postgres
-                Poll::Ready(Some(FrontendMessage::Raw(request))) => {
+                Poll::Ready(Some(request)) => {
                     Pin::new(&mut self.stream)
                         .start_send(request)
                         .map_err(Error::io)?;
                 }
-                Poll::Ready(Some(FrontendMessage::RecordNotices(notices))) => {
-                    self.notices = Some(notices)
-                }
                 // No more messages from the client, and no more responses to wait for.
                 // Send a terminate message to postgres
                 Poll::Ready(None) => {
                     trace!("poll_write: at eof, terminating");
-                    frontend::terminate(self.stream.write_buffer_mut());
+                    let mut request = BytesMut::new();
+                    frontend::terminate(&mut request);
+                    let request = FrontendMessage::Raw(request.freeze());
+
+                    Pin::new(&mut self.stream)
+                        .start_send(request)
+                        .map_err(Error::io)?;
 
                     trace!("poll_write: sent eof, closing");
                     trace!("poll_write: done");
@@ -228,10 +206,6 @@ where
         {
             Poll::Ready(()) => {
                 trace!("poll_flush: flushed");
-
-                // GC the write buffer if we managed to flush
-                gc_bytesmut(self.stream.write_buffer_mut());
-
                 Poll::Ready(Ok(()))
             }
             Poll::Pending => {
@@ -257,17 +231,34 @@ where
         }
     }
 
-    fn poll_message(&mut self, cx: &mut Context<'_>) -> Poll<Option<Result<Never, Error>>> {
+    /// Returns the value of a runtime parameter for this connection.
+    pub fn parameter(&self, name: &str) -> Option<&str> {
+        self.parameters.get(name).map(|s| &**s)
+    }
+
+    /// Polls for asynchronous messages from the server.
+    ///
+    /// The server can send notices as well as notifications asynchronously to the client. Applications that wish to
+    /// examine those messages should use this method to drive the connection rather than its `Future` implementation.
+    pub fn poll_message(
+        &mut self,
+        cx: &mut Context<'_>,
+    ) -> Poll<Option<Result<AsyncMessage, Error>>> {
         if self.state != State::Closing {
             // if the state is still active, try read from and write to postgres.
-            let Poll::Pending = self.poll_read(cx)?;
-            if self.poll_write(cx)?.is_ready() {
+            let message = self.poll_read(cx)?;
+            let closing = self.poll_write(cx)?;
+            if let Poll::Ready(()) = closing {
                 self.state = State::Closing;
             }
 
+            if let Poll::Ready(message) = message {
+                return Poll::Ready(Some(Ok(message)));
+            }
+
             // poll_read returned Pending.
-            // poll_write returned Pending or Ready(()).
-            // if poll_write returned Ready(()), then we are waiting to read more data from postgres.
+            // poll_write returned Pending or Ready(WriteReady::WaitingOnRead).
+            // if poll_write returned Ready(WriteReady::WaitingOnRead), then we are waiting to read more data from postgres.
             if self.state != State::Closing {
                 return Poll::Pending;
             }
@@ -289,9 +280,11 @@ where
     type Output = Result<(), Error>;
 
     fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Error>> {
-        match self.poll_message(cx)? {
-            Poll::Ready(None) => Poll::Ready(Ok(())),
-            Poll::Pending => Poll::Pending,
+        while let Some(message) = ready!(self.poll_message(cx)?) {
+            if let AsyncMessage::Notice(notice) = message {
+                info!("{}: {}", notice.severity(), notice.message());
+            }
         }
+        Poll::Ready(Ok(()))
     }
 }

libs/proxy/tokio-postgres2/src/error/mod.rs

@@ -452,16 +452,16 @@ impl Error {
         Error(Box::new(ErrorInner { kind, cause }))
     }
 
-    pub fn closed() -> Error {
+    pub(crate) fn closed() -> Error {
         Error::new(Kind::Closed, None)
     }
 
-    pub fn unexpected_message() -> Error {
+    pub(crate) fn unexpected_message() -> Error {
         Error::new(Kind::UnexpectedMessage, None)
     }
 
     #[allow(clippy::needless_pass_by_value)]
-    pub fn db(error: ErrorResponseBody) -> Error {
+    pub(crate) fn db(error: ErrorResponseBody) -> Error {
         match DbError::parse(&mut error.fields()) {
             Ok(e) => Error::new(Kind::Db, Some(Box::new(e))),
             Err(e) => Error::new(Kind::Parse, Some(Box::new(e))),
@@ -493,7 +493,7 @@ impl Error {
         Error::new(Kind::Tls, Some(e))
     }
 
-    pub fn io(e: io::Error) -> Error {
+    pub(crate) fn io(e: io::Error) -> Error {
         Error::new(Kind::Io, Some(Box::new(e)))
     }
 

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

@@ -6,7 +6,9 @@ use postgres_protocol2::message::backend::ReadyForQueryBody;
 pub use crate::cancel_token::{CancelToken, RawCancelToken};
 pub use crate::client::{Client, SocketConfig};
 pub use crate::config::Config;
+pub use crate::connect_raw::RawConnection;
 pub use crate::connection::Connection;
+use crate::error::DbError;
 pub use crate::error::Error;
 pub use crate::generic_client::GenericClient;
 pub use crate::query::RowStream;
@@ -48,8 +50,8 @@ mod cancel_token;
 mod client;
 mod codec;
 pub mod config;
-pub mod connect;
-pub mod connect_raw;
+mod connect;
+mod connect_raw;
 mod connect_socket;
 mod connect_tls;
 mod connection;
@@ -91,6 +93,21 @@ impl Notification {
     }
 }
 
+/// An asynchronous message from the server.
+#[allow(clippy::large_enum_variant)]
+#[derive(Debug, Clone)]
+#[non_exhaustive]
+pub enum AsyncMessage {
+    /// A notice.
+    ///
+    /// Notices use the same format as errors, but aren't "errors" per-se.
+    Notice(DbError),
+    /// A notification.
+    ///
+    /// Connections can subscribe to notifications with the `LISTEN` command.
+    Notification(Notification),
+}
+
 /// Message returned by the `SimpleQuery` stream.
 #[derive(Debug)]
 #[non_exhaustive]

libs/remote_storage/Cargo.toml

@@ -43,7 +43,7 @@ itertools.workspace = true
 sync_wrapper = { workspace = true, features = ["futures"] }
 
 byteorder = "1.4"
-rand.workspace = true
+rand = "0.8.5"
 
 [dev-dependencies]
 camino-tempfile.workspace = true

libs/remote_storage/src/simulate_failures.rs

@@ -81,7 +81,7 @@ impl UnreliableWrapper {
     ///
     fn attempt(&self, op: RemoteOp) -> anyhow::Result<u64> {
         let mut attempts = self.attempts.lock().unwrap();
-        let mut rng = rand::rng();
+        let mut rng = rand::thread_rng();
 
         match attempts.entry(op) {
             Entry::Occupied(mut e) => {
@@ -94,7 +94,7 @@ impl UnreliableWrapper {
                 /* BEGIN_HADRON */
                 // If there are more attempts to fail, fail the request by probability.
                 if (attempts_before_this < self.attempts_to_fail)
-                    && (rng.random_range(0..=100) < self.attempt_failure_probability)
+                    && (rng.gen_range(0..=100) < self.attempt_failure_probability)
                 {
                     let error =
                         anyhow::anyhow!("simulated failure of remote operation {:?}", e.key());

libs/remote_storage/tests/test_real_azure.rs

@@ -208,7 +208,7 @@ async fn create_azure_client(
         .as_millis();
 
     // because nanos can be the same for two threads so can millis, add randomness
-    let random = rand::rng().random::<u32>();
+    let random = rand::thread_rng().r#gen::<u32>();
 
     let remote_storage_config = RemoteStorageConfig {
         storage: RemoteStorageKind::AzureContainer(AzureConfig {

libs/remote_storage/tests/test_real_s3.rs

@@ -385,7 +385,7 @@ async fn create_s3_client(
         .as_millis();
 
     // because nanos can be the same for two threads so can millis, add randomness
-    let random = rand::rng().random::<u32>();
+    let random = rand::thread_rng().r#gen::<u32>();
 
     let remote_storage_config = RemoteStorageConfig {
         storage: RemoteStorageKind::AwsS3(S3Config {

libs/safekeeper_api/src/models.rs

@@ -301,12 +301,7 @@ pub struct PullTimelineRequest {
     pub tenant_id: TenantId,
     pub timeline_id: TimelineId,
     pub http_hosts: Vec<String>,
-    /// Membership configuration to switch to after pull.
-    /// It guarantees that if pull_timeline returns successfully, the timeline will
-    /// not be deleted by request with an older generation.
-    /// Storage controller always sets this field.
-    /// None is only allowed for manual pull_timeline requests.
-    pub mconf: Option<Configuration>,
+    pub ignore_tombstone: Option<bool>,
 }
 
 #[derive(Debug, Serialize, Deserialize)]

libs/tracing-utils/Cargo.toml

@@ -8,7 +8,7 @@ license.workspace = true
 hyper0.workspace = true
 opentelemetry = { workspace = true, features = ["trace"] }
 opentelemetry_sdk = { workspace = true, features = ["rt-tokio"] }
-opentelemetry-otlp = { workspace = true, default-features = false, features = ["http-proto", "trace", "http", "reqwest-blocking-client"] }
+opentelemetry-otlp = { workspace = true, default-features = false, features = ["http-proto", "trace", "http", "reqwest-client"] }
 opentelemetry-semantic-conventions.workspace = true
 tokio = { workspace = true, features = ["rt", "rt-multi-thread"] }
 tracing.workspace = true

libs/tracing-utils/src/lib.rs

@@ -1,5 +1,11 @@
 //! Helper functions to set up OpenTelemetry tracing.
 //!
+//! This comes in two variants, depending on whether you have a Tokio runtime available.
+//! If you do, call `init_tracing()`. It sets up the trace processor and exporter to use
+//! the current tokio runtime. If you don't have a runtime available, or you don't want
+//! to share the runtime with the tracing tasks, call `init_tracing_without_runtime()`
+//! instead. It sets up a dedicated single-threaded Tokio runtime for the tracing tasks.
+//!
 //! Example:
 //!
 //! ```rust,no_run
@@ -15,8 +21,7 @@
 //!         .with_writer(std::io::stderr);
 //!
 //!     // Initialize OpenTelemetry. Exports tracing spans as OpenTelemetry traces
-//!     let provider = tracing_utils::init_tracing("my_application", tracing_utils::ExportConfig::default());
-//!     let otlp_layer = provider.as_ref().map(tracing_utils::layer);
+//!     let otlp_layer = tracing_utils::init_tracing("my_application", tracing_utils::ExportConfig::default()).await;
 //!
 //!     // Put it all together
 //!     tracing_subscriber::registry()
@@ -31,18 +36,16 @@
 pub mod http;
 pub mod perf_span;
 
+use opentelemetry::KeyValue;
 use opentelemetry::trace::TracerProvider;
 use opentelemetry_otlp::WithExportConfig;
 pub use opentelemetry_otlp::{ExportConfig, Protocol};
-use opentelemetry_sdk::trace::SdkTracerProvider;
 use tracing::level_filters::LevelFilter;
 use tracing::{Dispatch, Subscriber};
 use tracing_subscriber::Layer;
 use tracing_subscriber::layer::SubscriberExt;
 use tracing_subscriber::registry::LookupSpan;
 
-pub type Provider = SdkTracerProvider;
-
 /// Set up OpenTelemetry exporter, using configuration from environment variables.
 ///
 /// `service_name` is set as the OpenTelemetry 'service.name' resource (see
@@ -67,7 +70,16 @@ pub type Provider = SdkTracerProvider;
 /// If you need some other setting, please test if it works first. And perhaps
 /// add a comment in the list above to save the effort of testing for the next
 /// person.
-pub fn init_tracing(service_name: &str, export_config: ExportConfig) -> Option<Provider> {
+///
+/// This doesn't block, but is marked as 'async' to hint that this must be called in
+/// asynchronous execution context.
+pub async fn init_tracing<S>(
+    service_name: &str,
+    export_config: ExportConfig,
+) -> Option<impl Layer<S>>
+where
+    S: Subscriber + for<'span> LookupSpan<'span>,
+{
     if std::env::var("OTEL_SDK_DISABLED") == Ok("true".to_string()) {
         return None;
     };
@@ -77,14 +89,52 @@ pub fn init_tracing(service_name: &str, export_config: ExportConfig) -> Option<P
     ))
 }
 
-pub fn layer<S>(p: &Provider) -> impl Layer<S>
+/// Like `init_tracing`, but creates a separate tokio Runtime for the tracing
+/// tasks.
+pub fn init_tracing_without_runtime<S>(
+    service_name: &str,
+    export_config: ExportConfig,
+) -> Option<impl Layer<S>>
 where
     S: Subscriber + for<'span> LookupSpan<'span>,
 {
-    tracing_opentelemetry::layer().with_tracer(p.tracer("global"))
+    if std::env::var("OTEL_SDK_DISABLED") == Ok("true".to_string()) {
+        return None;
+    };
+
+    // The opentelemetry batch processor and the OTLP exporter needs a Tokio
+    // runtime. Create a dedicated runtime for them. One thread should be
+    // enough.
+    //
+    // (Alternatively, instead of batching, we could use the "simple
+    // processor", which doesn't need Tokio, and use "reqwest-blocking"
+    // feature for the OTLP exporter, which also doesn't need Tokio.  However,
+    // batching is considered best practice, and also I have the feeling that
+    // the non-Tokio codepaths in the opentelemetry crate are less used and
+    // might be more buggy, so better to stay on the well-beaten path.)
+    //
+    // We leak the runtime so that it keeps running after we exit the
+    // function.
+    let runtime = Box::leak(Box::new(
+        tokio::runtime::Builder::new_multi_thread()
+            .enable_all()
+            .thread_name("otlp runtime thread")
+            .worker_threads(1)
+            .build()
+            .unwrap(),
+    ));
+    let _guard = runtime.enter();
+
+    Some(init_tracing_internal(
+        service_name.to_string(),
+        export_config,
+    ))
 }
 
-fn init_tracing_internal(service_name: String, export_config: ExportConfig) -> Provider {
+fn init_tracing_internal<S>(service_name: String, export_config: ExportConfig) -> impl Layer<S>
+where
+    S: Subscriber + for<'span> LookupSpan<'span>,
+{
     // Sets up exporter from the provided [`ExportConfig`] parameter.
     // If the endpoint is not specified, it is loaded from the
     // OTEL_EXPORTER_OTLP_ENDPOINT environment variable.
@@ -103,14 +153,22 @@ fn init_tracing_internal(service_name: String, export_config: ExportConfig) -> P
         opentelemetry_sdk::propagation::TraceContextPropagator::new(),
     );
 
-    Provider::builder()
-        .with_batch_exporter(exporter)
-        .with_resource(
-            opentelemetry_sdk::Resource::builder()
-                .with_service_name(service_name)
-                .build(),
-        )
+    let tracer = opentelemetry_sdk::trace::TracerProvider::builder()
+        .with_batch_exporter(exporter, opentelemetry_sdk::runtime::Tokio)
+        .with_resource(opentelemetry_sdk::Resource::new(vec![KeyValue::new(
+            opentelemetry_semantic_conventions::resource::SERVICE_NAME,
+            service_name,
+        )]))
         .build()
+        .tracer("global");
+
+    tracing_opentelemetry::layer().with_tracer(tracer)
+}
+
+// Shutdown trace pipeline gracefully, so that it has a chance to send any
+// pending traces before we exit.
+pub fn shutdown_tracing() {
+    opentelemetry::global::shutdown_tracer_provider();
 }
 
 pub enum OtelEnablement {
@@ -118,17 +176,17 @@ pub enum OtelEnablement {
     Enabled {
         service_name: String,
         export_config: ExportConfig,
+        runtime: &'static tokio::runtime::Runtime,
     },
 }
 
 pub struct OtelGuard {
-    provider: Provider,
     pub dispatch: Dispatch,
 }
 
 impl Drop for OtelGuard {
     fn drop(&mut self) {
-        _ = self.provider.shutdown();
+        shutdown_tracing();
     }
 }
 
@@ -141,19 +199,22 @@ impl Drop for OtelGuard {
 /// The lifetime of the guard should match taht of the application. On drop, it tears down the
 /// OTEL infra.
 pub fn init_performance_tracing(otel_enablement: OtelEnablement) -> Option<OtelGuard> {
-    match otel_enablement {
+    let otel_subscriber = match otel_enablement {
         OtelEnablement::Disabled => None,
         OtelEnablement::Enabled {
             service_name,
             export_config,
+            runtime,
         } => {
-            let provider = init_tracing(&service_name, export_config)?;
-
-            let otel_layer = layer(&provider).with_filter(LevelFilter::INFO);
+            let otel_layer = runtime
+                .block_on(init_tracing(&service_name, export_config))
+                .with_filter(LevelFilter::INFO);
             let otel_subscriber = tracing_subscriber::registry().with(otel_layer);
-            let dispatch = Dispatch::new(otel_subscriber);
+            let otel_dispatch = Dispatch::new(otel_subscriber);
 
-            Some(OtelGuard { dispatch, provider })
+            Some(otel_dispatch)
         }
-    }
+    };
+
+    otel_subscriber.map(|dispatch| OtelGuard { dispatch })
 }

libs/utils/src/id.rs

@@ -104,7 +104,7 @@ impl Id {
 
     pub fn generate() -> Self {
         let mut tli_buf = [0u8; 16];
-        rand::rng().fill(&mut tli_buf);
+        rand::thread_rng().fill(&mut tli_buf);
         Id::from(tli_buf)
     }
 

libs/utils/src/lsn.rs

@@ -364,37 +364,42 @@ impl MonotonicCounter<Lsn> for RecordLsn {
     }
 }
 
-/// Implements  [`rand::distr::uniform::UniformSampler`] so we can sample [`Lsn`]s.
+/// Implements  [`rand::distributions::uniform::UniformSampler`] so we can sample [`Lsn`]s.
 ///
 /// This is used by the `pagebench` pageserver benchmarking tool.
-pub struct LsnSampler(<u64 as rand::distr::uniform::SampleUniform>::Sampler);
+pub struct LsnSampler(<u64 as rand::distributions::uniform::SampleUniform>::Sampler);
 
-impl rand::distr::uniform::SampleUniform for Lsn {
+impl rand::distributions::uniform::SampleUniform for Lsn {
     type Sampler = LsnSampler;
 }
 
-impl rand::distr::uniform::UniformSampler for LsnSampler {
+impl rand::distributions::uniform::UniformSampler for LsnSampler {
     type X = Lsn;
 
-    fn new<B1, B2>(low: B1, high: B2) -> Result<Self, rand::distr::uniform::Error>
+    fn new<B1, B2>(low: B1, high: B2) -> Self
     where
-        B1: rand::distr::uniform::SampleBorrow<Self::X> + Sized,
-        B2: rand::distr::uniform::SampleBorrow<Self::X> + Sized,
+        B1: rand::distributions::uniform::SampleBorrow<Self::X> + Sized,
+        B2: rand::distributions::uniform::SampleBorrow<Self::X> + Sized,
     {
-        <u64 as rand::distr::uniform::SampleUniform>::Sampler::new(low.borrow().0, high.borrow().0)
-            .map(Self)
+        Self(
+            <u64 as rand::distributions::uniform::SampleUniform>::Sampler::new(
+                low.borrow().0,
+                high.borrow().0,
+            ),
+        )
     }
 
-    fn new_inclusive<B1, B2>(low: B1, high: B2) -> Result<Self, rand::distr::uniform::Error>
+    fn new_inclusive<B1, B2>(low: B1, high: B2) -> Self
     where
-        B1: rand::distr::uniform::SampleBorrow<Self::X> + Sized,
-        B2: rand::distr::uniform::SampleBorrow<Self::X> + Sized,
+        B1: rand::distributions::uniform::SampleBorrow<Self::X> + Sized,
+        B2: rand::distributions::uniform::SampleBorrow<Self::X> + Sized,
     {
-        <u64 as rand::distr::uniform::SampleUniform>::Sampler::new_inclusive(
-            low.borrow().0,
-            high.borrow().0,
+        Self(
+            <u64 as rand::distributions::uniform::SampleUniform>::Sampler::new_inclusive(
+                low.borrow().0,
+                high.borrow().0,
+            ),
         )
-        .map(Self)
     }
 
     fn sample<R: rand::prelude::Rng + ?Sized>(&self, rng: &mut R) -> Self::X {

libs/utils/src/shard.rs

@@ -25,12 +25,6 @@ pub struct ShardIndex {
     pub shard_count: ShardCount,
 }
 
-/// Stripe size as number of pages.
-///
-/// NB: don't implement Default, so callers don't lazily use it by mistake. See DEFAULT_STRIPE_SIZE.
-#[derive(Clone, Copy, Serialize, Deserialize, Eq, PartialEq, Debug)]
-pub struct ShardStripeSize(pub u32);
-
 /// Formatting helper, for generating the `shard_id` label in traces.
 pub struct ShardSlug<'a>(&'a TenantShardId);
 
@@ -183,12 +177,6 @@ impl std::fmt::Display for ShardCount {
     }
 }
 
-impl std::fmt::Display for ShardStripeSize {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        self.0.fmt(f)
-    }
-}
-
 impl std::fmt::Display for ShardSlug<'_> {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         write!(

libs/walproposer/src/api_bindings.rs

@@ -429,11 +429,9 @@ pub fn empty_shmem() -> crate::bindings::WalproposerShmemState {
     };
 
     let empty_wal_rate_limiter = crate::bindings::WalRateLimiter {
-        effective_max_wal_bytes_per_second: crate::bindings::pg_atomic_uint32 { value: 0 },
         should_limit: crate::bindings::pg_atomic_uint32 { value: 0 },
         sent_bytes: 0,
-        batch_start_time_us: crate::bindings::pg_atomic_uint64 { value: 0 },
-        batch_end_time_us: crate::bindings::pg_atomic_uint64 { value: 0 },
+        last_recorded_time_us: crate::bindings::pg_atomic_uint64 { value: 0 },
     };
 
     crate::bindings::WalproposerShmemState {

pageserver/benches/bench_layer_map.rs

@@ -11,8 +11,7 @@ use pageserver::tenant::layer_map::LayerMap;
 use pageserver::tenant::storage_layer::{LayerName, PersistentLayerDesc};
 use pageserver_api::key::Key;
 use pageserver_api::shard::TenantShardId;
-use rand::prelude::{SeedableRng, StdRng};
-use rand::seq::IndexedRandom;
+use rand::prelude::{SeedableRng, SliceRandom, StdRng};
 use utils::id::{TenantId, TimelineId};
 use utils::lsn::Lsn;
 

pageserver/client_grpc/src/client.rs

@@ -16,9 +16,10 @@ use crate::pool::{ChannelPool, ClientGuard, ClientPool, StreamGuard, StreamPool}
 use crate::retry::Retry;
 use crate::split::GetPageSplitter;
 use compute_api::spec::PageserverProtocol;
+use pageserver_api::shard::ShardStripeSize;
 use pageserver_page_api as page_api;
 use utils::id::{TenantId, TimelineId};
-use utils::shard::{ShardCount, ShardIndex, ShardNumber, ShardStripeSize};
+use utils::shard::{ShardCount, ShardIndex, ShardNumber};
 
 /// Max number of concurrent clients per channel (i.e. TCP connection). New channels will be spun up
 /// when full.
@@ -140,8 +141,8 @@ impl PageserverClient {
         if !old.count.is_unsharded() && shard_spec.stripe_size != old.stripe_size {
             return Err(anyhow!(
                 "can't change stripe size from {} to {}",
-                old.stripe_size.expect("always Some when sharded"),
-                shard_spec.stripe_size.expect("always Some when sharded")
+                old.stripe_size,
+                shard_spec.stripe_size
             ));
         }
 
@@ -156,6 +157,23 @@ impl PageserverClient {
         Ok(())
     }
 
+    /// Returns whether a relation exists.
+    #[instrument(skip_all, fields(rel=%req.rel, lsn=%req.read_lsn))]
+    pub async fn check_rel_exists(
+        &self,
+        req: page_api::CheckRelExistsRequest,
+    ) -> tonic::Result<page_api::CheckRelExistsResponse> {
+        debug!("sending request: {req:?}");
+        let resp = Self::with_retries(CALL_TIMEOUT, async |_| {
+            // Relation metadata is only available on shard 0.
+            let mut client = self.shards.load_full().get_zero().client().await?;
+            Self::with_timeout(REQUEST_TIMEOUT, client.check_rel_exists(req)).await
+        })
+        .await?;
+        debug!("received response: {resp:?}");
+        Ok(resp)
+    }
+
     /// Returns the total size of a database, as # of bytes.
     #[instrument(skip_all, fields(db_oid=%req.db_oid, lsn=%req.read_lsn))]
     pub async fn get_db_size(
@@ -231,15 +249,13 @@ impl PageserverClient {
         // Fast path: request is for a single shard.
         if let Some(shard_id) =
             GetPageSplitter::for_single_shard(&req, shards.count, shards.stripe_size)
-                .map_err(|err| tonic::Status::internal(err.to_string()))?
         {
             return Self::get_page_with_shard(req, shards.get(shard_id)?).await;
         }
 
         // Request spans multiple shards. Split it, dispatch concurrent per-shard requests, and
         // reassemble the responses.
-        let mut splitter = GetPageSplitter::split(req, shards.count, shards.stripe_size)
-            .map_err(|err| tonic::Status::internal(err.to_string()))?;
+        let mut splitter = GetPageSplitter::split(req, shards.count, shards.stripe_size);
 
         let mut shard_requests = FuturesUnordered::new();
         for (shard_id, shard_req) in splitter.drain_requests() {
@@ -249,14 +265,10 @@ impl PageserverClient {
         }
 
         while let Some((shard_id, shard_response)) = shard_requests.next().await.transpose()? {
-            splitter
-                .add_response(shard_id, shard_response)
-                .map_err(|err| tonic::Status::internal(err.to_string()))?;
+            splitter.add_response(shard_id, shard_response)?;
         }
 
-        splitter
-            .get_response()
-            .map_err(|err| tonic::Status::internal(err.to_string()))
+        splitter.get_response()
     }
 
     /// Fetches pages on the given shard. Does not retry internally.
@@ -384,14 +396,12 @@ pub struct ShardSpec {
     /// NB: this is 0 for unsharded tenants, following `ShardIndex::unsharded()` convention.
     count: ShardCount,
     /// The stripe size for these shards.
-    ///
-    /// INVARIANT: None for unsharded tenants, Some for sharded.
-    stripe_size: Option<ShardStripeSize>,
+    stripe_size: ShardStripeSize,
 }
 
 impl ShardSpec {
     /// Creates a new shard spec with the given URLs and stripe size. All shards must be given.
-    /// The stripe size must be Some for sharded tenants, or None for unsharded tenants.
+    /// The stripe size may be omitted for unsharded tenants.
     pub fn new(
         urls: HashMap<ShardIndex, String>,
         stripe_size: Option<ShardStripeSize>,
@@ -404,13 +414,11 @@ impl ShardSpec {
             n => ShardCount::new(n as u8),
         };
 
-        // Validate the stripe size.
+        // Determine the stripe size. It doesn't matter for unsharded tenants.
         if stripe_size.is_none() && !count.is_unsharded() {
             return Err(anyhow!("stripe size must be given for sharded tenants"));
         }
-        if stripe_size.is_some() && count.is_unsharded() {
-            return Err(anyhow!("stripe size can't be given for unsharded tenants"));
-        }
+        let stripe_size = stripe_size.unwrap_or_default();
 
         // Validate the shard spec.
         for (shard_id, url) in &urls {
@@ -450,10 +458,8 @@ struct Shards {
     ///
     /// NB: this is 0 for unsharded tenants, following `ShardIndex::unsharded()` convention.
     count: ShardCount,
-    /// The stripe size.
-    ///
-    /// INVARIANT: None for unsharded tenants, Some for sharded.
-    stripe_size: Option<ShardStripeSize>,
+    /// The stripe size. Only used for sharded tenants.
+    stripe_size: ShardStripeSize,
 }
 
 impl Shards {