Polish API handler and refresh OpenAPI spec

With this commit one can start compute with something like
```shell
cargo run --bin compute_ctl -- -i no-compute \
 -p http://localhost:9095 \
 -D compute_pgdata \
 -C "postgresql://cloud_admin@127.0.0.1:5434/postgres" \
 -b ./pg_install/v15/bin/postgres
```
and it will hang waiting for spec.

Then send one spec
```shell
curl -d "$(cat ./compute-spec.json)" http://localhost:3080/spec
```
Postgres will be started and configured.

Then reconfigure it with
```shell
curl -d "$(cat ./compute-spec-new.json)" http://localhost:3080/spec
```

Most of safeguards and comments are added. Some polishing especially
around HTTP API is still needed.

.github/PULL_REQUEST_TEMPLATE/release-pr.md

@@ -10,7 +10,6 @@
 <!-- List everything that should be done **before** release, any issues / setting changes / etc -->
 
 ### Checklist after release
-- [ ] Make sure instructions from PRs included in this release and labeled `manual_release_instructions` are executed (either by you or by people who wrote them).
 - [ ] Based on the merged commits write release notes and open a PR into `website` repo ([example](https://github.com/neondatabase/website/pull/219/files))
 - [ ] Check [#dev-production-stream](https://neondb.slack.com/archives/C03F5SM1N02) Slack channel
 - [ ] Check [stuck projects page](https://console.neon.tech/admin/projects?sort=last_active&order=desc&stuck=true)

.github/actions/allure-report/action.yml

@@ -45,12 +45,12 @@ runs:
       shell: bash -euxo pipefail {0}
       run: |
         if [ "${{ inputs.action }}" != "store" ] && [ "${{ inputs.action }}" != "generate" ]; then
-          echo >&2 "Unknown inputs.action type '${{ inputs.action }}'; allowed 'generate' or 'store' only"
+          echo 2>&1 "Unknown inputs.action type '${{ inputs.action }}'; allowed 'generate' or 'store' only"
           exit 1
         fi
 
         if [ -z "${{ inputs.test_selection }}" ] && [ "${{ inputs.action }}" == "store" ]; then
-          echo >&2 "inputs.test_selection must be set for 'store' action"
+          echo 2>&1 "inputs.test_selection must be set for 'store' action"
           exit 2
         fi
 

.github/actions/download/action.yml

@@ -37,7 +37,7 @@ runs:
             echo 'SKIPPED=true' >> $GITHUB_OUTPUT
             exit 0
           else
-            echo >&2 "Neither s3://${BUCKET}/${PREFIX}/${FILENAME} nor its version from previous attempts exist"
+            echo 2>&1 "Neither s3://${BUCKET}/${PREFIX}/${FILENAME} nor its version from previous attempts exist"
             exit 1
           fi
         fi

.github/actions/neon-branch-create/action.yml

@@ -58,7 +58,7 @@ runs:
         done
 
         if [ -z "${branch_id}" ] || [ "${branch_id}" == "null" ]; then
-          echo >&2 "Failed to create branch after 10 attempts, the latest response was: ${branch}"
+          echo 2>&1 "Failed to create branch after 10 attempts, the latest response was: ${branch}"
           exit 1
         fi
 
@@ -122,7 +122,7 @@ runs:
         done
 
         if [ -z "${password}" ] || [ "${password}" == "null" ]; then
-          echo >&2 "Failed to reset password after 10 attempts, the latest response was: ${reset_password}"
+          echo 2>&1 "Failed to reset password after 10 attempts, the latest response was: ${reset_password}"
           exit 1
         fi
 

.github/actions/neon-branch-delete/action.yml

@@ -48,7 +48,7 @@ runs:
         done
 
         if [ -z "${branch_id}" ] || [ "${branch_id}" == "null" ]; then
-          echo >&2 "Failed to delete branch after 10 attempts, the latest response was: ${deleted_branch}"
+          echo 2>&1 "Failed to delete branch after 10 attempts, the latest response was: ${deleted_branch}"
           exit 1
         fi
       env:

.github/actions/neon-project-create/action.yml

@@ -14,12 +14,6 @@ inputs:
   api_host:
     desctiption: 'Neon API host'
     default: console.stage.neon.tech
-  provisioner:
-    desctiption: 'k8s-pod or k8s-neonvm'
-    default: 'k8s-pod'
-  compute_units:
-    desctiption: '[Min, Max] compute units; Min and Max are used for k8s-neonvm with autoscaling, for k8s-pod values Min and Max should be equal'
-    default: '[1, 1]'
 
 outputs:
   dsn:
@@ -37,10 +31,6 @@ runs:
       # A shell without `set -x` to not to expose password/dsn in logs
       shell: bash -euo pipefail {0}
       run: |
-        if [ "${PROVISIONER}" == "k8s-pod" ] && [ "${MIN_CU}" != "${MAX_CU}" ]; then
-          echo >&2 "For k8s-pod provisioner MIN_CU should be equal to MAX_CU"
-        fi
-
         project=$(curl \
           "https://${API_HOST}/api/v2/projects" \
           --fail \
@@ -52,9 +42,6 @@ runs:
               \"name\": \"Created by actions/neon-project-create; GITHUB_RUN_ID=${GITHUB_RUN_ID}\",
               \"pg_version\": ${POSTGRES_VERSION},
               \"region_id\": \"${REGION_ID}\",
-              \"provisioner\": \"${PROVISIONER}\",
-              \"autoscaling_limit_min_cu\": ${MIN_CU},
-              \"autoscaling_limit_max_cu\": ${MAX_CU},
               \"settings\": { }
             }
           }")
@@ -75,6 +62,3 @@ runs:
         API_KEY: ${{ inputs.api_key }}
         REGION_ID: ${{ inputs.region_id }}
         POSTGRES_VERSION: ${{ inputs.postgres_version }}
-        PROVISIONER: ${{ inputs.provisioner }}
-        MIN_CU: ${{ fromJSON(inputs.compute_units)[0] }}
-        MAX_CU: ${{ fromJSON(inputs.compute_units)[1] }}

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

@@ -202,7 +202,7 @@ runs:
         prefix: latest
 
     - name: Create Allure report
-      if: ${{ !cancelled() }}
+      if: success() || failure()
       uses: ./.github/actions/allure-report
       with:
         action: store

.github/actions/upload/action.yml

@@ -23,7 +23,7 @@ runs:
         mkdir -p $(dirname $ARCHIVE)
 
         if [ -f ${ARCHIVE} ]; then
-          echo >&2 "File ${ARCHIVE} already exist. Something went wrong before"
+          echo 2>&1 "File ${ARCHIVE} already exist. Something went wrong before"
           exit 1
         fi
 
@@ -33,10 +33,10 @@ runs:
         elif [ -f ${SOURCE} ]; then
           time tar -cf ${ARCHIVE} --zstd ${SOURCE}
         elif ! ls ${SOURCE} > /dev/null 2>&1; then
-          echo >&2 "${SOURCE} does not exist"
+          echo 2>&1 "${SOURCE} does not exist"
           exit 2
         else
-          echo >&2 "${SOURCE} is neither a directory nor a file, do not know how to handle it"
+          echo 2>&1 "${SOURCE} is neither a directory nor a file, do not know how to handle it"
           exit 3
         fi
 

.github/ansible/scripts/init_pageserver.sh

@@ -3,8 +3,6 @@
 # fetch params from meta-data service
 INSTANCE_ID=$(curl -s http://169.254.169.254/latest/meta-data/instance-id)
 AZ_ID=$(curl -s http://169.254.169.254/latest/meta-data/placement/availability-zone)
-INSTANCE_TYPE=$(curl -s http://169.254.169.254/latest/meta-data/instance-type)
-DISK_SIZE=$(df -B1 /storage | tail -1 | awk '{print $2}')
 
 # store fqdn hostname in var
 HOST=$(hostname -f)
@@ -20,9 +18,7 @@ cat <<EOF | tee /tmp/payload
   "http_host": "${HOST}",
   "http_port": 9898,
   "active": false,
-  "availability_zone_id": "${AZ_ID}",
-  "disk_size": ${DISK_SIZE},
-  "instance_type": "${INSTANCE_TYPE}"
+  "availability_zone_id": "${AZ_ID}"
 }
 EOF
 

.github/helm-values/dev-us-east-2-beta.neon-proxy-scram.yaml

@@ -23,7 +23,6 @@ settings:
   authBackend: "console"
   authEndpoint: "http://neon-internal-api.aws.neon.build/management/api/v2"
   domain: "*.us-east-2.aws.neon.build"
-  extraDomains: ["*.us-east-2.postgres.zenith.tech", "*.us-east-2.retooldb-staging.com"]
   sentryEnvironment: "staging"
   wssPort: 8443
   metricCollectionEndpoint: "http://neon-internal-api.aws.neon.build/billing/api/v1/usage_events"

.github/helm-values/prod-ap-southeast-1-epsilon.neon-proxy-scram.yaml

@@ -24,7 +24,6 @@ settings:
   authBackend: "console"
   authEndpoint: "http://neon-internal-api.aws.neon.tech/management/api/v2"
   domain: "*.ap-southeast-1.aws.neon.tech"
-  extraDomains: ["*.ap-southeast-1.retooldb.com", "*.ap-southeast-1.postgres.vercel-storage.com"]
   sentryEnvironment: "production"
   wssPort: 8443
   metricCollectionEndpoint: "http://neon-internal-api.aws.neon.tech/billing/api/v1/usage_events"

.github/helm-values/prod-eu-central-1-gamma.neon-proxy-scram.yaml

@@ -24,7 +24,6 @@ settings:
   authBackend: "console"
   authEndpoint: "http://neon-internal-api.aws.neon.tech/management/api/v2"
   domain: "*.eu-central-1.aws.neon.tech"
-  extraDomains: ["*.eu-central-1.retooldb.com", "*.eu-central-1.postgres.vercel-storage.com"]
   sentryEnvironment: "production"
   wssPort: 8443
   metricCollectionEndpoint: "http://neon-internal-api.aws.neon.tech/billing/api/v1/usage_events"

.github/helm-values/prod-us-east-2-delta.neon-proxy-scram.yaml

@@ -24,7 +24,6 @@ settings:
   authBackend: "console"
   authEndpoint: "http://neon-internal-api.aws.neon.tech/management/api/v2"
   domain: "*.us-east-2.aws.neon.tech"
-  extraDomains: ["*.us-east-2.retooldb.com", "*.us-east-2.postgres.vercel-storage.com"]
   sentryEnvironment: "production"
   wssPort: 8443
   metricCollectionEndpoint: "http://neon-internal-api.aws.neon.tech/billing/api/v1/usage_events"

.github/helm-values/prod-us-west-2-eta.neon-proxy-scram.yaml

@@ -24,7 +24,6 @@ settings:
   authBackend: "console"
   authEndpoint: "http://neon-internal-api.aws.neon.tech/management/api/v2"
   domain: "*.us-west-2.aws.neon.tech"
-  extraDomains: ["*.us-west-2.retooldb.com", "*.us-west-2.postgres.vercel-storage.com"]
   sentryEnvironment: "production"
   wssPort: 8443
   metricCollectionEndpoint: "http://neon-internal-api.aws.neon.tech/billing/api/v1/usage_events"

.github/workflows/benchmarking.yml

@@ -30,7 +30,7 @@ defaults:
 
 concurrency:
   # Allow only one workflow per any non-`main` branch.
-  group: ${{ github.workflow }}-${{ github.ref_name }}-${{ github.ref_name == 'main' && github.sha || 'anysha' }}
+  group: ${{ github.workflow }}-${{ github.ref }}-${{ github.ref == 'refs/heads/main' && github.sha || 'anysha' }}
   cancel-in-progress: true
 
 jobs:
@@ -42,7 +42,7 @@ jobs:
       DEFAULT_PG_VERSION: 14
       TEST_OUTPUT: /tmp/test_output
       BUILD_TYPE: remote
-      SAVE_PERF_REPORT: ${{ github.event.inputs.save_perf_report || ( github.ref_name == 'main' ) }}
+      SAVE_PERF_REPORT: ${{ github.event.inputs.save_perf_report || ( github.ref == 'refs/heads/main' ) }}
       PLATFORM: "neon-staging"
 
     runs-on: [ self-hosted, us-east-2, x64 ]
@@ -92,7 +92,7 @@ jobs:
         api_key: ${{ secrets.NEON_STAGING_API_KEY }}
 
     - name: Create Allure report
-      if: ${{ !cancelled() }}
+      if: success() || failure()
       uses: ./.github/actions/allure-report
       with:
         action: generate
@@ -107,65 +107,25 @@ jobs:
       env:
         SLACK_BOT_TOKEN: ${{ secrets.SLACK_BOT_TOKEN }}
 
-  generate-matrices:
-    # Create matrices for the benchmarking jobs, so we run benchmarks on rds only once a week (on Saturday)
-    #
-    # Available platforms:
-    # - neon-captest-new: Freshly created project (1 CU)
-    # - neon-captest-freetier: Use freetier-sized compute (0.25 CU)
-    # - neon-captest-reuse: Reusing existing project
-    # - rds-aurora: Aurora Postgres Serverless v2 with autoscaling from 0.5 to 2 ACUs
-    # - rds-postgres: RDS Postgres db.m5.large instance (2 vCPU, 8 GiB) with gp3 EBS storage
-    runs-on: ubuntu-latest
-    outputs:
-      pgbench-compare-matrix: ${{ steps.pgbench-compare-matrix.outputs.matrix }}
-      olap-compare-matrix: ${{ steps.olap-compare-matrix.outputs.matrix }}
-
-    steps:
-    - name: Generate matrix for pgbench benchmark
-      id: pgbench-compare-matrix
-      run: |
-        matrix='{
-          "platform": [
-            "neon-captest-new",
-            "neon-captest-reuse"
-          ],
-          "db_size": [ "10gb" ],
-          "include": [
-            { "platform": "neon-captest-freetier", "db_size": "3gb"  },
-            { "platform": "neon-captest-new",      "db_size": "50gb" }
-          ]
-        }'
-
-        if [ "$(date +%A)" = "Saturday" ]; then
-          matrix=$(echo $matrix | jq '.include += [{ "platform": "rds-postgres", "db_size": "10gb"},
-                                                   { "platform": "rds-aurora",   "db_size": "50gb"}]')
-        fi
-
-        echo "matrix=$(echo $matrix | jq --compact-output '.')" >> $GITHUB_OUTPUT
-
-    - name: Generate matrix for OLAP benchmarks
-      id: olap-compare-matrix
-      run: |
-        matrix='{
-          "platform": [
-            "neon-captest-reuse"
-          ]
-        }'
-
-        if [ "$(date +%A)" = "Saturday" ]; then
-          matrix=$(echo $matrix | jq '.include += [{ "platform": "rds-postgres" },
-                                                   { "platform": "rds-aurora"   }]')
-        fi
-
-        echo "matrix=$(echo $matrix | jq --compact-output '.')" >> $GITHUB_OUTPUT
-
   pgbench-compare:
-    needs: [ generate-matrices ]
-
     strategy:
       fail-fast: false
-      matrix: ${{fromJson(needs.generate-matrices.outputs.pgbench-compare-matrix)}}
+      matrix:
+        # neon-captest-new: Run pgbench in a freshly created project
+        # neon-captest-reuse: Same, but reusing existing project
+        # neon-captest-prefetch: Same, with prefetching enabled (new project)
+        # rds-aurora: Aurora Postgres Serverless v2 with autoscaling from 0.5 to 2 ACUs
+        # rds-postgres: RDS Postgres db.m5.large instance (2 vCPU, 8 GiB) with gp3 EBS storage
+        platform: [ neon-captest-reuse, neon-captest-prefetch, rds-postgres ]
+        db_size: [ 10gb ]
+        runner: [ us-east-2 ]
+        include:
+          - platform: neon-captest-prefetch
+            db_size: 50gb
+            runner: us-east-2
+          - platform: rds-aurora
+            db_size: 50gb
+            runner: us-east-2
 
     env:
       TEST_PG_BENCH_DURATIONS_MATRIX: "60m"
@@ -174,10 +134,10 @@ jobs:
       DEFAULT_PG_VERSION: 14
       TEST_OUTPUT: /tmp/test_output
       BUILD_TYPE: remote
-      SAVE_PERF_REPORT: ${{ github.event.inputs.save_perf_report || ( github.ref_name == 'main' ) }}
+      SAVE_PERF_REPORT: ${{ github.event.inputs.save_perf_report || ( github.ref == 'refs/heads/main' ) }}
       PLATFORM: ${{ matrix.platform }}
 
-    runs-on: [ self-hosted, us-east-2, x64 ]
+    runs-on: [ self-hosted, "${{ matrix.runner }}", x64 ]
     container:
       image: 369495373322.dkr.ecr.eu-central-1.amazonaws.com/rust:pinned
       options: --init
@@ -200,14 +160,13 @@ jobs:
         echo "${POSTGRES_DISTRIB_DIR}/v${DEFAULT_PG_VERSION}/bin" >> $GITHUB_PATH
 
     - name: Create Neon Project
-      if: contains(fromJson('["neon-captest-new", "neon-captest-freetier"]'), matrix.platform)
+      if: contains(fromJson('["neon-captest-new", "neon-captest-prefetch"]'), matrix.platform)
       id: create-neon-project
       uses: ./.github/actions/neon-project-create
       with:
         region_id: ${{ github.event.inputs.region_id || 'aws-us-east-2' }}
         postgres_version: ${{ env.DEFAULT_PG_VERSION }}
         api_key: ${{ secrets.NEON_STAGING_API_KEY }}
-        compute_units: ${{ (matrix.platform == 'neon-captest-freetier' && '[0.25, 0.25]') || '[1, 1]' }}
 
     - name: Set up Connection String
       id: set-up-connstr
@@ -216,7 +175,7 @@ jobs:
           neon-captest-reuse)
             CONNSTR=${{ secrets.BENCHMARK_CAPTEST_CONNSTR }}
             ;;
-          neon-captest-new | neon-captest-freetier)
+          neon-captest-new | neon-captest-prefetch)
             CONNSTR=${{ steps.create-neon-project.outputs.dsn }}
             ;;
           rds-aurora)
@@ -226,7 +185,7 @@ jobs:
             CONNSTR=${{ secrets.BENCHMARK_RDS_POSTGRES_CONNSTR }}
             ;;
           *)
-            echo >&2 "Unknown PLATFORM=${PLATFORM}. Allowed only 'neon-captest-reuse', 'neon-captest-new', 'neon-captest-freetier', 'rds-aurora', or 'rds-postgres'"
+            echo 2>&1 "Unknown PLATFORM=${PLATFORM}. Allowed only 'neon-captest-reuse', 'neon-captest-new', 'neon-captest-prefetch', 'rds-aurora', or 'rds-postgres'"
             exit 1
             ;;
         esac
@@ -235,6 +194,17 @@ jobs:
 
         psql ${CONNSTR} -c "SELECT version();"
 
+    - name: Set database options
+      if: matrix.platform == 'neon-captest-prefetch'
+      run: |
+        DB_NAME=$(psql ${BENCHMARK_CONNSTR} --no-align --quiet -t -c "SELECT current_database()")
+
+        psql ${BENCHMARK_CONNSTR} -c "ALTER DATABASE ${DB_NAME} SET enable_seqscan_prefetch=on"
+        psql ${BENCHMARK_CONNSTR} -c "ALTER DATABASE ${DB_NAME} SET effective_io_concurrency=32"
+        psql ${BENCHMARK_CONNSTR} -c "ALTER DATABASE ${DB_NAME} SET maintenance_io_concurrency=32"
+      env:
+        BENCHMARK_CONNSTR: ${{ steps.set-up-connstr.outputs.connstr }}
+
     - name: Benchmark init
       uses: ./.github/actions/run-python-test-set
       with:
@@ -282,7 +252,7 @@ jobs:
         api_key: ${{ secrets.NEON_STAGING_API_KEY }}
 
     - name: Create Allure report
-      if: ${{ !cancelled() }}
+      if: success() || failure()
       uses: ./.github/actions/allure-report
       with:
         action: generate
@@ -305,19 +275,23 @@ jobs:
     #
     # *_CLICKBENCH_CONNSTR: Genuine ClickBench DB with ~100M rows
     # *_CLICKBENCH_10M_CONNSTR: DB with the first 10M rows of ClickBench DB
-    if: ${{ !cancelled() }}
-    needs: [ generate-matrices, pgbench-compare ]
+    if: success() || failure()
+    needs: [ pgbench-compare ]
 
     strategy:
       fail-fast: false
-      matrix: ${{ fromJson(needs.generate-matrices.outputs.olap-compare-matrix) }}
+      matrix:
+        # neon-captest-prefetch: We have pre-created projects with prefetch enabled
+        # rds-aurora: Aurora Postgres Serverless v2 with autoscaling from 0.5 to 2 ACUs
+        # rds-postgres: RDS Postgres db.m5.large instance (2 vCPU, 8 GiB) with gp3 EBS storage
+        platform: [ neon-captest-prefetch, rds-postgres, rds-aurora ]
 
     env:
       POSTGRES_DISTRIB_DIR: /tmp/neon/pg_install
       DEFAULT_PG_VERSION: 14
       TEST_OUTPUT: /tmp/test_output
       BUILD_TYPE: remote
-      SAVE_PERF_REPORT: ${{ github.event.inputs.save_perf_report || ( github.ref_name == 'main' ) }}
+      SAVE_PERF_REPORT: ${{ github.event.inputs.save_perf_report || ( github.ref == 'refs/heads/main' ) }}
       PLATFORM: ${{ matrix.platform }}
 
     runs-on: [ self-hosted, us-east-2, x64 ]
@@ -346,7 +320,7 @@ jobs:
       id: set-up-connstr
       run: |
         case "${PLATFORM}" in
-          neon-captest-reuse)
+          neon-captest-prefetch)
             CONNSTR=${{ secrets.BENCHMARK_CAPTEST_CLICKBENCH_10M_CONNSTR }}
             ;;
           rds-aurora)
@@ -356,7 +330,7 @@ jobs:
             CONNSTR=${{ secrets.BENCHMARK_RDS_POSTGRES_CLICKBENCH_10M_CONNSTR }}
             ;;
           *)
-            echo >&2 "Unknown PLATFORM=${PLATFORM}. Allowed only 'neon-captest-reuse', 'rds-aurora', or 'rds-postgres'"
+            echo 2>&1 "Unknown PLATFORM=${PLATFORM}. Allowed only 'neon-captest-prefetch', 'rds-aurora', or 'rds-postgres'"
             exit 1
             ;;
         esac
@@ -365,6 +339,17 @@ jobs:
 
         psql ${CONNSTR} -c "SELECT version();"
 
+    - name: Set database options
+      if: matrix.platform == 'neon-captest-prefetch'
+      run: |
+        DB_NAME=$(psql ${BENCHMARK_CONNSTR} --no-align --quiet -t -c "SELECT current_database()")
+
+        psql ${BENCHMARK_CONNSTR} -c "ALTER DATABASE ${DB_NAME} SET enable_seqscan_prefetch=on"
+        psql ${BENCHMARK_CONNSTR} -c "ALTER DATABASE ${DB_NAME} SET effective_io_concurrency=32"
+        psql ${BENCHMARK_CONNSTR} -c "ALTER DATABASE ${DB_NAME} SET maintenance_io_concurrency=32"
+      env:
+        BENCHMARK_CONNSTR: ${{ steps.set-up-connstr.outputs.connstr }}
+
     - name: ClickBench benchmark
       uses: ./.github/actions/run-python-test-set
       with:
@@ -379,7 +364,7 @@ jobs:
         BENCHMARK_CONNSTR: ${{ steps.set-up-connstr.outputs.connstr }}
 
     - name: Create Allure report
-      if: ${{ !cancelled() }}
+      if: success() || failure()
       uses: ./.github/actions/allure-report
       with:
         action: generate
@@ -401,19 +386,23 @@ jobs:
     # We might change it after https://github.com/neondatabase/neon/issues/2900.
     #
     # *_TPCH_S10_CONNSTR: DB generated with scale factor 10 (~10 GB)
-    if: ${{ !cancelled() }}
-    needs: [ generate-matrices, clickbench-compare ]
+    if: success() || failure()
+    needs: [ clickbench-compare ]
 
     strategy:
       fail-fast: false
-      matrix: ${{ fromJson(needs.generate-matrices.outputs.olap-compare-matrix) }}
+      matrix:
+        # neon-captest-prefetch: We have pre-created projects with prefetch enabled
+        # rds-aurora: Aurora Postgres Serverless v2 with autoscaling from 0.5 to 2 ACUs
+        # rds-postgres: RDS Postgres db.m5.large instance (2 vCPU, 8 GiB) with gp3 EBS storage
+        platform: [ neon-captest-prefetch, rds-postgres, rds-aurora ]
 
     env:
       POSTGRES_DISTRIB_DIR: /tmp/neon/pg_install
       DEFAULT_PG_VERSION: 14
       TEST_OUTPUT: /tmp/test_output
       BUILD_TYPE: remote
-      SAVE_PERF_REPORT: ${{ github.event.inputs.save_perf_report || ( github.ref_name == 'main' ) }}
+      SAVE_PERF_REPORT: ${{ github.event.inputs.save_perf_report || ( github.ref == 'refs/heads/main' ) }}
       PLATFORM: ${{ matrix.platform }}
 
     runs-on: [ self-hosted, us-east-2, x64 ]
@@ -442,7 +431,7 @@ jobs:
       id: set-up-connstr
       run: |
         case "${PLATFORM}" in
-          neon-captest-reuse)
+          neon-captest-prefetch)
             CONNSTR=${{ secrets.BENCHMARK_CAPTEST_TPCH_S10_CONNSTR }}
             ;;
           rds-aurora)
@@ -452,7 +441,7 @@ jobs:
             CONNSTR=${{ secrets.BENCHMARK_RDS_POSTGRES_TPCH_S10_CONNSTR }}
             ;;
           *)
-            echo >&2 "Unknown PLATFORM=${PLATFORM}. Allowed only 'neon-captest-reuse', 'rds-aurora', or 'rds-postgres'"
+            echo 2>&1 "Unknown PLATFORM=${PLATFORM}. Allowed only 'neon-captest-prefetch', 'rds-aurora', or 'rds-postgres'"
             exit 1
             ;;
         esac
@@ -461,6 +450,17 @@ jobs:
 
         psql ${CONNSTR} -c "SELECT version();"
 
+    - name: Set database options
+      if: matrix.platform == 'neon-captest-prefetch'
+      run: |
+        DB_NAME=$(psql ${BENCHMARK_CONNSTR} --no-align --quiet -t -c "SELECT current_database()")
+
+        psql ${BENCHMARK_CONNSTR} -c "ALTER DATABASE ${DB_NAME} SET enable_seqscan_prefetch=on"
+        psql ${BENCHMARK_CONNSTR} -c "ALTER DATABASE ${DB_NAME} SET effective_io_concurrency=32"
+        psql ${BENCHMARK_CONNSTR} -c "ALTER DATABASE ${DB_NAME} SET maintenance_io_concurrency=32"
+      env:
+        BENCHMARK_CONNSTR: ${{ steps.set-up-connstr.outputs.connstr }}
+
     - name: Run TPC-H benchmark
       uses: ./.github/actions/run-python-test-set
       with:
@@ -475,7 +475,7 @@ jobs:
         BENCHMARK_CONNSTR: ${{ steps.set-up-connstr.outputs.connstr }}
 
     - name: Create Allure report
-      if: ${{ !cancelled() }}
+      if: success() || failure()
       uses: ./.github/actions/allure-report
       with:
         action: generate
@@ -491,19 +491,23 @@ jobs:
         SLACK_BOT_TOKEN: ${{ secrets.SLACK_BOT_TOKEN }}
 
   user-examples-compare:
-    if: ${{ !cancelled() }}
-    needs: [ generate-matrices, tpch-compare ]
+    if: success() || failure()
+    needs: [ tpch-compare ]
 
     strategy:
       fail-fast: false
-      matrix: ${{ fromJson(needs.generate-matrices.outputs.olap-compare-matrix) }}
+      matrix:
+        # neon-captest-prefetch: We have pre-created projects with prefetch enabled
+        # rds-aurora: Aurora Postgres Serverless v2 with autoscaling from 0.5 to 2 ACUs
+        # rds-postgres: RDS Postgres db.m5.large instance (2 vCPU, 8 GiB) with gp3 EBS storage
+        platform: [ neon-captest-prefetch, rds-postgres, rds-aurora ]
 
     env:
       POSTGRES_DISTRIB_DIR: /tmp/neon/pg_install
       DEFAULT_PG_VERSION: 14
       TEST_OUTPUT: /tmp/test_output
       BUILD_TYPE: remote
-      SAVE_PERF_REPORT: ${{ github.event.inputs.save_perf_report || ( github.ref_name == 'main' ) }}
+      SAVE_PERF_REPORT: ${{ github.event.inputs.save_perf_report || ( github.ref == 'refs/heads/main' ) }}
       PLATFORM: ${{ matrix.platform }}
 
     runs-on: [ self-hosted, us-east-2, x64 ]
@@ -532,7 +536,7 @@ jobs:
       id: set-up-connstr
       run: |
         case "${PLATFORM}" in
-          neon-captest-reuse)
+          neon-captest-prefetch)
             CONNSTR=${{ secrets.BENCHMARK_USER_EXAMPLE_CAPTEST_CONNSTR }}
             ;;
           rds-aurora)
@@ -542,7 +546,7 @@ jobs:
             CONNSTR=${{ secrets.BENCHMARK_USER_EXAMPLE_RDS_POSTGRES_CONNSTR }}
             ;;
           *)
-            echo >&2 "Unknown PLATFORM=${PLATFORM}. Allowed only 'neon-captest-reuse', 'rds-aurora', or 'rds-postgres'"
+            echo 2>&1 "Unknown PLATFORM=${PLATFORM}. Allowed only 'neon-captest-prefetch', 'rds-aurora', or 'rds-postgres'"
             exit 1
             ;;
         esac
@@ -551,6 +555,17 @@ jobs:
 
         psql ${CONNSTR} -c "SELECT version();"
 
+    - name: Set database options
+      if: matrix.platform == 'neon-captest-prefetch'
+      run: |
+        DB_NAME=$(psql ${BENCHMARK_CONNSTR} --no-align --quiet -t -c "SELECT current_database()")
+
+        psql ${BENCHMARK_CONNSTR} -c "ALTER DATABASE ${DB_NAME} SET enable_seqscan_prefetch=on"
+        psql ${BENCHMARK_CONNSTR} -c "ALTER DATABASE ${DB_NAME} SET effective_io_concurrency=32"
+        psql ${BENCHMARK_CONNSTR} -c "ALTER DATABASE ${DB_NAME} SET maintenance_io_concurrency=32"
+      env:
+        BENCHMARK_CONNSTR: ${{ steps.set-up-connstr.outputs.connstr }}
+
     - name: Run user examples
       uses: ./.github/actions/run-python-test-set
       with:
@@ -565,7 +580,7 @@ jobs:
         BENCHMARK_CONNSTR: ${{ steps.set-up-connstr.outputs.connstr }}
 
     - name: Create Allure report
-      if: ${{ !cancelled() }}
+      if: success() || failure()
       uses: ./.github/actions/allure-report
       with:
         action: generate

.github/workflows/build_and_test.yml

@@ -13,7 +13,7 @@ defaults:
 
 concurrency:
   # Allow only one workflow per any non-`main` branch.
-  group: ${{ github.workflow }}-${{ github.ref_name }}-${{ github.ref_name == 'main' && github.sha || 'anysha' }}
+  group: ${{ github.workflow }}-${{ github.ref }}-${{ github.ref == 'refs/heads/main' && github.sha || 'anysha' }}
   cancel-in-progress: true
 
 env:
@@ -338,7 +338,6 @@ jobs:
           rerun_flaky: true
         env:
           TEST_RESULT_CONNSTR: ${{ secrets.REGRESS_TEST_RESULT_CONNSTR }}
-          CHECK_ONDISK_DATA_COMPATIBILITY: nonempty
 
       - name: Merge and upload coverage data
         if: matrix.build_type == 'debug'
@@ -368,7 +367,7 @@ jobs:
           build_type: ${{ matrix.build_type }}
           test_selection: performance
           run_in_parallel: false
-          save_perf_report: ${{ github.ref_name == 'main' }}
+          save_perf_report: ${{ github.ref == 'refs/heads/main' }}
         env:
           VIP_VAP_ACCESS_TOKEN: "${{ secrets.VIP_VAP_ACCESS_TOKEN }}"
           PERF_TEST_RESULT_CONNSTR: "${{ secrets.PERF_TEST_RESULT_CONNSTR }}"
@@ -410,8 +409,6 @@ jobs:
             steps.create-allure-report-release.outputs.report-url
           )
         with:
-          # Retry script for 5XX server errors: https://github.com/actions/github-script#retries
-          retries: 5
           script: |
             const reports = [{
               buildType: "debug",
@@ -1007,7 +1004,7 @@ jobs:
 
             S3_KEY=$(aws s3api list-objects-v2 --bucket ${BUCKET} --prefix ${OLD_PREFIX} | jq -r '.Contents[].Key' | grep ${FILENAME} | sort --version-sort | tail -1 || true)
             if [ -z "${S3_KEY}" ]; then
-              echo >&2 "Neither s3://${BUCKET}/${OLD_PREFIX}/${FILENAME} nor its version from previous attempts exist"
+              echo 2>&1 "Neither s3://${BUCKET}/${OLD_PREFIX}/${FILENAME} nor its version from previous attempts exist"
               exit 1
             fi
 

.github/workflows/neon_extra_builds.yml

@@ -12,7 +12,7 @@ defaults:
 
 concurrency:
   # Allow only one workflow per any non-`main` branch.
-  group: ${{ github.workflow }}-${{ github.ref_name }}-${{ github.ref_name == 'main' && github.sha || 'anysha' }}
+  group: ${{ github.workflow }}-${{ github.ref }}-${{ github.ref == 'refs/heads/main' && github.sha || 'anysha' }}
   cancel-in-progress: true
 
 env:

.github/workflows/pg_clients.yml

@@ -14,7 +14,7 @@ on:
 
 concurrency:
   # Allow only one workflow per any non-`main` branch.
-  group: ${{ github.workflow }}-${{ github.ref_name }}-${{ github.ref_name == 'main' && github.sha || 'anysha' }}
+  group: ${{ github.workflow }}-${{ github.ref }}-${{ github.ref == 'refs/heads/main' && github.sha || 'anysha' }}
   cancel-in-progress: true
 
 jobs:

Cargo.lock

@@ -879,7 +879,6 @@ dependencies = [
  "tracing-subscriber",
  "tracing-utils",
  "url",
- "utils",
  "workspace_hack",
 ]
 
@@ -1572,9 +1571,9 @@ checksum = "d2fabcfbdc87f4758337ca535fb41a6d701b65693ce38287d856d1674551ec9b"
 
 [[package]]
 name = "h2"
-version = "0.3.17"
+version = "0.3.15"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "66b91535aa35fea1523ad1b86cb6b53c28e0ae566ba4a460f4457e936cad7c6f"
+checksum = "5f9f29bc9dda355256b2916cf526ab02ce0aeaaaf2bad60d65ef3f12f11dd0f4"
 dependencies = [
  "bytes",
  "fnv",
@@ -2517,8 +2516,6 @@ dependencies = [
  "serde",
  "serde_json",
  "serde_with",
- "strum",
- "strum_macros",
  "utils",
  "workspace_hack",
 ]
@@ -2965,7 +2962,6 @@ dependencies = [
  "tokio-postgres",
  "tokio-postgres-rustls",
  "tokio-rustls",
- "tokio-util",
  "tracing",
  "tracing-opentelemetry",
  "tracing-subscriber",
@@ -3371,7 +3367,6 @@ dependencies = [
  "tempfile",
  "thiserror",
  "tokio",
- "tokio-io-timeout",
  "tokio-postgres",
  "toml_edit",
  "tracing",
@@ -3991,16 +3986,6 @@ dependencies = [
  "time-core",
 ]
 
-[[package]]
-name = "timeline_data_path"
-version = "0.1.0"
-dependencies = [
- "thiserror",
- "tokio",
- "utils",
- "workspace_hack",
-]
-
 [[package]]
 name = "tinytemplate"
 version = "1.2.1"
@@ -4574,7 +4559,6 @@ dependencies = [
  "byteorder",
  "bytes",
  "criterion",
- "either",
  "futures",
  "heapless",
  "hex",

Cargo.toml

@@ -101,7 +101,6 @@ test-context = "0.1"
 thiserror = "1.0"
 tls-listener = { version = "0.6", features = ["rustls", "hyper-h1"] }
 tokio = { version = "1.17", features = ["macros"] }
-tokio-io-timeout = "1.2.0"
 tokio-postgres-rustls = "0.9.0"
 tokio-rustls = "0.23"
 tokio-stream = "0.1"

Dockerfile.compute-node

@@ -38,7 +38,6 @@ RUN cd postgres && \
     echo 'trusted = true' >> /usr/local/pgsql/share/extension/insert_username.control && \
     echo 'trusted = true' >> /usr/local/pgsql/share/extension/intagg.control && \
     echo 'trusted = true' >> /usr/local/pgsql/share/extension/moddatetime.control && \
-    echo 'trusted = true' >> /usr/local/pgsql/share/extension/pg_stat_statements.control && \
     echo 'trusted = true' >> /usr/local/pgsql/share/extension/pgrowlocks.control && \
     echo 'trusted = true' >> /usr/local/pgsql/share/extension/pgstattuple.control && \
     echo 'trusted = true' >> /usr/local/pgsql/share/extension/refint.control && \
@@ -60,7 +59,6 @@ RUN apt update && \
 
 # SFCGAL > 1.3 requires CGAL > 5.2, Bullseye's libcgal-dev is 5.2
 RUN wget https://gitlab.com/Oslandia/SFCGAL/-/archive/v1.3.10/SFCGAL-v1.3.10.tar.gz -O SFCGAL.tar.gz && \
-    echo "4e39b3b2adada6254a7bdba6d297bb28e1a9835a9f879b74f37e2dab70203232 SFCGAL.tar.gz" | sha256sum --check && \
     mkdir sfcgal-src && cd sfcgal-src && tar xvzf ../SFCGAL.tar.gz --strip-components=1 -C . && \
     cmake . && make -j $(getconf _NPROCESSORS_ONLN) && \
     DESTDIR=/sfcgal make install -j $(getconf _NPROCESSORS_ONLN) && \
@@ -69,7 +67,6 @@ RUN wget https://gitlab.com/Oslandia/SFCGAL/-/archive/v1.3.10/SFCGAL-v1.3.10.tar
 ENV PATH "/usr/local/pgsql/bin:$PATH"
 
 RUN wget https://download.osgeo.org/postgis/source/postgis-3.3.2.tar.gz -O postgis.tar.gz && \
-    echo "9a2a219da005a1730a39d1959a1c7cec619b1efb009b65be80ffc25bad299068 postgis.tar.gz" | sha256sum --check && \
     mkdir postgis-src && cd postgis-src && tar xvzf ../postgis.tar.gz --strip-components=1 -C . && \
     ./autogen.sh && \
     ./configure --with-sfcgal=/usr/local/bin/sfcgal-config && \
@@ -86,7 +83,6 @@ RUN wget https://download.osgeo.org/postgis/source/postgis-3.3.2.tar.gz -O postg
     echo 'trusted = true' >> /usr/local/pgsql/share/extension/address_standardizer_data_us.control
 
 RUN wget https://github.com/pgRouting/pgrouting/archive/v3.4.2.tar.gz -O pgrouting.tar.gz && \
-    echo "cac297c07d34460887c4f3b522b35c470138760fe358e351ad1db4edb6ee306e pgrouting.tar.gz" | sha256sum --check && \
     mkdir pgrouting-src && cd pgrouting-src && tar xvzf ../pgrouting.tar.gz --strip-components=1 -C . && \
     mkdir build && \
     cd build && \
@@ -107,7 +103,6 @@ RUN apt update && \
     apt install -y ninja-build python3-dev libncurses5 binutils clang
 
 RUN wget https://github.com/plv8/plv8/archive/refs/tags/v3.1.5.tar.gz -O plv8.tar.gz && \
-    echo "1e108d5df639e4c189e1c5bdfa2432a521c126ca89e7e5a969d46899ca7bf106 plv8.tar.gz" | sha256sum --check && \
     mkdir plv8-src && cd plv8-src && tar xvzf ../plv8.tar.gz --strip-components=1 -C . && \
     export PATH="/usr/local/pgsql/bin:$PATH" && \
     make DOCKER=1 -j $(getconf _NPROCESSORS_ONLN) install && \
@@ -129,13 +124,11 @@ COPY --from=pg-build /usr/local/pgsql/ /usr/local/pgsql/
 # packaged cmake is too old
 RUN wget https://github.com/Kitware/CMake/releases/download/v3.24.2/cmake-3.24.2-linux-x86_64.sh \
       -q -O /tmp/cmake-install.sh \
-      && echo "739d372726cb23129d57a539ce1432453448816e345e1545f6127296926b6754 /tmp/cmake-install.sh" | sha256sum --check \
       && chmod u+x /tmp/cmake-install.sh \
       && /tmp/cmake-install.sh --skip-license --prefix=/usr/local/ \
       && rm /tmp/cmake-install.sh
 
 RUN wget https://github.com/uber/h3/archive/refs/tags/v4.1.0.tar.gz -O h3.tar.gz && \
-    echo "ec99f1f5974846bde64f4513cf8d2ea1b8d172d2218ab41803bf6a63532272bc h3.tar.gz" | sha256sum --check && \
     mkdir h3-src && cd h3-src && tar xvzf ../h3.tar.gz --strip-components=1 -C . && \
     mkdir build && cd build && \
     cmake .. -DCMAKE_BUILD_TYPE=Release && \
@@ -145,7 +138,6 @@ RUN wget https://github.com/uber/h3/archive/refs/tags/v4.1.0.tar.gz -O h3.tar.gz
     rm -rf build
 
 RUN wget https://github.com/zachasme/h3-pg/archive/refs/tags/v4.1.2.tar.gz -O h3-pg.tar.gz && \
-    echo "c135aa45999b2ad1326d2537c1cadef96d52660838e4ca371706c08fdea1a956 h3-pg.tar.gz" | sha256sum --check && \
     mkdir h3-pg-src && cd h3-pg-src && tar xvzf ../h3-pg.tar.gz --strip-components=1 -C . && \
     export PATH="/usr/local/pgsql/bin:$PATH" && \
     make -j $(getconf _NPROCESSORS_ONLN) && \
@@ -163,7 +155,6 @@ FROM build-deps AS unit-pg-build
 COPY --from=pg-build /usr/local/pgsql/ /usr/local/pgsql/
 
 RUN wget https://github.com/df7cb/postgresql-unit/archive/refs/tags/7.7.tar.gz -O postgresql-unit.tar.gz && \
-    echo "411d05beeb97e5a4abf17572bfcfbb5a68d98d1018918feff995f6ee3bb03e79 postgresql-unit.tar.gz" | sha256sum --check && \
     mkdir postgresql-unit-src && cd postgresql-unit-src && tar xvzf ../postgresql-unit.tar.gz --strip-components=1 -C . && \
     make -j $(getconf _NPROCESSORS_ONLN) PG_CONFIG=/usr/local/pgsql/bin/pg_config && \
     make -j $(getconf _NPROCESSORS_ONLN) install PG_CONFIG=/usr/local/pgsql/bin/pg_config && \
@@ -184,7 +175,6 @@ FROM build-deps AS vector-pg-build
 COPY --from=pg-build /usr/local/pgsql/ /usr/local/pgsql/
 
 RUN wget https://github.com/pgvector/pgvector/archive/refs/tags/v0.4.0.tar.gz -O pgvector.tar.gz && \
-    echo "b76cf84ddad452cc880a6c8c661d137ddd8679c000a16332f4f03ecf6e10bcc8 pgvector.tar.gz" | sha256sum --check && \
     mkdir pgvector-src && cd pgvector-src && tar xvzf ../pgvector.tar.gz --strip-components=1 -C . && \
     make -j $(getconf _NPROCESSORS_ONLN) PG_CONFIG=/usr/local/pgsql/bin/pg_config && \
     make -j $(getconf _NPROCESSORS_ONLN) install PG_CONFIG=/usr/local/pgsql/bin/pg_config && \
@@ -201,7 +191,6 @@ COPY --from=pg-build /usr/local/pgsql/ /usr/local/pgsql/
 
 # 9742dab1b2f297ad3811120db7b21451bca2d3c9 made on 13/11/2021
 RUN wget https://github.com/michelp/pgjwt/archive/9742dab1b2f297ad3811120db7b21451bca2d3c9.tar.gz -O pgjwt.tar.gz && \
-    echo "cfdefb15007286f67d3d45510f04a6a7a495004be5b3aecb12cda667e774203f pgjwt.tar.gz" | sha256sum --check && \
     mkdir pgjwt-src && cd pgjwt-src && tar xvzf ../pgjwt.tar.gz --strip-components=1 -C . && \
     make -j $(getconf _NPROCESSORS_ONLN) install PG_CONFIG=/usr/local/pgsql/bin/pg_config && \
     echo 'trusted = true' >> /usr/local/pgsql/share/extension/pgjwt.control
@@ -216,7 +205,6 @@ FROM build-deps AS hypopg-pg-build
 COPY --from=pg-build /usr/local/pgsql/ /usr/local/pgsql/
 
 RUN wget https://github.com/HypoPG/hypopg/archive/refs/tags/1.3.1.tar.gz -O hypopg.tar.gz && \
-    echo "e7f01ee0259dc1713f318a108f987663d60f3041948c2ada57a94b469565ca8e hypopg.tar.gz" | sha256sum --check && \
     mkdir hypopg-src && cd hypopg-src && tar xvzf ../hypopg.tar.gz --strip-components=1 -C . && \
     make -j $(getconf _NPROCESSORS_ONLN) PG_CONFIG=/usr/local/pgsql/bin/pg_config && \
     make -j $(getconf _NPROCESSORS_ONLN) install PG_CONFIG=/usr/local/pgsql/bin/pg_config && \
@@ -232,7 +220,6 @@ FROM build-deps AS pg-hashids-pg-build
 COPY --from=pg-build /usr/local/pgsql/ /usr/local/pgsql/
 
 RUN wget https://github.com/iCyberon/pg_hashids/archive/refs/tags/v1.2.1.tar.gz -O pg_hashids.tar.gz && \
-    echo "74576b992d9277c92196dd8d816baa2cc2d8046fe102f3dcd7f3c3febed6822a pg_hashids.tar.gz" | sha256sum --check && \
     mkdir pg_hashids-src && cd pg_hashids-src && tar xvzf ../pg_hashids.tar.gz --strip-components=1 -C . && \
     make -j $(getconf _NPROCESSORS_ONLN) PG_CONFIG=/usr/local/pgsql/bin/pg_config USE_PGXS=1 && \
     make -j $(getconf _NPROCESSORS_ONLN) install PG_CONFIG=/usr/local/pgsql/bin/pg_config USE_PGXS=1 && \
@@ -248,7 +235,6 @@ FROM build-deps AS rum-pg-build
 COPY --from=pg-build /usr/local/pgsql/ /usr/local/pgsql/
 
 RUN wget https://github.com/postgrespro/rum/archive/refs/tags/1.3.13.tar.gz -O rum.tar.gz && \
-    echo "6ab370532c965568df6210bd844ac6ba649f53055e48243525b0b7e5c4d69a7d rum.tar.gz" | sha256sum --check && \
     mkdir rum-src && cd rum-src && tar xvzf ../rum.tar.gz --strip-components=1 -C . && \
     make -j $(getconf _NPROCESSORS_ONLN) PG_CONFIG=/usr/local/pgsql/bin/pg_config USE_PGXS=1 && \
     make -j $(getconf _NPROCESSORS_ONLN) install PG_CONFIG=/usr/local/pgsql/bin/pg_config USE_PGXS=1 && \
@@ -264,28 +250,11 @@ FROM build-deps AS pgtap-pg-build
 COPY --from=pg-build /usr/local/pgsql/ /usr/local/pgsql/
 
 RUN wget https://github.com/theory/pgtap/archive/refs/tags/v1.2.0.tar.gz -O pgtap.tar.gz && \
-    echo "9c7c3de67ea41638e14f06da5da57bac6f5bd03fea05c165a0ec862205a5c052 pgtap.tar.gz" | sha256sum --check && \
     mkdir pgtap-src && cd pgtap-src && tar xvzf ../pgtap.tar.gz --strip-components=1 -C . && \
     make -j $(getconf _NPROCESSORS_ONLN) PG_CONFIG=/usr/local/pgsql/bin/pg_config && \
     make -j $(getconf _NPROCESSORS_ONLN) install PG_CONFIG=/usr/local/pgsql/bin/pg_config && \
     echo 'trusted = true' >> /usr/local/pgsql/share/extension/pgtap.control
 
-#########################################################################################
-#
-# Layer "ip4r-pg-build"
-# compile ip4r extension
-#
-#########################################################################################
-FROM build-deps AS ip4r-pg-build
-COPY --from=pg-build /usr/local/pgsql/ /usr/local/pgsql/
-
-RUN wget https://github.com/RhodiumToad/ip4r/archive/refs/tags/2.4.1.tar.gz -O ip4r.tar.gz && \
-    echo "78b9f0c1ae45c22182768fe892a32d533c82281035e10914111400bf6301c726 ip4r.tar.gz" | sha256sum --check && \
-    mkdir ip4r-src && cd ip4r-src && tar xvzf ../ip4r.tar.gz --strip-components=1 -C . && \
-    make -j $(getconf _NPROCESSORS_ONLN) PG_CONFIG=/usr/local/pgsql/bin/pg_config && \
-    make -j $(getconf _NPROCESSORS_ONLN) install PG_CONFIG=/usr/local/pgsql/bin/pg_config && \
-    echo 'trusted = true' >> /usr/local/pgsql/share/extension/ip4r.control
-
 #########################################################################################
 #
 # Layer "prefix-pg-build"
@@ -296,7 +265,6 @@ FROM build-deps AS prefix-pg-build
 COPY --from=pg-build /usr/local/pgsql/ /usr/local/pgsql/
 
 RUN wget https://github.com/dimitri/prefix/archive/refs/tags/v1.2.9.tar.gz -O prefix.tar.gz && \
-    echo "38d30a08d0241a8bbb8e1eb8f0152b385051665a8e621c8899e7c5068f8b511e prefix.tar.gz" | sha256sum --check && \
     mkdir prefix-src && cd prefix-src && tar xvzf ../prefix.tar.gz --strip-components=1 -C . && \
     make -j $(getconf _NPROCESSORS_ONLN) PG_CONFIG=/usr/local/pgsql/bin/pg_config && \
     make -j $(getconf _NPROCESSORS_ONLN) install PG_CONFIG=/usr/local/pgsql/bin/pg_config && \
@@ -312,7 +280,6 @@ FROM build-deps AS hll-pg-build
 COPY --from=pg-build /usr/local/pgsql/ /usr/local/pgsql/
 
 RUN wget https://github.com/citusdata/postgresql-hll/archive/refs/tags/v2.17.tar.gz -O hll.tar.gz && \
-    echo "9a18288e884f197196b0d29b9f178ba595b0dfc21fbf7a8699380e77fa04c1e9 hll.tar.gz" | sha256sum --check && \
     mkdir hll-src && cd hll-src && tar xvzf ../hll.tar.gz --strip-components=1 -C . && \
     make -j $(getconf _NPROCESSORS_ONLN) PG_CONFIG=/usr/local/pgsql/bin/pg_config && \
     make -j $(getconf _NPROCESSORS_ONLN) install PG_CONFIG=/usr/local/pgsql/bin/pg_config && \
@@ -328,68 +295,13 @@ FROM build-deps AS plpgsql-check-pg-build
 COPY --from=pg-build /usr/local/pgsql/ /usr/local/pgsql/
 
 RUN wget https://github.com/okbob/plpgsql_check/archive/refs/tags/v2.3.2.tar.gz -O plpgsql_check.tar.gz && \
-    echo "9d81167c4bbeb74eebf7d60147b21961506161addc2aee537f95ad8efeae427b plpgsql_check.tar.gz" | sha256sum --check && \
     mkdir plpgsql_check-src && cd plpgsql_check-src && tar xvzf ../plpgsql_check.tar.gz --strip-components=1 -C . && \
     make -j $(getconf _NPROCESSORS_ONLN) PG_CONFIG=/usr/local/pgsql/bin/pg_config USE_PGXS=1 && \
     make -j $(getconf _NPROCESSORS_ONLN) install PG_CONFIG=/usr/local/pgsql/bin/pg_config USE_PGXS=1 && \
     echo 'trusted = true' >> /usr/local/pgsql/share/extension/plpgsql_check.control
 
 #########################################################################################
-#
-# Layer "timescaledb-pg-build"
-# compile timescaledb extension
-#
-#########################################################################################
-FROM build-deps AS timescaledb-pg-build
-COPY --from=pg-build /usr/local/pgsql/ /usr/local/pgsql/
-
-ENV PATH "/usr/local/pgsql/bin:$PATH"
-
-RUN apt-get update && \
-    apt-get install -y cmake && \
-    wget https://github.com/timescale/timescaledb/archive/refs/tags/2.10.1.tar.gz -O timescaledb.tar.gz && \
-    echo "6fca72a6ed0f6d32d2b3523951ede73dc5f9b0077b38450a029a5f411fdb8c73 timescaledb.tar.gz" | sha256sum --check && \
-    mkdir timescaledb-src && cd timescaledb-src && tar xvzf ../timescaledb.tar.gz --strip-components=1 -C . && \
-    ./bootstrap -DSEND_TELEMETRY_DEFAULT:BOOL=OFF -DUSE_TELEMETRY:BOOL=OFF -DAPACHE_ONLY:BOOL=ON && \
-    cd build && \
-    make -j $(getconf _NPROCESSORS_ONLN) && \
-    make install -j $(getconf _NPROCESSORS_ONLN) && \
-    echo "trusted = true" >> /usr/local/pgsql/share/extension/timescaledb.control
-
-#########################################################################################
-#
-# Layer "pg-hint-plan-pg-build"
-# compile pg_hint_plan extension
-#
-#########################################################################################
-FROM build-deps AS pg-hint-plan-pg-build
-COPY --from=pg-build /usr/local/pgsql/ /usr/local/pgsql/
-
-ARG PG_VERSION
-ENV PATH "/usr/local/pgsql/bin:$PATH"
-
-RUN case "${PG_VERSION}" in \
-      "v14") \
-        export PG_HINT_PLAN_VERSION=14_1_4_1 \
-        export PG_HINT_PLAN_CHECKSUM=c3501becf70ead27f70626bce80ea401ceac6a77e2083ee5f3ff1f1444ec1ad1 \
-        ;; \
-      "v15") \
-        export PG_HINT_PLAN_VERSION=15_1_5_0 \
-        export PG_HINT_PLAN_CHECKSUM=564cbbf4820973ffece63fbf76e3c0af62c4ab23543142c7caaa682bc48918be \
-        ;; \
-      *) \
-        echo "Export the valid PG_HINT_PLAN_VERSION variable" && exit 1 \
-        ;; \
-    esac && \
-    wget https://github.com/ossc-db/pg_hint_plan/archive/refs/tags/REL${PG_HINT_PLAN_VERSION}.tar.gz -O pg_hint_plan.tar.gz && \
-    echo "${PG_HINT_PLAN_CHECKSUM} pg_hint_plan.tar.gz" | sha256sum --check && \
-    mkdir pg_hint_plan-src && cd pg_hint_plan-src && tar xvzf ../pg_hint_plan.tar.gz --strip-components=1 -C . && \
-    make -j $(getconf _NPROCESSORS_ONLN) && \
-    make install -j $(getconf _NPROCESSORS_ONLN) && \
-    echo "trusted = true" >> /usr/local/pgsql/share/extension/pg_hint_plan.control
-
-#########################################################################################
-#
+# 
 # Layer "rust extensions"
 # This layer is used to build `pgx` deps
 #
@@ -417,7 +329,7 @@ RUN curl -sSO https://static.rust-lang.org/rustup/dist/$(uname -m)-unknown-linux
 USER root
 
 #########################################################################################
-#
+# 
 # Layer "pg-jsonschema-pg-build"
 # Compile "pg_jsonschema" extension
 #
@@ -425,17 +337,15 @@ USER root
 
 FROM rust-extensions-build AS pg-jsonschema-pg-build
 
-# caeab60d70b2fd3ae421ec66466a3abbb37b7ee6 made on 06/03/2023
-# there is no release tag yet, but we need it due to the superuser fix in the control file, switch to git tag after release >= 0.1.5
+# there is no release tag yet, but we need it due to the superuser fix in the control file
 RUN wget https://github.com/supabase/pg_jsonschema/archive/caeab60d70b2fd3ae421ec66466a3abbb37b7ee6.tar.gz -O pg_jsonschema.tar.gz && \
-    echo "54129ce2e7ee7a585648dbb4cef6d73f795d94fe72f248ac01119992518469a4 pg_jsonschema.tar.gz" | sha256sum --check && \
     mkdir pg_jsonschema-src && cd pg_jsonschema-src && tar xvzf ../pg_jsonschema.tar.gz --strip-components=1 -C . && \
     sed -i 's/pgx = "0.7.1"/pgx = { version = "0.7.3", features = [ "unsafe-postgres" ] }/g' Cargo.toml && \
     cargo pgx install --release && \
     echo "trusted = true" >> /usr/local/pgsql/share/extension/pg_jsonschema.control
 
 #########################################################################################
-#
+# 
 # Layer "pg-graphql-pg-build"
 # Compile "pg_graphql" extension
 #
@@ -443,13 +353,11 @@ RUN wget https://github.com/supabase/pg_jsonschema/archive/caeab60d70b2fd3ae421e
 
 FROM rust-extensions-build AS pg-graphql-pg-build
 
-# b4988843647450a153439be367168ed09971af85 made on 22/02/2023 (from remove-pgx-contrib-spiext branch)
 # Currently pgx version bump to >= 0.7.2  causes "call to unsafe function" compliation errors in
 # pgx-contrib-spiext. There is a branch that removes that dependency, so use it. It is on the
 # same 1.1 version we've used before.
-RUN wget https://github.com/yrashk/pg_graphql/archive/b4988843647450a153439be367168ed09971af85.tar.gz -O pg_graphql.tar.gz && \
-    echo "0c7b0e746441b2ec24187d0e03555faf935c2159e2839bddd14df6dafbc8c9bd pg_graphql.tar.gz" | sha256sum --check && \
-    mkdir pg_graphql-src && cd pg_graphql-src && tar xvzf ../pg_graphql.tar.gz --strip-components=1 -C . && \
+RUN git clone -b remove-pgx-contrib-spiext --single-branch https://github.com/yrashk/pg_graphql && \
+    cd pg_graphql && \
     sed -i 's/pgx = "~0.7.1"/pgx = { version = "0.7.3", features = [ "unsafe-postgres" ] }/g' Cargo.toml && \
     sed -i 's/pgx-tests = "~0.7.1"/pgx-tests = "0.7.3"/g' Cargo.toml && \
     cargo pgx install --release && \
@@ -466,10 +374,8 @@ RUN wget https://github.com/yrashk/pg_graphql/archive/b4988843647450a153439be367
 
 FROM rust-extensions-build AS pg-tiktoken-pg-build
 
-# 801f84f08c6881c8aa30f405fafbf00eec386a72 made on 10/03/2023
-RUN wget https://github.com/kelvich/pg_tiktoken/archive/801f84f08c6881c8aa30f405fafbf00eec386a72.tar.gz -O pg_tiktoken.tar.gz && \
-    echo "52f60ac800993a49aa8c609961842b611b6b1949717b69ce2ec9117117e16e4a pg_tiktoken.tar.gz" | sha256sum --check && \
-    mkdir pg_tiktoken-src && cd pg_tiktoken-src && tar xvzf ../pg_tiktoken.tar.gz --strip-components=1 -C . && \
+RUN git clone --depth=1 --single-branch https://github.com/kelvich/pg_tiktoken && \
+    cd pg_tiktoken && \
     cargo pgx install --release && \
     echo "trusted = true" >> /usr/local/pgsql/share/extension/pg_tiktoken.control
 
@@ -495,12 +401,9 @@ COPY --from=hypopg-pg-build /usr/local/pgsql/ /usr/local/pgsql/
 COPY --from=pg-hashids-pg-build /usr/local/pgsql/ /usr/local/pgsql/
 COPY --from=rum-pg-build /usr/local/pgsql/ /usr/local/pgsql/
 COPY --from=pgtap-pg-build /usr/local/pgsql/ /usr/local/pgsql/
-COPY --from=ip4r-pg-build /usr/local/pgsql/ /usr/local/pgsql/
 COPY --from=prefix-pg-build /usr/local/pgsql/ /usr/local/pgsql/
 COPY --from=hll-pg-build /usr/local/pgsql/ /usr/local/pgsql/
 COPY --from=plpgsql-check-pg-build /usr/local/pgsql/ /usr/local/pgsql/
-COPY --from=timescaledb-pg-build /usr/local/pgsql/ /usr/local/pgsql/
-COPY --from=pg-hint-plan-pg-build /usr/local/pgsql/ /usr/local/pgsql/
 COPY pgxn/ pgxn/
 
 RUN make -j $(getconf _NPROCESSORS_ONLN) \

compute_tools/Cargo.toml

@@ -28,5 +28,4 @@ tracing-utils.workspace = true
 url.workspace = true
 
 compute_api.workspace = true
-utils.workspace = true
 workspace_hack.workspace = true

compute_tools/src/bin/compute_ctl.rs

@@ -34,24 +34,24 @@ use std::fs::File;
 use std::panic;
 use std::path::Path;
 use std::process::exit;
-use std::sync::{mpsc, Arc, Condvar, Mutex};
+use std::sync::{Arc, Condvar, Mutex};
 use std::{thread, time::Duration};
 
-use anyhow::{Context, Result};
+use anyhow::{anyhow, Context, Result};
 use chrono::Utc;
 use clap::Arg;
 use tracing::{error, info};
 use url::Url;
 
-use compute_api::responses::ComputeStatus;
+use compute_api::models::{ComputeMetrics, ComputeState, ComputeStatus};
 
-use compute_tools::compute::{ComputeNode, ComputeState, ParsedSpec};
+use compute_tools::compute::{ComputeNode, ComputeNodeInner, ParsedSpec};
 use compute_tools::configurator::launch_configurator;
 use compute_tools::http::api::launch_http_server;
 use compute_tools::logger::*;
 use compute_tools::monitor::launch_monitor;
 use compute_tools::params::*;
-use compute_tools::spec::*;
+use compute_tools::spec::get_spec_from_control_plane;
 
 fn main() -> Result<()> {
     init_tracing_and_logging(DEFAULT_LOG_LEVEL)?;
@@ -104,24 +104,37 @@ fn main() -> Result<()> {
         }
     };
 
-    let mut new_state = ComputeState::new();
-    let spec_set;
-    if let Some(spec) = spec {
-        let pspec = ParsedSpec::try_from(spec).map_err(|msg| anyhow::anyhow!(msg))?;
-        new_state.pspec = Some(pspec);
-        spec_set = true;
-    } else {
-        spec_set = false;
-    }
+    // Volatile compute state under mutex and condition variable to notify everyone
+    // who is interested in the state changes.
     let compute_node = ComputeNode {
         start_time: Utc::now(),
         connstr: Url::parse(connstr).context("cannot parse connstr as a URL")?,
         pgdata: pgdata.to_string(),
         pgbin: pgbin.to_string(),
         live_config_allowed,
-        state: Mutex::new(new_state),
-        state_changed: Condvar::new(),
+        inner: Mutex::new(ComputeNodeInner {
+            state: ComputeState {
+                status: ComputeStatus::Empty,
+                last_active: Utc::now(),
+                error: None,
+            },
+            spec: None,
+            metrics: ComputeMetrics::default(),
+        }),
+        state_changed: Condvar::new()
     };
+
+    // If we have a spec already, go immediately into Init state.
+    let spec_set = spec.is_some();
+    if let Some(spec) = spec {
+        let mut inner = compute_node.inner.lock().unwrap();
+
+        let parsed_spec = ParsedSpec::try_from(spec)
+            .map_err(|msg| anyhow!("error parsing compute spec: {msg}"))?;
+        inner.spec = Some(parsed_spec);
+        inner.state.status = ComputeStatus::Init;
+    }
+
     let compute = Arc::new(compute_node);
 
     // Launch http service first, so we were able to serve control-plane
@@ -129,27 +142,27 @@ fn main() -> Result<()> {
     let _http_handle = launch_http_server(&compute).expect("cannot launch http endpoint thread");
 
     if !spec_set {
-        // No spec provided, hang waiting for it.
+        // No spec was provided earlier, hang waiting for it.
         info!("no compute spec provided, waiting");
-        let mut state = compute.state.lock().unwrap();
-        while state.status != ComputeStatus::ConfigurationPending {
-            state = compute.state_changed.wait(state).unwrap();
 
-            if state.status == ComputeStatus::ConfigurationPending {
+        let mut inner = compute.inner.lock().unwrap();
+        while inner.state.status != ComputeStatus::ConfigurationPending {
+            inner = compute.state_changed.wait(inner).unwrap();
+
+            if inner.state.status == ComputeStatus::ConfigurationPending {
                 info!("got spec, continue configuration");
                 // Spec is already set by the http server handler.
+                inner.state.status = ComputeStatus::Init;
                 break;
             }
         }
-    }
+    };
 
-    // We got all we need, update the state.
-    let mut state = compute.state.lock().unwrap();
-    let pspec = state.pspec.as_ref().expect("spec must be set");
-    let startup_tracing_context = pspec.spec.startup_tracing_context.clone();
-    state.status = ComputeStatus::Init;
-    compute.state_changed.notify_all();
-    drop(state);
+    // We got the spec. Start up
+    let startup_tracing_context = {
+        let inner = compute.inner.lock().unwrap();
+        inner.spec.as_ref().unwrap().spec.startup_tracing_context.clone()
+    };
 
     // Extract OpenTelemetry context for the startup actions from the spec, and
     // attach it to the current tracing context.
@@ -186,10 +199,10 @@ fn main() -> Result<()> {
         Ok(pg) => Some(pg),
         Err(err) => {
             error!("could not start the compute node: {:?}", err);
-            let mut state = compute.state.lock().unwrap();
-            state.error = Some(format!("{:?}", err));
-            state.status = ComputeStatus::Failed;
-            drop(state);
+            let mut inner = compute.inner.lock().unwrap();
+            inner.state.error = Some(format!("{:?}", err));
+            inner.state.status = ComputeStatus::Failed;
+            drop(inner);
             delay_exit = true;
             None
         }
@@ -219,25 +232,10 @@ fn main() -> Result<()> {
         thread::sleep(Duration::from_secs(30));
     }
 
-    // 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");
-    }
+    // Shutdown trace pipeline gracefully, so that it has a chance to send any
+    // pending traces before we exit.
+    tracing_utils::shutdown_tracing();
 
     info!("shutting down");
     exit(exit_code.unwrap_or(1))

compute_tools/src/compute.rs

@@ -26,10 +26,8 @@ use chrono::{DateTime, Utc};
 use postgres::{Client, NoTls};
 use tokio_postgres;
 use tracing::{info, instrument, warn};
-use utils::id::{TenantId, TimelineId};
-use utils::lsn::Lsn;
 
-use compute_api::responses::{ComputeMetrics, ComputeStatus};
+use compute_api::models::{ComputeMetrics, ComputeState, ComputeStatus};
 use compute_api::spec::ComputeSpec;
 
 use crate::checker::create_writability_check_data;
@@ -44,66 +42,47 @@ pub struct ComputeNode {
     pub connstr: url::Url,
     pub pgdata: String,
     pub pgbin: String,
-    /// We should only allow live re- / configuration of the compute node if
-    /// it uses 'pull model', i.e. it can go to control-plane and fetch
-    /// the latest configuration. Otherwise, there could be a case:
-    /// - we start compute with some spec provided as argument
-    /// - we push new spec and it does reconfiguration
-    /// - but then something happens and compute pod / VM is destroyed,
-    ///   so k8s controller starts it again with the **old** spec
-    /// and the same for empty computes:
-    /// - we started compute without any spec
-    /// - we push spec and it does configuration
-    /// - but then it is restarted without any spec again
+    // We only allow live re- / configuration of the compute node if
+    // it uses 'pull model', i.e. it can go to control-plane and fetch
+    // the latest configuration. Otherwise, there could be a case:
+    // - we start compute with some spec provided as argument
+    // - we push new spec and it does reconfiguration
+    // - but then something happens and compute pod / VM is destroyed,
+    //   so k8s controller starts it again with the **old** spec
     pub live_config_allowed: bool,
+
     /// Volatile part of the `ComputeNode`, which should be used under `Mutex`.
-    /// To allow HTTP API server to serving status requests, while configuration
-    /// is in progress, lock should be held only for short periods of time to do
-    /// read/write, not the whole configuration process.
-    pub state: Mutex<ComputeState>,
-    /// `Condvar` to allow notifying waiters about state changes.
+    /// Coupled with `Condvar` to allow notifying HTTP API and configurator
+    /// thread about state changes. To allow HTTP API server to serving status
+    /// requests, while configuration is in progress, lock should be held only
+    /// for short periods of time to do read/write, not the whole configuration
+    /// process.
+    pub inner: Mutex<ComputeNodeInner>,
     pub state_changed: Condvar,
 }
 
-#[derive(Clone, Debug)]
-pub struct ComputeState {
-    pub status: ComputeStatus,
-    /// Timestamp of the last Postgres activity
-    pub last_active: DateTime<Utc>,
-    pub error: Option<String>,
-    pub pspec: Option<ParsedSpec>,
+pub struct ComputeNodeInner {
+    pub state: ComputeState,
+
+    pub spec: Option<ParsedSpec>,
+
     pub metrics: ComputeMetrics,
 }
 
-impl ComputeState {
-    pub fn new() -> Self {
-        Self {
-            status: ComputeStatus::Empty,
-            last_active: Utc::now(),
-            error: None,
-            pspec: None,
-            metrics: ComputeMetrics::default(),
-        }
-    }
-}
-
-impl Default for ComputeState {
-    fn default() -> Self {
-        Self::new()
-    }
-}
-
-#[derive(Clone, Debug)]
+#[derive(Clone)]
 pub struct ParsedSpec {
     pub spec: ComputeSpec,
-    pub tenant_id: TenantId,
-    pub timeline_id: TimelineId,
+
+    // extra fields extracted from 'spec'.
+    pub tenant: String,
+    pub timeline: String,
     pub pageserver_connstr: String,
     pub storage_auth_token: Option<String>,
 }
 
 impl TryFrom<ComputeSpec> for ParsedSpec {
     type Error = String;
+
     fn try_from(spec: ComputeSpec) -> Result<Self, String> {
         let pageserver_connstr = spec
             .cluster
@@ -111,40 +90,44 @@ impl TryFrom<ComputeSpec> for ParsedSpec {
             .find("neon.pageserver_connstring")
             .ok_or("pageserver connstr should be provided")?;
         let storage_auth_token = spec.storage_auth_token.clone();
-        let tenant_id: TenantId = spec
+        let tenant = spec
             .cluster
             .settings
             .find("neon.tenant_id")
-            .ok_or("tenant id should be provided")
-            .map(|s| TenantId::from_str(&s))?
-            .or(Err("invalid tenant id"))?;
-        let timeline_id: TimelineId = spec
+            .ok_or("tenant id should be provided")?;
+        let timeline = spec
             .cluster
             .settings
             .find("neon.timeline_id")
-            .ok_or("timeline id should be provided")
-            .map(|s| TimelineId::from_str(&s))?
-            .or(Err("invalid timeline id"))?;
+            .ok_or("tenant id should be provided")?;
 
         Ok(ParsedSpec {
             spec,
             pageserver_connstr,
             storage_auth_token,
-            tenant_id,
-            timeline_id,
+            tenant,
+            timeline,
         })
     }
 }
 
 impl ComputeNode {
     pub fn set_status(&self, status: ComputeStatus) {
-        let mut state = self.state.lock().unwrap();
-        state.status = status;
+        let mut inner = self.inner.lock().unwrap();
+        inner.state.status = status;
         self.state_changed.notify_all();
     }
 
     pub fn get_status(&self) -> ComputeStatus {
-        self.state.lock().unwrap().status
+        self.inner.lock().unwrap().state.status
+    }
+
+    pub fn get_state(&self) -> ComputeState {
+        self.inner.lock().unwrap().state.clone()
+    }
+
+    pub fn get_metrics(&self) -> ComputeMetrics {
+        self.inner.lock().unwrap().metrics.clone()
     }
 
     // Remove `pgdata` directory and create it again with right permissions.
@@ -160,9 +143,8 @@ impl ComputeNode {
 
     // Get basebackup from the libpq connection to pageserver using `connstr` and
     // unarchive it to `pgdata` directory overriding all its previous content.
-    #[instrument(skip(self, compute_state))]
-    fn get_basebackup(&self, compute_state: &ComputeState, lsn: Lsn) -> Result<()> {
-        let spec = compute_state.pspec.as_ref().expect("spec must be set");
+    #[instrument(skip(self, spec))]
+    fn get_basebackup(&self, spec: &ParsedSpec, lsn: &str) -> Result<()> {
         let start_time = Utc::now();
 
         let mut config = postgres::Config::from_str(&spec.pageserver_connstr)?;
@@ -178,8 +160,8 @@ impl ComputeNode {
 
         let mut client = config.connect(NoTls)?;
         let basebackup_cmd = match lsn {
-            Lsn(0) => format!("basebackup {} {}", spec.tenant_id, spec.timeline_id), // First start of the compute
-            _ => format!("basebackup {} {} {}", spec.tenant_id, spec.timeline_id, lsn),
+            "0/0" => format!("basebackup {} {}", &spec.tenant, &spec.timeline), // First start of the compute
+            _ => format!("basebackup {} {} {}", &spec.tenant, &spec.timeline, lsn),
         };
         let copyreader = client.copy_out(basebackup_cmd.as_str())?;
 
@@ -192,7 +174,7 @@ impl ComputeNode {
         ar.set_ignore_zeros(true);
         ar.unpack(&self.pgdata)?;
 
-        self.state.lock().unwrap().metrics.basebackup_ms = Utc::now()
+        self.inner.lock().unwrap().metrics.basebackup_ms = Utc::now()
             .signed_duration_since(start_time)
             .to_std()
             .unwrap()
@@ -203,7 +185,7 @@ impl ComputeNode {
     // Run `postgres` in a special mode with `--sync-safekeepers` argument
     // and return the reported LSN back to the caller.
     #[instrument(skip(self, storage_auth_token))]
-    fn sync_safekeepers(&self, storage_auth_token: Option<String>) -> Result<Lsn> {
+    fn sync_safekeepers(&self, storage_auth_token: Option<String>) -> Result<String> {
         let start_time = Utc::now();
 
         let sync_handle = Command::new(&self.pgbin)
@@ -234,42 +216,41 @@ impl ComputeNode {
             );
         }
 
-        self.state.lock().unwrap().metrics.sync_safekeepers_ms = Utc::now()
+        self.inner.lock().unwrap().metrics.sync_safekeepers_ms = Utc::now()
             .signed_duration_since(start_time)
             .to_std()
             .unwrap()
             .as_millis() as u64;
 
-        let lsn = Lsn::from_str(String::from_utf8(sync_output.stdout)?.trim())?;
+        let lsn = String::from(String::from_utf8(sync_output.stdout)?.trim());
 
         Ok(lsn)
     }
 
     /// Do all the preparations like PGDATA directory creation, configuration,
     /// safekeepers sync, basebackup, etc.
-    #[instrument(skip(self, compute_state))]
-    pub fn prepare_pgdata(&self, compute_state: &ComputeState) -> Result<()> {
-        let pspec = compute_state.pspec.as_ref().expect("spec must be set");
+    #[instrument(skip(self, spec))]
+    fn prepare_pgdata(&self, spec: &ParsedSpec) -> Result<()> {
         let pgdata_path = Path::new(&self.pgdata);
 
         // Remove/create an empty pgdata directory and put configuration there.
         self.create_pgdata()?;
-        config::write_postgres_conf(&pgdata_path.join("postgresql.conf"), &pspec.spec)?;
+        config::write_postgres_conf(&pgdata_path.join("postgresql.conf"), &spec.spec)?;
 
         info!("starting safekeepers syncing");
         let lsn = self
-            .sync_safekeepers(pspec.storage_auth_token.clone())
+            .sync_safekeepers(spec.storage_auth_token.clone())
             .with_context(|| "failed to sync safekeepers")?;
         info!("safekeepers synced at LSN {}", lsn);
 
         info!(
             "getting basebackup@{} from pageserver {}",
-            lsn, &pspec.pageserver_connstr
+            lsn, &spec.pageserver_connstr
         );
-        self.get_basebackup(compute_state, lsn).with_context(|| {
+        self.get_basebackup(spec, &lsn).with_context(|| {
             format!(
                 "failed to get basebackup@{} from pageserver {}",
-                lsn, &pspec.pageserver_connstr
+                lsn, &spec.pageserver_connstr
             )
         })?;
 
@@ -305,8 +286,8 @@ impl ComputeNode {
     }
 
     /// Do initial configuration of the already started Postgres.
-    #[instrument(skip(self, compute_state))]
-    pub fn apply_config(&self, compute_state: &ComputeState) -> Result<()> {
+    #[instrument(skip(self, spec))]
+    fn apply_config(&self, spec: &ParsedSpec) -> Result<()> {
         // If connection fails,
         // it may be the old node with `zenith_admin` superuser.
         //
@@ -337,20 +318,19 @@ impl ComputeNode {
         };
 
         // Proceed with post-startup configuration. Note, that order of operations is important.
-        let spec = &compute_state.pspec.as_ref().expect("spec must be set").spec;
-        handle_roles(spec, &mut client)?;
-        handle_databases(spec, &mut client)?;
-        handle_role_deletions(spec, self.connstr.as_str(), &mut client)?;
-        handle_grants(spec, self.connstr.as_str(), &mut client)?;
+        handle_roles(&spec.spec, &mut client)?;
+        handle_databases(&spec.spec, &mut client)?;
+        handle_role_deletions(&spec.spec, self.connstr.as_str(), &mut client)?;
+        handle_grants(&spec.spec, self.connstr.as_str(), &mut client)?;
         create_writability_check_data(&mut client)?;
-        handle_extensions(spec, &mut client)?;
+        handle_extensions(&spec.spec, &mut client)?;
 
         // 'Close' connection
         drop(client);
 
         info!(
             "finished configuration of compute for project {}",
-            spec.cluster.cluster_id
+            spec.spec.cluster.cluster_id
         );
 
         Ok(())
@@ -369,7 +349,10 @@ impl ComputeNode {
     /// as it's used to reconfigure a previously started and configured Postgres node.
     #[instrument(skip(self))]
     pub fn reconfigure(&self) -> Result<()> {
-        let spec = self.state.lock().unwrap().pspec.clone().unwrap().spec;
+        let spec = {
+            let inner = self.inner.lock().unwrap();
+            inner.spec.as_ref().expect("cannot start_compute without spec").spec.clone()
+        };
 
         // Write new config
         let pgdata_path = Path::new(&self.pgdata);
@@ -400,33 +383,39 @@ impl ComputeNode {
 
     #[instrument(skip(self))]
     pub fn start_compute(&self) -> Result<std::process::Child> {
-        let compute_state = self.state.lock().unwrap().clone();
-        let spec = compute_state.pspec.as_ref().expect("spec must be set");
+        let spec = self
+            .inner
+            .lock()
+            .unwrap()
+            .spec
+            .as_ref()
+            .expect("cannot start_compute without spec")
+            .clone();
         info!(
             "starting compute for project {}, operation {}, tenant {}, timeline {}",
             spec.spec.cluster.cluster_id,
-            spec.spec.operation_uuid.as_deref().unwrap_or("None"),
-            spec.tenant_id,
-            spec.timeline_id,
+            spec.spec.operation_uuid.as_ref().unwrap(),
+            spec.tenant,
+            spec.timeline,
         );
 
-        self.prepare_pgdata(&compute_state)?;
+        self.prepare_pgdata(&spec)?;
 
         let start_time = Utc::now();
 
         let pg = self.start_postgres(spec.storage_auth_token.clone())?;
 
-        self.apply_config(&compute_state)?;
+        self.apply_config(&spec)?;
 
         let startup_end_time = Utc::now();
         {
-            let mut state = self.state.lock().unwrap();
-            state.metrics.config_ms = startup_end_time
+            let mut inner = self.inner.lock().unwrap();
+            inner.metrics.config_ms = startup_end_time
                 .signed_duration_since(start_time)
                 .to_std()
                 .unwrap()
                 .as_millis() as u64;
-            state.metrics.total_startup_ms = startup_end_time
+            inner.metrics.total_startup_ms = startup_end_time
                 .signed_duration_since(self.start_time)
                 .to_std()
                 .unwrap()

compute_tools/src/configurator.rs

@@ -4,22 +4,21 @@ use std::thread;
 use anyhow::Result;
 use tracing::{error, info, instrument};
 
-use compute_api::responses::ComputeStatus;
-
 use crate::compute::ComputeNode;
+use compute_api::models::ComputeStatus;
 
 #[instrument(skip(compute))]
 fn configurator_main_loop(compute: &Arc<ComputeNode>) {
     info!("waiting for reconfiguration requests");
     loop {
-        let state = compute.state.lock().unwrap();
-        let mut state = compute.state_changed.wait(state).unwrap();
+        let inner = compute.inner.lock().unwrap();
+        let mut inner = compute.state_changed.wait(inner).unwrap();
 
-        if state.status == ComputeStatus::ConfigurationPending {
+        if inner.state.status == ComputeStatus::ConfigurationPending {
             info!("got configuration request");
-            state.status = ComputeStatus::Configuration;
+            inner.state.status = ComputeStatus::Configuration;
             compute.state_changed.notify_all();
-            drop(state);
+            drop(inner);
 
             let mut new_status = ComputeStatus::Failed;
             if let Err(e) = compute.reconfigure() {
@@ -30,14 +29,14 @@ fn configurator_main_loop(compute: &Arc<ComputeNode>) {
             }
 
             // XXX: used to test that API is blocking
-            // std::thread::sleep(std::time::Duration::from_millis(10000));
+            // std::thread::sleep(std::time::Duration::from_millis(2000));
 
             compute.set_status(new_status);
-        } else if state.status == ComputeStatus::Failed {
+        } else if inner.state.status == ComputeStatus::Failed {
             info!("compute node is now in Failed state, exiting");
             break;
         } else {
-            info!("woken up for compute status: {:?}, sleeping", state.status);
+            info!("woken up for compute status: {:?}, sleeping", inner.state.status);
         }
     }
 }

compute_tools/src/http/api.rs

@@ -3,33 +3,54 @@ use std::net::SocketAddr;
 use std::sync::Arc;
 use std::thread;
 
-use crate::compute::{ComputeNode, ComputeState, ParsedSpec};
-use compute_api::requests::ConfigurationRequest;
-use compute_api::responses::{ComputeStatus, ComputeStatusResponse, GenericAPIError};
+use crate::compute::{ComputeNode, ParsedSpec};
+use crate::http::models::{ConfigurationRequest, GenericAPIError};
+use compute_api::models::ComputeStatus;
 
 use anyhow::Result;
 use hyper::service::{make_service_fn, service_fn};
 use hyper::{Body, Method, Request, Response, Server, StatusCode};
 use num_cpus;
 use serde_json;
-use tokio::task;
 use tracing::{error, info};
 use tracing_utils::http::OtelName;
 
-fn status_response_from_state(state: &ComputeState) -> ComputeStatusResponse {
-    ComputeStatusResponse {
-        tenant: state
-            .pspec
-            .as_ref()
-            .map(|pspec| pspec.tenant_id.to_string()),
-        timeline: state
-            .pspec
-            .as_ref()
-            .map(|pspec| pspec.timeline_id.to_string()),
-        status: state.status,
-        last_active: state.last_active,
-        error: state.error.clone(),
+async fn handle_spec_request(req: Request<Body>, compute: &Arc<ComputeNode>) -> Result<(), (String, StatusCode)> {
+    if !compute.live_config_allowed {
+        return Err(("live reconfiguration is not allowed for this compute node".to_string(), StatusCode::PRECONDITION_FAILED));
     }
+
+    let body_bytes = hyper::body::to_bytes(req.into_body()).await.unwrap();
+    let spec_raw = String::from_utf8(body_bytes.to_vec()).unwrap();
+
+    let request = serde_json::from_str::<ConfigurationRequest>(&spec_raw)
+        .map_err(|err| (format!("could not parse request json: {err}"), StatusCode::BAD_REQUEST))?;
+    let spec = ParsedSpec::try_from(request.spec)
+        .map_err(|err| (format!("could not parse spec: {err}"), StatusCode::BAD_REQUEST))?;
+
+    let mut inner = compute.inner.lock().unwrap();
+    if !(inner.state.status == ComputeStatus::Empty
+         || inner.state.status == ComputeStatus::Running)
+    {
+        return Err((format!(
+            "invalid compute status for reconfiguration request: {}",
+            serde_json::to_string(&inner.state).unwrap()
+        ), StatusCode::PRECONDITION_FAILED));
+    }
+    inner.spec = Some(spec);
+    inner.state.status = ComputeStatus::ConfigurationPending;
+    compute.state_changed.notify_all();
+    info!("set new spec and notified configurator");
+
+    while inner.state.status != ComputeStatus::Running {
+        inner = compute.state_changed.wait(inner).unwrap();
+        info!(
+            "waiting for compute to become Running, current status: {:?}",
+            inner.state.status
+        );
+    }
+    drop(inner);
+    Ok(())
 }
 
 // Service function to handle all available routes.
@@ -43,16 +64,15 @@ async fn routes(req: Request<Body>, compute: &Arc<ComputeNode>) -> Response<Body
         // Serialized compute state.
         (&Method::GET, "/status") => {
             info!("serving /status GET request");
-            let state = compute.state.lock().unwrap();
-            let status_response = status_response_from_state(&state);
-            Response::new(Body::from(serde_json::to_string(&status_response).unwrap()))
+            let state = compute.get_state();
+            Response::new(Body::from(serde_json::to_string(&state).unwrap()))
         }
 
         // Startup metrics in JSON format. Keep /metrics reserved for a possible
         // future use for Prometheus metrics format.
         (&Method::GET, "/metrics.json") => {
             info!("serving /metrics.json GET request");
-            let metrics = compute.state.lock().unwrap().metrics.clone();
+            let metrics = compute.get_metrics();
             Response::new(Body::from(serde_json::to_string(&metrics).unwrap()))
         }
 
@@ -100,18 +120,19 @@ async fn routes(req: Request<Body>, compute: &Arc<ComputeNode>) -> Response<Body
             ))
         }
 
-        // Accept spec in JSON format and request compute configuration. If
-        // anything goes wrong after we set the compute status to `ConfigurationPending`
-        // and update compute state with new spec, we basically leave compute
-        // in the potentially wrong state. That said, it's control-plane's
-        // responsibility to watch compute state after reconfiguration request
-        // and to clean restart in case of errors.
+        // Accept spec in JSON format and request compute configuration from
+        // the configurator thread. If anything goes wrong after we set the
+        // compute state to `ConfigurationPending` and / or sent spec to the
+        // configurator thread, we basically leave compute in the potentially
+        // wrong state. That said, it's control-plane's responsibility to
+        // watch compute state after reconfiguration request and to clean
+        // restart in case of errors.
         (&Method::POST, "/configure") => {
             info!("serving /configure POST request");
-            match handle_configure_request(req, compute).await {
-                Ok(msg) => Response::new(Body::from(msg)),
-                Err((msg, code)) => {
-                    error!("error handling /configure request: {msg}");
+            match handle_spec_request(req, compute).await {
+                Ok(()) => Response::new(Body::from("ok")),
+                Err((msg, code) ) => {
+                    error!("error handling /spec request: {msg}");
                     render_json_error(&msg, code)
                 }
             }
@@ -126,84 +147,6 @@ async fn routes(req: Request<Body>, compute: &Arc<ComputeNode>) -> Response<Body
     }
 }
 
-async fn handle_configure_request(
-    req: Request<Body>,
-    compute: &Arc<ComputeNode>,
-) -> Result<String, (String, StatusCode)> {
-    if !compute.live_config_allowed {
-        return Err((
-            "live configuration is not allowed for this compute node".to_string(),
-            StatusCode::PRECONDITION_FAILED,
-        ));
-    }
-
-    let body_bytes = hyper::body::to_bytes(req.into_body()).await.unwrap();
-    let spec_raw = String::from_utf8(body_bytes.to_vec()).unwrap();
-    if let Ok(request) = serde_json::from_str::<ConfigurationRequest>(&spec_raw) {
-        let spec = request.spec;
-
-        let parsed_spec = match ParsedSpec::try_from(spec) {
-            Ok(ps) => ps,
-            Err(msg) => return Err((msg, StatusCode::PRECONDITION_FAILED)),
-        };
-
-        // XXX: wrap state update under lock in code blocks. Otherwise,
-        // we will try to `Send` `mut state` into the spawned thread
-        // bellow, which will cause error:
-        // ```
-        // error: future cannot be sent between threads safely
-        // ```
-        {
-            let mut state = compute.state.lock().unwrap();
-            if state.status != ComputeStatus::Empty && state.status != ComputeStatus::Running {
-                let msg = format!(
-                    "invalid compute status for configuration request: {:?}",
-                    state.status.clone()
-                );
-                return Err((msg, StatusCode::PRECONDITION_FAILED));
-            }
-            state.pspec = Some(parsed_spec);
-            state.status = ComputeStatus::ConfigurationPending;
-            compute.state_changed.notify_all();
-            drop(state);
-            info!("set new spec and notified waiters");
-        }
-
-        // Spawn a blocking thread to wait for compute to become Running.
-        // This is needed to do not block the main pool of workers and
-        // be able to serve other requests while some particular request
-        // is waiting for compute to finish configuration.
-        let c = compute.clone();
-        task::spawn_blocking(move || {
-            let mut state = c.state.lock().unwrap();
-            while state.status != ComputeStatus::Running {
-                state = c.state_changed.wait(state).unwrap();
-                info!(
-                    "waiting for compute to become Running, current status: {:?}",
-                    state.status
-                );
-
-                if state.status == ComputeStatus::Failed {
-                    let err = state.error.as_ref().map_or("unknown error", |x| x);
-                    let msg = format!("compute configuration failed: {:?}", err);
-                    return Err((msg, StatusCode::INTERNAL_SERVER_ERROR));
-                }
-            }
-
-            Ok(())
-        })
-        .await
-        .unwrap()?;
-
-        // Return current compute state if everything went well.
-        let state = compute.state.lock().unwrap().clone();
-        let status_response = status_response_from_state(&state);
-        Ok(serde_json::to_string(&status_response).unwrap())
-    } else {
-        Err(("invalid spec".to_string(), StatusCode::BAD_REQUEST))
-    }
-}
-
 fn render_json_error(e: &str, status: StatusCode) -> Response<Body> {
     let error = GenericAPIError {
         error: e.to_string(),
@@ -254,7 +197,6 @@ async fn serve(state: Arc<ComputeNode>) {
 /// Launch a separate Hyper HTTP API server thread and return its `JoinHandle`.
 pub fn launch_http_server(state: &Arc<ComputeNode>) -> Result<thread::JoinHandle<()>> {
     let state = Arc::clone(state);
-
     Ok(thread::Builder::new()
         .name("http-endpoint".into())
         .spawn(move || serve(state))?)

compute_tools/src/http/mod.rs

@@ -1 +1,2 @@
 pub mod api;
+pub mod models;

compute_tools/src/http/models.rs

@@ -1,10 +1,7 @@
-//! Structs representing the JSON formats used in the compute_ctl's HTTP API.
+use serde::{Deserialize, Serialize};
 
-use crate::spec::ComputeSpec;
-use serde::Deserialize;
+use compute_api::spec::ComputeSpec;
 
-/// Request of the /configure API
-///
 /// We now pass only `spec` in the configuration request, but later we can
 /// extend it and something like `restart: bool` or something else. So put
 /// `spec` into a struct initially to be more flexible in the future.
@@ -12,3 +9,8 @@ use serde::Deserialize;
 pub struct ConfigurationRequest {
     pub spec: ComputeSpec,
 }
+
+#[derive(Serialize, Debug)]
+pub struct GenericAPIError {
+    pub error: String,
+}

compute_tools/src/http/openapi_spec.yaml

@@ -89,11 +89,12 @@ paths:
     post:
       tags:
       - Configure
-      summary: Perform compute node configuration.
+      summary: Request compute node configuration.
       description: |
         This is a blocking API endpoint, i.e. it blocks waiting until
         compute is finished configuration and is in `Running` state.
-        Optional non-blocking mode could be added later.
+        Optional non-blocking mode could be added later. Currently,
+        it's also assumed that reconfiguration doesn't require restart.
       operationId: configureCompute
       requestBody:
         description: Configuration request.
@@ -131,14 +132,6 @@ paths:
             application/json:
               schema:
                 $ref: "#/components/schemas/GenericError"
-        500:
-          description: |
-            Compute configuration request was processed, but error
-            occurred. Compute will likely shutdown soon.
-          content:
-            application/json:
-              schema:
-                $ref: "#/components/schemas/GenericError"
 
 components:
   securitySchemes:

compute_tools/src/monitor.rs

@@ -46,7 +46,7 @@ fn watch_compute_activity(compute: &ComputeNode) {
                             AND usename != 'cloud_admin';", // XXX: find a better way to filter other monitors?
                         &[],
                     );
-                let mut last_active = compute.state.lock().unwrap().last_active;
+                let mut last_active = compute.inner.lock().unwrap().state.last_active;
 
                 if let Ok(backs) = backends {
                     let mut idle_backs: Vec<DateTime<Utc>> = vec![];
@@ -87,9 +87,9 @@ fn watch_compute_activity(compute: &ComputeNode) {
                 }
 
                 // Update the last activity in the shared state if we got a more recent one.
-                let mut state = compute.state.lock().unwrap();
-                if last_active > state.last_active {
-                    state.last_active = last_active;
+                let mut inner = compute.inner.lock().unwrap();
+                if last_active > inner.state.last_active {
+                    inner.state.last_active = last_active;
                     debug!("set the last compute activity time to: {}", last_active);
                 }
             }

compute_tools/src/spec.rs

@@ -1,7 +1,7 @@
 use std::path::Path;
 use std::str::FromStr;
 
-use anyhow::{anyhow, bail, Result};
+use anyhow::Result;
 use postgres::config::Config;
 use postgres::{Client, NoTls};
 use tracing::{info, info_span, instrument, span_enabled, warn, Level};
@@ -10,7 +10,6 @@ use crate::config;
 use crate::params::PG_HBA_ALL_MD5;
 use crate::pg_helpers::*;
 
-use compute_api::responses::ControlPlaneSpecResponse;
 use compute_api::spec::{ComputeSpec, Database, PgIdent, Role};
 
 /// Request spec from the control-plane by compute_id. If `NEON_CONSOLE_JWT`
@@ -27,19 +26,13 @@ pub fn get_spec_from_control_plane(base_uri: &str, compute_id: &str) -> Result<C
     // - network error, then retry
     // - no spec for compute yet, then wait
     // - compute id is unknown or any other error, then bail out
-    let resp: ControlPlaneSpecResponse = reqwest::blocking::Client::new()
+    let spec = reqwest::blocking::Client::new()
         .get(cp_uri)
         .header("Authorization", jwt)
-        .send()
-        .map_err(|e| anyhow!("could not send spec request to control plane: {}", e))?
-        .json()
-        .map_err(|e| anyhow!("could not get compute spec from control plane: {}", e))?;
+        .send()?
+        .json()?;
 
-    if let Some(spec) = resp.spec {
-        Ok(spec)
-    } else {
-        bail!("could not get compute spec from control plane")
-    }
+    Ok(spec)
 }
 
 /// It takes cluster specification and does the following:

control_plane/src/bin/neon_local.rs

@@ -108,7 +108,6 @@ fn main() -> Result<()> {
             "pageserver" => handle_pageserver(sub_args, &env),
             "safekeeper" => handle_safekeeper(sub_args, &env),
             "endpoint" => handle_endpoint(sub_args, &env),
-            "pg" => bail!("'pg' subcommand has been renamed to 'endpoint'"),
             _ => bail!("unexpected subcommand {sub_name}"),
         };
 
@@ -1077,13 +1076,6 @@ fn cli() -> Command {
                 )
 
         )
-        // Obsolete old name for 'endpoint'. We now just print an error if it's used.
-        .subcommand(
-            Command::new("pg")
-                .hide(true)
-                .arg(Arg::new("ignore-rest").allow_hyphen_values(true).num_args(0..).required(false))
-                .trailing_var_arg(true)
-        )
         .subcommand(
             Command::new("start")
                 .about("Start page server and safekeepers")

libs/compute_api/src/lib.rs

@@ -1,3 +1,2 @@
-pub mod requests;
-pub mod responses;
+pub mod models;
 pub mod spec;

libs/compute_api/src/models.rs

@@ -1,28 +1,10 @@
 //! Structs representing the JSON formats used in the compute_ctl's HTTP API.
-
 use chrono::{DateTime, Utc};
-use serde::{Deserialize, Serialize, Serializer};
-
-use crate::spec::ComputeSpec;
-
-#[derive(Serialize, Debug)]
-pub struct GenericAPIError {
-    pub error: String,
-}
+use serde::{Serialize, Serializer};
 
 /// Response of the /status API
-#[derive(Serialize, Debug)]
-#[serde(rename_all = "snake_case")]
-pub struct ComputeStatusResponse {
-    pub tenant: Option<String>,
-    pub timeline: Option<String>,
-    pub status: ComputeStatus,
-    #[serde(serialize_with = "rfc3339_serialize")]
-    pub last_active: DateTime<Utc>,
-    pub error: Option<String>,
-}
-
-#[derive(Serialize)]
+///
+#[derive(Clone, Serialize)]
 #[serde(rename_all = "snake_case")]
 pub struct ComputeState {
     pub status: ComputeStatus,
@@ -32,25 +14,25 @@ pub struct ComputeState {
     pub error: Option<String>,
 }
 
-#[derive(Serialize, Clone, Copy, Debug, PartialEq, Eq)]
+#[derive(Serialize, Clone, Copy, PartialEq, Eq, Debug)]
 #[serde(rename_all = "snake_case")]
 pub enum ComputeStatus {
-    // Spec wasn't provided at start, waiting for it to be
+    // Spec wasn't provided as start, waiting for it to be
     // provided by control-plane.
     Empty,
-    // Compute configuration was requested.
-    ConfigurationPending,
     // Compute node has spec and initial startup and
     // configuration is in progress.
     Init,
     // Compute is configured and running.
     Running,
-    // New spec is being applied.
-    Configuration,
     // Either startup or configuration failed,
     // compute will exit soon or is waiting for
     // control-plane to terminate it.
     Failed,
+    // Control-plane requested reconfiguration.
+    ConfigurationPending,
+    // New spec is being applied.
+    Configuration,
 }
 
 fn rfc3339_serialize<S>(x: &DateTime<Utc>, s: S) -> Result<S::Ok, S::Error>
@@ -61,18 +43,10 @@ where
 }
 
 /// Response of the /metrics.json API
-#[derive(Clone, Debug, Default, Serialize)]
+#[derive(Clone, Default, Serialize)]
 pub struct ComputeMetrics {
     pub sync_safekeepers_ms: u64,
     pub basebackup_ms: u64,
     pub config_ms: u64,
     pub total_startup_ms: u64,
 }
-
-/// Response of the `/computes/{compute_id}/spec` control-plane API.
-/// This is not actually a compute API response, so consider moving
-/// to a different place.
-#[derive(Deserialize, Debug)]
-pub struct ControlPlaneSpecResponse {
-    pub spec: Option<ComputeSpec>,
-}

libs/compute_api/src/spec.rs

@@ -12,13 +12,10 @@ pub type PgIdent = String;
 
 /// Cluster spec or configuration represented as an optional number of
 /// delta operations + final cluster state description.
-#[derive(Clone, Debug, Default, Deserialize)]
+#[derive(Clone, Debug, Deserialize)]
 pub struct ComputeSpec {
     pub format_version: f32,
-
-    // The control plane also includes a 'timestamp' field in the JSON document,
-    // but we don't use it for anything. Serde will ignore missing fields when
-    // deserializing it.
+    pub timestamp: String,
     pub operation_uuid: Option<String>,
     /// Expected cluster state at the end of transition process.
     pub cluster: Cluster,
@@ -29,7 +26,7 @@ pub struct ComputeSpec {
     pub startup_tracing_context: Option<HashMap<String, String>>,
 }
 
-#[derive(Clone, Debug, Default, Deserialize)]
+#[derive(Clone, Debug, Deserialize)]
 pub struct Cluster {
     pub cluster_id: String,
     pub name: String,

libs/pageserver_api/Cargo.toml

@@ -7,7 +7,6 @@ license.workspace = true
 [dependencies]
 serde.workspace = true
 serde_with.workspace = true
-serde_json.workspace = true
 const_format.workspace = true
 anyhow.workspace = true
 bytes.workspace = true
@@ -15,7 +14,6 @@ byteorder.workspace = true
 utils.workspace = true
 postgres_ffi.workspace = true
 enum-map.workspace = true
-strum.workspace = true
-strum_macros.workspace = true
+serde_json.workspace = true
 
 workspace_hack.workspace = true

libs/pageserver_api/src/models.rs

@@ -7,7 +7,6 @@ use std::{
 use byteorder::{BigEndian, ReadBytesExt};
 use serde::{Deserialize, Serialize};
 use serde_with::{serde_as, DisplayFromStr};
-use strum_macros;
 use utils::{
     history_buffer::HistoryBufferWithDropCounter,
     id::{NodeId, TenantId, TimelineId},
@@ -19,23 +18,11 @@ use anyhow::bail;
 use bytes::{BufMut, Bytes, BytesMut};
 
 /// A state of a tenant in pageserver's memory.
-#[derive(
-    Clone,
-    PartialEq,
-    Eq,
-    serde::Serialize,
-    serde::Deserialize,
-    strum_macros::Display,
-    strum_macros::EnumString,
-    strum_macros::EnumVariantNames,
-    strum_macros::AsRefStr,
-    strum_macros::IntoStaticStr,
-)]
-#[serde(tag = "slug", content = "data")]
+#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
 pub enum TenantState {
-    /// This tenant is being loaded from local disk
+    // This tenant is being loaded from local disk
     Loading,
-    /// This tenant is being downloaded from cloud storage.
+    // This tenant is being downloaded from cloud storage.
     Attaching,
     /// Tenant is fully operational
     Active,
@@ -44,7 +31,15 @@ pub enum TenantState {
     Stopping,
     /// A tenant is recognized by the pageserver, but can no longer be used for
     /// any operations, because it failed to be activated.
-    Broken { reason: String, backtrace: String },
+    Broken,
+}
+
+pub mod state {
+    pub const LOADING: &str = "loading";
+    pub const ATTACHING: &str = "attaching";
+    pub const ACTIVE: &str = "active";
+    pub const STOPPING: &str = "stopping";
+    pub const BROKEN: &str = "broken";
 }
 
 impl TenantState {
@@ -54,26 +49,17 @@ impl TenantState {
             Self::Attaching => true,
             Self::Active => false,
             Self::Stopping => false,
-            Self::Broken { .. } => false,
+            Self::Broken => false,
         }
     }
 
-    pub fn broken_from_reason(reason: String) -> Self {
-        let backtrace_str: String = format!("{}", std::backtrace::Backtrace::force_capture());
-        Self::Broken {
-            reason,
-            backtrace: backtrace_str,
-        }
-    }
-}
-
-impl std::fmt::Debug for TenantState {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+    pub fn as_str(&self) -> &'static str {
         match self {
-            Self::Broken { reason, backtrace } if !reason.is_empty() => {
-                write!(f, "Broken due to: {reason}. Backtrace:\n{backtrace}")
-            }
-            _ => write!(f, "{self}"),
+            TenantState::Loading => state::LOADING,
+            TenantState::Attaching => state::ATTACHING,
+            TenantState::Active => state::ACTIVE,
+            TenantState::Stopping => state::STOPPING,
+            TenantState::Broken => state::BROKEN,
         }
     }
 }
@@ -629,7 +615,6 @@ impl PagestreamBeMessage {
 #[cfg(test)]
 mod tests {
     use bytes::Buf;
-    use serde_json::json;
 
     use super::*;
 
@@ -680,57 +665,4 @@ mod tests {
             assert!(msg == reconstructed);
         }
     }
-
-    #[test]
-    fn test_tenantinfo_serde() {
-        // Test serialization/deserialization of TenantInfo
-        let original_active = TenantInfo {
-            id: TenantId::generate(),
-            state: TenantState::Active,
-            current_physical_size: Some(42),
-            has_in_progress_downloads: Some(false),
-        };
-        let expected_active = json!({
-            "id": original_active.id.to_string(),
-            "state": {
-                "slug": "Active",
-            },
-            "current_physical_size": 42,
-            "has_in_progress_downloads": false,
-        });
-
-        let original_broken = TenantInfo {
-            id: TenantId::generate(),
-            state: TenantState::Broken {
-                reason: "reason".into(),
-                backtrace: "backtrace info".into(),
-            },
-            current_physical_size: Some(42),
-            has_in_progress_downloads: Some(false),
-        };
-        let expected_broken = json!({
-            "id": original_broken.id.to_string(),
-            "state": {
-                "slug": "Broken",
-                "data": {
-                    "backtrace": "backtrace info",
-                    "reason": "reason",
-                }
-            },
-            "current_physical_size": 42,
-            "has_in_progress_downloads": false,
-        });
-
-        assert_eq!(
-            serde_json::to_value(&original_active).unwrap(),
-            expected_active
-        );
-
-        assert_eq!(
-            serde_json::to_value(&original_broken).unwrap(),
-            expected_broken
-        );
-        assert!(format!("{:?}", &original_broken.state).contains("reason"));
-        assert!(format!("{:?}", &original_broken.state).contains("backtrace info"));
-    }
 }

libs/postgres_backend/src/lib.rs

@@ -54,7 +54,7 @@ pub fn is_expected_io_error(e: &io::Error) -> bool {
     use io::ErrorKind::*;
     matches!(
         e.kind(),
-        ConnectionRefused | ConnectionAborted | ConnectionReset | TimedOut
+        ConnectionRefused | ConnectionAborted | ConnectionReset
     )
 }
 
@@ -320,17 +320,9 @@ impl<IO: AsyncRead + AsyncWrite + Unpin> PostgresBackend<IO> {
         if let ProtoState::Closed = self.state {
             Ok(None)
         } else {
-            match self.framed.read_message().await {
-                Ok(m) => {
-                    trace!("read msg {:?}", m);
-                    Ok(m)
-                }
-                Err(e) => {
-                    // remember not to try to read anymore
-                    self.state = ProtoState::Closed;
-                    Err(e)
-                }
-            }
+            let m = self.framed.read_message().await?;
+            trace!("read msg {:?}", m);
+            Ok(m)
         }
     }
 
@@ -501,10 +493,7 @@ impl<IO: AsyncRead + AsyncWrite + Unpin> PostgresBackend<IO> {
             MaybeWriteOnly::Full(framed) => {
                 let (reader, writer) = framed.split();
                 self.framed = MaybeWriteOnly::WriteOnly(writer);
-                Ok(PostgresBackendReader {
-                    reader,
-                    closed: false,
-                })
+                Ok(PostgresBackendReader(reader))
             }
             MaybeWriteOnly::WriteOnly(_) => {
                 anyhow::bail!("PostgresBackend is already split")
@@ -521,12 +510,8 @@ impl<IO: AsyncRead + AsyncWrite + Unpin> PostgresBackend<IO> {
                 anyhow::bail!("PostgresBackend is not split")
             }
             MaybeWriteOnly::WriteOnly(writer) => {
-                let joined = Framed::unsplit(reader.reader, writer);
+                let joined = Framed::unsplit(reader.0, writer);
                 self.framed = MaybeWriteOnly::Full(joined);
-                // if reader encountered connection error, do not attempt reading anymore
-                if reader.closed {
-                    self.state = ProtoState::Closed;
-                }
                 Ok(())
             }
             MaybeWriteOnly::Broken => panic!("unsplit on framed in invalid state"),
@@ -812,25 +797,15 @@ impl<IO: AsyncRead + AsyncWrite + Unpin> PostgresBackend<IO> {
     }
 }
 
-pub struct PostgresBackendReader<IO> {
-    reader: FramedReader<MaybeTlsStream<IO>>,
-    closed: bool, // true if received error closing the connection
-}
+pub struct PostgresBackendReader<IO>(FramedReader<MaybeTlsStream<IO>>);
 
 impl<IO: AsyncRead + AsyncWrite + Unpin> PostgresBackendReader<IO> {
     /// Read full message or return None if connection is cleanly closed with no
     /// unprocessed data.
     pub async fn read_message(&mut self) -> Result<Option<FeMessage>, ConnectionError> {
-        match self.reader.read_message().await {
-            Ok(m) => {
-                trace!("read msg {:?}", m);
-                Ok(m)
-            }
-            Err(e) => {
-                self.closed = true;
-                Err(e)
-            }
-        }
+        let m = self.0.read_message().await?;
+        trace!("read msg {:?}", m);
+        Ok(m)
     }
 
     /// Get CopyData contents of the next message in COPY stream or error
@@ -948,7 +923,7 @@ pub enum CopyStreamHandlerEnd {
     #[error("EOF on COPY stream")]
     EOF,
     /// The connection was lost
-    #[error("connection error: {0}")]
+    #[error(transparent)]
     Disconnected(#[from] ConnectionError),
     /// Some other error
     #[error(transparent)]

libs/pq_proto/src/lib.rs

@@ -293,9 +293,6 @@ impl FeStartupPacket {
         // We shouldn't advance `buf` as probably full message is not there yet,
         // so can't directly use Bytes::get_u32 etc.
         let len = (&buf[0..4]).read_u32::<BigEndian>().unwrap() as usize;
-        // The proposed replacement is `!(4..=MAX_STARTUP_PACKET_LENGTH).contains(&len)`
-        // which is less readable
-        #[allow(clippy::manual_range_contains)]
         if len < 4 || len > MAX_STARTUP_PACKET_LENGTH {
             return Err(ProtocolError::Protocol(format!(
                 "invalid startup packet message length {}",

libs/remote_storage/src/lib.rs

@@ -13,6 +13,7 @@ use std::{
     collections::HashMap,
     fmt::Debug,
     num::{NonZeroU32, NonZeroUsize},
+    ops::Deref,
     path::{Path, PathBuf},
     pin::Pin,
     sync::Arc,
@@ -89,7 +90,7 @@ pub trait RemoteStorage: Send + Sync + 'static {
     /// Streams the local file contents into remote into the remote storage entry.
     async fn upload(
         &self,
-        from: impl io::AsyncRead + Unpin + Send + Sync + 'static,
+        data: Box<(dyn io::AsyncRead + Unpin + Send + Sync + 'static)>,
         // S3 PUT request requires the content length to be specified,
         // otherwise it starts to fail with the concurrent connection count increasing.
         data_size_bytes: usize,
@@ -160,67 +161,14 @@ pub enum GenericRemoteStorage {
     Unreliable(Arc<UnreliableWrapper>),
 }
 
-impl GenericRemoteStorage {
-    pub async fn list_prefixes(
-        &self,
-        prefix: Option<&RemotePath>,
-    ) -> Result<Vec<RemotePath>, DownloadError> {
-        match self {
-            Self::LocalFs(s) => s.list_prefixes(prefix).await,
-            Self::AwsS3(s) => s.list_prefixes(prefix).await,
-            Self::Unreliable(s) => s.list_prefixes(prefix).await,
-        }
-    }
+impl Deref for GenericRemoteStorage {
+    type Target = dyn RemoteStorage;
 
-    pub async fn upload(
-        &self,
-        from: impl io::AsyncRead + Unpin + Send + Sync + 'static,
-        data_size_bytes: usize,
-        to: &RemotePath,
-        metadata: Option<StorageMetadata>,
-    ) -> anyhow::Result<()> {
+    fn deref(&self) -> &Self::Target {
         match self {
-            Self::LocalFs(s) => s.upload(from, data_size_bytes, to, metadata).await,
-            Self::AwsS3(s) => s.upload(from, data_size_bytes, to, metadata).await,
-            Self::Unreliable(s) => s.upload(from, data_size_bytes, to, metadata).await,
-        }
-    }
-
-    pub async fn download(&self, from: &RemotePath) -> Result<Download, DownloadError> {
-        match self {
-            Self::LocalFs(s) => s.download(from).await,
-            Self::AwsS3(s) => s.download(from).await,
-            Self::Unreliable(s) => s.download(from).await,
-        }
-    }
-
-    pub async fn download_byte_range(
-        &self,
-        from: &RemotePath,
-        start_inclusive: u64,
-        end_exclusive: Option<u64>,
-    ) -> Result<Download, DownloadError> {
-        match self {
-            Self::LocalFs(s) => {
-                s.download_byte_range(from, start_inclusive, end_exclusive)
-                    .await
-            }
-            Self::AwsS3(s) => {
-                s.download_byte_range(from, start_inclusive, end_exclusive)
-                    .await
-            }
-            Self::Unreliable(s) => {
-                s.download_byte_range(from, start_inclusive, end_exclusive)
-                    .await
-            }
-        }
-    }
-
-    pub async fn delete(&self, path: &RemotePath) -> anyhow::Result<()> {
-        match self {
-            Self::LocalFs(s) => s.delete(path).await,
-            Self::AwsS3(s) => s.delete(path).await,
-            Self::Unreliable(s) => s.delete(path).await,
+            GenericRemoteStorage::LocalFs(local_fs) => local_fs,
+            GenericRemoteStorage::AwsS3(s3_bucket) => s3_bucket.as_ref(),
+            GenericRemoteStorage::Unreliable(s) => s.as_ref(),
         }
     }
 }
@@ -251,7 +199,7 @@ impl GenericRemoteStorage {
     /// this path is used for the remote object id conversion only.
     pub async fn upload_storage_object(
         &self,
-        from: impl tokio::io::AsyncRead + Unpin + Send + Sync + 'static,
+        from: Box<dyn tokio::io::AsyncRead + Unpin + Send + Sync + 'static>,
         from_size_bytes: usize,
         to: &RemotePath,
     ) -> anyhow::Result<()> {

libs/remote_storage/src/local_fs.rs

@@ -118,7 +118,7 @@ impl RemoteStorage for LocalFs {
 
     async fn upload(
         &self,
-        data: impl io::AsyncRead + Unpin + Send + Sync + 'static,
+        data: Box<(dyn io::AsyncRead + Unpin + Send + Sync + 'static)>,
         data_size_bytes: usize,
         to: &RemotePath,
         metadata: Option<StorageMetadata>,

libs/remote_storage/src/s3_bucket.rs

@@ -343,7 +343,7 @@ impl RemoteStorage for S3Bucket {
 
     async fn upload(
         &self,
-        from: impl io::AsyncRead + Unpin + Send + Sync + 'static,
+        from: Box<(dyn io::AsyncRead + Unpin + Send + Sync + 'static)>,
         from_size_bytes: usize,
         to: &RemotePath,
         metadata: Option<StorageMetadata>,

libs/remote_storage/src/simulate_failures.rs

@@ -84,7 +84,7 @@ impl RemoteStorage for UnreliableWrapper {
 
     async fn upload(
         &self,
-        data: impl tokio::io::AsyncRead + Unpin + Send + Sync + 'static,
+        data: Box<(dyn tokio::io::AsyncRead + Unpin + Send + Sync + 'static)>,
         // S3 PUT request requires the content length to be specified,
         // otherwise it starts to fail with the concurrent connection count increasing.
         data_size_bytes: usize,

libs/timeline_data_path/Cargo.toml

@@ -1,13 +0,0 @@
-[package]
-name = "timeline_data_path"
-version = "0.1.0"
-edition.workspace = true
-license.workspace = true
-
-# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
-
-[dependencies]
-utils.workspace = true
-workspace_hack.workspace = true
-tokio.workspace = true
-thiserror.workspace = true

libs/timeline_data_path/src/lib.rs

@@ -1,396 +0,0 @@
-//! The Timeline's core data path.
-//!
-//! # Overview
-//!
-//! This crate implements the core data path of a Timeline inside Pageserver:
-//!
-//! 1. WAL records from `walreceiver`, via in-memory layers, into persistent L0 layers.
-//! 1. `GetPage@LSN`: retrieval of WAL records and page images for feeding into WAL redo.
-//! 1. Data re-shuffeling through compaction (TODO).
-//! 1. Page image creation & garbage collection through GC (TODO).
-//!
-//! This crate assumes the following concepts, but is fully generic over their implementation:
-//!
-//! - **Delta Records**: data is written into the system in the form of self-descriptive deltas.
-//!   For the Pageserver use case, these deltas are derived from Postgres WAL records.
-//! - **Page Numbers**: Delta Records always affect a single key.
-//!   That key is called page number, because, in the Pageserver use case, the Postgres table page numbers are the keys.
-//! - **LSN**: When writing Delta Records into the system, they are associated with a monotonically increasing LSN.
-//!   Subsequently written Delta Records must have increasing LSNs.
-//! - **Page Images**: Delta Records for a given page can be used to reconstruct the page. Think of it like squashing diffs.
-//!   - When sorting the Delta Records for a given key by their LSN, any prefix of that sorting can be squashed into a page image.
-//!   - Delta Records following such a squash can be squashed into that page image.
-//!   - In Pageserver, WAL redo implements the (pure) function of squashing.
-//! - **In-Memory Layer**: an object that represents an "unfinished" L0 layer file, holding Delta Records in insertion order.
-//!   "Unfinished" means that we're still writing Delta Records to that file.
-//! - **Historic Layer**: an object that represents a "finished" layer file, at any compaction level.
-//!   Such objects reside on disk and/or in remote storage.
-//!   They may contain Delta Records, Page Images, or a mixture thereof. It doesn't matter.
-//! - **HistoricStuff**: an efficient lookup data structure to find the list of Historic Layer objects
-//!   that hold the Delta Records / PageImages required to reconstruct a Page Image at a given LSN.
-//!
-//! # API
-//!
-//! The core idea is that of a specialized single-producer multi-consumer structure,
-//! embodied by a Read-end and a Write-end.
-//!
-//! The Write-end is used to push new `DeltaRecord @ LSN`s into the system.
-//! In Pageserver, this is used by the `WalReceiver`.
-//!
-//! The Read-end provides the `GetPage@LSN` API.
-//! In the current iteration, we actually return something called `ReconstructWork`.
-//! I.e., we leave the work of reading the values from the layers, and the WAL redo invocation to the caller.
-//! Find rationale for this design in the *Scope* section.
-//!
-//! ## Immutability
-//!
-//! The traits defined by this crate assume immutable data structures that are multi-versioned.
-//!
-//! As an example for what "immutable" means, take the case where we add a new Historic Layer to HistoricStuff.
-//! Traditionally, one would use shared mutable state, i.e. `Arc<RwLock<...>>`.
-//! To insert the new Historic Layer, we would acquire the RwLock in write mode and modify a lookup data structure to accomodate the new layer.
-//! The Read-ends would use RwLock in read mode to read from the data structure.
-//!
-//! Conversely, with  *immutable data structures*, writers create new version (aka *snapshots*) of the lookup data structure.
-//! New reads on the Read-ends will use the new snapshot, but old ongoing reads would use the old version(s).
-//! An efficient implementation would likely share the Historic Layer objects, e.g., using `Arc`.
-//! And maybe there's internally mutable state inside the layer objects, e.g., to track residence (i.e., *on-demand downloaded* vs *evicted*).
-//! But the important point is that there's no synchronization / lock-holding at any higher level, except when grabbing a reference to the snapshot (Read-end), or when publishing a new snapshot (Write-end).
-//!
-//! ## Scope
-//!
-//! The following concerns are considered implementation details from the perspective of this crate:
-//!
-//! - **Layer File Persistence**: `HistoricStuff::make_historic` is responsible for this.
-//! - **Reading Layer Files**: the `ReconstructWork` that the Read-end returns from `GetPage@LSN` requests contains the list of layers to consult.
-//!   The crate consumer is responsible for reading the layers & doing WAL redo.
-//!   Likely the implementation of `HistoricStuff` plays a role here, because it is responsible for persisting the layer files.
-//! - **Layer Eviction & On-Demand Download**: this is just an aspect of the above.
-//!   The crate consumer can choose to implement eviction & on-demand download however they wish.
-//!   The only requirement is that the Historic Layers don't change their contents, i.e., they always returnt he same reconstruct values for the same lookup.
-//!   - For example, a `LayerCache` modoule or service could take care of layer uploads, eviction, and on-demand downloads.
-//!     Initially, the `layer cache` can be local-only.
-//!     But in the future, it can be multi-machine / clustered pagesevers / aka "sharding".
-//!
-//! # Example
-//!
-//! The [`new`] function is the entrypoint to this crate.
-//!
-//! See the test cases for how it is used.
-
-use std::{marker::PhantomData, time::Duration};
-
-use utils::seqwait::{self, Advance, SeqWait, Wait};
-
-#[cfg(test)]
-mod tests;
-
-/// Collection of types / type bounds used by Read-end and Write-end.
-///
-/// See the [`crate`]-level docs's *Concepts* section to learn about
-/// the meaning of each associated `type`.
-///
-/// # Usage
-///
-/// Define a zero-sized-type and impl this Trait for it.
-/// Then use that zero-sized-type as the single generic argument to [`new`]
-/// and almost all types declared in this crate.
-///
-/// It might feel a bit weird, but, the alternative is to have umpteen generic
-/// types per `impl` with repetitive trait bounds.
-///
-/// Search the test cases for an example of how this can be used to improve testability.
-pub trait Types {
-    type Key: Copy;
-    type Lsn: Ord + Copy;
-    type LsnCounter: seqwait::MonotonicCounter<Self::Lsn> + Copy;
-    type DeltaRecord;
-    type HistoricLayer;
-    type InMemoryLayer: InMemoryLayer<Types = Self> + Clone;
-    type HistoricStuff: HistoricStuff<Types = Self> + Clone;
-    type GetReconstructPathError: std::error::Error;
-}
-
-/// Error returned by [`InMemoryLayer::put`].
-#[derive(thiserror::Error)]
-pub struct InMemoryLayerPutError<DeltaRecord> {
-    delta: DeltaRecord,
-    kind: InMemoryLayerPutErrorKind,
-}
-
-/// Part of [`InMemoryLayerPutError`].
-#[derive(Debug)]
-pub enum InMemoryLayerPutErrorKind {
-    LayerFull,
-    AlreadyHaveRecordForKeyAndLsn,
-}
-
-impl<DeltaRecord> std::fmt::Debug for InMemoryLayerPutError<DeltaRecord> {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        f.debug_struct("InMemoryLayerPutError")
-            // would require DeltaRecord to impl Debug
-            //         .field("delta", &self.delta)
-            .field("kind", &self.kind)
-            .finish()
-    }
-}
-
-/// An in-memory layer. See [`crate`] docs for details on this concept.
-pub trait InMemoryLayer: std::fmt::Debug + Default + Clone {
-    type Types: Types;
-    fn put(
-        &mut self,
-        key: <Self::Types as Types>::Key,
-        lsn: <Self::Types as Types>::Lsn,
-        delta: <Self::Types as Types>::DeltaRecord,
-    ) -> Result<Self, InMemoryLayerPutError<<Self::Types as Types>::DeltaRecord>>;
-    fn get(
-        &self,
-        key: <Self::Types as Types>::Key,
-        lsn: <Self::Types as Types>::Lsn,
-    ) -> Vec<<Self::Types as Types>::DeltaRecord>;
-}
-
-/// The manager of [`Types::HistoricLayer`]s.
-pub trait HistoricStuff {
-    type Types: Types;
-    fn get_reconstruct_path(
-        &self,
-        key: <Self::Types as Types>::Key,
-        lsn: <Self::Types as Types>::Lsn,
-    ) -> Result<
-        Vec<<Self::Types as Types>::HistoricLayer>,
-        <Self::Types as Types>::GetReconstructPathError,
-    >;
-    /// Produce a new version of `self` that includes the given inmem layer.
-    fn make_historic(&self, inmem: <Self::Types as Types>::InMemoryLayer) -> Self;
-}
-
-/// A snapshot of the data. See [`crate`]-level docs section on *immutability* for details.
-struct Snapshot<T: Types> {
-    _types: PhantomData<T>,
-    inmem: Option<T::InMemoryLayer>,
-    historic: T::HistoricStuff,
-}
-
-impl<T: Types> Clone for Snapshot<T> {
-    fn clone(&self) -> Self {
-        Self {
-            _types: self._types.clone(),
-            inmem: self.inmem.clone(),
-            historic: self.historic.clone(),
-        }
-    }
-}
-
-/// The Read-end. See [`crate`]-level docs for details.
-pub struct Reader<T: Types> {
-    wait: Wait<T::LsnCounter, T::Lsn, Snapshot<T>>,
-}
-
-/// The Write-end. See [`crate`]-level docs for details.
-pub struct Writer<T: Types> {
-    advance: Advance<T::LsnCounter, T::Lsn, Snapshot<T>>,
-}
-
-/// Setup a pair of Read-end and Write-End. This is the entrypoint to this crate.
-///
-/// The idea is that the caller loads the arguments from persistent state that `HistoricStuff` wrote at an earlier point in time.
-pub fn new<T: Types>(lsn: T::LsnCounter, historic: T::HistoricStuff) -> (Reader<T>, Writer<T>) {
-    let state = Snapshot {
-        _types: PhantomData::<T>::default(),
-        inmem: None,
-        historic: historic,
-    };
-    let (wait, advance) = SeqWait::new(lsn, state).split_spmc();
-    let reader = Reader { wait };
-    let read_writer = Writer { advance };
-    (reader, read_writer)
-}
-
-/// Error returned by the get-page operations.
-#[derive(Debug, thiserror::Error)]
-pub enum GetError<T: Types> {
-    #[error(transparent)]
-    SeqWait(seqwait::SeqWaitError),
-    #[error(transparent)]
-    GetReconstructPath(T::GetReconstructPathError),
-}
-
-/// Self-contained set of objects required to reconstruct a page image for the given `key` @ `lsn`.
-///
-/// This is returned by the `get` methods of [`Reader`] and [`Writer`].
-///
-/// To reconstruct the page image, stack up (top to bottom) `inmem_records` plus all records found for `key` and `lsn` along the `historic_path` until an initial page image is found.
-/// Then feed that stack to WAL-redo to get the page image.
-///
-/// See [`crate`]-level docs on *scope* for why we don't return page images from these functions.
-pub struct ReconstructWork<T: Types> {
-    pub key: T::Key,
-    pub lsn: T::Lsn,
-    pub inmem_records: Vec<T::DeltaRecord>,
-    pub historic_path: Vec<T::HistoricLayer>,
-}
-
-impl<T: Types> Reader<T> {
-    /// This is the `GetPage@LSN` operation.
-    ///
-    /// See the [`crate`]-level docs for why we return [`ReconstructWork`] instead of a Page Image here.
-    pub async fn get(&self, key: T::Key, lsn: T::Lsn) -> Result<ReconstructWork<T>, GetError<T>> {
-        // XXX dedup with Writer::get_nowait
-        let state = self.wait.wait_for(lsn).await.map_err(GetError::SeqWait)?;
-        let inmem_records = state
-            .inmem
-            .as_ref()
-            .map(|iml| iml.get(key, lsn))
-            .unwrap_or_default();
-        let historic_path = state
-            .historic
-            .get_reconstruct_path(key, lsn)
-            .map_err(GetError::GetReconstructPath)?;
-        Ok(ReconstructWork {
-            key,
-            lsn,
-            inmem_records,
-            historic_path,
-        })
-    }
-}
-
-/// Error returned by the `put` operation.
-#[derive(thiserror::Error)]
-pub struct PutError<T: Types> {
-    /// The `delta` record which we failed to `put`.
-    pub delta: T::DeltaRecord,
-    /// Description of what went wrong.
-    pub kind: PutErrorKind,
-}
-
-/// Part of [`PutError`].
-#[derive(Debug)]
-pub enum PutErrorKind {
-    AlreadyHaveInMemoryRecordForKeyAndLsn,
-}
-
-impl<T: Types> std::fmt::Debug for PutError<T> {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        f.debug_struct("PutError")
-            // would need to require Debug for DeltaRecord
-            // .field("delta", &self.delta)
-            .field("kind", &self.kind)
-            .finish()
-    }
-}
-
-impl<T: Types> Writer<T> {
-    /// Insert data into the system.
-    pub async fn put(
-        &mut self,
-        key: T::Key,
-        lsn: T::Lsn,
-        delta: T::DeltaRecord,
-    ) -> Result<(), PutError<T>> {
-        let (_snapshot_lsn, snapshot) = self.advance.get_current_data();
-        // TODO ensure snapshot_lsn <= lsn?
-        let mut inmem = snapshot
-            .inmem
-            .unwrap_or_else(|| T::InMemoryLayer::default());
-        // XXX: use the Advance as witness and only allow witness to access inmem in write mode
-        match inmem.put(key, lsn, delta) {
-            Ok(new_inmem) => {
-                let new_snapshot = Snapshot {
-                    _types: PhantomData,
-                    inmem: Some(new_inmem),
-                    historic: snapshot.historic,
-                };
-                self.advance.advance(lsn, Some(new_snapshot));
-            }
-            Err(InMemoryLayerPutError {
-                delta,
-                kind: InMemoryLayerPutErrorKind::AlreadyHaveRecordForKeyAndLsn,
-            }) => {
-                return Err(PutError {
-                    delta,
-                    kind: PutErrorKind::AlreadyHaveInMemoryRecordForKeyAndLsn,
-                });
-            }
-            Err(InMemoryLayerPutError {
-                delta,
-                kind: InMemoryLayerPutErrorKind::LayerFull,
-            }) => {
-                let new_historic = snapshot.historic.make_historic(inmem);
-                let mut new_inmem = T::InMemoryLayer::default();
-                let new_inmem = new_inmem
-                    .put(key, lsn, delta)
-                    .expect("put into default inmem layer must not fail");
-                let new_state = Snapshot {
-                    _types: PhantomData::<T>::default(),
-                    inmem: Some(new_inmem),
-                    historic: new_historic,
-                };
-                self.advance.advance(lsn, Some(new_state));
-            }
-        }
-        Ok(())
-    }
-
-    /// Force flushing of the current in-memory layer.
-    ///
-    /// Usually, flushing happens only if the in-memory layer is full.
-    /// Use this API to make it happen in other circumstances (shutdown, periodic ticker, etc.).
-    pub async fn force_flush(&mut self) -> tokio::io::Result<()> {
-        let (snapshot_lsn, snapshot) = self.advance.get_current_data();
-        let Snapshot {
-            _types,
-            inmem,
-            historic,
-        } = snapshot;
-        // XXX: use the Advance as witness and only allow witness to access inmem in "write" mode
-        let Some(inmem) = inmem else {
-            // nothing to do
-            return Ok(());
-        };
-        let new_historic = historic.make_historic(inmem);
-        let new_snapshot = Snapshot {
-            _types: PhantomData::<T>::default(),
-            inmem: None,
-            historic: new_historic,
-        };
-        self.advance.advance(snapshot_lsn, Some(new_snapshot)); // TODO: should fail if we're past snapshot_lsn
-        Ok(())
-    }
-
-    /// `get` at the given LSN, without blocking.
-    ///
-    /// Fails with a timeout error if the `lsn` isn't there yet.
-    /// That makes sense because the only way we'd stop waiting is by a `self.put()`.
-    /// But concurrent `put()` is forbidden.
-    pub async fn get_nowait(
-        &self,
-        key: T::Key,
-        lsn: T::Lsn,
-    ) -> Result<ReconstructWork<T>, GetError<T>> {
-        // XXX dedup with Reader::get
-        let state = self
-            .advance
-            .wait_for_timeout(lsn, Duration::from_secs(0))
-            // The await is never going to block because we pass from_secs(0).
-            .await
-            .map_err(GetError::SeqWait)?;
-        let inmem_records = state
-            .inmem
-            .as_ref()
-            .map(|iml| iml.get(key, lsn))
-            .unwrap_or_default();
-        let historic_path = state
-            .historic
-            .get_reconstruct_path(key, lsn)
-            .map_err(GetError::GetReconstructPath)?;
-        Ok(ReconstructWork {
-            key,
-            lsn,
-            inmem_records,
-            historic_path,
-        })
-    }
-}

libs/timeline_data_path/src/tests.rs

@@ -1,170 +0,0 @@
-use std::collections::{btree_map::Entry, BTreeMap};
-use std::sync::Arc;
-use utils::seqwait;
-
-/// The ZST for which we impl the `super::Types` type collection trait.
-struct TestTypes;
-
-impl super::Types for TestTypes {
-    type Key = usize;
-
-    type Lsn = usize;
-
-    type LsnCounter = UsizeCounter;
-
-    type DeltaRecord = &'static str;
-
-    type HistoricLayer = Arc<TestHistoricLayer>;
-
-    type InMemoryLayer = TestInMemoryLayer;
-
-    type HistoricStuff = TestHistoricStuff;
-}
-
-/// For testing, our in-memory layer is a simple hashmap.
-#[derive(Clone, Default, Debug)]
-struct TestInMemoryLayer {
-    by_key: BTreeMap<usize, BTreeMap<usize, &'static str>>,
-}
-
-/// For testing, our historic layers are just in-memory layer objects with `frozen==true`.
-struct TestHistoricLayer(TestInMemoryLayer);
-
-/// This is the data structure that impls the `HistoricStuff` trait.
-#[derive(Default, Clone)]
-struct TestHistoricStuff {
-    by_key: BTreeMap<usize, BTreeMap<usize, Arc<TestHistoricLayer>>>,
-}
-
-/// `seqwait::MonotonicCounter` impl
-#[derive(Copy, Clone)]
-pub struct UsizeCounter(usize);
-
-// Our testing impl of HistoricStuff references the frozen InMemoryLayer objects
-// from all the (key,lsn) entries that it covers.
-// This mimics the (much more efficient) search tree in the real impl.
-impl super::HistoricStuff for TestHistoricStuff {
-    type Types = TestTypes;
-    fn get_reconstruct_path(
-        &self,
-        key: usize,
-        lsn: usize,
-    ) -> Result<Vec<Arc<TestHistoricLayer>>, super::GetReconstructPathError> {
-        let Some(bk) = self.by_key.get(&key) else {
-                return Ok(vec![]);
-            };
-        Ok(bk.range(..=lsn).rev().map(|(_, l)| Arc::clone(l)).collect())
-    }
-
-    fn make_historic(&self, inmem: TestInMemoryLayer) -> Self {
-        // For the purposes of testing, just turn the inmemory layer historic through the type system
-        let historic = Arc::new(TestHistoricLayer(inmem));
-        // Deep-copy
-        let mut copy = self.by_key.clone();
-        // Add the references to `inmem` to the deep-copied struct
-        for (k, v) in historic.0.by_key.iter() {
-            for (lsn, _deltas) in v.into_iter() {
-                let by_key = copy.entry(*k).or_default();
-                let overwritten = by_key.insert(*lsn, historic.clone());
-                assert!(matches!(overwritten, None), "layers must not overlap");
-            }
-        }
-        Self { by_key: copy }
-    }
-}
-
-impl super::InMemoryLayer for TestInMemoryLayer {
-    type Types = TestTypes;
-
-    fn put(
-        &mut self,
-        key: usize,
-        lsn: usize,
-        delta: &'static str,
-    ) -> Result<Self, super::InMemoryLayerPutError<&'static str>> {
-        let mut clone = self.clone();
-        drop(self);
-        let by_key = clone.by_key.entry(key).or_default();
-        match by_key.entry(lsn) {
-            Entry::Occupied(_record) => {
-                return Err(super::InMemoryLayerPutError {
-                    delta,
-                    kind: super::InMemoryLayerPutErrorKind::AlreadyHaveRecordForKeyAndLsn,
-                });
-            }
-            Entry::Vacant(vacant) => vacant.insert(delta),
-        };
-        Ok(clone)
-    }
-
-    fn get(&self, key: usize, lsn: usize) -> Vec<&'static str> {
-        let by_key = match self.by_key.get(&key) {
-            Some(by_key) => by_key,
-            None => return vec![],
-        };
-        by_key
-            .range(..=lsn)
-            .map(|(_, v)| v)
-            .rev()
-            .cloned()
-            .collect()
-    }
-}
-
-impl UsizeCounter {
-    pub fn new(inital: usize) -> Self {
-        UsizeCounter(inital)
-    }
-}
-
-impl seqwait::MonotonicCounter<usize> for UsizeCounter {
-    fn cnt_advance(&mut self, new_val: usize) {
-        assert!(self.0 < new_val);
-        self.0 = new_val;
-    }
-
-    fn cnt_value(&self) -> usize {
-        self.0
-    }
-}
-
-#[test]
-fn basic() {
-    let lm = TestHistoricStuff::default();
-
-    let (r, mut rw) = super::new::<TestTypes>(UsizeCounter::new(0), lm);
-
-    let r = Arc::new(r);
-    let r2 = Arc::clone(&r);
-
-    let rt = tokio::runtime::Builder::new_current_thread()
-        .enable_all()
-        .build()
-        .unwrap();
-
-    let read_jh = rt.spawn(async move { r.get(0, 10).await });
-
-    let mut rw = rt.block_on(async move {
-        rw.put(0, 1, "foo").await.unwrap();
-        rw.put(1, 1, "bar").await.unwrap();
-        rw.put(0, 10, "baz").await.unwrap();
-        rw
-    });
-
-    let read_res = rt.block_on(read_jh).unwrap().unwrap();
-    assert!(
-        read_res.historic_path.is_empty(),
-        "we have pushed less than needed for flush"
-    );
-    assert_eq!(read_res.inmem_records, vec!["baz", "foo"]);
-
-    let rw = rt.block_on(async move {
-        rw.put(0, 11, "blup").await.unwrap();
-        rw
-    });
-    let read_res = rt.block_on(async move { r2.get(0, 11).await.unwrap() });
-    assert_eq!(read_res.historic_path.len(), 0);
-    assert_eq!(read_res.inmem_records, vec!["blup", "baz", "foo"]);
-
-    drop(rw);
-}

libs/utils/Cargo.toml

@@ -37,7 +37,6 @@ uuid = { version = "1.2", features = ["v4", "serde"] }
 
 metrics.workspace = true
 workspace_hack.workspace = true
-either.workspace = true
 
 [dev-dependencies]
 byteorder.workspace = true

libs/utils/src/seqwait.rs

@@ -1,13 +1,12 @@
 #![warn(missing_docs)]
 
-use either::Either;
 use std::cmp::{Eq, Ordering, PartialOrd};
 use std::collections::BinaryHeap;
 use std::fmt::Debug;
 use std::mem;
-use std::sync::{Arc, Mutex};
+use std::sync::Mutex;
 use std::time::Duration;
-use tokio::sync::oneshot::{channel, Receiver, Sender};
+use tokio::sync::watch::{channel, Receiver, Sender};
 use tokio::time::timeout;
 
 /// An error happened while waiting for a number
@@ -37,48 +36,45 @@ pub trait MonotonicCounter<V> {
 }
 
 /// Internal components of a `SeqWait`
-struct SeqWaitInt<S, V, T>
+struct SeqWaitInt<S, V>
 where
     S: MonotonicCounter<V>,
     V: Ord,
-    T: Clone,
 {
-    waiters: BinaryHeap<Waiter<V, T>>,
+    waiters: BinaryHeap<Waiter<V>>,
     current: S,
     shutdown: bool,
-    data: T,
 }
 
-struct Waiter<V, T>
+struct Waiter<T>
 where
-    V: Ord,
-    T: Clone,
+    T: Ord,
 {
-    wake_num: V,             // wake me when this number arrives ...
-    wake_channel: Sender<T>, // ... by sending a message to this channel
+    wake_num: T,              // wake me when this number arrives ...
+    wake_channel: Sender<()>, // ... by sending a message to this channel
 }
 
 // BinaryHeap is a max-heap, and we want a min-heap. Reverse the ordering here
 // to get that.
-impl<V: Ord, T: Clone> PartialOrd for Waiter<V, T> {
+impl<T: Ord> PartialOrd for Waiter<T> {
     fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
         other.wake_num.partial_cmp(&self.wake_num)
     }
 }
 
-impl<V: Ord, T: Clone> Ord for Waiter<V, T> {
+impl<T: Ord> Ord for Waiter<T> {
     fn cmp(&self, other: &Self) -> Ordering {
         other.wake_num.cmp(&self.wake_num)
     }
 }
 
-impl<V: Ord, T: Clone> PartialEq for Waiter<V, T> {
+impl<T: Ord> PartialEq for Waiter<T> {
     fn eq(&self, other: &Self) -> bool {
         other.wake_num == self.wake_num
     }
 }
 
-impl<V: Ord, T: Clone> Eq for Waiter<V, T> {}
+impl<T: Ord> Eq for Waiter<T> {}
 
 /// A tool for waiting on a sequence number
 ///
@@ -96,28 +92,25 @@ impl<V: Ord, T: Clone> Eq for Waiter<V, T> {}
 ///
 /// <S> means Storage, <V> is type of counter that this storage exposes.
 ///
-pub struct SeqWait<S, V, T>
+pub struct SeqWait<S, V>
 where
     S: MonotonicCounter<V>,
     V: Ord,
-    T: Clone,
 {
-    internal: Mutex<SeqWaitInt<S, V, T>>,
+    internal: Mutex<SeqWaitInt<S, V>>,
 }
 
-impl<S, V, T> SeqWait<S, V, T>
+impl<S, V> SeqWait<S, V>
 where
     S: MonotonicCounter<V> + Copy,
     V: Ord + Copy,
-    T: Clone,
 {
     /// Create a new `SeqWait`, initialized to a particular number
-    pub fn new(starting_num: S, data: T) -> Self {
+    pub fn new(starting_num: S) -> Self {
         let internal = SeqWaitInt {
             waiters: BinaryHeap::new(),
             current: starting_num,
             shutdown: false,
-            data,
         };
         SeqWait {
             internal: Mutex::new(internal),
@@ -151,13 +144,10 @@ where
     ///
     /// This call won't complete until someone has called `advance`
     /// with a number greater than or equal to the one we're waiting for.
-    pub async fn wait_for(&self, num: V) -> Result<T, SeqWaitError> {
-        match self.queue_for_wait(num, false) {
-            Ok(Either::Left(data)) => Ok(data),
-            Ok(Either::Right(rx)) => match rx.await {
-                Err(_) => Err(SeqWaitError::Shutdown),
-                Ok(data) => Ok(data),
-            },
+    pub async fn wait_for(&self, num: V) -> Result<(), SeqWaitError> {
+        match self.queue_for_wait(num) {
+            Ok(None) => Ok(()),
+            Ok(Some(mut rx)) => rx.changed().await.map_err(|_| SeqWaitError::Shutdown),
             Err(e) => Err(e),
         }
     }
@@ -169,18 +159,15 @@ where
     ///
     /// If that hasn't happened after the specified timeout duration,
     /// [`SeqWaitError::Timeout`] will be returned.
-    ///
-    /// Pass `timeout_duration.is_zero() == true` to guarantee that the
-    /// future that is this function will never await.
     pub async fn wait_for_timeout(
         &self,
         num: V,
         timeout_duration: Duration,
-    ) -> Result<T, SeqWaitError> {
-        match self.queue_for_wait(num, timeout_duration.is_zero()) {
-            Ok(Either::Left(data)) => Ok(data),
-            Ok(Either::Right(rx)) => match timeout(timeout_duration, rx).await {
-                Ok(Ok(data)) => Ok(data),
+    ) -> Result<(), SeqWaitError> {
+        match self.queue_for_wait(num) {
+            Ok(None) => Ok(()),
+            Ok(Some(mut rx)) => match timeout(timeout_duration, rx.changed()).await {
+                Ok(Ok(())) => Ok(()),
                 Ok(Err(_)) => Err(SeqWaitError::Shutdown),
                 Err(_) => Err(SeqWaitError::Timeout),
             },
@@ -190,50 +177,41 @@ where
 
     /// Register and return a channel that will be notified when a number arrives,
     /// or None, if it has already arrived.
-    fn queue_for_wait(&self, num: V, nowait: bool) -> Result<Either<T, Receiver<T>>, SeqWaitError> {
+    fn queue_for_wait(&self, num: V) -> Result<Option<Receiver<()>>, SeqWaitError> {
         let mut internal = self.internal.lock().unwrap();
         if internal.current.cnt_value() >= num {
-            return Ok(Either::Left(internal.data.clone()));
+            return Ok(None);
         }
         if internal.shutdown {
             return Err(SeqWaitError::Shutdown);
         }
-        if nowait {
-            return Err(SeqWaitError::Timeout);
-        }
 
         // Create a new channel.
-        let (tx, rx) = channel();
+        let (tx, rx) = channel(());
         internal.waiters.push(Waiter {
             wake_num: num,
             wake_channel: tx,
         });
         // Drop the lock as we exit this scope.
-        Ok(Either::Right(rx))
+        Ok(Some(rx))
     }
 
     /// Announce a new number has arrived
     ///
     /// All waiters at this value or below will be woken.
     ///
-    /// If `new_data` is Some(), it will update the internal data,
-    /// even if `num` is smaller than the internal counter.
-    /// It will not cause a wake-up though, in this case.
-    ///
     /// Returns the old number.
-    pub fn advance(&self, num: V, new_data: Option<T>) -> V {
+    pub fn advance(&self, num: V) -> V {
         let old_value;
-        let (wake_these, with_data) = {
+        let wake_these = {
             let mut internal = self.internal.lock().unwrap();
-            if let Some(new_data) = new_data {
-                internal.data = new_data;
-            }
 
             old_value = internal.current.cnt_value();
             if old_value >= num {
                 return old_value;
             }
             internal.current.cnt_advance(num);
+
             // Pop all waiters <= num from the heap. Collect them in a vector, and
             // wake them up after releasing the lock.
             let mut wake_these = Vec::new();
@@ -243,13 +221,13 @@ where
                 }
                 wake_these.push(internal.waiters.pop().unwrap().wake_channel);
             }
-            (wake_these, internal.data.clone())
+            wake_these
         };
 
         for tx in wake_these {
             // This can fail if there are no receivers.
             // We don't care; discard the error.
-            let _ = tx.send(with_data.clone());
+            let _ = tx.send(());
         }
         old_value
     }
@@ -258,106 +236,6 @@ where
     pub fn load(&self) -> S {
         self.internal.lock().unwrap().current
     }
-
-    /// Split the seqwait into a part than can only do wait,
-    /// and another part that can do advance + wait.
-    ///
-    /// The wait-only part can be cloned, the advance part cannot be cloned.
-    /// This provides a single-producer multi-consumer scheme.
-    pub fn split_spmc(self) -> (Wait<S, V, T>, Advance<S, V, T>) {
-        let inner = Arc::new(self);
-        let w = Wait {
-            inner: inner.clone(),
-        };
-        let a = Advance { inner };
-        (w, a)
-    }
-}
-
-/// See [`SeqWait::split_spmc`].
-pub struct Wait<S, V, T>
-where
-    S: MonotonicCounter<V> + Copy,
-    V: Ord + Copy,
-    T: Clone,
-{
-    inner: Arc<SeqWait<S, V, T>>,
-}
-
-/// See [`SeqWait::split_spmc`].
-pub struct Advance<S, V, T>
-where
-    S: MonotonicCounter<V> + Copy,
-    V: Ord + Copy,
-    T: Clone,
-{
-    inner: Arc<SeqWait<S, V, T>>,
-}
-
-impl<S, V, T> Wait<S, V, T>
-where
-    S: MonotonicCounter<V> + Copy,
-    V: Ord + Copy,
-    T: Clone,
-{
-    /// See [`SeqWait::wait_for`].
-    pub async fn wait_for(&self, num: V) -> Result<T, SeqWaitError> {
-        self.inner.wait_for(num).await
-    }
-
-    /// See [`SeqWait::wait_for_timeout`].
-    pub async fn wait_for_timeout(
-        &self,
-        num: V,
-        timeout_duration: Duration,
-    ) -> Result<T, SeqWaitError> {
-        self.inner.wait_for_timeout(num, timeout_duration).await
-    }
-}
-
-impl<S, V, T> Advance<S, V, T>
-where
-    S: MonotonicCounter<V> + Copy,
-    V: Ord + Copy,
-    T: Clone,
-{
-    /// See [`SeqWait::advance`].
-    pub fn advance(&self, num: V, new_data: Option<T>) -> V {
-        self.inner.advance(num, new_data)
-    }
-
-    /// See [`SeqWait::wait_for`].
-    pub async fn wait_for(&self, num: V) -> Result<T, SeqWaitError> {
-        self.inner.wait_for(num).await
-    }
-
-    /// See [`SeqWait::wait_for_timeout`].
-    pub async fn wait_for_timeout(
-        &self,
-        num: V,
-        timeout_duration: Duration,
-    ) -> Result<T, SeqWaitError> {
-        self.inner.wait_for_timeout(num, timeout_duration).await
-    }
-
-    /// Get a `Clone::clone` of the current data inside the seqwait.
-    pub fn get_current_data(&self) -> (V, T) {
-        let inner = self.inner.internal.lock().unwrap();
-        (inner.current.cnt_value(), inner.data.clone())
-    }
-}
-
-impl<S, V, T> Clone for Wait<S, V, T>
-where
-    S: MonotonicCounter<V> + Copy,
-    V: Ord + Copy,
-    T: Clone,
-{
-    fn clone(&self) -> Self {
-        Self {
-            inner: self.inner.clone(),
-        }
-    }
 }
 
 #[cfg(test)]
@@ -378,12 +256,12 @@ mod tests {
 
     #[tokio::test]
     async fn seqwait() {
-        let seq = Arc::new(SeqWait::new(0, ()));
+        let seq = Arc::new(SeqWait::new(0));
         let seq2 = Arc::clone(&seq);
         let seq3 = Arc::clone(&seq);
         let jh1 = tokio::task::spawn(async move {
             seq2.wait_for(42).await.expect("wait_for 42");
-            let old = seq2.advance(100, None);
+            let old = seq2.advance(100);
             assert_eq!(old, 99);
             seq2.wait_for_timeout(999, Duration::from_millis(100))
                 .await
@@ -394,12 +272,12 @@ mod tests {
             seq3.wait_for(0).await.expect("wait_for 0");
         });
         tokio::time::sleep(Duration::from_millis(200)).await;
-        let old = seq.advance(99, None);
+        let old = seq.advance(99);
         assert_eq!(old, 0);
         seq.wait_for(100).await.expect("wait_for 100");
 
         // Calling advance with a smaller value is a no-op
-        assert_eq!(seq.advance(98, None), 100);
+        assert_eq!(seq.advance(98), 100);
         assert_eq!(seq.load(), 100);
 
         jh1.await.unwrap();
@@ -410,7 +288,7 @@ mod tests {
 
     #[tokio::test]
     async fn seqwait_timeout() {
-        let seq = Arc::new(SeqWait::new(0, ()));
+        let seq = Arc::new(SeqWait::new(0));
         let seq2 = Arc::clone(&seq);
         let jh = tokio::task::spawn(async move {
             let timeout = Duration::from_millis(1);
@@ -420,104 +298,10 @@ mod tests {
         tokio::time::sleep(Duration::from_millis(200)).await;
         // This will attempt to wake, but nothing will happen
         // because the waiter already dropped its Receiver.
-        let old = seq.advance(99, None);
+        let old = seq.advance(99);
         assert_eq!(old, 0);
         jh.await.unwrap();
 
         seq.shutdown();
     }
-
-    #[tokio::test]
-    async fn data_basic() {
-        let seq = Arc::new(SeqWait::new(0, "a"));
-        let seq2 = Arc::clone(&seq);
-        let jh = tokio::task::spawn(async move {
-            let data = seq.wait_for(2).await.unwrap();
-            assert_eq!(data, "b");
-        });
-        seq2.advance(1, Some("x"));
-        seq2.advance(2, Some("b"));
-        jh.await.unwrap();
-    }
-
-    #[test]
-    fn data_always_most_recent() {
-        let rt = tokio::runtime::Builder::new_current_thread()
-            .build()
-            .unwrap();
-
-        let seq = Arc::new(SeqWait::new(0, "a"));
-        let seq2 = Arc::clone(&seq);
-
-        let jh = rt.spawn(async move {
-            let data = seq.wait_for(2).await.unwrap();
-            assert_eq!(data, "d");
-        });
-
-        // jh is not running until we poll it, thanks to current thread runtime
-
-        rt.block_on(async move {
-            seq2.advance(2, Some("b"));
-            seq2.advance(3, Some("c"));
-            seq2.advance(4, Some("d"));
-        });
-
-        rt.block_on(jh).unwrap();
-    }
-
-    #[tokio::test]
-    async fn split_spmc_api_surface() {
-        let seq = SeqWait::new(0, 1);
-        let (w, a) = seq.split_spmc();
-
-        let _ = w.wait_for(1);
-        let _ = w.wait_for_timeout(0, Duration::from_secs(10));
-        let _ = w.clone();
-
-        let _ = a.advance(1, None);
-        let _ = a.wait_for(1);
-        let _ = a.wait_for_timeout(0, Duration::from_secs(10));
-
-        // TODO would be nice to have must-not-compile tests for Advance not being clonable.
-    }
-
-    #[tokio::test]
-    async fn new_data_same_lsn() {
-        let seq = Arc::new(SeqWait::new(0, "a"));
-
-        seq.advance(1, Some("b"));
-        let data = seq.wait_for(1).await.unwrap();
-        assert_eq!(data, "b", "the regular case where lsn and data advance");
-
-        seq.advance(1, Some("c"));
-        let data = seq.wait_for(1).await.unwrap();
-        assert_eq!(
-            data, "c",
-            "no lsn advance still gives new data for old lsn wait_for's"
-        );
-
-        let (start_wait_for_sender, start_wait_for_receiver) = tokio::sync::oneshot::channel();
-        // ensure we don't wake waiters for data-only change
-        let jh = tokio::spawn({
-            let seq = seq.clone();
-            async move {
-                start_wait_for_receiver.await.unwrap();
-                match tokio::time::timeout(Duration::from_secs(2), seq.wait_for(2)).await {
-                    Ok(_) => {
-                        assert!(
-                            false,
-                            "advance should not wake waiters if data changes but LSN doesn't"
-                        );
-                    }
-                    Err(_) => {
-                        // Good, we weren't woken up.
-                    }
-                }
-            }
-        });
-
-        seq.advance(1, Some("d"));
-        start_wait_for_sender.send(()).unwrap();
-        jh.await.unwrap();
-    }
 }

libs/utils/src/serde_percent.rs

@@ -11,14 +11,6 @@ use serde::{Deserialize, Serialize};
 pub struct Percent(#[serde(deserialize_with = "deserialize_pct_0_to_100")] u8);
 
 impl Percent {
-    pub const fn new(pct: u8) -> Option<Self> {
-        if pct <= 100 {
-            Some(Percent(pct))
-        } else {
-            None
-        }
-    }
-
     pub fn get(&self) -> u8 {
         self.0
     }

pageserver/benches/bench_layer_map.rs

@@ -33,7 +33,7 @@ fn build_layer_map(filename_dump: PathBuf) -> LayerMap<LayerDescriptor> {
         min_lsn = min(min_lsn, lsn_range.start);
         max_lsn = max(max_lsn, Lsn(lsn_range.end.0 - 1));
 
-        updates.insert_historic(Arc::new(layer)).unwrap();
+        updates.insert_historic(Arc::new(layer));
     }
 
     println!("min: {min_lsn}, max: {max_lsn}");
@@ -215,7 +215,7 @@ fn bench_sequential(c: &mut Criterion) {
             is_incremental: false,
             short_id: format!("Layer {}", i),
         };
-        updates.insert_historic(Arc::new(layer)).unwrap();
+        updates.insert_historic(Arc::new(layer));
     }
     updates.flush();
     println!("Finished layer map init in {:?}", now.elapsed());

pageserver/src/disk_usage_eviction_task.rs

@@ -639,7 +639,7 @@ mod filesystem_level_usage {
                 ),
                 (
                     "max_usage_pct",
-                    usage_pct >= self.config.max_usage_pct.get() as u64,
+                    usage_pct > self.config.max_usage_pct.get() as u64,
                 ),
             ];
 
@@ -686,43 +686,4 @@ mod filesystem_level_usage {
             avail_bytes,
         })
     }
-
-    #[test]
-    fn max_usage_pct_pressure() {
-        use super::Usage as _;
-        use std::time::Duration;
-        use utils::serde_percent::Percent;
-
-        let mut usage = Usage {
-            config: &DiskUsageEvictionTaskConfig {
-                max_usage_pct: Percent::new(85).unwrap(),
-                min_avail_bytes: 0,
-                period: Duration::MAX,
-                #[cfg(feature = "testing")]
-                mock_statvfs: None,
-            },
-            total_bytes: 100_000,
-            avail_bytes: 0,
-        };
-
-        assert!(usage.has_pressure(), "expected pressure at 100%");
-
-        usage.add_available_bytes(14_000);
-        assert!(usage.has_pressure(), "expected pressure at 86%");
-
-        usage.add_available_bytes(999);
-        assert!(usage.has_pressure(), "expected pressure at 85.001%");
-
-        usage.add_available_bytes(1);
-        assert!(usage.has_pressure(), "expected pressure at precisely 85%");
-
-        usage.add_available_bytes(1);
-        assert!(!usage.has_pressure(), "no pressure at 84.999%");
-
-        usage.add_available_bytes(999);
-        assert!(!usage.has_pressure(), "no pressure at 84%");
-
-        usage.add_available_bytes(16_000);
-        assert!(!usage.has_pressure());
-    }
 }

pageserver/src/http/openapi_spec.yml

@@ -829,9 +829,12 @@ components:
       type: object
       required:
         - id
+        - state
       properties:
         id:
           type: string
+        state:
+          type: string
         current_physical_size:
           type: integer
         has_in_progress_downloads:

pageserver/src/http/routes.rs

@@ -465,7 +465,7 @@ async fn tenant_list_handler(request: Request<Body>) -> Result<Response<Body>, A
         .iter()
         .map(|(id, state)| TenantInfo {
             id: *id,
-            state: state.clone(),
+            state: *state,
             current_physical_size: None,
             has_in_progress_downloads: Some(state.has_in_progress_downloads()),
         })
@@ -490,7 +490,7 @@ async fn tenant_status(request: Request<Body>) -> Result<Response<Body>, ApiErro
         let state = tenant.current_state();
         Ok(TenantInfo {
             id: tenant_id,
-            state: state.clone(),
+            state,
             current_physical_size: Some(current_physical_size),
             has_in_progress_downloads: Some(state.has_in_progress_downloads()),
         })
@@ -931,7 +931,7 @@ async fn handle_tenant_break(r: Request<Body>) -> Result<Response<Body>, ApiErro
         .await
         .map_err(|_| ApiError::Conflict(String::from("no active tenant found")))?;
 
-    tenant.set_broken("broken from test".to_owned());
+    tenant.set_broken("broken from test");
 
     json_response(StatusCode::OK, ())
 }

pageserver/src/metrics.rs

@@ -6,8 +6,7 @@ use metrics::{
     UIntGauge, UIntGaugeVec,
 };
 use once_cell::sync::Lazy;
-use pageserver_api::models::TenantState;
-use strum::VariantNames;
+use pageserver_api::models::state;
 use utils::id::{TenantId, TimelineId};
 
 /// Prometheus histogram buckets (in seconds) for operations in the critical
@@ -148,6 +147,15 @@ static CURRENT_LOGICAL_SIZE: Lazy<UIntGaugeVec> = Lazy::new(|| {
     .expect("failed to define current logical size metric")
 });
 
+// Metrics collected on tenant states.
+const TENANT_STATE_OPTIONS: &[&str] = &[
+    state::LOADING,
+    state::ATTACHING,
+    state::ACTIVE,
+    state::STOPPING,
+    state::BROKEN,
+];
+
 pub static TENANT_STATE_METRIC: Lazy<UIntGaugeVec> = Lazy::new(|| {
     register_uint_gauge_vec!(
         "pageserver_tenant_states_count",
@@ -699,7 +707,7 @@ impl Drop for TimelineMetrics {
 pub fn remove_tenant_metrics(tenant_id: &TenantId) {
     let tid = tenant_id.to_string();
     let _ = TENANT_SYNTHETIC_SIZE_METRIC.remove_label_values(&[&tid]);
-    for state in TenantState::VARIANTS {
+    for state in TENANT_STATE_OPTIONS {
         let _ = TENANT_STATE_METRIC.remove_label_values(&[&tid, state]);
     }
 }

pageserver/src/tenant.rs

@@ -177,9 +177,9 @@ impl UninitializedTimeline<'_> {
     ///
     /// The new timeline is initialized in Active state, and its background jobs are
     /// started
-    pub fn initialize(self, ctx: &RequestContext) -> anyhow::Result<Arc<Timeline>> {
+    pub fn initialize(self, _ctx: &RequestContext) -> anyhow::Result<Arc<Timeline>> {
         let mut timelines = self.owning_tenant.timelines.lock().unwrap();
-        self.initialize_with_lock(ctx, &mut timelines, true, true)
+        self.initialize_with_lock(&mut timelines, true, true)
     }
 
     /// Like `initialize`, but the caller is already holding lock on Tenant::timelines.
@@ -189,7 +189,6 @@ impl UninitializedTimeline<'_> {
     /// been initialized.
     fn initialize_with_lock(
         mut self,
-        ctx: &RequestContext,
         timelines: &mut HashMap<TimelineId, Arc<Timeline>>,
         load_layer_map: bool,
         activate: bool,
@@ -230,9 +229,7 @@ impl UninitializedTimeline<'_> {
                 new_timeline.maybe_spawn_flush_loop();
 
                 if activate {
-                    new_timeline
-                        .activate(ctx)
-                        .context("initializing timeline activation")?;
+                    new_timeline.activate();
                 }
             }
         }
@@ -267,10 +264,7 @@ impl UninitializedTimeline<'_> {
             .await
             .context("Failed to flush after basebackup import")?;
 
-        // Initialize without loading the layer map. We started with an empty layer map, and already
-        // updated it for the layers that we created during the import.
-        let mut timelines = self.owning_tenant.timelines.lock().unwrap();
-        self.initialize_with_lock(ctx, &mut timelines, false, true)
+        self.initialize(ctx)
     }
 
     fn raw_timeline(&self) -> anyhow::Result<&Arc<Timeline>> {
@@ -475,7 +469,7 @@ impl Tenant {
         local_metadata: Option<TimelineMetadata>,
         ancestor: Option<Arc<Timeline>>,
         first_save: bool,
-        ctx: &RequestContext,
+        _ctx: &RequestContext,
     ) -> anyhow::Result<()> {
         let tenant_id = self.tenant_id;
 
@@ -510,7 +504,7 @@ impl Tenant {
             // Do not start walreceiver here. We do need loaded layer map for reconcile_with_remote
             // But we shouldnt start walreceiver before we have all the data locally, because working walreceiver
             // will ingest data which may require looking at the layers which are not yet available locally
-            match timeline.initialize_with_lock(ctx, &mut timelines_accessor, true, false) {
+            match timeline.initialize_with_lock(&mut timelines_accessor, true, false) {
                 Ok(new_timeline) => new_timeline,
                 Err(e) => {
                     error!("Failed to initialize timeline {tenant_id}/{timeline_id}: {e:?}");
@@ -622,7 +616,7 @@ impl Tenant {
                 match tenant_clone.attach(ctx).await {
                     Ok(_) => {}
                     Err(e) => {
-                        tenant_clone.set_broken(e.to_string());
+                        tenant_clone.set_broken(&e.to_string());
                         error!("error attaching tenant: {:?}", e);
                     }
                 }
@@ -635,7 +629,7 @@ impl Tenant {
     ///
     /// Background task that downloads all data for a tenant and brings it to Active state.
     ///
-    #[instrument(skip_all, fields(tenant_id=%self.tenant_id))]
+    #[instrument(skip(self, ctx), fields(tenant_id=%self.tenant_id))]
     async fn attach(self: &Arc<Tenant>, ctx: RequestContext) -> anyhow::Result<()> {
         // Create directory with marker file to indicate attaching state.
         // The load_local_tenants() function in tenant::mgr relies on the marker file
@@ -756,7 +750,7 @@ impl Tenant {
 
         // Start background operations and open the tenant for business.
         // The loops will shut themselves down when they notice that the tenant is inactive.
-        self.activate(&ctx)?;
+        self.activate()?;
 
         info!("Done");
 
@@ -830,10 +824,7 @@ impl Tenant {
     pub fn create_broken_tenant(conf: &'static PageServerConf, tenant_id: TenantId) -> Arc<Tenant> {
         let wal_redo_manager = Arc::new(PostgresRedoManager::new(conf, tenant_id));
         Arc::new(Tenant::new(
-            TenantState::Broken {
-                reason: "create_broken_tenant".into(),
-                backtrace: String::new(),
-            },
+            TenantState::Broken,
             conf,
             TenantConfOpt::default(),
             wal_redo_manager,
@@ -894,7 +885,7 @@ impl Tenant {
                 match tenant_clone.load(&ctx).await {
                     Ok(()) => {}
                     Err(err) => {
-                        tenant_clone.set_broken(err.to_string());
+                        tenant_clone.set_broken(&err.to_string());
                         error!("could not load tenant {tenant_id}: {err:?}");
                     }
                 }
@@ -1031,7 +1022,7 @@ impl Tenant {
 
         // Start background operations and open the tenant for business.
         // The loops will shut themselves down when they notice that the tenant is inactive.
-        self.activate(ctx)?;
+        self.activate()?;
 
         info!("Done");
 
@@ -1367,7 +1358,12 @@ impl Tenant {
 
         // Stop the walreceiver first.
         debug!("waiting for wal receiver to shutdown");
-        timeline.walreceiver.stop().await;
+        task_mgr::shutdown_tasks(
+            Some(TaskKind::WalReceiverManager),
+            Some(self.tenant_id),
+            Some(timeline_id),
+        )
+        .await;
         debug!("wal receiver shutdown confirmed");
 
         info!("waiting for timeline tasks to shutdown");
@@ -1446,7 +1442,7 @@ impl Tenant {
     }
 
     pub fn current_state(&self) -> TenantState {
-        self.state.borrow().clone()
+        *self.state.borrow()
     }
 
     pub fn is_active(&self) -> bool {
@@ -1454,18 +1450,18 @@ impl Tenant {
     }
 
     /// Changes tenant status to active, unless shutdown was already requested.
-    fn activate(&self, ctx: &RequestContext) -> anyhow::Result<()> {
+    fn activate(&self) -> anyhow::Result<()> {
         let mut result = Ok(());
         self.state.send_modify(|current_state| {
-            match &*current_state {
+            match *current_state {
                 TenantState::Active => {
                     // activate() was called on an already Active tenant. Shouldn't happen.
                     result = Err(anyhow::anyhow!("Tenant is already active"));
                 }
-                TenantState::Broken { reason, .. } => {
+                TenantState::Broken => {
                     // This shouldn't happen either
                     result = Err(anyhow::anyhow!(
-                        "Could not activate tenant because it is in broken state due to: {reason}",
+                        "Could not activate tenant because it is in broken state"
                     ));
                 }
                 TenantState::Stopping => {
@@ -1488,23 +1484,7 @@ impl Tenant {
                     tasks::start_background_loops(self.tenant_id);
 
                     for timeline in not_broken_timelines {
-                        match timeline
-                            .activate(ctx)
-                            .context("timeline activation for activating tenant")
-                        {
-                            Ok(()) => {}
-                            Err(e) => {
-                                error!(
-                                    "Failed to activate timeline {}: {:#}",
-                                    timeline.timeline_id, e
-                                );
-                                timeline.set_state(TimelineState::Broken);
-                                *current_state = TenantState::broken_from_reason(format!(
-                                    "failed to activate timeline {}: {}",
-                                    timeline.timeline_id, e
-                                ));
-                            }
-                        }
+                        timeline.activate();
                     }
                 }
             }
@@ -1515,7 +1495,7 @@ impl Tenant {
     /// Change tenant status to Stopping, to mark that it is being shut down
     pub fn set_stopping(&self) {
         self.state.send_modify(|current_state| {
-            match current_state {
+            match *current_state {
                 TenantState::Active | TenantState::Loading | TenantState::Attaching => {
                     *current_state = TenantState::Stopping;
 
@@ -1531,8 +1511,8 @@ impl Tenant {
                         timeline.set_state(TimelineState::Stopping);
                     }
                 }
-                TenantState::Broken { reason, .. } => {
-                    info!("Cannot set tenant to Stopping state, it is in Broken state due to: {reason}");
+                TenantState::Broken => {
+                    info!("Cannot set tenant to Stopping state, it is already in Broken state");
                 }
                 TenantState::Stopping => {
                     // The tenant was detached, or system shutdown was requested, while we were
@@ -1543,7 +1523,7 @@ impl Tenant {
         });
     }
 
-    pub fn set_broken(&self, reason: String) {
+    pub fn set_broken(&self, reason: &str) {
         self.state.send_modify(|current_state| {
             match *current_state {
                 TenantState::Active => {
@@ -1551,24 +1531,24 @@ impl Tenant {
                     // while loading or attaching a tenant. A tenant that has already been
                     // activated should never be marked as broken. We cope with it the best
                     // we can, but it shouldn't happen.
+                    *current_state = TenantState::Broken;
                     warn!("Changing Active tenant to Broken state, reason: {}", reason);
-                    *current_state = TenantState::broken_from_reason(reason);
                 }
-                TenantState::Broken { .. } => {
+                TenantState::Broken => {
                     // This shouldn't happen either
                     warn!("Tenant is already in Broken state");
                 }
                 TenantState::Stopping => {
                     // This shouldn't happen either
+                    *current_state = TenantState::Broken;
                     warn!(
                         "Marking Stopping tenant as Broken state, reason: {}",
                         reason
                     );
-                    *current_state = TenantState::broken_from_reason(reason);
                 }
                 TenantState::Loading | TenantState::Attaching => {
                     info!("Setting tenant as Broken state, reason: {}", reason);
-                    *current_state = TenantState::broken_from_reason(reason);
+                    *current_state = TenantState::Broken;
                 }
             }
         });
@@ -1581,7 +1561,7 @@ impl Tenant {
     pub async fn wait_to_become_active(&self) -> anyhow::Result<()> {
         let mut receiver = self.state.subscribe();
         loop {
-            let current_state = receiver.borrow_and_update().clone();
+            let current_state = *receiver.borrow_and_update();
             match current_state {
                 TenantState::Loading | TenantState::Attaching => {
                     // in these states, there's a chance that we can reach ::Active
@@ -1590,12 +1570,12 @@ impl Tenant {
                 TenantState::Active { .. } => {
                     return Ok(());
                 }
-                TenantState::Broken { .. } | TenantState::Stopping => {
+                TenantState::Broken | TenantState::Stopping => {
                     // There's no chance the tenant can transition back into ::Active
                     anyhow::bail!(
                         "Tenant {} will not become active. Current state: {:?}",
                         self.tenant_id,
-                        &current_state,
+                        current_state,
                     );
                 }
             }
@@ -1776,23 +1756,21 @@ impl Tenant {
         let (state, mut rx) = watch::channel(state);
 
         tokio::spawn(async move {
-            let mut current_state: &'static str = From::from(&*rx.borrow_and_update());
+            let current_state = *rx.borrow_and_update();
             let tid = tenant_id.to_string();
             TENANT_STATE_METRIC
-                .with_label_values(&[&tid, current_state])
+                .with_label_values(&[&tid, current_state.as_str()])
                 .inc();
             loop {
                 match rx.changed().await {
                     Ok(()) => {
-                        let new_state: &'static str = From::from(&*rx.borrow_and_update());
+                        let new_state = *rx.borrow();
                         TENANT_STATE_METRIC
-                            .with_label_values(&[&tid, current_state])
+                            .with_label_values(&[&tid, current_state.as_str()])
                             .dec();
                         TENANT_STATE_METRIC
-                            .with_label_values(&[&tid, new_state])
+                            .with_label_values(&[&tid, new_state.as_str()])
                             .inc();
-
-                        current_state = new_state;
                     }
                     Err(_sender_dropped_error) => {
                         info!("Tenant dropped the state updates sender, quitting waiting for tenant state change");
@@ -2115,7 +2093,7 @@ impl Tenant {
         src_timeline: &Arc<Timeline>,
         dst_id: TimelineId,
         start_lsn: Option<Lsn>,
-        ctx: &RequestContext,
+        _ctx: &RequestContext,
     ) -> anyhow::Result<Arc<Timeline>> {
         let src_id = src_timeline.timeline_id;
 
@@ -2208,7 +2186,7 @@ impl Tenant {
                 false,
                 Some(Arc::clone(src_timeline)),
             )?
-            .initialize_with_lock(ctx, &mut timelines, true, true)?;
+            .initialize_with_lock(&mut timelines, true, true)?;
         drop(timelines);
 
         // Root timeline gets its layers during creation and uploads them along with the metadata.
@@ -2319,11 +2297,9 @@ impl Tenant {
                 )
             })?;
 
-        // Initialize the timeline without loading the layer map, because we already updated the layer
-        // map above, when we imported the datadir.
         let timeline = {
             let mut timelines = self.timelines.lock().unwrap();
-            raw_timeline.initialize_with_lock(ctx, &mut timelines, false, true)?
+            raw_timeline.initialize_with_lock(&mut timelines, false, true)?
         };
 
         info!(

pageserver/src/tenant/layer_map.rs

@@ -52,7 +52,7 @@ use crate::metrics::NUM_ONDISK_LAYERS;
 use crate::repository::Key;
 use crate::tenant::storage_layer::InMemoryLayer;
 use crate::tenant::storage_layer::Layer;
-use anyhow::{bail, Result};
+use anyhow::Result;
 use std::collections::VecDeque;
 use std::ops::Range;
 use std::sync::Arc;
@@ -126,7 +126,7 @@ where
     ///
     /// Insert an on-disk layer.
     ///
-    pub fn insert_historic(&mut self, layer: Arc<L>) -> anyhow::Result<()> {
+    pub fn insert_historic(&mut self, layer: Arc<L>) {
         self.layer_map.insert_historic_noflush(layer)
     }
 
@@ -274,22 +274,17 @@ where
     ///
     /// Helper function for BatchedUpdates::insert_historic
     ///
-    pub(self) fn insert_historic_noflush(&mut self, layer: Arc<L>) -> anyhow::Result<()> {
-        let key = historic_layer_coverage::LayerKey::from(&*layer);
-        if self.historic.contains(&key) {
-            bail!(
-                "Attempt to insert duplicate layer {} in layer map",
-                layer.short_id()
-            );
-        }
-        self.historic.insert(key, Arc::clone(&layer));
+    pub(self) fn insert_historic_noflush(&mut self, layer: Arc<L>) {
+        self.historic.insert(
+            historic_layer_coverage::LayerKey::from(&*layer),
+            Arc::clone(&layer),
+        );
 
         if Self::is_l0(&layer) {
             self.l0_delta_layers.push(layer);
         }
 
         NUM_ONDISK_LAYERS.inc();
-        Ok(())
     }
 
     ///
@@ -843,7 +838,7 @@ mod tests {
 
             let expected_in_counts = (1, usize::from(expected_l0));
 
-            map.batch_update().insert_historic(remote.clone()).unwrap();
+            map.batch_update().insert_historic(remote.clone());
             assert_eq!(count_layer_in(&map, &remote), expected_in_counts);
 
             let replaced = map

pageserver/src/tenant/layer_map/historic_layer_coverage.rs

@@ -417,14 +417,6 @@ impl<Value: Clone> BufferedHistoricLayerCoverage<Value> {
         }
     }
 
-    pub fn contains(&self, layer_key: &LayerKey) -> bool {
-        match self.buffer.get(layer_key) {
-            Some(None) => false,                         // layer remove was buffered
-            Some(_) => true,                             // layer insert was buffered
-            None => self.layers.contains_key(layer_key), // no buffered ops for this layer
-        }
-    }
-
     pub fn insert(&mut self, layer_key: LayerKey, value: Value) {
         self.buffer.insert(layer_key, Some(value));
     }

pageserver/src/tenant/mgr.rs

@@ -537,7 +537,7 @@ where
             Some(tenant) => match tenant.current_state() {
                 TenantState::Attaching
                 | TenantState::Loading
-                | TenantState::Broken { .. }
+                | TenantState::Broken
                 | TenantState::Active => tenant.set_stopping(),
                 TenantState::Stopping => return Err(TenantStateError::IsStopping(tenant_id)),
             },
@@ -565,7 +565,7 @@ where
             let tenants_accessor = TENANTS.read().await;
             match tenants_accessor.get(&tenant_id) {
                 Some(tenant) => {
-                    tenant.set_broken(e.to_string());
+                    tenant.set_broken(&e.to_string());
                 }
                 None => {
                     warn!("Tenant {tenant_id} got removed from memory");

pageserver/src/tenant/tasks.rs

@@ -209,7 +209,7 @@ async fn wait_for_active_tenant(
         loop {
             match tenant_state_updates.changed().await {
                 Ok(()) => {
-                    let new_state = &*tenant_state_updates.borrow();
+                    let new_state = *tenant_state_updates.borrow();
                     match new_state {
                         TenantState::Active => {
                             debug!("Tenant state changed to active, continuing the task loop");

pageserver/src/tenant/timeline.rs

@@ -14,7 +14,6 @@ use pageserver_api::models::{
     DownloadRemoteLayersTaskState, LayerMapInfo, LayerResidenceStatus, TimelineState,
 };
 use remote_storage::GenericRemoteStorage;
-use storage_broker::BrokerClientChannel;
 use tokio::sync::{oneshot, watch, Semaphore, TryAcquireError};
 use tokio_util::sync::CancellationToken;
 use tracing::*;
@@ -31,7 +30,7 @@ use std::sync::atomic::{AtomicI64, Ordering as AtomicOrdering};
 use std::sync::{Arc, Mutex, MutexGuard, RwLock, Weak};
 use std::time::{Duration, Instant, SystemTime};
 
-use crate::broker_client::{get_broker_client, is_broker_client_initialized};
+use crate::broker_client::is_broker_client_initialized;
 use crate::context::{DownloadBehavior, RequestContext};
 use crate::tenant::remote_timeline_client::{self, index::LayerFileMetadata};
 use crate::tenant::storage_layer::{
@@ -72,10 +71,10 @@ use crate::walredo::WalRedoManager;
 use crate::METADATA_FILE_NAME;
 use crate::ZERO_PAGE;
 use crate::{is_temporary, task_mgr};
+use walreceiver::spawn_connection_manager_task;
 
 pub(super) use self::eviction_task::EvictionTaskTenantState;
 use self::eviction_task::EvictionTaskTimelineState;
-use self::walreceiver::{WalReceiver, WalReceiverConf};
 
 use super::layer_map::BatchedUpdates;
 use super::remote_timeline_client::index::IndexPart;
@@ -145,7 +144,7 @@ pub struct Timeline {
     // 'last_record_lsn.load().prev'. It's used to set the xl_prev pointer of the
     // first WAL record when the node is started up. But here, we just
     // keep track of it.
-    last_record_lsn: SeqWait<RecordLsn, Lsn, ()>,
+    last_record_lsn: SeqWait<RecordLsn, Lsn>,
 
     // All WAL records have been processed and stored durably on files on
     // local disk, up to this LSN. On crash and restart, we need to re-process
@@ -215,7 +214,6 @@ pub struct Timeline {
     /// or None if WAL receiver has not received anything for this timeline
     /// yet.
     pub last_received_wal: Mutex<Option<WalReceiverInfo>>,
-    pub walreceiver: WalReceiver,
 
     /// Relation size cache
     pub rel_size_cache: RwLock<HashMap<RelTag, (Lsn, BlockNumber)>>,
@@ -868,18 +866,10 @@ impl Timeline {
         Ok(())
     }
 
-    pub fn activate(self: &Arc<Self>, ctx: &RequestContext) -> anyhow::Result<()> {
-        if is_broker_client_initialized() {
-            self.launch_wal_receiver(ctx, get_broker_client().clone())?;
-        } else if cfg!(test) {
-            info!("not launching WAL receiver because broker client hasn't been initialized");
-        } else {
-            anyhow::bail!("broker client not initialized");
-        }
-
+    pub fn activate(self: &Arc<Self>) {
         self.set_state(TimelineState::Active);
+        self.launch_wal_receiver();
         self.launch_eviction_task();
-        Ok(())
     }
 
     pub fn set_state(&self, new_state: TimelineState) {
@@ -1230,31 +1220,7 @@ impl Timeline {
         let (layer_flush_start_tx, _) = tokio::sync::watch::channel(0);
         let (layer_flush_done_tx, _) = tokio::sync::watch::channel((0, Ok(())));
 
-        let tenant_conf_guard = tenant_conf.read().unwrap();
-        let wal_connect_timeout = tenant_conf_guard
-            .walreceiver_connect_timeout
-            .unwrap_or(conf.default_tenant_conf.walreceiver_connect_timeout);
-        let lagging_wal_timeout = tenant_conf_guard
-            .lagging_wal_timeout
-            .unwrap_or(conf.default_tenant_conf.lagging_wal_timeout);
-        let max_lsn_wal_lag = tenant_conf_guard
-            .max_lsn_wal_lag
-            .unwrap_or(conf.default_tenant_conf.max_lsn_wal_lag);
-        drop(tenant_conf_guard);
-
         Arc::new_cyclic(|myself| {
-            let walreceiver = WalReceiver::new(
-                TenantTimelineId::new(tenant_id, timeline_id),
-                Weak::clone(myself),
-                WalReceiverConf {
-                    wal_connect_timeout,
-                    lagging_wal_timeout,
-                    max_lsn_wal_lag,
-                    auth_token: crate::config::SAFEKEEPER_AUTH_TOKEN.get().cloned(),
-                    availability_zone: conf.availability_zone.clone(),
-                },
-            );
-
             let mut result = Timeline {
                 conf,
                 tenant_conf,
@@ -1265,18 +1231,14 @@ impl Timeline {
                 layers: RwLock::new(LayerMap::default()),
 
                 walredo_mgr,
-                walreceiver,
 
                 remote_client: remote_client.map(Arc::new),
 
                 // initialize in-memory 'last_record_lsn' from 'disk_consistent_lsn'.
-                last_record_lsn: SeqWait::new(
-                    RecordLsn {
-                        last: disk_consistent_lsn,
-                        prev: metadata.prev_record_lsn().unwrap_or(Lsn(0)),
-                    },
-                    (),
-                ),
+                last_record_lsn: SeqWait::new(RecordLsn {
+                    last: disk_consistent_lsn,
+                    prev: metadata.prev_record_lsn().unwrap_or(Lsn(0)),
+                }),
                 disk_consistent_lsn: AtomicLsn::new(disk_consistent_lsn.0),
 
                 last_freeze_at: AtomicLsn::new(disk_consistent_lsn.0),
@@ -1388,17 +1350,44 @@ impl Timeline {
         *flush_loop_state = FlushLoopState::Running;
     }
 
-    pub(super) fn launch_wal_receiver(
-        &self,
-        ctx: &RequestContext,
-        broker_client: BrokerClientChannel,
-    ) -> anyhow::Result<()> {
+    pub(super) fn launch_wal_receiver(self: &Arc<Self>) {
+        if !is_broker_client_initialized() {
+            if cfg!(test) {
+                info!("not launching WAL receiver because broker client hasn't been initialized");
+                return;
+            } else {
+                panic!("broker client not initialized");
+            }
+        }
+
         info!(
             "launching WAL receiver for timeline {} of tenant {}",
             self.timeline_id, self.tenant_id
         );
-        self.walreceiver.start(ctx, broker_client)?;
-        Ok(())
+        let tenant_conf_guard = self.tenant_conf.read().unwrap();
+        let lagging_wal_timeout = tenant_conf_guard
+            .lagging_wal_timeout
+            .unwrap_or(self.conf.default_tenant_conf.lagging_wal_timeout);
+        let walreceiver_connect_timeout = tenant_conf_guard
+            .walreceiver_connect_timeout
+            .unwrap_or(self.conf.default_tenant_conf.walreceiver_connect_timeout);
+        let max_lsn_wal_lag = tenant_conf_guard
+            .max_lsn_wal_lag
+            .unwrap_or(self.conf.default_tenant_conf.max_lsn_wal_lag);
+        drop(tenant_conf_guard);
+        let self_clone = Arc::clone(self);
+        let background_ctx =
+            // XXX: this is a detached_child. Plumb through the ctx from call sites.
+            RequestContext::todo_child(TaskKind::WalReceiverManager, DownloadBehavior::Error);
+        spawn_connection_manager_task(
+            self_clone,
+            walreceiver_connect_timeout,
+            lagging_wal_timeout,
+            max_lsn_wal_lag,
+            crate::config::SAFEKEEPER_AUTH_TOKEN.get().cloned(),
+            self.conf.availability_zone.clone(),
+            background_ctx,
+        );
     }
 
     ///
@@ -1449,7 +1438,7 @@ impl Timeline {
 
                 trace!("found layer {}", layer.path().display());
                 total_physical_size += file_size;
-                updates.insert_historic(Arc::new(layer))?;
+                updates.insert_historic(Arc::new(layer));
                 num_layers += 1;
             } else if let Some(deltafilename) = DeltaFileName::parse_str(&fname) {
                 // Create a DeltaLayer struct for each delta file.
@@ -1481,7 +1470,7 @@ impl Timeline {
 
                 trace!("found layer {}", layer.path().display());
                 total_physical_size += file_size;
-                updates.insert_historic(Arc::new(layer))?;
+                updates.insert_historic(Arc::new(layer));
                 num_layers += 1;
             } else if fname == METADATA_FILE_NAME || fname.ends_with(".old") {
                 // ignore these
@@ -1555,7 +1544,7 @@ impl Timeline {
             // remote index file?
             // If so, rename_to_backup those files & replace their local layer with
             // a RemoteLayer in the layer map so that we re-download them on-demand.
-            if let Some(local_layer) = &local_layer {
+            if let Some(local_layer) = local_layer {
                 let local_layer_path = local_layer
                     .local_path()
                     .expect("caller must ensure that local_layers only contains local layers");
@@ -1580,6 +1569,7 @@ impl Timeline {
                         anyhow::bail!("could not rename file {local_layer_path:?}: {err:?}");
                     } else {
                         self.metrics.resident_physical_size_gauge.sub(local_size);
+                        updates.remove_historic(local_layer);
                         // fall-through to adding the remote layer
                     }
                 } else {
@@ -1615,11 +1605,7 @@ impl Timeline {
                     );
                     let remote_layer = Arc::new(remote_layer);
 
-                    if let Some(local_layer) = &local_layer {
-                        updates.replace_historic(local_layer, remote_layer)?;
-                    } else {
-                        updates.insert_historic(remote_layer)?;
-                    }
+                    updates.insert_historic(remote_layer);
                 }
                 LayerFileName::Delta(deltafilename) => {
                     // Create a RemoteLayer for the delta file.
@@ -1643,11 +1629,7 @@ impl Timeline {
                         LayerAccessStats::for_loading_layer(LayerResidenceStatus::Evicted),
                     );
                     let remote_layer = Arc::new(remote_layer);
-                    if let Some(local_layer) = &local_layer {
-                        updates.replace_historic(local_layer, remote_layer)?;
-                    } else {
-                        updates.insert_historic(remote_layer)?;
-                    }
+                    updates.insert_historic(remote_layer);
                 }
             }
         }
@@ -2423,7 +2405,7 @@ impl Timeline {
         assert!(new_lsn.is_aligned());
 
         self.metrics.last_record_gauge.set(new_lsn.0 as i64);
-        self.last_record_lsn.advance(new_lsn, None);
+        self.last_record_lsn.advance(new_lsn);
     }
 
     fn freeze_inmem_layer(&self, write_lock_held: bool) {
@@ -2694,7 +2676,7 @@ impl Timeline {
             .write()
             .unwrap()
             .batch_update()
-            .insert_historic(Arc::new(new_delta))?;
+            .insert_historic(Arc::new(new_delta));
 
         // update the timeline's physical size
         let sz = new_delta_path.metadata()?.len();
@@ -2899,7 +2881,7 @@ impl Timeline {
             self.metrics
                 .resident_physical_size_gauge
                 .add(metadata.len());
-            updates.insert_historic(Arc::new(l))?;
+            updates.insert_historic(Arc::new(l));
         }
         updates.flush();
         drop(layers);
@@ -3332,7 +3314,7 @@ impl Timeline {
 
             new_layer_paths.insert(new_delta_path, LayerFileMetadata::new(metadata.len()));
             let x: Arc<dyn PersistentLayer + 'static> = Arc::new(l);
-            updates.insert_historic(x)?;
+            updates.insert_historic(x);
         }
 
         // Now that we have reshuffled the data to set of new delta layers, we can

pageserver/src/tenant/timeline/walreceiver.rs

@@ -23,133 +23,14 @@
 mod connection_manager;
 mod walreceiver_connection;
 
-use crate::context::{DownloadBehavior, RequestContext};
-use crate::task_mgr::{self, TaskKind, WALRECEIVER_RUNTIME};
-use crate::tenant::timeline::walreceiver::connection_manager::{
-    connection_manager_loop_step, ConnectionManagerState,
-};
+use crate::task_mgr::WALRECEIVER_RUNTIME;
 
-use anyhow::Context;
 use std::future::Future;
-use std::num::NonZeroU64;
-use std::ops::ControlFlow;
-use std::sync::atomic::{self, AtomicBool};
-use std::sync::{Arc, Weak};
-use std::time::Duration;
-use storage_broker::BrokerClientChannel;
-use tokio::select;
 use tokio::sync::watch;
 use tokio_util::sync::CancellationToken;
 use tracing::*;
 
-use utils::id::TenantTimelineId;
-
-use super::Timeline;
-
-#[derive(Clone)]
-pub struct WalReceiverConf {
-    /// The timeout on the connection to safekeeper for WAL streaming.
-    pub wal_connect_timeout: Duration,
-    /// The timeout to use to determine when the current connection is "stale" and reconnect to the other one.
-    pub lagging_wal_timeout: Duration,
-    /// The Lsn lag to use to determine when the current connection is lagging to much behind and reconnect to the other one.
-    pub max_lsn_wal_lag: NonZeroU64,
-    pub auth_token: Option<Arc<String>>,
-    pub availability_zone: Option<String>,
-}
-
-pub struct WalReceiver {
-    timeline: TenantTimelineId,
-    timeline_ref: Weak<Timeline>,
-    conf: WalReceiverConf,
-    started: AtomicBool,
-}
-
-impl WalReceiver {
-    pub fn new(
-        timeline: TenantTimelineId,
-        timeline_ref: Weak<Timeline>,
-        conf: WalReceiverConf,
-    ) -> Self {
-        Self {
-            timeline,
-            timeline_ref,
-            conf,
-            started: AtomicBool::new(false),
-        }
-    }
-
-    pub fn start(
-        &self,
-        ctx: &RequestContext,
-        mut broker_client: BrokerClientChannel,
-    ) -> anyhow::Result<()> {
-        if self.started.load(atomic::Ordering::Acquire) {
-            anyhow::bail!("Wal receiver is already started");
-        }
-
-        let timeline = self.timeline_ref.upgrade().with_context(|| {
-            format!("walreceiver start on a dropped timeline {}", self.timeline)
-        })?;
-
-        let tenant_id = timeline.tenant_id;
-        let timeline_id = timeline.timeline_id;
-        let walreceiver_ctx =
-            ctx.detached_child(TaskKind::WalReceiverManager, DownloadBehavior::Error);
-
-        let wal_receiver_conf = self.conf.clone();
-        task_mgr::spawn(
-            WALRECEIVER_RUNTIME.handle(),
-            TaskKind::WalReceiverManager,
-            Some(tenant_id),
-            Some(timeline_id),
-            &format!("walreceiver for timeline {tenant_id}/{timeline_id}"),
-            false,
-            async move {
-                info!("WAL receiver manager started, connecting to broker");
-                let mut connection_manager_state = ConnectionManagerState::new(
-                    timeline,
-                    wal_receiver_conf,
-                );
-                loop {
-                    select! {
-                        _ = task_mgr::shutdown_watcher() => {
-                            info!("WAL receiver shutdown requested, shutting down");
-                            connection_manager_state.shutdown().await;
-                            return Ok(());
-                        },
-                        loop_step_result = connection_manager_loop_step(
-                            &mut broker_client,
-                            &mut connection_manager_state,
-                            &walreceiver_ctx,
-                        ) => match loop_step_result {
-                            ControlFlow::Continue(()) => continue,
-                            ControlFlow::Break(()) => {
-                                info!("Connection manager loop ended, shutting down");
-                                connection_manager_state.shutdown().await;
-                                return Ok(());
-                            }
-                        },
-                    }
-                }
-            }.instrument(info_span!(parent: None, "wal_connection_manager", tenant = %tenant_id, timeline = %timeline_id))
-        );
-
-        self.started.store(true, atomic::Ordering::Release);
-
-        Ok(())
-    }
-
-    pub async fn stop(&self) {
-        task_mgr::shutdown_tasks(
-            Some(TaskKind::WalReceiverManager),
-            Some(self.timeline.tenant_id),
-            Some(self.timeline.timeline_id),
-        )
-        .await;
-        self.started.store(false, atomic::Ordering::Release);
-    }
-}
+pub use connection_manager::spawn_connection_manager_task;
 
 /// A handle of an asynchronous task.
 /// The task has a channel that it can use to communicate its lifecycle events in a certain form, see [`TaskEvent`]
@@ -158,26 +39,26 @@ impl WalReceiver {
 /// Note that the communication happens via the `watch` channel, that does not accumulate the events, replacing the old one with the never one on submission.
 /// That may lead to certain events not being observed by the listener.
 #[derive(Debug)]
-struct TaskHandle<E> {
+pub struct TaskHandle<E> {
     join_handle: Option<tokio::task::JoinHandle<anyhow::Result<()>>>,
     events_receiver: watch::Receiver<TaskStateUpdate<E>>,
     cancellation: CancellationToken,
 }
 
-enum TaskEvent<E> {
+pub enum TaskEvent<E> {
     Update(TaskStateUpdate<E>),
     End(anyhow::Result<()>),
 }
 
 #[derive(Debug, Clone)]
-enum TaskStateUpdate<E> {
+pub enum TaskStateUpdate<E> {
     Started,
     Progress(E),
 }
 
 impl<E: Clone> TaskHandle<E> {
     /// Initializes the task, starting it immediately after the creation.
-    fn spawn<Fut>(
+    pub fn spawn<Fut>(
         task: impl FnOnce(watch::Sender<TaskStateUpdate<E>>, CancellationToken) -> Fut + Send + 'static,
     ) -> Self
     where
@@ -250,7 +131,7 @@ impl<E: Clone> TaskHandle<E> {
     }
 
     /// Aborts current task, waiting for it to finish.
-    async fn shutdown(self) {
+    pub async fn shutdown(self) {
         if let Some(jh) = self.join_handle {
             self.cancellation.cancel();
             match jh.await {

pageserver/src/tenant/timeline/walreceiver/connection_manager.rs

@@ -11,9 +11,11 @@
 
 use std::{collections::HashMap, num::NonZeroU64, ops::ControlFlow, sync::Arc, time::Duration};
 
-use super::{TaskStateUpdate, WalReceiverConf};
+use super::TaskStateUpdate;
+use crate::broker_client::get_broker_client;
 use crate::context::{DownloadBehavior, RequestContext};
-use crate::task_mgr::TaskKind;
+use crate::task_mgr::WALRECEIVER_RUNTIME;
+use crate::task_mgr::{self, TaskKind};
 use crate::tenant::Timeline;
 use anyhow::Context;
 use chrono::{NaiveDateTime, Utc};
@@ -36,17 +38,75 @@ use utils::{
 
 use super::{walreceiver_connection::WalConnectionStatus, TaskEvent, TaskHandle};
 
+/// Spawns the loop to take care of the timeline's WAL streaming connection.
+pub fn spawn_connection_manager_task(
+    timeline: Arc<Timeline>,
+    wal_connect_timeout: Duration,
+    lagging_wal_timeout: Duration,
+    max_lsn_wal_lag: NonZeroU64,
+    auth_token: Option<Arc<String>>,
+    availability_zone: Option<String>,
+    ctx: RequestContext,
+) {
+    let mut broker_client = get_broker_client().clone();
+
+    let tenant_id = timeline.tenant_id;
+    let timeline_id = timeline.timeline_id;
+
+    task_mgr::spawn(
+        WALRECEIVER_RUNTIME.handle(),
+        TaskKind::WalReceiverManager,
+        Some(tenant_id),
+        Some(timeline_id),
+        &format!("walreceiver for timeline {tenant_id}/{timeline_id}"),
+        false,
+        async move {
+            info!("WAL receiver manager started, connecting to broker");
+            let mut walreceiver_state = WalreceiverState::new(
+                timeline,
+                wal_connect_timeout,
+                lagging_wal_timeout,
+                max_lsn_wal_lag,
+                auth_token,
+                availability_zone,
+            );
+            loop {
+                select! {
+                    _ = task_mgr::shutdown_watcher() => {
+                        info!("WAL receiver shutdown requested, shutting down");
+                        walreceiver_state.shutdown().await;
+                        return Ok(());
+                    },
+                    loop_step_result = connection_manager_loop_step(
+                        &mut broker_client,
+                        &mut walreceiver_state,
+                        &ctx,
+                    ) => match loop_step_result {
+                        ControlFlow::Continue(()) => continue,
+                        ControlFlow::Break(()) => {
+                            info!("Connection manager loop ended, shutting down");
+                            walreceiver_state.shutdown().await;
+                            return Ok(());
+                        }
+                    },
+                }
+            }
+        }
+        .instrument(
+            info_span!(parent: None, "wal_connection_manager", tenant = %tenant_id, timeline = %timeline_id),
+        ),
+    );
+}
+
 /// Attempts to subscribe for timeline updates, pushed by safekeepers into the broker.
 /// Based on the updates, desides whether to start, keep or stop a WAL receiver task.
 /// If storage broker subscription is cancelled, exits.
-pub(super) async fn connection_manager_loop_step(
+async fn connection_manager_loop_step(
     broker_client: &mut BrokerClientChannel,
-    connection_manager_state: &mut ConnectionManagerState,
+    walreceiver_state: &mut WalreceiverState,
     ctx: &RequestContext,
 ) -> ControlFlow<(), ()> {
-    let mut timeline_state_updates = connection_manager_state
-        .timeline
-        .subscribe_for_state_updates();
+    let mut timeline_state_updates = walreceiver_state.timeline.subscribe_for_state_updates();
 
     match wait_for_active_timeline(&mut timeline_state_updates).await {
         ControlFlow::Continue(()) => {}
@@ -57,8 +117,8 @@ pub(super) async fn connection_manager_loop_step(
     }
 
     let id = TenantTimelineId {
-        tenant_id: connection_manager_state.timeline.tenant_id,
-        timeline_id: connection_manager_state.timeline.timeline_id,
+        tenant_id: walreceiver_state.timeline.tenant_id,
+        timeline_id: walreceiver_state.timeline.timeline_id,
     };
 
     // Subscribe to the broker updates. Stream shares underlying TCP connection
@@ -68,7 +128,7 @@ pub(super) async fn connection_manager_loop_step(
     info!("Subscribed for broker timeline updates");
 
     loop {
-        let time_until_next_retry = connection_manager_state.time_until_next_retry();
+        let time_until_next_retry = walreceiver_state.time_until_next_retry();
 
         // These things are happening concurrently:
         //
@@ -81,12 +141,12 @@ pub(super) async fn connection_manager_loop_step(
         //  - timeline state changes to something that does not allow walreceiver to run concurrently
         select! {
             Some(wal_connection_update) = async {
-                match connection_manager_state.wal_connection.as_mut() {
+                match walreceiver_state.wal_connection.as_mut() {
                     Some(wal_connection) => Some(wal_connection.connection_task.next_task_event().await),
                     None => None,
                 }
             } => {
-                let wal_connection = connection_manager_state.wal_connection.as_mut()
+                let wal_connection = walreceiver_state.wal_connection.as_mut()
                     .expect("Should have a connection, as checked by the corresponding select! guard");
                 match wal_connection_update {
                     TaskEvent::Update(TaskStateUpdate::Started) => {},
@@ -96,7 +156,7 @@ pub(super) async fn connection_manager_loop_step(
                             // from this safekeeper. This is good enough to clean unsuccessful
                             // retries history and allow reconnecting to this safekeeper without
                             // sleeping for a long time.
-                            connection_manager_state.wal_connection_retries.remove(&wal_connection.sk_id);
+                            walreceiver_state.wal_connection_retries.remove(&wal_connection.sk_id);
                         }
                         wal_connection.status = new_status;
                     }
@@ -105,7 +165,7 @@ pub(super) async fn connection_manager_loop_step(
                             Ok(()) => debug!("WAL receiving task finished"),
                             Err(e) => error!("wal receiver task finished with an error: {e:?}"),
                         }
-                        connection_manager_state.drop_old_connection(false).await;
+                        walreceiver_state.drop_old_connection(false).await;
                     },
                 }
             },
@@ -113,7 +173,7 @@ pub(super) async fn connection_manager_loop_step(
             // Got a new update from the broker
             broker_update = broker_subscription.message() => {
                 match broker_update {
-                    Ok(Some(broker_update)) => connection_manager_state.register_timeline_update(broker_update),
+                    Ok(Some(broker_update)) => walreceiver_state.register_timeline_update(broker_update),
                     Err(e) => {
                         error!("broker subscription failed: {e}");
                         return ControlFlow::Continue(());
@@ -127,12 +187,12 @@ pub(super) async fn connection_manager_loop_step(
 
             new_event = async {
                 loop {
-                    if connection_manager_state.timeline.current_state() == TimelineState::Loading {
+                    if walreceiver_state.timeline.current_state() == TimelineState::Loading {
                         warn!("wal connection manager should only be launched after timeline has become active");
                     }
                     match timeline_state_updates.changed().await {
                         Ok(()) => {
-                            let new_state = connection_manager_state.timeline.current_state();
+                            let new_state = walreceiver_state.timeline.current_state();
                             match new_state {
                                 // we're already active as walreceiver, no need to reactivate
                                 TimelineState::Active => continue,
@@ -174,9 +234,9 @@ pub(super) async fn connection_manager_loop_step(
             } => debug!("Waking up for the next retry after waiting for {time_until_next_retry:?}"),
         }
 
-        if let Some(new_candidate) = connection_manager_state.next_connection_candidate() {
+        if let Some(new_candidate) = walreceiver_state.next_connection_candidate() {
             info!("Switching to new connection candidate: {new_candidate:?}");
-            connection_manager_state
+            walreceiver_state
                 .change_connection(new_candidate, ctx)
                 .await
         }
@@ -254,17 +314,25 @@ const WALCONNECTION_RETRY_MAX_BACKOFF_SECONDS: f64 = 15.0;
 const WALCONNECTION_RETRY_BACKOFF_MULTIPLIER: f64 = 1.5;
 
 /// All data that's needed to run endless broker loop and keep the WAL streaming connection alive, if possible.
-pub(super) struct ConnectionManagerState {
+struct WalreceiverState {
     id: TenantTimelineId,
+
     /// Use pageserver data about the timeline to filter out some of the safekeepers.
     timeline: Arc<Timeline>,
-    conf: WalReceiverConf,
+    /// The timeout on the connection to safekeeper for WAL streaming.
+    wal_connect_timeout: Duration,
+    /// The timeout to use to determine when the current connection is "stale" and reconnect to the other one.
+    lagging_wal_timeout: Duration,
+    /// The Lsn lag to use to determine when the current connection is lagging to much behind and reconnect to the other one.
+    max_lsn_wal_lag: NonZeroU64,
     /// Current connection to safekeeper for WAL streaming.
     wal_connection: Option<WalConnection>,
     /// Info about retries and unsuccessful attempts to connect to safekeepers.
     wal_connection_retries: HashMap<NodeId, RetryInfo>,
     /// Data about all timelines, available for connection, fetched from storage broker, grouped by their corresponding safekeeper node id.
     wal_stream_candidates: HashMap<NodeId, BrokerSkTimeline>,
+    auth_token: Option<Arc<String>>,
+    availability_zone: Option<String>,
 }
 
 /// Current connection data.
@@ -307,8 +375,15 @@ struct BrokerSkTimeline {
     latest_update: NaiveDateTime,
 }
 
-impl ConnectionManagerState {
-    pub(super) fn new(timeline: Arc<Timeline>, conf: WalReceiverConf) -> Self {
+impl WalreceiverState {
+    fn new(
+        timeline: Arc<Timeline>,
+        wal_connect_timeout: Duration,
+        lagging_wal_timeout: Duration,
+        max_lsn_wal_lag: NonZeroU64,
+        auth_token: Option<Arc<String>>,
+        availability_zone: Option<String>,
+    ) -> Self {
         let id = TenantTimelineId {
             tenant_id: timeline.tenant_id,
             timeline_id: timeline.timeline_id,
@@ -316,10 +391,14 @@ impl ConnectionManagerState {
         Self {
             id,
             timeline,
-            conf,
+            wal_connect_timeout,
+            lagging_wal_timeout,
+            max_lsn_wal_lag,
             wal_connection: None,
             wal_stream_candidates: HashMap::new(),
             wal_connection_retries: HashMap::new(),
+            auth_token,
+            availability_zone,
         }
     }
 
@@ -328,7 +407,7 @@ impl ConnectionManagerState {
         self.drop_old_connection(true).await;
 
         let id = self.id;
-        let connect_timeout = self.conf.wal_connect_timeout;
+        let connect_timeout = self.wal_connect_timeout;
         let timeline = Arc::clone(&self.timeline);
         let ctx = ctx.detached_child(
             TaskKind::WalReceiverConnectionHandler,
@@ -484,7 +563,7 @@ impl ConnectionManagerState {
                     (now - existing_wal_connection.status.latest_connection_update).to_std()
                 {
                     // Drop connection if we haven't received keepalive message for a while.
-                    if latest_interaciton > self.conf.wal_connect_timeout {
+                    if latest_interaciton > self.wal_connect_timeout {
                         return Some(NewWalConnectionCandidate {
                             safekeeper_id: new_sk_id,
                             wal_source_connconf: new_wal_source_connconf,
@@ -494,7 +573,7 @@ impl ConnectionManagerState {
                                     existing_wal_connection.status.latest_connection_update,
                                 ),
                                 check_time: now,
-                                threshold: self.conf.wal_connect_timeout,
+                                threshold: self.wal_connect_timeout,
                             },
                         });
                     }
@@ -510,7 +589,7 @@ impl ConnectionManagerState {
                     // Check if the new candidate has much more WAL than the current one.
                     match new_commit_lsn.0.checked_sub(current_commit_lsn.0) {
                         Some(new_sk_lsn_advantage) => {
-                            if new_sk_lsn_advantage >= self.conf.max_lsn_wal_lag.get() {
+                            if new_sk_lsn_advantage >= self.max_lsn_wal_lag.get() {
                                 return Some(NewWalConnectionCandidate {
                                     safekeeper_id: new_sk_id,
                                     wal_source_connconf: new_wal_source_connconf,
@@ -518,16 +597,16 @@ impl ConnectionManagerState {
                                     reason: ReconnectReason::LaggingWal {
                                         current_commit_lsn,
                                         new_commit_lsn,
-                                        threshold: self.conf.max_lsn_wal_lag,
+                                        threshold: self.max_lsn_wal_lag,
                                     },
                                 });
                             }
                             // If we have a candidate with the same commit_lsn as the current one, which is in the same AZ as pageserver,
                             // and the current one is not, switch to the new one.
-                            if self.conf.availability_zone.is_some()
+                            if self.availability_zone.is_some()
                                 && existing_wal_connection.availability_zone
-                                    != self.conf.availability_zone
-                                && self.conf.availability_zone == new_availability_zone
+                                    != self.availability_zone
+                                && self.availability_zone == new_availability_zone
                             {
                                 return Some(NewWalConnectionCandidate {
                                     safekeeper_id: new_sk_id,
@@ -598,7 +677,7 @@ impl ConnectionManagerState {
                 if let Some(waiting_for_new_lsn_since) = waiting_for_new_lsn_since {
                     if let Ok(waiting_for_new_wal) = (now - waiting_for_new_lsn_since).to_std() {
                         if candidate_commit_lsn > current_commit_lsn
-                            && waiting_for_new_wal > self.conf.lagging_wal_timeout
+                            && waiting_for_new_wal > self.lagging_wal_timeout
                         {
                             return Some(NewWalConnectionCandidate {
                                 safekeeper_id: new_sk_id,
@@ -612,7 +691,7 @@ impl ConnectionManagerState {
                                         existing_wal_connection.status.latest_wal_update,
                                     ),
                                     check_time: now,
-                                    threshold: self.conf.lagging_wal_timeout,
+                                    threshold: self.lagging_wal_timeout,
                                 },
                             });
                         }
@@ -678,11 +757,11 @@ impl ConnectionManagerState {
                 match wal_stream_connection_config(
                     self.id,
                     info.safekeeper_connstr.as_ref(),
-                    match &self.conf.auth_token {
+                    match &self.auth_token {
                         None => None,
                         Some(x) => Some(x),
                     },
-                    self.conf.availability_zone.as_deref(),
+                    self.availability_zone.as_deref(),
                 ) {
                     Ok(connstr) => Some((*sk_id, info, connstr)),
                     Err(e) => {
@@ -696,7 +775,7 @@ impl ConnectionManagerState {
     /// Remove candidates which haven't sent broker updates for a while.
     fn cleanup_old_candidates(&mut self) {
         let mut node_ids_to_remove = Vec::with_capacity(self.wal_stream_candidates.len());
-        let lagging_wal_timeout = self.conf.lagging_wal_timeout;
+        let lagging_wal_timeout = self.lagging_wal_timeout;
 
         self.wal_stream_candidates.retain(|node_id, broker_info| {
             if let Ok(time_since_latest_broker_update) =
@@ -720,7 +799,7 @@ impl ConnectionManagerState {
         }
     }
 
-    pub(super) async fn shutdown(mut self) {
+    async fn shutdown(mut self) {
         if let Some(wal_connection) = self.wal_connection.take() {
             wal_connection.connection_task.shutdown().await;
         }
@@ -824,7 +903,7 @@ mod tests {
         let mut state = dummy_state(&harness).await;
         let now = Utc::now().naive_utc();
 
-        let lagging_wal_timeout = chrono::Duration::from_std(state.conf.lagging_wal_timeout)?;
+        let lagging_wal_timeout = chrono::Duration::from_std(state.lagging_wal_timeout)?;
         let delay_over_threshold = now - lagging_wal_timeout - lagging_wal_timeout;
 
         state.wal_connection = None;
@@ -835,7 +914,7 @@ mod tests {
             (
                 NodeId(3),
                 dummy_broker_sk_timeline(
-                    1 + state.conf.max_lsn_wal_lag.get(),
+                    1 + state.max_lsn_wal_lag.get(),
                     "delay_over_threshold",
                     delay_over_threshold,
                 ),
@@ -869,7 +948,7 @@ mod tests {
             streaming_lsn: Some(Lsn(current_lsn)),
         };
 
-        state.conf.max_lsn_wal_lag = NonZeroU64::new(100).unwrap();
+        state.max_lsn_wal_lag = NonZeroU64::new(100).unwrap();
         state.wal_connection = Some(WalConnection {
             started_at: now,
             sk_id: connected_sk_id,
@@ -887,7 +966,7 @@ mod tests {
             (
                 connected_sk_id,
                 dummy_broker_sk_timeline(
-                    current_lsn + state.conf.max_lsn_wal_lag.get() * 2,
+                    current_lsn + state.max_lsn_wal_lag.get() * 2,
                     DUMMY_SAFEKEEPER_HOST,
                     now,
                 ),
@@ -899,7 +978,7 @@ mod tests {
             (
                 NodeId(2),
                 dummy_broker_sk_timeline(
-                    current_lsn + state.conf.max_lsn_wal_lag.get() / 2,
+                    current_lsn + state.max_lsn_wal_lag.get() / 2,
                     "not_enough_advanced_lsn",
                     now,
                 ),
@@ -924,11 +1003,7 @@ mod tests {
         state.wal_connection = None;
         state.wal_stream_candidates = HashMap::from([(
             NodeId(0),
-            dummy_broker_sk_timeline(
-                1 + state.conf.max_lsn_wal_lag.get(),
-                DUMMY_SAFEKEEPER_HOST,
-                now,
-            ),
+            dummy_broker_sk_timeline(1 + state.max_lsn_wal_lag.get(), DUMMY_SAFEKEEPER_HOST, now),
         )]);
 
         let only_candidate = state
@@ -1026,7 +1101,7 @@ mod tests {
         let now = Utc::now().naive_utc();
 
         let connected_sk_id = NodeId(0);
-        let new_lsn = Lsn(current_lsn.0 + state.conf.max_lsn_wal_lag.get() + 1);
+        let new_lsn = Lsn(current_lsn.0 + state.max_lsn_wal_lag.get() + 1);
 
         let connection_status = WalConnectionStatus {
             is_connected: true,
@@ -1071,7 +1146,7 @@ mod tests {
             ReconnectReason::LaggingWal {
                 current_commit_lsn: current_lsn,
                 new_commit_lsn: new_lsn,
-                threshold: state.conf.max_lsn_wal_lag
+                threshold: state.max_lsn_wal_lag
             },
             "Should select bigger WAL safekeeper if it starts to lag enough"
         );
@@ -1090,7 +1165,7 @@ mod tests {
         let current_lsn = Lsn(100_000).align();
         let now = Utc::now().naive_utc();
 
-        let wal_connect_timeout = chrono::Duration::from_std(state.conf.wal_connect_timeout)?;
+        let wal_connect_timeout = chrono::Duration::from_std(state.wal_connect_timeout)?;
         let time_over_threshold =
             Utc::now().naive_utc() - wal_connect_timeout - wal_connect_timeout;
 
@@ -1133,7 +1208,7 @@ mod tests {
                 ..
             } => {
                 assert_eq!(last_keep_alive, Some(time_over_threshold));
-                assert_eq!(threshold, state.conf.lagging_wal_timeout);
+                assert_eq!(threshold, state.lagging_wal_timeout);
             }
             unexpected => panic!("Unexpected reason: {unexpected:?}"),
         }
@@ -1153,7 +1228,7 @@ mod tests {
         let new_lsn = Lsn(100_100).align();
         let now = Utc::now().naive_utc();
 
-        let lagging_wal_timeout = chrono::Duration::from_std(state.conf.lagging_wal_timeout)?;
+        let lagging_wal_timeout = chrono::Duration::from_std(state.lagging_wal_timeout)?;
         let time_over_threshold =
             Utc::now().naive_utc() - lagging_wal_timeout - lagging_wal_timeout;
 
@@ -1200,7 +1275,7 @@ mod tests {
                 assert_eq!(current_commit_lsn, current_lsn);
                 assert_eq!(candidate_commit_lsn, new_lsn);
                 assert_eq!(last_wal_interaction, Some(time_over_threshold));
-                assert_eq!(threshold, state.conf.lagging_wal_timeout);
+                assert_eq!(threshold, state.lagging_wal_timeout);
             }
             unexpected => panic!("Unexpected reason: {unexpected:?}"),
         }
@@ -1214,29 +1289,27 @@ mod tests {
 
     const DUMMY_SAFEKEEPER_HOST: &str = "safekeeper_connstr";
 
-    async fn dummy_state(harness: &TenantHarness<'_>) -> ConnectionManagerState {
+    async fn dummy_state(harness: &TenantHarness<'_>) -> WalreceiverState {
         let (tenant, ctx) = harness.load().await;
         let timeline = tenant
             .create_empty_timeline(TIMELINE_ID, Lsn(0), crate::DEFAULT_PG_VERSION, &ctx)
             .expect("Failed to create an empty timeline for dummy wal connection manager");
         let timeline = timeline.initialize(&ctx).unwrap();
 
-        ConnectionManagerState {
+        WalreceiverState {
             id: TenantTimelineId {
                 tenant_id: harness.tenant_id,
                 timeline_id: TIMELINE_ID,
             },
             timeline,
-            conf: WalReceiverConf {
-                wal_connect_timeout: Duration::from_secs(1),
-                lagging_wal_timeout: Duration::from_secs(1),
-                max_lsn_wal_lag: NonZeroU64::new(1024 * 1024).unwrap(),
-                auth_token: None,
-                availability_zone: None,
-            },
+            wal_connect_timeout: Duration::from_secs(1),
+            lagging_wal_timeout: Duration::from_secs(1),
+            max_lsn_wal_lag: NonZeroU64::new(1024 * 1024).unwrap(),
             wal_connection: None,
             wal_stream_candidates: HashMap::new(),
             wal_connection_retries: HashMap::new(),
+            auth_token: None,
+            availability_zone: None,
         }
     }
 
@@ -1248,7 +1321,7 @@ mod tests {
 
         let harness = TenantHarness::create("switch_to_same_availability_zone")?;
         let mut state = dummy_state(&harness).await;
-        state.conf.availability_zone = test_az.clone();
+        state.availability_zone = test_az.clone();
         let current_lsn = Lsn(100_000).align();
         let now = Utc::now().naive_utc();
 

pageserver/src/tenant/timeline/walreceiver/walreceiver_connection.rs

@@ -42,7 +42,7 @@ use utils::lsn::Lsn;
 
 /// Status of the connection.
 #[derive(Debug, Clone, Copy)]
-pub(super) struct WalConnectionStatus {
+pub struct WalConnectionStatus {
     /// If we were able to initiate a postgres connection, this means that safekeeper process is at least running.
     pub is_connected: bool,
     /// Defines a healthy connection as one on which pageserver received WAL from safekeeper
@@ -60,7 +60,7 @@ pub(super) struct WalConnectionStatus {
 
 /// Open a connection to the given safekeeper and receive WAL, sending back progress
 /// messages as we go.
-pub(super) async fn handle_walreceiver_connection(
+pub async fn handle_walreceiver_connection(
     timeline: Arc<Timeline>,
     wal_source_connconf: PgConnectionConfig,
     events_sender: watch::Sender<TaskStateUpdate<WalConnectionStatus>>,

pgxn/neon_walredo/seccomp.c

@@ -9,14 +9,6 @@
  * To prevent this, it has been decided to limit possible interactions
  * with the outside world using the Secure Computing BPF mode.
  *
- * This code is intended to support both x86_64 and aarch64. The latter
- * doesn't implement some syscalls like open and select. We allow both
- * select (absent on aarch64) and pselect6 (present on both architectures)
- * We call select(2) through libc, and the libc wrapper calls select or pselect6
- * depending on the architecture. You can check which syscalls are present on
- * different architectures with the `scmp_sys_resolver` tool from the
- * seccomp package.
- *
  * We use this mode to disable all syscalls not in the allowlist. This
  * approach has its pros & cons:
  *
@@ -81,6 +73,8 @@
  *    I suspect that certain libc functions might involve slightly
  *    different syscalls, e.g. select/pselect6/pselect6_time64/whatever.
  *
+ *  - Test on any arch other than amd64 to see if it works there.
+ *
  *-------------------------------------------------------------------------
  */
 
@@ -128,10 +122,9 @@ seccomp_load_rules(PgSeccompRule *rules, int count)
 
 	/*
 	 * First, check that open of a well-known file works.
-	 * XXX: We use raw syscall() to call the very openat() which is
-	 * present both on x86_64 and on aarch64.
+	 * XXX: We use raw syscall() to call the very open().
 	 */
-	fd = syscall(SCMP_SYS(openat), AT_FDCWD, "/dev/null", O_RDONLY, 0);
+	fd = syscall(SCMP_SYS(open), "/dev/null", O_RDONLY, 0);
 	if (seccomp_test_sighandler_done)
 		ereport(FATAL,
 				(errcode(ERRCODE_SYSTEM_ERROR),
@@ -142,15 +135,15 @@ seccomp_load_rules(PgSeccompRule *rules, int count)
 				 errmsg("seccomp: could not open /dev/null for seccomp testing: %m")));
 	close((int) fd);
 
-	/* Set a trap on openat() to test seccomp bpf */
-	rule = PG_SCMP(openat, SCMP_ACT_TRAP);
+	/* Set a trap on open() to test seccomp bpf */
+	rule = PG_SCMP(open, SCMP_ACT_TRAP);
 	if (do_seccomp_load_rules(&rule, 1, SCMP_ACT_ALLOW) != 0)
 		ereport(FATAL,
 				(errcode(ERRCODE_SYSTEM_ERROR),
 				 errmsg("seccomp: could not load test trap")));
 
-	/* Finally, check that openat() now raises SIGSYS */
-	(void) syscall(SCMP_SYS(openat), AT_FDCWD, "/dev/null", O_RDONLY, 0);
+	/* Finally, check that open() now raises SIGSYS */
+	(void) syscall(SCMP_SYS(open), "/dev/null", O_RDONLY, 0);
 	if (!seccomp_test_sighandler_done)
 		ereport(FATAL,
 				(errcode(ERRCODE_SYSTEM_ERROR),
@@ -231,7 +224,7 @@ seccomp_test_sighandler(int signum, siginfo_t *info, void *cxt pg_attribute_unus
 		die(1, DIE_PREFIX "bad signal number\n");
 
 	/* TODO: maybe somehow extract the hardcoded syscall number */
-	if (info->si_syscall != SCMP_SYS(openat))
+	if (info->si_syscall != SCMP_SYS(open))
 		die(1, DIE_PREFIX "bad syscall number\n");
 
 #undef DIE_PREFIX

proxy/Cargo.toml

@@ -64,7 +64,6 @@ webpki-roots.workspace = true
 x509-parser.workspace = true
 
 workspace_hack.workspace = true
-tokio-util.workspace = true
 
 [dev-dependencies]
 rcgen.workspace = true

proxy/src/auth/backend/hacks.rs

@@ -53,7 +53,7 @@ pub async fn password_hack(
         .await?;
 
     info!(project = &payload.project, "received missing parameter");
-    creds.project = Some(payload.project);
+    creds.project = Some(payload.project.into());
 
     let mut node = api.wake_compute(extra, creds).await?;
     node.config.password(payload.password);

proxy/src/auth/credentials.rs

@@ -2,7 +2,7 @@
 
 use crate::error::UserFacingError;
 use pq_proto::StartupMessageParams;
-use std::collections::HashSet;
+use std::borrow::Cow;
 use thiserror::Error;
 use tracing::info;
 
@@ -19,10 +19,11 @@ pub enum ClientCredsParseError {
     InconsistentProjectNames { domain: String, option: String },
 
     #[error(
-        "Common name inferred from SNI ('{}') is not known",
-        .cn,
+        "SNI ('{}') inconsistently formatted with respect to common name ('{}'). \
+         SNI should be formatted as '<project-name>.{}'.",
+        .sni, .cn, .cn,
     )]
-    UnknownCommonName { cn: String },
+    InconsistentSni { sni: String, cn: String },
 
     #[error("Project name ('{0}') must contain only alphanumeric characters and hyphen.")]
     MalformedProjectName(String),
@@ -36,7 +37,7 @@ impl UserFacingError for ClientCredsParseError {}
 pub struct ClientCredentials<'a> {
     pub user: &'a str,
     // TODO: this is a severe misnomer! We should think of a new name ASAP.
-    pub project: Option<String>,
+    pub project: Option<Cow<'a, str>>,
 }
 
 impl ClientCredentials<'_> {
@@ -50,7 +51,7 @@ impl<'a> ClientCredentials<'a> {
     pub fn parse(
         params: &'a StartupMessageParams,
         sni: Option<&str>,
-        common_names: Option<HashSet<String>>,
+        common_name: Option<&str>,
     ) -> Result<Self, ClientCredsParseError> {
         use ClientCredsParseError::*;
 
@@ -59,43 +60,37 @@ impl<'a> ClientCredentials<'a> {
         let user = get_param("user")?;
 
         // Project name might be passed via PG's command-line options.
-        let project_option = params
-            .options_raw()
-            .and_then(|mut options| options.find_map(|opt| opt.strip_prefix("project=")))
-            .map(|name| name.to_string());
+        let project_option = params.options_raw().and_then(|mut options| {
+            options
+                .find_map(|opt| opt.strip_prefix("project="))
+                .map(Cow::Borrowed)
+        });
 
-        let project_from_domain = if let Some(sni_str) = sni {
-            if let Some(cn) = common_names {
-                let common_name_from_sni = sni_str.split_once('.').map(|(_, domain)| domain);
-
-                let project = common_name_from_sni
-                    .and_then(|domain| {
-                        if cn.contains(domain) {
-                            subdomain_from_sni(sni_str, domain)
-                        } else {
-                            None
-                        }
+        // Alternative project name is in fact a subdomain from SNI.
+        // NOTE: we do not consider SNI if `common_name` is missing.
+        let project_domain = sni
+            .zip(common_name)
+            .map(|(sni, cn)| {
+                subdomain_from_sni(sni, cn)
+                    .ok_or_else(|| InconsistentSni {
+                        sni: sni.into(),
+                        cn: cn.into(),
                     })
-                    .ok_or_else(|| UnknownCommonName {
-                        cn: common_name_from_sni.unwrap_or("").into(),
-                    })?;
+                    .map(Cow::<'static, str>::Owned)
+            })
+            .transpose()?;
 
-                Some(project)
-            } else {
-                None
-            }
-        } else {
-            None
-        };
-
-        let project = match (project_option, project_from_domain) {
+        let project = match (project_option, project_domain) {
             // Invariant: if we have both project name variants, they should match.
             (Some(option), Some(domain)) if option != domain => {
-                Some(Err(InconsistentProjectNames { domain, option }))
+                Some(Err(InconsistentProjectNames {
+                    domain: domain.into(),
+                    option: option.into(),
+                }))
             }
             // Invariant: project name may not contain certain characters.
             (a, b) => a.or(b).map(|name| match project_name_valid(&name) {
-                false => Err(MalformedProjectName(name)),
+                false => Err(MalformedProjectName(name.into())),
                 true => Ok(name),
             }),
         }
@@ -154,9 +149,9 @@ mod tests {
         let options = StartupMessageParams::new([("user", "john_doe")]);
 
         let sni = Some("foo.localhost");
-        let common_names = Some(["localhost".into()].into());
+        let common_name = Some("localhost");
 
-        let creds = ClientCredentials::parse(&options, sni, common_names)?;
+        let creds = ClientCredentials::parse(&options, sni, common_name)?;
         assert_eq!(creds.user, "john_doe");
         assert_eq!(creds.project.as_deref(), Some("foo"));
 
@@ -182,41 +177,24 @@ mod tests {
         let options = StartupMessageParams::new([("user", "john_doe"), ("options", "project=baz")]);
 
         let sni = Some("baz.localhost");
-        let common_names = Some(["localhost".into()].into());
+        let common_name = Some("localhost");
 
-        let creds = ClientCredentials::parse(&options, sni, common_names)?;
+        let creds = ClientCredentials::parse(&options, sni, common_name)?;
         assert_eq!(creds.user, "john_doe");
         assert_eq!(creds.project.as_deref(), Some("baz"));
 
         Ok(())
     }
 
-    #[test]
-    fn parse_multi_common_names() -> anyhow::Result<()> {
-        let options = StartupMessageParams::new([("user", "john_doe")]);
-
-        let common_names = Some(["a.com".into(), "b.com".into()].into());
-        let sni = Some("p1.a.com");
-        let creds = ClientCredentials::parse(&options, sni, common_names)?;
-        assert_eq!(creds.project.as_deref(), Some("p1"));
-
-        let common_names = Some(["a.com".into(), "b.com".into()].into());
-        let sni = Some("p1.b.com");
-        let creds = ClientCredentials::parse(&options, sni, common_names)?;
-        assert_eq!(creds.project.as_deref(), Some("p1"));
-
-        Ok(())
-    }
-
     #[test]
     fn parse_projects_different() {
         let options =
             StartupMessageParams::new([("user", "john_doe"), ("options", "project=first")]);
 
         let sni = Some("second.localhost");
-        let common_names = Some(["localhost".into()].into());
+        let common_name = Some("localhost");
 
-        let err = ClientCredentials::parse(&options, sni, common_names).expect_err("should fail");
+        let err = ClientCredentials::parse(&options, sni, common_name).expect_err("should fail");
         match err {
             InconsistentProjectNames { domain, option } => {
                 assert_eq!(option, "first");
@@ -231,12 +209,13 @@ mod tests {
         let options = StartupMessageParams::new([("user", "john_doe")]);
 
         let sni = Some("project.localhost");
-        let common_names = Some(["example.com".into()].into());
+        let common_name = Some("example.com");
 
-        let err = ClientCredentials::parse(&options, sni, common_names).expect_err("should fail");
+        let err = ClientCredentials::parse(&options, sni, common_name).expect_err("should fail");
         match err {
-            UnknownCommonName { cn } => {
-                assert_eq!(cn, "localhost");
+            InconsistentSni { sni, cn } => {
+                assert_eq!(sni, "project.localhost");
+                assert_eq!(cn, "example.com");
             }
             _ => panic!("bad error: {err:?}"),
         }

proxy/src/config.rs

@@ -1,12 +1,6 @@
 use crate::auth;
-use anyhow::{bail, ensure, Context, Ok};
-use rustls::sign;
-use std::{
-    collections::{HashMap, HashSet},
-    str::FromStr,
-    sync::Arc,
-    time::Duration,
-};
+use anyhow::{bail, ensure, Context};
+use std::{str::FromStr, sync::Arc, time::Duration};
 
 pub struct ProxyConfig {
     pub tls_config: Option<TlsConfig>,
@@ -22,7 +16,7 @@ pub struct MetricCollectionConfig {
 
 pub struct TlsConfig {
     pub config: Arc<rustls::ServerConfig>,
-    pub common_names: Option<HashSet<String>>,
+    pub common_name: Option<String>,
 }
 
 impl TlsConfig {
@@ -32,37 +26,28 @@ impl TlsConfig {
 }
 
 /// Configure TLS for the main endpoint.
-pub fn configure_tls(
-    key_path: &str,
-    cert_path: &str,
-    certs_dir: Option<&String>,
-) -> anyhow::Result<TlsConfig> {
-    let mut cert_resolver = CertResolver::new();
+pub fn configure_tls(key_path: &str, cert_path: &str) -> anyhow::Result<TlsConfig> {
+    let key = {
+        let key_bytes = std::fs::read(key_path).context("TLS key file")?;
+        let mut keys = rustls_pemfile::pkcs8_private_keys(&mut &key_bytes[..])
+            .context(format!("Failed to read TLS keys at '{key_path}'"))?;
 
-    // add default certificate
-    cert_resolver.add_cert(key_path, cert_path, true)?;
+        ensure!(keys.len() == 1, "keys.len() = {} (should be 1)", keys.len());
+        keys.pop().map(rustls::PrivateKey).unwrap()
+    };
 
-    // add extra certificates
-    if let Some(certs_dir) = certs_dir {
-        for entry in std::fs::read_dir(certs_dir)? {
-            let entry = entry?;
-            let path = entry.path();
-            if path.is_dir() {
-                // file names aligned with default cert-manager names
-                let key_path = path.join("tls.key");
-                let cert_path = path.join("tls.crt");
-                if key_path.exists() && cert_path.exists() {
-                    cert_resolver.add_cert(
-                        &key_path.to_string_lossy(),
-                        &cert_path.to_string_lossy(),
-                        false,
-                    )?;
-                }
-            }
-        }
-    }
+    let cert_chain_bytes = std::fs::read(cert_path)
+        .context(format!("Failed to read TLS cert file at '{cert_path}.'"))?;
 
-    let common_names = cert_resolver.get_common_names();
+    let cert_chain = {
+        rustls_pemfile::certs(&mut &cert_chain_bytes[..])
+            .context(format!(
+                "Failed to read TLS certificate chain from bytes from file at '{cert_path}'."
+            ))?
+            .into_iter()
+            .map(rustls::Certificate)
+            .collect()
+    };
 
     let config = rustls::ServerConfig::builder()
         .with_safe_default_cipher_suites()
@@ -70,136 +55,27 @@ pub fn configure_tls(
         // allow TLS 1.2 to be compatible with older client libraries
         .with_protocol_versions(&[&rustls::version::TLS13, &rustls::version::TLS12])?
         .with_no_client_auth()
-        .with_cert_resolver(Arc::new(cert_resolver))
+        .with_single_cert(cert_chain, key)?
         .into();
 
+    // determine common name from tls-cert (-c server.crt param).
+    // used in asserting project name formatting invariant.
+    let common_name = {
+        let pem = x509_parser::pem::parse_x509_pem(&cert_chain_bytes)
+            .context(format!(
+                "Failed to parse PEM object from bytes from file at '{cert_path}'."
+            ))?
+            .1;
+        let common_name = pem.parse_x509()?.subject().to_string();
+        common_name.strip_prefix("CN=*.").map(|s| s.to_string())
+    };
+
     Ok(TlsConfig {
         config,
-        common_names: Some(common_names),
+        common_name,
     })
 }
 
-struct CertResolver {
-    certs: HashMap<String, Arc<rustls::sign::CertifiedKey>>,
-    default: Option<Arc<rustls::sign::CertifiedKey>>,
-}
-
-impl CertResolver {
-    fn new() -> Self {
-        Self {
-            certs: HashMap::new(),
-            default: None,
-        }
-    }
-
-    fn add_cert(
-        &mut self,
-        key_path: &str,
-        cert_path: &str,
-        is_default: bool,
-    ) -> anyhow::Result<()> {
-        let priv_key = {
-            let key_bytes = std::fs::read(key_path).context("TLS key file")?;
-            let mut keys = rustls_pemfile::pkcs8_private_keys(&mut &key_bytes[..])
-                .context(format!("Failed to read TLS keys at '{key_path}'"))?;
-
-            ensure!(keys.len() == 1, "keys.len() = {} (should be 1)", keys.len());
-            keys.pop().map(rustls::PrivateKey).unwrap()
-        };
-
-        let key = sign::any_supported_type(&priv_key).context("invalid private key")?;
-
-        let cert_chain_bytes = std::fs::read(cert_path)
-            .context(format!("Failed to read TLS cert file at '{cert_path}.'"))?;
-
-        let cert_chain = {
-            rustls_pemfile::certs(&mut &cert_chain_bytes[..])
-                .context(format!(
-                    "Failed to read TLS certificate chain from bytes from file at '{cert_path}'."
-                ))?
-                .into_iter()
-                .map(rustls::Certificate)
-                .collect()
-        };
-
-        let common_name = {
-            let pem = x509_parser::pem::parse_x509_pem(&cert_chain_bytes)
-                .context(format!(
-                    "Failed to parse PEM object from bytes from file at '{cert_path}'."
-                ))?
-                .1;
-            let common_name = pem.parse_x509()?.subject().to_string();
-
-            // We only use non-wildcard certificates in link proxy so it seems okay to treat them the same as
-            // wildcard ones as we don't use SNI there. That treatment only affects certificate selection, so
-            // verify-full will still check wildcard match. Old coding here just ignored non-wildcard common names
-            // and passed None instead, which blows up number of cases downstream code should handle. Proper coding
-            // here should better avoid Option for common_names, and do wildcard-based certificate selection instead
-            // of cutting off '*.' parts.
-            if common_name.starts_with("CN=*.") {
-                common_name.strip_prefix("CN=*.").map(|s| s.to_string())
-            } else {
-                common_name.strip_prefix("CN=").map(|s| s.to_string())
-            }
-        }
-        .context(format!(
-            "Failed to parse common name from certificate at '{cert_path}'."
-        ))?;
-
-        let cert = Arc::new(rustls::sign::CertifiedKey::new(cert_chain, key));
-
-        if is_default {
-            self.default = Some(cert.clone());
-        }
-
-        self.certs.insert(common_name, cert);
-
-        Ok(())
-    }
-
-    fn get_common_names(&self) -> HashSet<String> {
-        self.certs.keys().map(|s| s.to_string()).collect()
-    }
-}
-
-impl rustls::server::ResolvesServerCert for CertResolver {
-    fn resolve(
-        &self,
-        _client_hello: rustls::server::ClientHello,
-    ) -> Option<Arc<rustls::sign::CertifiedKey>> {
-        // loop here and cut off more and more subdomains until we find
-        // a match to get a proper wildcard support. OTOH, we now do not
-        // use nested domains, so keep this simple for now.
-        //
-        // With the current coding foo.com will match *.foo.com and that
-        // repeats behavior of the old code.
-        if let Some(mut sni_name) = _client_hello.server_name() {
-            loop {
-                if let Some(cert) = self.certs.get(sni_name) {
-                    return Some(cert.clone());
-                }
-                if let Some((_, rest)) = sni_name.split_once('.') {
-                    sni_name = rest;
-                } else {
-                    return None;
-                }
-            }
-        } else {
-            // No SNI, use the default certificate, otherwise we can't get to
-            // options parameter which can be used to set endpoint name too.
-            // That means that non-SNI flow will not work for CNAME domains in
-            // verify-full mode.
-            //
-            // If that will be a problem we can:
-            //
-            // a) Instead of multi-cert approach use single cert with extra
-            //    domains listed in Subject Alternative Name (SAN).
-            // b) Deploy separate proxy instances for extra domains.
-            self.default.as_ref().cloned()
-        }
-    }
-}
-
 /// Helper for cmdline cache options parsing.
 pub struct CacheOptions {
     /// Max number of entries.

proxy/src/http/websocket.rs

@@ -22,7 +22,6 @@ use tokio::{
     io::{self, AsyncBufRead, AsyncRead, AsyncWrite, ReadBuf},
     net::TcpListener,
 };
-use tokio_util::sync::CancellationToken;
 use tracing::{error, info, info_span, warn, Instrument};
 use utils::http::{error::ApiError, json::json_response};
 
@@ -189,7 +188,6 @@ async fn ws_handler(
 pub async fn task_main(
     config: &'static ProxyConfig,
     ws_listener: TcpListener,
-    cancellation_token: CancellationToken,
 ) -> anyhow::Result<()> {
     scopeguard::defer! {
         info!("websocket server has shut down");
@@ -233,7 +231,6 @@ pub async fn task_main(
 
     hyper::Server::builder(accept::from_stream(tls_listener))
         .serve(make_svc)
-        .with_graceful_shutdown(cancellation_token.cancelled())
         .await?;
 
     Ok(())

proxy/src/main.rs

@@ -28,7 +28,6 @@ use config::ProxyConfig;
 use futures::FutureExt;
 use std::{borrow::Cow, future::Future, net::SocketAddr};
 use tokio::{net::TcpListener, task::JoinError};
-use tokio_util::sync::CancellationToken;
 use tracing::{info, warn};
 use utils::{project_git_version, sentry_init::init_sentry};
 
@@ -67,48 +66,39 @@ async fn main() -> anyhow::Result<()> {
     let proxy_address: SocketAddr = args.get_one::<String>("proxy").unwrap().parse()?;
     info!("Starting proxy on {proxy_address}");
     let proxy_listener = TcpListener::bind(proxy_address).await?;
-    let cancellation_token = CancellationToken::new();
 
-    let mut client_tasks = vec![tokio::spawn(proxy::task_main(
-        config,
-        proxy_listener,
-        cancellation_token.clone(),
-    ))];
+    let mut tasks = vec![
+        tokio::spawn(handle_signals()),
+        tokio::spawn(http::server::task_main(http_listener)),
+        tokio::spawn(proxy::task_main(config, proxy_listener)),
+        tokio::spawn(console::mgmt::task_main(mgmt_listener)),
+    ];
 
     if let Some(wss_address) = args.get_one::<String>("wss") {
         let wss_address: SocketAddr = wss_address.parse()?;
         info!("Starting wss on {wss_address}");
         let wss_listener = TcpListener::bind(wss_address).await?;
 
-        client_tasks.push(tokio::spawn(http::websocket::task_main(
+        tasks.push(tokio::spawn(http::websocket::task_main(
             config,
             wss_listener,
-            cancellation_token.clone(),
         )));
     }
 
-    let mut tasks = vec![
-        tokio::spawn(handle_signals(cancellation_token)),
-        tokio::spawn(http::server::task_main(http_listener)),
-        tokio::spawn(console::mgmt::task_main(mgmt_listener)),
-    ];
-
     if let Some(metrics_config) = &config.metric_collection {
         tasks.push(tokio::spawn(metrics::task_main(metrics_config)));
     }
 
-    let tasks = futures::future::try_join_all(tasks.into_iter().map(flatten_err));
-    let client_tasks = futures::future::try_join_all(client_tasks.into_iter().map(flatten_err));
-    tokio::select! {
-        // We are only expecting an error from these forever tasks
-        res = tasks => { res?; },
-        res = client_tasks => { res?; },
-    }
+    // This combinator will block until either all tasks complete or
+    // one of them finishes with an error (others will be cancelled).
+    let tasks = tasks.into_iter().map(flatten_err);
+    let _: Vec<()> = futures::future::try_join_all(tasks).await?;
+
     Ok(())
 }
 
 /// Handle unix signals appropriately.
-async fn handle_signals(token: CancellationToken) -> anyhow::Result<()> {
+async fn handle_signals() -> anyhow::Result<()> {
     use tokio::signal::unix::{signal, SignalKind};
 
     let mut hangup = signal(SignalKind::hangup())?;
@@ -126,9 +116,11 @@ async fn handle_signals(token: CancellationToken) -> anyhow::Result<()> {
                 warn!("received SIGINT, exiting immediately");
                 bail!("interrupted");
             }
+            // TODO: Don't accept new proxy connections.
+            // TODO: Shut down once all exisiting connections have been closed.
             _ = terminate.recv() => {
-                warn!("received SIGTERM, shutting down once all existing connections have closed");
-                token.cancel();
+                warn!("received SIGTERM, exiting immediately");
+                bail!("terminated");
             }
         }
     }
@@ -140,11 +132,7 @@ fn build_config(args: &clap::ArgMatches) -> anyhow::Result<&'static ProxyConfig>
         args.get_one::<String>("tls-key"),
         args.get_one::<String>("tls-cert"),
     ) {
-        (Some(key_path), Some(cert_path)) => Some(config::configure_tls(
-            key_path,
-            cert_path,
-            args.get_one::<String>("certs-dir"),
-        )?),
+        (Some(key_path), Some(cert_path)) => Some(config::configure_tls(key_path, cert_path)?),
         (None, None) => None,
         _ => bail!("either both or neither tls-key and tls-cert must be specified"),
     };
@@ -266,12 +254,6 @@ fn cli() -> clap::Command {
                 .alias("ssl-cert") // backwards compatibility
                 .help("path to TLS cert for client postgres connections"),
         )
-        // tls-key and tls-cert are for backwards compatibility, we can put all certs in one dir
-        .arg(
-            Arg::new("certs-dir")
-                .long("certs-dir")
-                .help("path to directory with TLS certificates for client postgres connections"),
-        )
         .arg(
             Arg::new("metric-collection-endpoint")
                 .long("metric-collection-endpoint")

proxy/src/metrics.rs

@@ -5,7 +5,7 @@ use chrono::{DateTime, Utc};
 use consumption_metrics::{idempotency_key, Event, EventChunk, EventType, CHUNK_SIZE};
 use serde::Serialize;
 use std::collections::HashMap;
-use tracing::{error, info, instrument, trace, warn};
+use tracing::{debug, error, info, instrument, trace};
 
 const PROXY_IO_BYTES_PER_CLIENT: &str = "proxy_io_bytes_per_client";
 
@@ -84,14 +84,10 @@ fn gather_proxy_io_bytes_per_client() -> Vec<(Ids, (u64, DateTime<Utc>))> {
 
                     let value = ms.get_counter().get_value() as u64;
 
-                    // Report if the metric value is suspiciously large
-                    if value > (1u64 << 40) {
-                        warn!(
-                            "potentially abnormal counter value: branch_id {} endpoint_id {} val: {}",
-                            branch_id, endpoint_id, value
-                        );
-                    }
-
+                    debug!(
+                        "branch_id {} endpoint_id {} val: {}",
+                        branch_id, endpoint_id, value
+                    );
                     current_metrics.push((
                         Ids {
                             endpoint_id: endpoint_id.to_string(),
@@ -128,15 +124,11 @@ async fn collect_metrics_iteration(
             let mut value = *curr_val;
 
             if let Some((prev_val, prev_time)) = cached_metrics.get(curr_key) {
-                // Only send metrics updates if the metric has increased
-                if curr_val > prev_val {
+                // Only send metrics updates if the metric has changed
+                if curr_val - prev_val > 0 {
                     value = curr_val - prev_val;
                     start_time = *prev_time;
                 } else {
-                    if curr_val < prev_val {
-                        error!("proxy_io_bytes_per_client metric value decreased from {} to {} for key {:?}",
-                        prev_val, curr_val, curr_key);
-                    }
                     return None;
                 }
             };
@@ -197,7 +189,7 @@ async fn collect_metrics_iteration(
                     })
                     // update cached value (add delta) and time
                     .and_modify(|e| {
-                        e.0 = e.0.saturating_add(send_metric.value);
+                        e.0 += send_metric.value;
                         e.1 = stop_time
                     })
                     // cache new metric

proxy/src/proxy.rs

@@ -17,7 +17,6 @@ use once_cell::sync::Lazy;
 use pq_proto::{BeMessage as Be, FeStartupPacket, StartupMessageParams};
 use std::sync::Arc;
 use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt};
-use tokio_util::sync::CancellationToken;
 use tracing::{error, info, warn};
 use utils::measured_stream::MeasuredStream;
 
@@ -64,7 +63,6 @@ static NUM_BYTES_PROXIED_COUNTER: Lazy<IntCounterVec> = Lazy::new(|| {
 pub async fn task_main(
     config: &'static ProxyConfig,
     listener: tokio::net::TcpListener,
-    cancellation_token: CancellationToken,
 ) -> anyhow::Result<()> {
     scopeguard::defer! {
         info!("proxy has shut down");
@@ -74,48 +72,29 @@ pub async fn task_main(
     // will be inherited by all accepted client sockets.
     socket2::SockRef::from(&listener).set_keepalive(true)?;
 
-    let mut connections = tokio::task::JoinSet::new();
     let cancel_map = Arc::new(CancelMap::default());
-
     loop {
-        tokio::select! {
-            accept_result = listener.accept() => {
-                let (socket, peer_addr) = accept_result?;
-                info!("accepted postgres client connection from {peer_addr}");
+        let (socket, peer_addr) = listener.accept().await?;
+        info!("accepted postgres client connection from {peer_addr}");
 
-                let session_id = uuid::Uuid::new_v4();
-                let cancel_map = Arc::clone(&cancel_map);
-                connections.spawn(
-                    async move {
-                        info!("spawned a task for {peer_addr}");
+        let session_id = uuid::Uuid::new_v4();
+        let cancel_map = Arc::clone(&cancel_map);
+        tokio::spawn(
+            async move {
+                info!("spawned a task for {peer_addr}");
 
-                        socket
-                            .set_nodelay(true)
-                            .context("failed to set socket option")?;
+                socket
+                    .set_nodelay(true)
+                    .context("failed to set socket option")?;
 
-                        handle_client(config, &cancel_map, session_id, socket).await
-                    }
-                    .unwrap_or_else(|e| {
-                        // Acknowledge that the task has finished with an error.
-                        error!("per-client task finished with an error: {e:#}");
-                    }),
-                );
+                handle_client(config, &cancel_map, session_id, socket).await
             }
-            _ = cancellation_token.cancelled() => {
-                drop(listener);
-                break;
-            }
-        }
+            .unwrap_or_else(|e| {
+                // Acknowledge that the task has finished with an error.
+                error!("per-client task finished with an error: {e:#}");
+            }),
+        );
     }
-    // Drain connections
-    while let Some(res) = connections.join_next().await {
-        if let Err(e) = res {
-            if !e.is_panic() && !e.is_cancelled() {
-                warn!("unexpected error from joined connection task: {e:?}");
-            }
-        }
-    }
-    Ok(())
 }
 
 // TODO(tech debt): unite this with its twin below.
@@ -145,11 +124,11 @@ pub async fn handle_ws_client(
 
     // Extract credentials which we're going to use for auth.
     let creds = {
-        let common_names = tls.and_then(|tls| tls.common_names.clone());
+        let common_name = tls.and_then(|tls| tls.common_name.as_deref());
         let result = config
             .auth_backend
             .as_ref()
-            .map(|_| auth::ClientCredentials::parse(&params, hostname, common_names))
+            .map(|_| auth::ClientCredentials::parse(&params, hostname, common_name))
             .transpose();
 
         async { result }.or_else(|e| stream.throw_error(e)).await?
@@ -184,11 +163,11 @@ async fn handle_client(
     // Extract credentials which we're going to use for auth.
     let creds = {
         let sni = stream.get_ref().sni_hostname();
-        let common_names = tls.and_then(|tls| tls.common_names.clone());
+        let common_name = tls.and_then(|tls| tls.common_name.as_deref());
         let result = config
             .auth_backend
             .as_ref()
-            .map(|_| auth::ClientCredentials::parse(&params, sni, common_names))
+            .map(|_| auth::ClientCredentials::parse(&params, sni, common_name))
             .transpose();
 
         async { result }.or_else(|e| stream.throw_error(e)).await?

proxy/src/proxy/tests.rs

@@ -54,11 +54,9 @@ fn generate_tls_config<'a>(
             .with_single_cert(vec![cert], key)?
             .into();
 
-        let common_names = Some([common_name.to_owned()].iter().cloned().collect());
-
         TlsConfig {
             config,
-            common_names,
+            common_name: Some(common_name.to_string()),
         }
     };
 

run_clippy.sh

@@ -8,7 +8,13 @@
 # warnings and errors right in the editor.
 # In vscode, this setting is Rust-analyzer>Check On Save:Command
 
+# manual-range-contains wants
+#   !(4..=MAX_STARTUP_PACKET_LENGTH).contains(&len)
+# instead of
+#   len < 4 || len > MAX_STARTUP_PACKET_LENGTH
+# , let's disagree.
+
 # * `-A unknown_lints` – do not warn about unknown lint suppressions
 #                        that people with newer toolchains might use
 # * `-D warnings`      - fail on any warnings (`cargo` returns non-zero exit status)
-cargo clippy --locked --all --all-targets --all-features -- -A unknown_lints -D warnings
+cargo clippy --locked --all --all-targets --all-features -- -A unknown_lints -A clippy::manual-range-contains -D warnings

safekeeper/Cargo.toml

@@ -30,7 +30,6 @@ serde_with.workspace = true
 signal-hook.workspace = true
 thiserror.workspace = true
 tokio = { workspace = true, features = ["fs"] }
-tokio-io-timeout.workspace = true
 tokio-postgres.workspace = true
 toml_edit.workspace = true
 tracing.workspace = true

safekeeper/src/timeline.rs

@@ -674,8 +674,7 @@ impl Timeline {
             bail!(TimelineError::Cancelled(self.ttid));
         }
 
-        let mut state = self.write_shared_state();
-        state.sk.inmem.backup_lsn = max(state.sk.inmem.backup_lsn, backup_lsn);
+        self.write_shared_state().sk.inmem.backup_lsn = backup_lsn;
         // we should check whether to shut down offloader, but this will be done
         // soon by peer communication anyway.
         Ok(())

safekeeper/src/wal_backup.rs

@@ -323,8 +323,7 @@ impl WalBackupTask {
             }
 
             match backup_lsn_range(
-                &self.timeline,
-                &mut backup_lsn,
+                backup_lsn,
                 commit_lsn,
                 self.wal_seg_size,
                 &self.timeline_dir,
@@ -332,7 +331,13 @@ impl WalBackupTask {
             )
             .await
             {
-                Ok(()) => {
+                Ok(backup_lsn_result) => {
+                    backup_lsn = backup_lsn_result;
+                    let res = self.timeline.set_wal_backup_lsn(backup_lsn_result);
+                    if let Err(e) = res {
+                        error!("failed to set wal_backup_lsn: {}", e);
+                        return;
+                    }
                     retry_attempt = 0;
                 }
                 Err(e) => {
@@ -349,25 +354,20 @@ impl WalBackupTask {
 }
 
 pub async fn backup_lsn_range(
-    timeline: &Arc<Timeline>,
-    backup_lsn: &mut Lsn,
+    start_lsn: Lsn,
     end_lsn: Lsn,
     wal_seg_size: usize,
     timeline_dir: &Path,
     workspace_dir: &Path,
-) -> Result<()> {
-    let start_lsn = *backup_lsn;
+) -> Result<Lsn> {
+    let mut res = start_lsn;
     let segments = get_segments(start_lsn, end_lsn, wal_seg_size);
     for s in &segments {
         backup_single_segment(s, timeline_dir, workspace_dir)
             .await
             .with_context(|| format!("offloading segno {}", s.seg_no))?;
 
-        let new_backup_lsn = s.end_lsn;
-        timeline
-            .set_wal_backup_lsn(new_backup_lsn)
-            .context("setting wal_backup_lsn")?;
-        *backup_lsn = new_backup_lsn;
+        res = s.end_lsn;
     }
     info!(
         "offloaded segnos {:?} up to {}, previous backup_lsn {}",
@@ -375,7 +375,7 @@ pub async fn backup_lsn_range(
         end_lsn,
         start_lsn,
     );
-    Ok(())
+    Ok(res)
 }
 
 async fn backup_single_segment(

safekeeper/src/wal_service.rs

@@ -4,9 +4,8 @@
 //!
 use anyhow::{Context, Result};
 use postgres_backend::QueryError;
-use std::{future, thread, time::Duration};
+use std::{future, thread};
 use tokio::net::TcpStream;
-use tokio_io_timeout::TimeoutReader;
 use tracing::*;
 use utils::measured_stream::MeasuredStream;
 
@@ -68,52 +67,41 @@ fn handle_socket(
     let runtime = tokio::runtime::Builder::new_current_thread()
         .enable_all()
         .build()?;
+    let local = tokio::task::LocalSet::new();
 
     socket.set_nodelay(true)?;
     let peer_addr = socket.peer_addr()?;
 
-    // TimeoutReader wants async runtime during creation.
-    runtime.block_on(async move {
-        // Set timeout on reading from the socket. It prevents hanged up connection
-        // if client suddenly disappears. Note that TCP_KEEPALIVE is not enabled by
-        // default, and tokio doesn't provide ability to set it out of the box.
-        let mut socket = TimeoutReader::new(socket);
-        let wal_service_timeout = Duration::from_secs(60 * 10);
-        socket.set_timeout(Some(wal_service_timeout));
-        // pin! is here because TimeoutReader (due to storing sleep future inside)
-        // is not Unpin, and all pgbackend/framed/tokio dependencies require stream
-        // to be Unpin. Which is reasonable, as indeed something like TimeoutReader
-        // shouldn't be moved.
-        tokio::pin!(socket);
+    let traffic_metrics = TrafficMetrics::new();
+    if let Some(current_az) = conf.availability_zone.as_deref() {
+        traffic_metrics.set_sk_az(current_az);
+    }
 
-        let traffic_metrics = TrafficMetrics::new();
-        if let Some(current_az) = conf.availability_zone.as_deref() {
-            traffic_metrics.set_sk_az(current_az);
-        }
+    let socket = MeasuredStream::new(
+        socket,
+        |cnt| {
+            traffic_metrics.observe_read(cnt);
+        },
+        |cnt| {
+            traffic_metrics.observe_write(cnt);
+        },
+    );
 
-        let socket = MeasuredStream::new(
-            socket,
-            |cnt| {
-                traffic_metrics.observe_read(cnt);
-            },
-            |cnt| {
-                traffic_metrics.observe_write(cnt);
-            },
-        );
+    let auth_type = match conf.auth {
+        None => AuthType::Trust,
+        Some(_) => AuthType::NeonJWT,
+    };
+    let mut conn_handler =
+        SafekeeperPostgresHandler::new(conf, conn_id, Some(traffic_metrics.clone()));
+    let pgbackend = PostgresBackend::new_from_io(socket, peer_addr, auth_type, None)?;
+    // libpq protocol between safekeeper and walproposer / pageserver
+    // We don't use shutdown.
+    local.block_on(
+        &runtime,
+        pgbackend.run(&mut conn_handler, future::pending::<()>),
+    )?;
 
-        let auth_type = match conf.auth {
-            None => AuthType::Trust,
-            Some(_) => AuthType::NeonJWT,
-        };
-        let mut conn_handler =
-            SafekeeperPostgresHandler::new(conf, conn_id, Some(traffic_metrics.clone()));
-        let pgbackend = PostgresBackend::new_from_io(socket, peer_addr, auth_type, None)?;
-        // libpq protocol between safekeeper and walproposer / pageserver
-        // We don't use shutdown.
-        pgbackend
-            .run(&mut conn_handler, future::pending::<()>)
-            .await
-    })
+    Ok(())
 }
 
 /// Unique WAL service connection ids are logged in spans for observability.

test_runner/fixtures/neon_fixtures.py

@@ -14,9 +14,9 @@ import tempfile
 import textwrap
 import time
 import uuid
+from collections import defaultdict
 from contextlib import closing, contextmanager
 from dataclasses import dataclass, field
-from datetime import datetime
 from enum import Flag, auto
 from functools import cached_property
 from itertools import chain, product
@@ -43,12 +43,11 @@ from psycopg2.extensions import make_dsn, parse_dsn
 from typing_extensions import Literal
 
 from fixtures.log_helper import log
-from fixtures.pageserver.http import PageserverHttpClient
-from fixtures.pageserver.utils import wait_for_last_record_lsn, wait_for_upload
+from fixtures.metrics import Metrics, parse_metrics
 from fixtures.types import Lsn, TenantId, TimelineId
 from fixtures.utils import (
     ATTACHMENT_NAME_REGEX,
-    allure_add_grafana_links,
+    Fn,
     allure_attach_from_dir,
     get_self_dir,
     subprocess_capture,
@@ -1127,6 +1126,538 @@ def neon_env_builder(
         yield builder
 
 
+class PageserverApiException(Exception):
+    def __init__(self, message, status_code: int):
+        super().__init__(message)
+        self.status_code = status_code
+
+
+class PageserverHttpClient(requests.Session):
+    def __init__(self, port: int, is_testing_enabled_or_skip: Fn, auth_token: Optional[str] = None):
+        super().__init__()
+        self.port = port
+        self.auth_token = auth_token
+        self.is_testing_enabled_or_skip = is_testing_enabled_or_skip
+
+        if auth_token is not None:
+            self.headers["Authorization"] = f"Bearer {auth_token}"
+
+    def verbose_error(self, res: requests.Response):
+        try:
+            res.raise_for_status()
+        except requests.RequestException as e:
+            try:
+                msg = res.json()["msg"]
+            except:  # noqa: E722
+                msg = ""
+            raise PageserverApiException(msg, res.status_code) from e
+
+    def check_status(self):
+        self.get(f"http://localhost:{self.port}/v1/status").raise_for_status()
+
+    def configure_failpoints(self, config_strings: Tuple[str, str] | List[Tuple[str, str]]):
+        self.is_testing_enabled_or_skip()
+
+        if isinstance(config_strings, tuple):
+            pairs = [config_strings]
+        else:
+            pairs = config_strings
+
+        log.info(f"Requesting config failpoints: {repr(pairs)}")
+
+        res = self.put(
+            f"http://localhost:{self.port}/v1/failpoints",
+            json=[{"name": name, "actions": actions} for name, actions in pairs],
+        )
+        log.info(f"Got failpoints request response code {res.status_code}")
+        self.verbose_error(res)
+        res_json = res.json()
+        assert res_json is None
+        return res_json
+
+    def tenant_list(self) -> List[Dict[Any, Any]]:
+        res = self.get(f"http://localhost:{self.port}/v1/tenant")
+        self.verbose_error(res)
+        res_json = res.json()
+        assert isinstance(res_json, list)
+        return res_json
+
+    def tenant_create(self, new_tenant_id: Optional[TenantId] = None) -> TenantId:
+        res = self.post(
+            f"http://localhost:{self.port}/v1/tenant",
+            json={
+                "new_tenant_id": str(new_tenant_id) if new_tenant_id else None,
+            },
+        )
+        self.verbose_error(res)
+        if res.status_code == 409:
+            raise Exception(f"could not create tenant: already exists for id {new_tenant_id}")
+        new_tenant_id = res.json()
+        assert isinstance(new_tenant_id, str)
+        return TenantId(new_tenant_id)
+
+    def tenant_attach(self, tenant_id: TenantId):
+        res = self.post(f"http://localhost:{self.port}/v1/tenant/{tenant_id}/attach")
+        self.verbose_error(res)
+
+    def tenant_detach(self, tenant_id: TenantId, detach_ignored=False):
+        params = {}
+        if detach_ignored:
+            params["detach_ignored"] = "true"
+
+        res = self.post(f"http://localhost:{self.port}/v1/tenant/{tenant_id}/detach", params=params)
+        self.verbose_error(res)
+
+    def tenant_load(self, tenant_id: TenantId):
+        res = self.post(f"http://localhost:{self.port}/v1/tenant/{tenant_id}/load")
+        self.verbose_error(res)
+
+    def tenant_ignore(self, tenant_id: TenantId):
+        res = self.post(f"http://localhost:{self.port}/v1/tenant/{tenant_id}/ignore")
+        self.verbose_error(res)
+
+    def tenant_status(self, tenant_id: TenantId) -> Dict[Any, Any]:
+        res = self.get(f"http://localhost:{self.port}/v1/tenant/{tenant_id}")
+        self.verbose_error(res)
+        res_json = res.json()
+        assert isinstance(res_json, dict)
+        return res_json
+
+    def tenant_config(self, tenant_id: TenantId) -> TenantConfig:
+        res = self.get(f"http://localhost:{self.port}/v1/tenant/{tenant_id}/config")
+        self.verbose_error(res)
+        return TenantConfig.from_json(res.json())
+
+    def set_tenant_config(self, tenant_id: TenantId, config: dict[str, Any]):
+        assert "tenant_id" not in config.keys()
+        res = self.put(
+            f"http://localhost:{self.port}/v1/tenant/config",
+            json={**config, "tenant_id": str(tenant_id)},
+        )
+        self.verbose_error(res)
+
+    def patch_tenant_config_client_side(
+        self,
+        tenant_id: TenantId,
+        inserts: Optional[Dict[str, Any]] = None,
+        removes: Optional[List[str]] = None,
+    ):
+        current = self.tenant_config(tenant_id).tenant_specific_overrides
+        if inserts is not None:
+            current.update(inserts)
+        if removes is not None:
+            for key in removes:
+                del current[key]
+        self.set_tenant_config(tenant_id, current)
+
+    def tenant_size(self, tenant_id: TenantId) -> int:
+        return self.tenant_size_and_modelinputs(tenant_id)[0]
+
+    def tenant_size_and_modelinputs(self, tenant_id: TenantId) -> Tuple[int, Dict[str, Any]]:
+        """
+        Returns the tenant size, together with the model inputs as the second tuple item.
+        """
+        res = self.get(f"http://localhost:{self.port}/v1/tenant/{tenant_id}/synthetic_size")
+        self.verbose_error(res)
+        res = res.json()
+        assert isinstance(res, dict)
+        assert TenantId(res["id"]) == tenant_id
+        size = res["size"]
+        assert type(size) == int
+        inputs = res["inputs"]
+        assert type(inputs) is dict
+        return (size, inputs)
+
+    def tenant_size_debug(self, tenant_id: TenantId) -> str:
+        """
+        Returns the tenant size debug info, as an HTML string
+        """
+        res = self.get(
+            f"http://localhost:{self.port}/v1/tenant/{tenant_id}/synthetic_size",
+            headers={"Accept": "text/html"},
+        )
+        return res.text
+
+    def timeline_list(
+        self,
+        tenant_id: TenantId,
+        include_non_incremental_logical_size: bool = False,
+        include_timeline_dir_layer_file_size_sum: bool = False,
+    ) -> List[Dict[str, Any]]:
+        params = {}
+        if include_non_incremental_logical_size:
+            params["include-non-incremental-logical-size"] = "true"
+        if include_timeline_dir_layer_file_size_sum:
+            params["include-timeline-dir-layer-file-size-sum"] = "true"
+
+        res = self.get(
+            f"http://localhost:{self.port}/v1/tenant/{tenant_id}/timeline", params=params
+        )
+        self.verbose_error(res)
+        res_json = res.json()
+        assert isinstance(res_json, list)
+        return res_json
+
+    def timeline_create(
+        self,
+        tenant_id: TenantId,
+        new_timeline_id: Optional[TimelineId] = None,
+        ancestor_timeline_id: Optional[TimelineId] = None,
+        ancestor_start_lsn: Optional[Lsn] = None,
+    ) -> Dict[Any, Any]:
+        res = self.post(
+            f"http://localhost:{self.port}/v1/tenant/{tenant_id}/timeline",
+            json={
+                "new_timeline_id": str(new_timeline_id) if new_timeline_id else None,
+                "ancestor_start_lsn": str(ancestor_start_lsn) if ancestor_start_lsn else None,
+                "ancestor_timeline_id": str(ancestor_timeline_id) if ancestor_timeline_id else None,
+            },
+        )
+        self.verbose_error(res)
+        if res.status_code == 409:
+            raise Exception(f"could not create timeline: already exists for id {new_timeline_id}")
+
+        res_json = res.json()
+        assert isinstance(res_json, dict)
+        return res_json
+
+    def timeline_detail(
+        self,
+        tenant_id: TenantId,
+        timeline_id: TimelineId,
+        include_non_incremental_logical_size: bool = False,
+        include_timeline_dir_layer_file_size_sum: bool = False,
+        **kwargs,
+    ) -> Dict[Any, Any]:
+        params = {}
+        if include_non_incremental_logical_size:
+            params["include-non-incremental-logical-size"] = "true"
+        if include_timeline_dir_layer_file_size_sum:
+            params["include-timeline-dir-layer-file-size-sum"] = "true"
+
+        res = self.get(
+            f"http://localhost:{self.port}/v1/tenant/{tenant_id}/timeline/{timeline_id}",
+            params=params,
+            **kwargs,
+        )
+        self.verbose_error(res)
+        res_json = res.json()
+        assert isinstance(res_json, dict)
+        return res_json
+
+    def timeline_delete(self, tenant_id: TenantId, timeline_id: TimelineId):
+        res = self.delete(
+            f"http://localhost:{self.port}/v1/tenant/{tenant_id}/timeline/{timeline_id}"
+        )
+        self.verbose_error(res)
+        res_json = res.json()
+        assert res_json is None
+
+    def timeline_gc(
+        self, tenant_id: TenantId, timeline_id: TimelineId, gc_horizon: Optional[int]
+    ) -> dict[str, Any]:
+        self.is_testing_enabled_or_skip()
+
+        log.info(
+            f"Requesting GC: tenant {tenant_id}, timeline {timeline_id}, gc_horizon {repr(gc_horizon)}"
+        )
+        res = self.put(
+            f"http://localhost:{self.port}/v1/tenant/{tenant_id}/timeline/{timeline_id}/do_gc",
+            json={"gc_horizon": gc_horizon},
+        )
+        log.info(f"Got GC request response code: {res.status_code}")
+        self.verbose_error(res)
+        res_json = res.json()
+        assert res_json is not None
+        assert isinstance(res_json, dict)
+        return res_json
+
+    def timeline_compact(self, tenant_id: TenantId, timeline_id: TimelineId):
+        self.is_testing_enabled_or_skip()
+
+        log.info(f"Requesting compact: tenant {tenant_id}, timeline {timeline_id}")
+        res = self.put(
+            f"http://localhost:{self.port}/v1/tenant/{tenant_id}/timeline/{timeline_id}/compact"
+        )
+        log.info(f"Got compact request response code: {res.status_code}")
+        self.verbose_error(res)
+        res_json = res.json()
+        assert res_json is None
+
+    def timeline_get_lsn_by_timestamp(
+        self, tenant_id: TenantId, timeline_id: TimelineId, timestamp
+    ):
+        log.info(
+            f"Requesting lsn by timestamp {timestamp}, tenant {tenant_id}, timeline {timeline_id}"
+        )
+        res = self.get(
+            f"http://localhost:{self.port}/v1/tenant/{tenant_id}/timeline/{timeline_id}/get_lsn_by_timestamp?timestamp={timestamp}",
+        )
+        self.verbose_error(res)
+        res_json = res.json()
+        return res_json
+
+    def timeline_checkpoint(self, tenant_id: TenantId, timeline_id: TimelineId):
+        self.is_testing_enabled_or_skip()
+
+        log.info(f"Requesting checkpoint: tenant {tenant_id}, timeline {timeline_id}")
+        res = self.put(
+            f"http://localhost:{self.port}/v1/tenant/{tenant_id}/timeline/{timeline_id}/checkpoint"
+        )
+        log.info(f"Got checkpoint request response code: {res.status_code}")
+        self.verbose_error(res)
+        res_json = res.json()
+        assert res_json is None
+
+    def timeline_spawn_download_remote_layers(
+        self,
+        tenant_id: TenantId,
+        timeline_id: TimelineId,
+        max_concurrent_downloads: int,
+    ) -> dict[str, Any]:
+        body = {
+            "max_concurrent_downloads": max_concurrent_downloads,
+        }
+        res = self.post(
+            f"http://localhost:{self.port}/v1/tenant/{tenant_id}/timeline/{timeline_id}/download_remote_layers",
+            json=body,
+        )
+        self.verbose_error(res)
+        res_json = res.json()
+        assert res_json is not None
+        assert isinstance(res_json, dict)
+        return res_json
+
+    def timeline_poll_download_remote_layers_status(
+        self,
+        tenant_id: TenantId,
+        timeline_id: TimelineId,
+        spawn_response: dict[str, Any],
+        poll_state=None,
+    ) -> None | dict[str, Any]:
+        res = self.get(
+            f"http://localhost:{self.port}/v1/tenant/{tenant_id}/timeline/{timeline_id}/download_remote_layers",
+        )
+        self.verbose_error(res)
+        res_json = res.json()
+        assert res_json is not None
+        assert isinstance(res_json, dict)
+
+        # assumption in this API client here is that nobody else spawns the task
+        assert res_json["task_id"] == spawn_response["task_id"]
+
+        if poll_state is None or res_json["state"] == poll_state:
+            return res_json
+        return None
+
+    def timeline_download_remote_layers(
+        self,
+        tenant_id: TenantId,
+        timeline_id: TimelineId,
+        max_concurrent_downloads: int,
+        errors_ok=False,
+        at_least_one_download=True,
+    ):
+        res = self.timeline_spawn_download_remote_layers(
+            tenant_id, timeline_id, max_concurrent_downloads
+        )
+        while True:
+            completed = self.timeline_poll_download_remote_layers_status(
+                tenant_id, timeline_id, res, poll_state="Completed"
+            )
+            if not completed:
+                time.sleep(0.1)
+                continue
+            if not errors_ok:
+                assert completed["failed_download_count"] == 0
+            if at_least_one_download:
+                assert completed["successful_download_count"] > 0
+            return completed
+
+    def get_metrics_str(self) -> str:
+        """You probably want to use get_metrics() instead."""
+        res = self.get(f"http://localhost:{self.port}/metrics")
+        self.verbose_error(res)
+        return res.text
+
+    def get_metrics(self) -> Metrics:
+        res = self.get_metrics_str()
+        return parse_metrics(res)
+
+    def get_timeline_metric(
+        self, tenant_id: TenantId, timeline_id: TimelineId, metric_name: str
+    ) -> float:
+        metrics = self.get_metrics()
+        return metrics.query_one(
+            metric_name,
+            filter={
+                "tenant_id": str(tenant_id),
+                "timeline_id": str(timeline_id),
+            },
+        ).value
+
+    def get_remote_timeline_client_metric(
+        self,
+        metric_name: str,
+        tenant_id: TenantId,
+        timeline_id: TimelineId,
+        file_kind: str,
+        op_kind: str,
+    ) -> Optional[float]:
+        metrics = self.get_metrics()
+        matches = metrics.query_all(
+            name=metric_name,
+            filter={
+                "tenant_id": str(tenant_id),
+                "timeline_id": str(timeline_id),
+                "file_kind": str(file_kind),
+                "op_kind": str(op_kind),
+            },
+        )
+        if len(matches) == 0:
+            value = None
+        elif len(matches) == 1:
+            value = matches[0].value
+            assert value is not None
+        else:
+            assert len(matches) < 2, "above filter should uniquely identify metric"
+        return value
+
+    def get_metric_value(
+        self, name: str, filter: Optional[Dict[str, str]] = None
+    ) -> Optional[float]:
+        metrics = self.get_metrics()
+        results = metrics.query_all(name, filter=filter)
+        if not results:
+            log.info(f'could not find metric "{name}"')
+            return None
+        assert len(results) == 1, f"metric {name} with given filters is not unique, got: {results}"
+        return results[0].value
+
+    def layer_map_info(
+        self,
+        tenant_id: TenantId,
+        timeline_id: TimelineId,
+    ) -> LayerMapInfo:
+        res = self.get(
+            f"http://localhost:{self.port}/v1/tenant/{tenant_id}/timeline/{timeline_id}/layer/",
+        )
+        self.verbose_error(res)
+        return LayerMapInfo.from_json(res.json())
+
+    def download_layer(self, tenant_id: TenantId, timeline_id: TimelineId, layer_name: str):
+        res = self.get(
+            f"http://localhost:{self.port}/v1/tenant/{tenant_id}/timeline/{timeline_id}/layer/{layer_name}",
+        )
+        self.verbose_error(res)
+
+        assert res.status_code == 200
+
+    def evict_layer(self, tenant_id: TenantId, timeline_id: TimelineId, layer_name: str):
+        res = self.delete(
+            f"http://localhost:{self.port}/v1/tenant/{tenant_id}/timeline/{timeline_id}/layer/{layer_name}",
+        )
+        self.verbose_error(res)
+
+        assert res.status_code == 200
+
+    def evict_all_layers(self, tenant_id: TenantId, timeline_id: TimelineId):
+        info = self.layer_map_info(tenant_id, timeline_id)
+        for layer in info.historic_layers:
+            self.evict_layer(tenant_id, timeline_id, layer.layer_file_name)
+
+    def disk_usage_eviction_run(self, request: dict[str, Any]):
+        res = self.put(
+            f"http://localhost:{self.port}/v1/disk_usage_eviction/run",
+            json=request,
+        )
+        self.verbose_error(res)
+        return res.json()
+
+    def tenant_break(self, tenant_id: TenantId):
+        res = self.put(f"http://localhost:{self.port}/v1/tenant/{tenant_id}/break")
+        self.verbose_error(res)
+
+
+@dataclass
+class TenantConfig:
+    tenant_specific_overrides: Dict[str, Any]
+    effective_config: Dict[str, Any]
+
+    @classmethod
+    def from_json(cls, d: Dict[str, Any]) -> TenantConfig:
+        return TenantConfig(
+            tenant_specific_overrides=d["tenant_specific_overrides"],
+            effective_config=d["effective_config"],
+        )
+
+
+@dataclass
+class LayerMapInfo:
+    in_memory_layers: List[InMemoryLayerInfo]
+    historic_layers: List[HistoricLayerInfo]
+
+    @classmethod
+    def from_json(cls, d: Dict[str, Any]) -> LayerMapInfo:
+        info = LayerMapInfo(in_memory_layers=[], historic_layers=[])
+
+        json_in_memory_layers = d["in_memory_layers"]
+        assert isinstance(json_in_memory_layers, List)
+        for json_in_memory_layer in json_in_memory_layers:
+            info.in_memory_layers.append(InMemoryLayerInfo.from_json(json_in_memory_layer))
+
+        json_historic_layers = d["historic_layers"]
+        assert isinstance(json_historic_layers, List)
+        for json_historic_layer in json_historic_layers:
+            info.historic_layers.append(HistoricLayerInfo.from_json(json_historic_layer))
+
+        return info
+
+    def kind_count(self) -> Dict[str, int]:
+        counts: Dict[str, int] = defaultdict(int)
+        for inmem_layer in self.in_memory_layers:
+            counts[inmem_layer.kind] += 1
+        for hist_layer in self.historic_layers:
+            counts[hist_layer.kind] += 1
+        return counts
+
+
+@dataclass
+class InMemoryLayerInfo:
+    kind: str
+    lsn_start: str
+    lsn_end: Optional[str]
+
+    @classmethod
+    def from_json(cls, d: Dict[str, Any]) -> InMemoryLayerInfo:
+        return InMemoryLayerInfo(
+            kind=d["kind"],
+            lsn_start=d["lsn_start"],
+            lsn_end=d.get("lsn_end"),
+        )
+
+
+@dataclass(frozen=True)
+class HistoricLayerInfo:
+    kind: str
+    layer_file_name: str
+    layer_file_size: Optional[int]
+    lsn_start: str
+    lsn_end: Optional[str]
+    remote: bool
+
+    @classmethod
+    def from_json(cls, d: Dict[str, Any]) -> HistoricLayerInfo:
+        return HistoricLayerInfo(
+            kind=d["kind"],
+            layer_file_name=d["layer_file_name"],
+            layer_file_size=d.get("layer_file_size"),
+            lsn_start=d["lsn_start"],
+            lsn_end=d.get("lsn_end"),
+            remote=d["remote"],
+        )
+
+
 @dataclass
 class PageserverPort:
     pg: int
@@ -1913,16 +2444,10 @@ def remote_pg(
     connstr = os.getenv("BENCHMARK_CONNSTR")
     if connstr is None:
         raise ValueError("no connstr provided, use BENCHMARK_CONNSTR environment variable")
-    start_ms = int(datetime.utcnow().timestamp() * 1000)
+
     with RemotePostgres(pg_bin, connstr) as remote_pg:
         yield remote_pg
 
-    end_ms = int(datetime.utcnow().timestamp() * 1000)
-    host = parse_dsn(connstr).get("host", "")
-    if host.endswith(".neon.build"):
-        # Add 10s margin to the start and end times
-        allure_add_grafana_links(host, start_ms - 10_000, end_ms + 10_000)
-
 
 class PSQL:
     """
@@ -2041,17 +2566,6 @@ class NeonProxy(PgProtocol):
         self._wait_until_ready()
         return self
 
-    # Sends SIGTERM to the proxy if it has been started
-    def terminate(self):
-        if self._popen:
-            self._popen.terminate()
-
-    # Waits for proxy to exit if it has been opened with a default timeout of
-    # two seconds. Raises subprocess.TimeoutExpired if the proxy does not exit in time.
-    def wait_for_exit(self, timeout=2):
-        if self._popen:
-            self._popen.wait(timeout=2)
-
     @backoff.on_exception(backoff.expo, requests.exceptions.RequestException, max_time=10)
     def _wait_until_ready(self):
         requests.get(f"http://{self.host}:{self.http_port}/v1/status")
@@ -2885,6 +3399,151 @@ def check_restored_datadir_content(
     assert (mismatch, error) == ([], [])
 
 
+def wait_until(number_of_iterations: int, interval: float, func):
+    """
+    Wait until 'func' returns successfully, without exception. Returns the
+    last return value from the function.
+    """
+    last_exception = None
+    for i in range(number_of_iterations):
+        try:
+            res = func()
+        except Exception as e:
+            log.info("waiting for %s iteration %s failed", func, i + 1)
+            last_exception = e
+            time.sleep(interval)
+            continue
+        return res
+    raise Exception("timed out while waiting for %s" % func) from last_exception
+
+
+def wait_while(number_of_iterations: int, interval: float, func):
+    """
+    Wait until 'func' returns false, or throws an exception.
+    """
+    for i in range(number_of_iterations):
+        try:
+            if not func():
+                return
+            log.info("waiting for %s iteration %s failed", func, i + 1)
+            time.sleep(interval)
+            continue
+        except Exception:
+            return
+    raise Exception("timed out while waiting for %s" % func)
+
+
+def assert_tenant_status(
+    pageserver_http_client: PageserverHttpClient, tenant: TenantId, expected_status: str
+):
+    tenant_status = pageserver_http_client.tenant_status(tenant)
+    log.info(f"tenant_status: {tenant_status}")
+    assert tenant_status["state"] == expected_status, tenant_status
+
+
+def tenant_exists(ps_http: PageserverHttpClient, tenant_id: TenantId):
+    tenants = ps_http.tenant_list()
+    matching = [t for t in tenants if TenantId(t["id"]) == tenant_id]
+    assert len(matching) < 2
+    if len(matching) == 0:
+        return None
+    return matching[0]
+
+
+def remote_consistent_lsn(
+    pageserver_http_client: PageserverHttpClient, tenant: TenantId, timeline: TimelineId
+) -> Lsn:
+    detail = pageserver_http_client.timeline_detail(tenant, timeline)
+
+    if detail["remote_consistent_lsn"] is None:
+        # No remote information at all. This happens right after creating
+        # a timeline, before any part of it has been uploaded to remote
+        # storage yet.
+        return Lsn(0)
+    else:
+        lsn_str = detail["remote_consistent_lsn"]
+        assert isinstance(lsn_str, str)
+        return Lsn(lsn_str)
+
+
+def wait_for_upload(
+    pageserver_http_client: PageserverHttpClient,
+    tenant: TenantId,
+    timeline: TimelineId,
+    lsn: Lsn,
+):
+    """waits for local timeline upload up to specified lsn"""
+    for i in range(20):
+        current_lsn = remote_consistent_lsn(pageserver_http_client, tenant, timeline)
+        if current_lsn >= lsn:
+            log.info("wait finished")
+            return
+        log.info(
+            "waiting for remote_consistent_lsn to reach {}, now {}, iteration {}".format(
+                lsn, current_lsn, i + 1
+            )
+        )
+        time.sleep(1)
+    raise Exception(
+        "timed out while waiting for remote_consistent_lsn to reach {}, was {}".format(
+            lsn, current_lsn
+        )
+    )
+
+
+# Does not use `wait_until` for debugging purposes
+def wait_until_tenant_state(
+    pageserver_http: PageserverHttpClient,
+    tenant_id: TenantId,
+    expected_state: str,
+    iterations: int,
+) -> bool:
+    for _ in range(iterations):
+        try:
+            tenant = pageserver_http.tenant_status(tenant_id=tenant_id)
+            log.debug(f"Tenant {tenant_id} data: {tenant}")
+            if tenant["state"] == expected_state:
+                return True
+        except Exception as e:
+            log.debug(f"Tenant {tenant_id} state retrieval failure: {e}")
+
+        time.sleep(1)
+
+    raise Exception(f"Tenant {tenant_id} did not become {expected_state} in {iterations} seconds")
+
+
+def last_record_lsn(
+    pageserver_http_client: PageserverHttpClient, tenant: TenantId, timeline: TimelineId
+) -> Lsn:
+    detail = pageserver_http_client.timeline_detail(tenant, timeline)
+
+    lsn_str = detail["last_record_lsn"]
+    assert isinstance(lsn_str, str)
+    return Lsn(lsn_str)
+
+
+def wait_for_last_record_lsn(
+    pageserver_http_client: PageserverHttpClient,
+    tenant: TenantId,
+    timeline: TimelineId,
+    lsn: Lsn,
+) -> Lsn:
+    """waits for pageserver to catch up to a certain lsn, returns the last observed lsn."""
+    for i in range(10):
+        current_lsn = last_record_lsn(pageserver_http_client, tenant, timeline)
+        if current_lsn >= lsn:
+            return current_lsn
+        log.info(
+            "waiting for last_record_lsn to reach {}, now {}, iteration {}".format(
+                lsn, current_lsn, i + 1
+            )
+        )
+        time.sleep(1)
+    raise Exception(
+        "timed out while waiting for last_record_lsn to reach {}, was {}".format(lsn, current_lsn)
+    )
+
+
 def wait_for_last_flush_lsn(
     env: NeonEnv, endpoint: Endpoint, tenant: TenantId, timeline: TimelineId
 ) -> Lsn:
@@ -2946,3 +3605,23 @@ def wait_for_sk_commit_lsn_to_reach_remote_storage(
     ps_http.timeline_checkpoint(tenant_id, timeline_id)
     wait_for_upload(ps_http, tenant_id, timeline_id, lsn)
     return lsn
+
+
+def wait_for_upload_queue_empty(
+    pageserver: NeonPageserver, tenant_id: TenantId, timeline_id: TimelineId
+):
+    ps_http = pageserver.http_client()
+    while True:
+        all_metrics = ps_http.get_metrics()
+        tl = all_metrics.query_all(
+            "pageserver_remote_timeline_client_calls_unfinished",
+            {
+                "tenant_id": str(tenant_id),
+                "timeline_id": str(timeline_id),
+            },
+        )
+        assert len(tl) > 0
+        log.info(f"upload queue for {tenant_id}/{timeline_id}: {tl}")
+        if all(m.value == 0 for m in tl):
+            return
+        time.sleep(0.2)

test_runner/fixtures/pageserver/http.py

@@ -1,545 +0,0 @@
-from __future__ import annotations
-
-import time
-from collections import defaultdict
-from dataclasses import dataclass
-from typing import Any, Dict, List, Optional, Tuple
-
-import requests
-
-from fixtures.log_helper import log
-from fixtures.metrics import Metrics, parse_metrics
-from fixtures.types import Lsn, TenantId, TimelineId
-from fixtures.utils import Fn
-
-
-class PageserverApiException(Exception):
-    def __init__(self, message, status_code: int):
-        super().__init__(message)
-        self.status_code = status_code
-
-
-@dataclass
-class InMemoryLayerInfo:
-    kind: str
-    lsn_start: str
-    lsn_end: Optional[str]
-
-    @classmethod
-    def from_json(cls, d: Dict[str, Any]) -> InMemoryLayerInfo:
-        return InMemoryLayerInfo(
-            kind=d["kind"],
-            lsn_start=d["lsn_start"],
-            lsn_end=d.get("lsn_end"),
-        )
-
-
-@dataclass(frozen=True)
-class HistoricLayerInfo:
-    kind: str
-    layer_file_name: str
-    layer_file_size: Optional[int]
-    lsn_start: str
-    lsn_end: Optional[str]
-    remote: bool
-
-    @classmethod
-    def from_json(cls, d: Dict[str, Any]) -> HistoricLayerInfo:
-        return HistoricLayerInfo(
-            kind=d["kind"],
-            layer_file_name=d["layer_file_name"],
-            layer_file_size=d.get("layer_file_size"),
-            lsn_start=d["lsn_start"],
-            lsn_end=d.get("lsn_end"),
-            remote=d["remote"],
-        )
-
-
-@dataclass
-class LayerMapInfo:
-    in_memory_layers: List[InMemoryLayerInfo]
-    historic_layers: List[HistoricLayerInfo]
-
-    @classmethod
-    def from_json(cls, d: Dict[str, Any]) -> LayerMapInfo:
-        info = LayerMapInfo(in_memory_layers=[], historic_layers=[])
-
-        json_in_memory_layers = d["in_memory_layers"]
-        assert isinstance(json_in_memory_layers, List)
-        for json_in_memory_layer in json_in_memory_layers:
-            info.in_memory_layers.append(InMemoryLayerInfo.from_json(json_in_memory_layer))
-
-        json_historic_layers = d["historic_layers"]
-        assert isinstance(json_historic_layers, List)
-        for json_historic_layer in json_historic_layers:
-            info.historic_layers.append(HistoricLayerInfo.from_json(json_historic_layer))
-
-        return info
-
-    def kind_count(self) -> Dict[str, int]:
-        counts: Dict[str, int] = defaultdict(int)
-        for inmem_layer in self.in_memory_layers:
-            counts[inmem_layer.kind] += 1
-        for hist_layer in self.historic_layers:
-            counts[hist_layer.kind] += 1
-        return counts
-
-
-@dataclass
-class TenantConfig:
-    tenant_specific_overrides: Dict[str, Any]
-    effective_config: Dict[str, Any]
-
-    @classmethod
-    def from_json(cls, d: Dict[str, Any]) -> TenantConfig:
-        return TenantConfig(
-            tenant_specific_overrides=d["tenant_specific_overrides"],
-            effective_config=d["effective_config"],
-        )
-
-
-class PageserverHttpClient(requests.Session):
-    def __init__(self, port: int, is_testing_enabled_or_skip: Fn, auth_token: Optional[str] = None):
-        super().__init__()
-        self.port = port
-        self.auth_token = auth_token
-        self.is_testing_enabled_or_skip = is_testing_enabled_or_skip
-
-        if auth_token is not None:
-            self.headers["Authorization"] = f"Bearer {auth_token}"
-
-    def verbose_error(self, res: requests.Response):
-        try:
-            res.raise_for_status()
-        except requests.RequestException as e:
-            try:
-                msg = res.json()["msg"]
-            except:  # noqa: E722
-                msg = ""
-            raise PageserverApiException(msg, res.status_code) from e
-
-    def check_status(self):
-        self.get(f"http://localhost:{self.port}/v1/status").raise_for_status()
-
-    def configure_failpoints(self, config_strings: Tuple[str, str] | List[Tuple[str, str]]):
-        self.is_testing_enabled_or_skip()
-
-        if isinstance(config_strings, tuple):
-            pairs = [config_strings]
-        else:
-            pairs = config_strings
-
-        log.info(f"Requesting config failpoints: {repr(pairs)}")
-
-        res = self.put(
-            f"http://localhost:{self.port}/v1/failpoints",
-            json=[{"name": name, "actions": actions} for name, actions in pairs],
-        )
-        log.info(f"Got failpoints request response code {res.status_code}")
-        self.verbose_error(res)
-        res_json = res.json()
-        assert res_json is None
-        return res_json
-
-    def tenant_list(self) -> List[Dict[Any, Any]]:
-        res = self.get(f"http://localhost:{self.port}/v1/tenant")
-        self.verbose_error(res)
-        res_json = res.json()
-        assert isinstance(res_json, list)
-        return res_json
-
-    def tenant_create(self, new_tenant_id: Optional[TenantId] = None) -> TenantId:
-        res = self.post(
-            f"http://localhost:{self.port}/v1/tenant",
-            json={
-                "new_tenant_id": str(new_tenant_id) if new_tenant_id else None,
-            },
-        )
-        self.verbose_error(res)
-        if res.status_code == 409:
-            raise Exception(f"could not create tenant: already exists for id {new_tenant_id}")
-        new_tenant_id = res.json()
-        assert isinstance(new_tenant_id, str)
-        return TenantId(new_tenant_id)
-
-    def tenant_attach(self, tenant_id: TenantId):
-        res = self.post(f"http://localhost:{self.port}/v1/tenant/{tenant_id}/attach")
-        self.verbose_error(res)
-
-    def tenant_detach(self, tenant_id: TenantId, detach_ignored=False):
-        params = {}
-        if detach_ignored:
-            params["detach_ignored"] = "true"
-
-        res = self.post(f"http://localhost:{self.port}/v1/tenant/{tenant_id}/detach", params=params)
-        self.verbose_error(res)
-
-    def tenant_load(self, tenant_id: TenantId):
-        res = self.post(f"http://localhost:{self.port}/v1/tenant/{tenant_id}/load")
-        self.verbose_error(res)
-
-    def tenant_ignore(self, tenant_id: TenantId):
-        res = self.post(f"http://localhost:{self.port}/v1/tenant/{tenant_id}/ignore")
-        self.verbose_error(res)
-
-    def tenant_status(self, tenant_id: TenantId) -> Dict[Any, Any]:
-        res = self.get(f"http://localhost:{self.port}/v1/tenant/{tenant_id}")
-        self.verbose_error(res)
-        res_json = res.json()
-        assert isinstance(res_json, dict)
-        return res_json
-
-    def tenant_config(self, tenant_id: TenantId) -> TenantConfig:
-        res = self.get(f"http://localhost:{self.port}/v1/tenant/{tenant_id}/config")
-        self.verbose_error(res)
-        return TenantConfig.from_json(res.json())
-
-    def set_tenant_config(self, tenant_id: TenantId, config: dict[str, Any]):
-        assert "tenant_id" not in config.keys()
-        res = self.put(
-            f"http://localhost:{self.port}/v1/tenant/config",
-            json={**config, "tenant_id": str(tenant_id)},
-        )
-        self.verbose_error(res)
-
-    def patch_tenant_config_client_side(
-        self,
-        tenant_id: TenantId,
-        inserts: Optional[Dict[str, Any]] = None,
-        removes: Optional[List[str]] = None,
-    ):
-        current = self.tenant_config(tenant_id).tenant_specific_overrides
-        if inserts is not None:
-            current.update(inserts)
-        if removes is not None:
-            for key in removes:
-                del current[key]
-        self.set_tenant_config(tenant_id, current)
-
-    def tenant_size(self, tenant_id: TenantId) -> int:
-        return self.tenant_size_and_modelinputs(tenant_id)[0]
-
-    def tenant_size_and_modelinputs(self, tenant_id: TenantId) -> Tuple[int, Dict[str, Any]]:
-        """
-        Returns the tenant size, together with the model inputs as the second tuple item.
-        """
-        res = self.get(f"http://localhost:{self.port}/v1/tenant/{tenant_id}/synthetic_size")
-        self.verbose_error(res)
-        res = res.json()
-        assert isinstance(res, dict)
-        assert TenantId(res["id"]) == tenant_id
-        size = res["size"]
-        assert type(size) == int
-        inputs = res["inputs"]
-        assert type(inputs) is dict
-        return (size, inputs)
-
-    def tenant_size_debug(self, tenant_id: TenantId) -> str:
-        """
-        Returns the tenant size debug info, as an HTML string
-        """
-        res = self.get(
-            f"http://localhost:{self.port}/v1/tenant/{tenant_id}/synthetic_size",
-            headers={"Accept": "text/html"},
-        )
-        return res.text
-
-    def timeline_list(
-        self,
-        tenant_id: TenantId,
-        include_non_incremental_logical_size: bool = False,
-        include_timeline_dir_layer_file_size_sum: bool = False,
-    ) -> List[Dict[str, Any]]:
-        params = {}
-        if include_non_incremental_logical_size:
-            params["include-non-incremental-logical-size"] = "true"
-        if include_timeline_dir_layer_file_size_sum:
-            params["include-timeline-dir-layer-file-size-sum"] = "true"
-
-        res = self.get(
-            f"http://localhost:{self.port}/v1/tenant/{tenant_id}/timeline", params=params
-        )
-        self.verbose_error(res)
-        res_json = res.json()
-        assert isinstance(res_json, list)
-        return res_json
-
-    def timeline_create(
-        self,
-        tenant_id: TenantId,
-        new_timeline_id: Optional[TimelineId] = None,
-        ancestor_timeline_id: Optional[TimelineId] = None,
-        ancestor_start_lsn: Optional[Lsn] = None,
-    ) -> Dict[Any, Any]:
-        res = self.post(
-            f"http://localhost:{self.port}/v1/tenant/{tenant_id}/timeline",
-            json={
-                "new_timeline_id": str(new_timeline_id) if new_timeline_id else None,
-                "ancestor_start_lsn": str(ancestor_start_lsn) if ancestor_start_lsn else None,
-                "ancestor_timeline_id": str(ancestor_timeline_id) if ancestor_timeline_id else None,
-            },
-        )
-        self.verbose_error(res)
-        if res.status_code == 409:
-            raise Exception(f"could not create timeline: already exists for id {new_timeline_id}")
-
-        res_json = res.json()
-        assert isinstance(res_json, dict)
-        return res_json
-
-    def timeline_detail(
-        self,
-        tenant_id: TenantId,
-        timeline_id: TimelineId,
-        include_non_incremental_logical_size: bool = False,
-        include_timeline_dir_layer_file_size_sum: bool = False,
-        **kwargs,
-    ) -> Dict[Any, Any]:
-        params = {}
-        if include_non_incremental_logical_size:
-            params["include-non-incremental-logical-size"] = "true"
-        if include_timeline_dir_layer_file_size_sum:
-            params["include-timeline-dir-layer-file-size-sum"] = "true"
-
-        res = self.get(
-            f"http://localhost:{self.port}/v1/tenant/{tenant_id}/timeline/{timeline_id}",
-            params=params,
-            **kwargs,
-        )
-        self.verbose_error(res)
-        res_json = res.json()
-        assert isinstance(res_json, dict)
-        return res_json
-
-    def timeline_delete(self, tenant_id: TenantId, timeline_id: TimelineId):
-        res = self.delete(
-            f"http://localhost:{self.port}/v1/tenant/{tenant_id}/timeline/{timeline_id}"
-        )
-        self.verbose_error(res)
-        res_json = res.json()
-        assert res_json is None
-
-    def timeline_gc(
-        self, tenant_id: TenantId, timeline_id: TimelineId, gc_horizon: Optional[int]
-    ) -> dict[str, Any]:
-        self.is_testing_enabled_or_skip()
-
-        log.info(
-            f"Requesting GC: tenant {tenant_id}, timeline {timeline_id}, gc_horizon {repr(gc_horizon)}"
-        )
-        res = self.put(
-            f"http://localhost:{self.port}/v1/tenant/{tenant_id}/timeline/{timeline_id}/do_gc",
-            json={"gc_horizon": gc_horizon},
-        )
-        log.info(f"Got GC request response code: {res.status_code}")
-        self.verbose_error(res)
-        res_json = res.json()
-        assert res_json is not None
-        assert isinstance(res_json, dict)
-        return res_json
-
-    def timeline_compact(self, tenant_id: TenantId, timeline_id: TimelineId):
-        self.is_testing_enabled_or_skip()
-
-        log.info(f"Requesting compact: tenant {tenant_id}, timeline {timeline_id}")
-        res = self.put(
-            f"http://localhost:{self.port}/v1/tenant/{tenant_id}/timeline/{timeline_id}/compact"
-        )
-        log.info(f"Got compact request response code: {res.status_code}")
-        self.verbose_error(res)
-        res_json = res.json()
-        assert res_json is None
-
-    def timeline_get_lsn_by_timestamp(
-        self, tenant_id: TenantId, timeline_id: TimelineId, timestamp
-    ):
-        log.info(
-            f"Requesting lsn by timestamp {timestamp}, tenant {tenant_id}, timeline {timeline_id}"
-        )
-        res = self.get(
-            f"http://localhost:{self.port}/v1/tenant/{tenant_id}/timeline/{timeline_id}/get_lsn_by_timestamp?timestamp={timestamp}",
-        )
-        self.verbose_error(res)
-        res_json = res.json()
-        return res_json
-
-    def timeline_checkpoint(self, tenant_id: TenantId, timeline_id: TimelineId):
-        self.is_testing_enabled_or_skip()
-
-        log.info(f"Requesting checkpoint: tenant {tenant_id}, timeline {timeline_id}")
-        res = self.put(
-            f"http://localhost:{self.port}/v1/tenant/{tenant_id}/timeline/{timeline_id}/checkpoint"
-        )
-        log.info(f"Got checkpoint request response code: {res.status_code}")
-        self.verbose_error(res)
-        res_json = res.json()
-        assert res_json is None
-
-    def timeline_spawn_download_remote_layers(
-        self,
-        tenant_id: TenantId,
-        timeline_id: TimelineId,
-        max_concurrent_downloads: int,
-    ) -> dict[str, Any]:
-        body = {
-            "max_concurrent_downloads": max_concurrent_downloads,
-        }
-        res = self.post(
-            f"http://localhost:{self.port}/v1/tenant/{tenant_id}/timeline/{timeline_id}/download_remote_layers",
-            json=body,
-        )
-        self.verbose_error(res)
-        res_json = res.json()
-        assert res_json is not None
-        assert isinstance(res_json, dict)
-        return res_json
-
-    def timeline_poll_download_remote_layers_status(
-        self,
-        tenant_id: TenantId,
-        timeline_id: TimelineId,
-        spawn_response: dict[str, Any],
-        poll_state=None,
-    ) -> None | dict[str, Any]:
-        res = self.get(
-            f"http://localhost:{self.port}/v1/tenant/{tenant_id}/timeline/{timeline_id}/download_remote_layers",
-        )
-        self.verbose_error(res)
-        res_json = res.json()
-        assert res_json is not None
-        assert isinstance(res_json, dict)
-
-        # assumption in this API client here is that nobody else spawns the task
-        assert res_json["task_id"] == spawn_response["task_id"]
-
-        if poll_state is None or res_json["state"] == poll_state:
-            return res_json
-        return None
-
-    def timeline_download_remote_layers(
-        self,
-        tenant_id: TenantId,
-        timeline_id: TimelineId,
-        max_concurrent_downloads: int,
-        errors_ok=False,
-        at_least_one_download=True,
-    ):
-        res = self.timeline_spawn_download_remote_layers(
-            tenant_id, timeline_id, max_concurrent_downloads
-        )
-        while True:
-            completed = self.timeline_poll_download_remote_layers_status(
-                tenant_id, timeline_id, res, poll_state="Completed"
-            )
-            if not completed:
-                time.sleep(0.1)
-                continue
-            if not errors_ok:
-                assert completed["failed_download_count"] == 0
-            if at_least_one_download:
-                assert completed["successful_download_count"] > 0
-            return completed
-
-    def get_metrics_str(self) -> str:
-        """You probably want to use get_metrics() instead."""
-        res = self.get(f"http://localhost:{self.port}/metrics")
-        self.verbose_error(res)
-        return res.text
-
-    def get_metrics(self) -> Metrics:
-        res = self.get_metrics_str()
-        return parse_metrics(res)
-
-    def get_timeline_metric(
-        self, tenant_id: TenantId, timeline_id: TimelineId, metric_name: str
-    ) -> float:
-        metrics = self.get_metrics()
-        return metrics.query_one(
-            metric_name,
-            filter={
-                "tenant_id": str(tenant_id),
-                "timeline_id": str(timeline_id),
-            },
-        ).value
-
-    def get_remote_timeline_client_metric(
-        self,
-        metric_name: str,
-        tenant_id: TenantId,
-        timeline_id: TimelineId,
-        file_kind: str,
-        op_kind: str,
-    ) -> Optional[float]:
-        metrics = self.get_metrics()
-        matches = metrics.query_all(
-            name=metric_name,
-            filter={
-                "tenant_id": str(tenant_id),
-                "timeline_id": str(timeline_id),
-                "file_kind": str(file_kind),
-                "op_kind": str(op_kind),
-            },
-        )
-        if len(matches) == 0:
-            value = None
-        elif len(matches) == 1:
-            value = matches[0].value
-            assert value is not None
-        else:
-            assert len(matches) < 2, "above filter should uniquely identify metric"
-        return value
-
-    def get_metric_value(
-        self, name: str, filter: Optional[Dict[str, str]] = None
-    ) -> Optional[float]:
-        metrics = self.get_metrics()
-        results = metrics.query_all(name, filter=filter)
-        if not results:
-            log.info(f'could not find metric "{name}"')
-            return None
-        assert len(results) == 1, f"metric {name} with given filters is not unique, got: {results}"
-        return results[0].value
-
-    def layer_map_info(
-        self,
-        tenant_id: TenantId,
-        timeline_id: TimelineId,
-    ) -> LayerMapInfo:
-        res = self.get(
-            f"http://localhost:{self.port}/v1/tenant/{tenant_id}/timeline/{timeline_id}/layer/",
-        )
-        self.verbose_error(res)
-        return LayerMapInfo.from_json(res.json())
-
-    def download_layer(self, tenant_id: TenantId, timeline_id: TimelineId, layer_name: str):
-        res = self.get(
-            f"http://localhost:{self.port}/v1/tenant/{tenant_id}/timeline/{timeline_id}/layer/{layer_name}",
-        )
-        self.verbose_error(res)
-
-        assert res.status_code == 200
-
-    def evict_layer(self, tenant_id: TenantId, timeline_id: TimelineId, layer_name: str):
-        res = self.delete(
-            f"http://localhost:{self.port}/v1/tenant/{tenant_id}/timeline/{timeline_id}/layer/{layer_name}",
-        )
-        self.verbose_error(res)
-
-        assert res.status_code == 200
-
-    def evict_all_layers(self, tenant_id: TenantId, timeline_id: TimelineId):
-        info = self.layer_map_info(tenant_id, timeline_id)
-        for layer in info.historic_layers:
-            self.evict_layer(tenant_id, timeline_id, layer.layer_file_name)
-
-    def disk_usage_eviction_run(self, request: dict[str, Any]):
-        res = self.put(
-            f"http://localhost:{self.port}/v1/disk_usage_eviction/run",
-            json=request,
-        )
-        self.verbose_error(res)
-        return res.json()
-
-    def tenant_break(self, tenant_id: TenantId):
-        res = self.put(f"http://localhost:{self.port}/v1/tenant/{tenant_id}/break")
-        self.verbose_error(res)

test_runner/fixtures/pageserver/utils.py

@@ -1,157 +0,0 @@
-import time
-from typing import Optional
-
-from fixtures.log_helper import log
-from fixtures.pageserver.http import PageserverHttpClient
-from fixtures.types import Lsn, TenantId, TimelineId
-
-
-def assert_tenant_state(
-    pageserver_http: PageserverHttpClient,
-    tenant: TenantId,
-    expected_state: str,
-    message: Optional[str] = None,
-):
-    tenant_status = pageserver_http.tenant_status(tenant)
-    log.info(f"tenant_status: {tenant_status}")
-    assert tenant_status["state"]["slug"] == expected_state, message or tenant_status
-
-
-def tenant_exists(pageserver_http: PageserverHttpClient, tenant_id: TenantId):
-    tenants = pageserver_http.tenant_list()
-    matching = [t for t in tenants if TenantId(t["id"]) == tenant_id]
-    assert len(matching) < 2
-    if len(matching) == 0:
-        return None
-    return matching[0]
-
-
-def remote_consistent_lsn(
-    pageserver_http: PageserverHttpClient, tenant: TenantId, timeline: TimelineId
-) -> Lsn:
-    detail = pageserver_http.timeline_detail(tenant, timeline)
-
-    if detail["remote_consistent_lsn"] is None:
-        # No remote information at all. This happens right after creating
-        # a timeline, before any part of it has been uploaded to remote
-        # storage yet.
-        return Lsn(0)
-    else:
-        lsn_str = detail["remote_consistent_lsn"]
-        assert isinstance(lsn_str, str)
-        return Lsn(lsn_str)
-
-
-def wait_for_upload(
-    pageserver_http: PageserverHttpClient,
-    tenant: TenantId,
-    timeline: TimelineId,
-    lsn: Lsn,
-):
-    """waits for local timeline upload up to specified lsn"""
-    for i in range(20):
-        current_lsn = remote_consistent_lsn(pageserver_http, tenant, timeline)
-        if current_lsn >= lsn:
-            log.info("wait finished")
-            return
-        log.info(
-            "waiting for remote_consistent_lsn to reach {}, now {}, iteration {}".format(
-                lsn, current_lsn, i + 1
-            )
-        )
-        time.sleep(1)
-    raise Exception(
-        "timed out while waiting for remote_consistent_lsn to reach {}, was {}".format(
-            lsn, current_lsn
-        )
-    )
-
-
-def wait_until_tenant_state(
-    pageserver_http: PageserverHttpClient,
-    tenant_id: TenantId,
-    expected_state: str,
-    iterations: int,
-    period: float = 1.0,
-) -> bool:
-    """
-    Does not use `wait_until` for debugging purposes
-    """
-    for _ in range(iterations):
-        try:
-            tenant = pageserver_http.tenant_status(tenant_id=tenant_id)
-            log.debug(f"Tenant {tenant_id} data: {tenant}")
-            if tenant["state"]["slug"] == expected_state:
-                return True
-        except Exception as e:
-            log.debug(f"Tenant {tenant_id} state retrieval failure: {e}")
-
-        time.sleep(period)
-
-    raise Exception(f"Tenant {tenant_id} did not become {expected_state} in {iterations} seconds")
-
-
-def wait_until_tenant_active(
-    pageserver_http: PageserverHttpClient,
-    tenant_id: TenantId,
-    iterations: int = 30,
-    period: float = 1.0,
-):
-    wait_until_tenant_state(
-        pageserver_http,
-        tenant_id,
-        expected_state="Active",
-        iterations=iterations,
-        period=period,
-    )
-
-
-def last_record_lsn(
-    pageserver_http_client: PageserverHttpClient, tenant: TenantId, timeline: TimelineId
-) -> Lsn:
-    detail = pageserver_http_client.timeline_detail(tenant, timeline)
-
-    lsn_str = detail["last_record_lsn"]
-    assert isinstance(lsn_str, str)
-    return Lsn(lsn_str)
-
-
-def wait_for_last_record_lsn(
-    pageserver_http: PageserverHttpClient,
-    tenant: TenantId,
-    timeline: TimelineId,
-    lsn: Lsn,
-) -> Lsn:
-    """waits for pageserver to catch up to a certain lsn, returns the last observed lsn."""
-    for i in range(10):
-        current_lsn = last_record_lsn(pageserver_http, tenant, timeline)
-        if current_lsn >= lsn:
-            return current_lsn
-        log.info(
-            "waiting for last_record_lsn to reach {}, now {}, iteration {}".format(
-                lsn, current_lsn, i + 1
-            )
-        )
-        time.sleep(1)
-    raise Exception(
-        "timed out while waiting for last_record_lsn to reach {}, was {}".format(lsn, current_lsn)
-    )
-
-
-def wait_for_upload_queue_empty(
-    pageserver_http: PageserverHttpClient, tenant_id: TenantId, timeline_id: TimelineId
-):
-    while True:
-        all_metrics = pageserver_http.get_metrics()
-        tl = all_metrics.query_all(
-            "pageserver_remote_timeline_client_calls_unfinished",
-            {
-                "tenant_id": str(tenant_id),
-                "timeline_id": str(timeline_id),
-            },
-        )
-        assert len(tl) > 0
-        log.info(f"upload queue for {tenant_id}/{timeline_id}: {tl}")
-        if all(m.value == 0 for m in tl):
-            return
-        time.sleep(0.2)

test_runner/fixtures/utils.py

@@ -1,5 +1,4 @@
 import contextlib
-import json
 import os
 import re
 import subprocess
@@ -7,7 +6,6 @@ import tarfile
 import time
 from pathlib import Path
 from typing import Any, Callable, Dict, List, Tuple, TypeVar
-from urllib.parse import urlencode
 
 import allure
 from psycopg2.extensions import cursor
@@ -186,46 +184,6 @@ def allure_attach_from_dir(dir: Path):
             allure.attach.file(source, name, attachment_type, extension)
 
 
-DATASOURCE_ID = "xHHYY0dVz"
-
-
-def allure_add_grafana_links(host: str, start_ms: int, end_ms: int):
-    """Add links to server logs in Grafana to Allure report"""
-    # We expect host to be in format like ep-divine-night-159320.us-east-2.aws.neon.build
-    endpoint_id, region_id, _ = host.split(".", 2)
-
-    expressions = {
-        "compute logs": f'{{app="compute-node-{endpoint_id}", neon_region="{region_id}"}}',
-        "k8s events": f'{{job="integrations/kubernetes/eventhandler"}} |~ "name=compute-node-{endpoint_id}-"',
-        "console logs": f'{{neon_service="console", neon_region="{region_id}"}} | json | endpoint_id = "{endpoint_id}"',
-        "proxy logs": f'{{neon_service="proxy-scram", neon_region="{region_id}"}}',
-    }
-
-    params: Dict[str, Any] = {
-        "datasource": DATASOURCE_ID,
-        "queries": [
-            {
-                "expr": "<PUT AN EXPRESSION HERE>",
-                "refId": "A",
-                "datasource": {"type": "loki", "uid": DATASOURCE_ID},
-                "editorMode": "code",
-                "queryType": "range",
-            }
-        ],
-        "range": {
-            "from": str(start_ms),
-            "to": str(end_ms),
-        },
-    }
-    for name, expr in expressions.items():
-        params["queries"][0]["expr"] = expr
-        query_string = urlencode({"orgId": 1, "left": json.dumps(params)})
-        link = f"https://neonprod.grafana.net/explore?{query_string}"
-
-        allure.dynamic.link(link, name=name)
-        log.info(f"{name}: {link}")
-
-
 def start_in_background(
     command: list[str], cwd: Path, log_file_name: str, is_started: Fn
 ) -> subprocess.Popen[bytes]:
@@ -278,19 +236,3 @@ def wait_until(number_of_iterations: int, interval: float, func: Fn):
             continue
         return res
     raise Exception("timed out while waiting for %s" % func) from last_exception
-
-
-def wait_while(number_of_iterations: int, interval: float, func):
-    """
-    Wait until 'func' returns false, or throws an exception.
-    """
-    for i in range(number_of_iterations):
-        try:
-            if not func():
-                return
-            log.info("waiting for %s iteration %s failed", func, i + 1)
-            time.sleep(interval)
-            continue
-        except Exception:
-            return
-    raise Exception("timed out while waiting for %s" % func)

test_runner/performance/test_branch_creation.py

@@ -10,7 +10,7 @@ import pytest
 from fixtures.benchmark_fixture import MetricReport
 from fixtures.compare_fixtures import NeonCompare
 from fixtures.log_helper import log
-from fixtures.pageserver.utils import wait_for_last_record_lsn
+from fixtures.neon_fixtures import wait_for_last_record_lsn
 from fixtures.types import Lsn
 
 

test_runner/regress/test_auth.py

@@ -1,8 +1,7 @@
 from contextlib import closing
 
 import pytest
-from fixtures.neon_fixtures import NeonEnvBuilder, PgProtocol
-from fixtures.pageserver.http import PageserverApiException
+from fixtures.neon_fixtures import NeonEnvBuilder, PageserverApiException, PgProtocol
 from fixtures.types import TenantId
 
 

test_runner/regress/test_compatibility.py

@@ -11,11 +11,12 @@ from fixtures.log_helper import log
 from fixtures.neon_fixtures import (
     NeonCli,
     NeonEnvBuilder,
+    PageserverHttpClient,
     PgBin,
     PortDistributor,
+    wait_for_last_record_lsn,
+    wait_for_upload,
 )
-from fixtures.pageserver.http import PageserverHttpClient
-from fixtures.pageserver.utils import wait_for_last_record_lsn, wait_for_upload
 from fixtures.types import Lsn
 from pytest import FixtureRequest
 
@@ -34,15 +35,9 @@ from pytest import FixtureRequest
 # - check_neon_works performs the test itself, feel free to add more checks there.
 #
 
-check_ondisk_data_compatibility_if_enabled = pytest.mark.skipif(
-    os.environ.get("CHECK_ONDISK_DATA_COMPATIBILITY") is None,
-    reason="CHECK_ONDISK_DATA_COMPATIBILITY env is not set",
-)
-
 
 # Note: if renaming this test, don't forget to update a reference to it in a workflow file:
 # "Upload compatibility snapshot" step in .github/actions/run-python-test-set/action.yml
-@check_ondisk_data_compatibility_if_enabled
 @pytest.mark.xdist_group("compatibility")
 @pytest.mark.order(before="test_forward_compatibility")
 def test_create_snapshot(neon_env_builder: NeonEnvBuilder, pg_bin: PgBin, test_output_dir: Path):
@@ -89,7 +84,6 @@ def test_create_snapshot(neon_env_builder: NeonEnvBuilder, pg_bin: PgBin, test_o
     # Directory `test_output_dir / "compatibility_snapshot_pg14"` is uploaded to S3 in a workflow, keep the name in sync with it
 
 
-@check_ondisk_data_compatibility_if_enabled
 @pytest.mark.xdist_group("compatibility")
 @pytest.mark.order(after="test_create_snapshot")
 def test_backward_compatibility(
@@ -147,7 +141,6 @@ def test_backward_compatibility(
     ), "Breaking changes are allowed by ALLOW_BACKWARD_COMPATIBILITY_BREAKAGE, but the test has passed without any breakage"
 
 
-@check_ondisk_data_compatibility_if_enabled
 @pytest.mark.xdist_group("compatibility")
 @pytest.mark.order(after="test_create_snapshot")
 def test_forward_compatibility(

test_runner/regress/test_disk_usage_eviction.py

@@ -11,14 +11,14 @@ from fixtures.neon_fixtures import (
     LocalFsStorage,
     NeonEnv,
     NeonEnvBuilder,
+    PageserverHttpClient,
     PgBin,
     RemoteStorageKind,
     wait_for_last_flush_lsn,
+    wait_for_upload_queue_empty,
+    wait_until,
 )
-from fixtures.pageserver.http import PageserverHttpClient
-from fixtures.pageserver.utils import wait_for_upload_queue_empty
 from fixtures.types import Lsn, TenantId, TimelineId
-from fixtures.utils import wait_until
 
 GLOBAL_LRU_LOG_LINE = "tenant_min_resident_size-respecting LRU would not relieve pressure, evicting more following global LRU policy"
 
@@ -138,7 +138,7 @@ def eviction_env(request, neon_env_builder: NeonEnvBuilder, pg_bin: PgBin) -> Ev
     # remove the initial tenant
     ## why wait for upload queue? => https://github.com/neondatabase/neon/issues/3865
     assert env.initial_timeline
-    wait_for_upload_queue_empty(pageserver_http, env.initial_tenant, env.initial_timeline)
+    wait_for_upload_queue_empty(env.pageserver, env.initial_tenant, env.initial_timeline)
     pageserver_http.tenant_detach(env.initial_tenant)
     assert isinstance(env.remote_storage, LocalFsStorage)
     tenant_remote_storage = env.remote_storage.root / "tenants" / str(env.initial_tenant)
@@ -182,7 +182,7 @@ def eviction_env(request, neon_env_builder: NeonEnvBuilder, pg_bin: PgBin) -> Ev
     # after stopping the safekeepers, we know that no new WAL will be coming in
     for tenant_id, timeline_id in timelines:
         pageserver_http.timeline_checkpoint(tenant_id, timeline_id)
-        wait_for_upload_queue_empty(pageserver_http, tenant_id, timeline_id)
+        wait_for_upload_queue_empty(env.pageserver, tenant_id, timeline_id)
         tl_info = pageserver_http.timeline_detail(tenant_id, timeline_id)
         assert tl_info["last_record_lsn"] == tl_info["disk_consistent_lsn"]
         assert tl_info["disk_consistent_lsn"] == tl_info["remote_consistent_lsn"]

test_runner/regress/test_import.py

@@ -13,8 +13,9 @@ from fixtures.neon_fixtures import (
     NeonEnv,
     NeonEnvBuilder,
     PgBin,
+    wait_for_last_record_lsn,
+    wait_for_upload,
 )
-from fixtures.pageserver.utils import wait_for_last_record_lsn, wait_for_upload
 from fixtures.types import Lsn, TenantId, TimelineId
 from fixtures.utils import subprocess_capture
 

test_runner/regress/test_layer_eviction.py

@@ -6,9 +6,10 @@ from fixtures.neon_fixtures import (
     NeonEnvBuilder,
     RemoteStorageKind,
     wait_for_last_flush_lsn,
+    wait_for_last_record_lsn,
     wait_for_sk_commit_lsn_to_reach_remote_storage,
+    wait_for_upload,
 )
-from fixtures.pageserver.utils import wait_for_last_record_lsn, wait_for_upload
 from fixtures.types import Lsn, TenantId, TimelineId
 from fixtures.utils import query_scalar
 

test_runner/regress/test_neon_cli.py

@@ -5,8 +5,8 @@ from fixtures.neon_fixtures import (
     DEFAULT_BRANCH_NAME,
     NeonEnv,
     NeonEnvBuilder,
+    PageserverHttpClient,
 )
-from fixtures.pageserver.http import PageserverHttpClient
 from fixtures.types import TenantId, TimelineId
 
 

test_runner/regress/test_normal_work.py

@@ -1,7 +1,6 @@
 import pytest
 from fixtures.log_helper import log
-from fixtures.neon_fixtures import NeonEnv, NeonEnvBuilder
-from fixtures.pageserver.http import PageserverHttpClient
+from fixtures.neon_fixtures import NeonEnv, NeonEnvBuilder, PageserverHttpClient
 
 
 def check_tenant(env: NeonEnv, pageserver_http: PageserverHttpClient):

test_runner/regress/test_ondemand_download.py

@@ -10,20 +10,20 @@ import pytest
 from fixtures.log_helper import log
 from fixtures.neon_fixtures import (
     NeonEnvBuilder,
+    PageserverApiException,
+    PageserverHttpClient,
     RemoteStorageKind,
+    assert_tenant_status,
     available_remote_storages,
     wait_for_last_flush_lsn,
-    wait_for_sk_commit_lsn_to_reach_remote_storage,
-)
-from fixtures.pageserver.http import PageserverApiException, PageserverHttpClient
-from fixtures.pageserver.utils import (
-    assert_tenant_state,
     wait_for_last_record_lsn,
+    wait_for_sk_commit_lsn_to_reach_remote_storage,
     wait_for_upload,
+    wait_until,
     wait_until_tenant_state,
 )
 from fixtures.types import Lsn
-from fixtures.utils import query_scalar, wait_until
+from fixtures.utils import query_scalar
 
 
 def get_num_downloaded_layers(client: PageserverHttpClient, tenant_id, timeline_id):
@@ -239,7 +239,7 @@ def test_ondemand_download_timetravel(
     ##### Second start, restore the data and ensure it's the same
     env.pageserver.start()
 
-    wait_until(10, 0.2, lambda: assert_tenant_state(client, tenant_id, "Active"))
+    wait_until(10, 0.2, lambda: assert_tenant_status(client, tenant_id, "Active"))
 
     # The current_physical_size reports the sum of layers loaded in the layer
     # map, regardless of where the layer files are located. So even though we
@@ -392,7 +392,7 @@ def test_download_remote_layers_api(
         ]
     )
 
-    wait_until(10, 0.2, lambda: assert_tenant_state(client, tenant_id, "Active"))
+    wait_until(10, 0.2, lambda: assert_tenant_status(client, tenant_id, "Active"))
 
     ###### Phase 1: exercise download error code path
     assert (

test_runner/regress/test_pageserver_api.py

@@ -6,8 +6,8 @@ from fixtures.neon_fixtures import (
     DEFAULT_BRANCH_NAME,
     NeonEnv,
     NeonEnvBuilder,
+    PageserverHttpClient,
 )
-from fixtures.pageserver.http import PageserverHttpClient
 from fixtures.types import Lsn, TenantId, TimelineId
 from fixtures.utils import wait_until
 

test_runner/regress/test_pageserver_restart.py

@@ -59,7 +59,7 @@ def test_pageserver_restart(neon_env_builder: NeonEnvBuilder):
     client = env.pageserver.http_client()
     tenant_status = client.tenant_status(env.initial_tenant)
     log.info("Tenant status : %s", tenant_status)
-    assert tenant_status["state"]["slug"] == "Loading"
+    assert tenant_status["state"] == "Loading"
 
     # Try to read. This waits until the loading finishes, and then return normally.
     cur.execute("SELECT count(*) FROM foo")

test_runner/regress/test_proxy.py

@@ -1,5 +1,3 @@
-import subprocess
-
 import psycopg2
 import pytest
 from fixtures.neon_fixtures import PSQL, NeonProxy, VanillaPostgres
@@ -136,19 +134,3 @@ def test_forward_params_to_client(static_proxy: NeonProxy):
             for name, value in cur.fetchall():
                 # Check that proxy has forwarded this parameter.
                 assert conn.get_parameter_status(name) == value
-
-
-@pytest.mark.timeout(5)
-def test_close_on_connections_exit(static_proxy: NeonProxy):
-    # Open two connections, send SIGTERM, then ensure that proxy doesn't exit
-    # until after connections close.
-    with static_proxy.connect(options="project=irrelevant"), static_proxy.connect(
-        options="project=irrelevant"
-    ):
-        static_proxy.terminate()
-        with pytest.raises(subprocess.TimeoutExpired):
-            static_proxy.wait_for_exit(timeout=2)
-        # Ensure we don't accept any more connections
-        with pytest.raises(psycopg2.OperationalError):
-            static_proxy.connect(options="project=irrelevant")
-    static_proxy.wait_for_exit()

test_runner/regress/test_read_trace.py

@@ -1,7 +1,6 @@
 from contextlib import closing
 
-from fixtures.neon_fixtures import NeonEnvBuilder
-from fixtures.pageserver.utils import wait_for_last_record_lsn
+from fixtures.neon_fixtures import NeonEnvBuilder, wait_for_last_record_lsn
 from fixtures.types import Lsn, TenantId, TimelineId
 from fixtures.utils import query_scalar
 

test_runner/regress/test_readonly_node.py

@@ -1,7 +1,6 @@
 import pytest
 from fixtures.log_helper import log
-from fixtures.neon_fixtures import NeonEnv
-from fixtures.pageserver.utils import wait_for_last_record_lsn
+from fixtures.neon_fixtures import NeonEnv, wait_for_last_record_lsn
 from fixtures.types import Lsn
 from fixtures.utils import query_scalar
 

test_runner/regress/test_remote_storage.py

@@ -13,15 +13,13 @@ from fixtures.log_helper import log
 from fixtures.neon_fixtures import (
     LocalFsStorage,
     NeonEnvBuilder,
+    PageserverApiException,
+    PageserverHttpClient,
     RemoteStorageKind,
     available_remote_storages,
     wait_for_last_flush_lsn,
-)
-from fixtures.pageserver.http import PageserverApiException, PageserverHttpClient
-from fixtures.pageserver.utils import (
     wait_for_last_record_lsn,
     wait_for_upload,
-    wait_until_tenant_active,
     wait_until_tenant_state,
 )
 from fixtures.types import Lsn, TenantId, TimelineId
@@ -174,10 +172,15 @@ def test_remote_storage_backup_and_restore(
         client.tenant_attach(tenant_id)
     log.info("waiting for tenant to become active. this should be quick with on-demand download")
 
-    wait_until_tenant_active(
-        pageserver_http=client,
-        tenant_id=tenant_id,
-        iterations=5,
+    def tenant_active():
+        all_states = client.tenant_list()
+        [tenant] = [t for t in all_states if TenantId(t["id"]) == tenant_id]
+        assert tenant["state"] == "Active"
+
+    wait_until(
+        number_of_iterations=5,
+        interval=1,
+        func=tenant_active,
     )
 
     detail = client.timeline_detail(tenant_id, timeline_id)
@@ -354,7 +357,12 @@ def test_remote_storage_upload_queue_retries(
 
     client.tenant_attach(tenant_id)
 
-    wait_until_tenant_active(client, tenant_id)
+    def tenant_active():
+        all_states = client.tenant_list()
+        [tenant] = [t for t in all_states if TenantId(t["id"]) == tenant_id]
+        assert tenant["state"] == "Active"
+
+    wait_until(30, 1, tenant_active)
 
     log.info("restarting postgres to validate")
     endpoint = env.endpoints.create_start("main", tenant_id=tenant_id)
@@ -489,7 +497,12 @@ def test_remote_timeline_client_calls_started_metric(
 
     client.tenant_attach(tenant_id)
 
-    wait_until_tenant_active(client, tenant_id)
+    def tenant_active():
+        all_states = client.tenant_list()
+        [tenant] = [t for t in all_states if TenantId(t["id"]) == tenant_id]
+        assert tenant["state"] == "Active"
+
+    wait_until(30, 1, tenant_active)
 
     log.info("restarting postgres to validate")
     endpoint = env.endpoints.create_start("main", tenant_id=tenant_id)

test_runner/regress/test_tenant_conf.py

@@ -6,8 +6,9 @@ from fixtures.neon_fixtures import (
     LocalFsStorage,
     NeonEnvBuilder,
     RemoteStorageKind,
+    assert_tenant_status,
+    wait_for_upload,
 )
-from fixtures.pageserver.utils import assert_tenant_state, wait_for_upload
 from fixtures.types import Lsn
 from fixtures.utils import wait_until
 
@@ -274,7 +275,7 @@ def test_creating_tenant_conf_after_attach(neon_env_builder: NeonEnvBuilder):
     wait_until(
         number_of_iterations=5,
         interval=1,
-        func=lambda: assert_tenant_state(http_client, tenant_id, "Active"),
+        func=lambda: assert_tenant_status(http_client, tenant_id, "Active"),
     )
 
     env.neon_cli.config_tenant(tenant_id, {"gc_horizon": "1000000"})

test_runner/regress/test_tenant_detach.py

@@ -10,17 +10,17 @@ from fixtures.neon_fixtures import (
     Endpoint,
     NeonEnv,
     NeonEnvBuilder,
+    PageserverApiException,
+    PageserverHttpClient,
     RemoteStorageKind,
     available_remote_storages,
-)
-from fixtures.pageserver.http import PageserverApiException, PageserverHttpClient
-from fixtures.pageserver.utils import (
     wait_for_last_record_lsn,
     wait_for_upload,
+    wait_until,
     wait_until_tenant_state,
 )
 from fixtures.types import Lsn, TenantId, TimelineId
-from fixtures.utils import query_scalar, wait_until
+from fixtures.utils import query_scalar
 
 
 def do_gc_target(