allow pushing fake SystemTimes for refresh_gc_info

this should theoretically make the situation reproducable by first:

0. configure local pageserver to never run gc or compaction, allow
   access to tenant via local storage
1. `curl -X POST localhost:9898/v1/tenant/$tid/attach`
2. `curl -X PUT --data "{ \"tenant_id\": \"$tid\", \"gc_horizon\": 87772208 }" --header 'content-type: application/json' localhost:9898/v1/tenant/config`
3. `curl -X POST localhost:9898/add_forced_now?now=2023-01-24T04:58:17.319972Z`
4. `curl -X POST localhost:9898/add_forced_now?now=2023-01-24T04:59:32.436Z`
5. `curl -X PUT --data '{}' --header 'content-type: application/json' localhost:9898/v1/tenant/$tid/timeline/$ttid/do_gc`
    - this now uses now from step 3
6. `curl localhost:9898/v1/tenant/$tid/size`
    - this now uses now from step 4

this comes quite close, but doesn't end up downloading the same file.

.config/hakari.toml

@@ -4,7 +4,7 @@
 hakari-package = "workspace_hack"
 
 # Format for `workspace-hack = ...` lines in other Cargo.tomls. Requires cargo-hakari 0.9.8 or above.
-dep-format-version = "2"
+dep-format-version = "3"
 
 # Setting workspace.resolver = "2" in the root Cargo.toml is HIGHLY recommended.
 # Hakari works much better with the new feature resolver.

.github/PULL_REQUEST_TEMPLATE/pull_request_template.md

@@ -0,0 +1,10 @@
+## Describe your changes
+
+## Issue ticket number and link
+
+## Checklist before requesting a review
+- [ ] I have performed a self-review of my code.
+- [ ] If it is a core feature, I have added thorough tests.
+- [ ] Do we need to implement analytics? if so did you add the relevant metrics to the dashboard?
+- [ ] If this PR requires public announcement, mark it with /release-notes label and add several sentences in this section.
+

.github/ansible/prod.ap-southeast-1.hosts.yaml

@@ -6,6 +6,8 @@ storage:
     broker_endpoint: http://storage-broker-lb.epsilon.ap-southeast-1.internal.aws.neon.tech:50051
     pageserver_config_stub:
       pg_distrib_dir: /usr/local
+      metric_collection_endpoint: http://console-release.local/billing/api/v1/usage_events
+      metric_collection_interval: 10min
       remote_storage:
         bucket_name: "{{ bucket_name }}"
         bucket_region: "{{ bucket_region }}"

.github/ansible/prod.eu-central-1.hosts.yaml

@@ -6,6 +6,8 @@ storage:
     broker_endpoint: http://storage-broker-lb.gamma.eu-central-1.internal.aws.neon.tech:50051
     pageserver_config_stub:
       pg_distrib_dir: /usr/local
+      metric_collection_endpoint: http://console-release.local/billing/api/v1/usage_events
+      metric_collection_interval: 10min
       remote_storage:
         bucket_name: "{{ bucket_name }}"
         bucket_region: "{{ bucket_region }}"

.github/ansible/prod.us-east-2.hosts.yaml

@@ -6,6 +6,8 @@ storage:
     broker_endpoint: http://storage-broker-lb.delta.us-east-2.internal.aws.neon.tech:50051
     pageserver_config_stub:
       pg_distrib_dir: /usr/local
+      metric_collection_endpoint: http://console-release.local/billing/api/v1/usage_events
+      metric_collection_interval: 10min
       remote_storage:
         bucket_name: "{{ bucket_name }}"
         bucket_region: "{{ bucket_region }}"
@@ -34,4 +36,4 @@ storage:
           ansible_host:  i-06d113fb73bfddeb0
         safekeeper-2.us-east-2.aws.neon.tech:
           ansible_host:  i-09f66c8e04afff2e8
-          
+

.github/ansible/prod.us-west-2.hosts.yaml

@@ -6,6 +6,8 @@ storage:
     broker_endpoint: http://storage-broker-lb.eta.us-west-2.internal.aws.neon.tech:50051
     pageserver_config_stub:
       pg_distrib_dir: /usr/local
+      metric_collection_endpoint: http://console-release.local/billing/api/v1/usage_events
+      metric_collection_interval: 10min
       remote_storage:
         bucket_name: "{{ bucket_name }}"
         bucket_region: "{{ bucket_region }}"
@@ -25,6 +27,8 @@ storage:
           ansible_host: i-0d9f6dfae0e1c780d 
         pageserver-1.us-west-2.aws.neon.tech:
           ansible_host: i-0c834be1dddba8b3f
+        pageserver-2.us-west-2.aws.neon.tech:
+          ansible_host: i-051642d372c0a4f32
 
     safekeepers:
       hosts:

.github/ansible/production.hosts.yaml

@@ -7,6 +7,8 @@ storage:
     broker_endpoint: http://storage-broker.prod.local:50051
     pageserver_config_stub:
       pg_distrib_dir: /usr/local
+      metric_collection_endpoint: http://console-release.local/billing/api/v1/usage_events
+      metric_collection_interval: 10min
       remote_storage:
         bucket_name: "{{ bucket_name }}"
         bucket_region: "{{ bucket_region }}"

.github/ansible/staging.eu-west-1.hosts.yaml

@@ -18,7 +18,7 @@ storage:
     ansible_aws_ssm_region: eu-west-1
     ansible_aws_ssm_bucket_name: neon-dev-storage-eu-west-1
     console_region_id: aws-eu-west-1
-    sentry_environment: development
+    sentry_environment: staging
 
   children:
     pageservers:

.github/ansible/staging.us-east-2.hosts.yaml

@@ -18,7 +18,7 @@ storage:
     ansible_aws_ssm_region: us-east-2
     ansible_aws_ssm_bucket_name: neon-staging-storage-us-east-2
     console_region_id: aws-us-east-2
-    sentry_environment: development
+    sentry_environment: staging
 
   children:
     pageservers:
@@ -27,6 +27,10 @@ storage:
           ansible_host: i-0c3e70929edb5d691
         pageserver-1.us-east-2.aws.neon.build:
           ansible_host: i-0565a8b4008aa3f40
+        pageserver-2.us-east-2.aws.neon.build:
+          ansible_host: i-01e31cdf7e970586a
+        pageserver-3.us-east-2.aws.neon.build:
+          ansible_host: i-0602a0291365ef7cc
 
     safekeepers:
       hosts:

.github/helm-values/dev-eu-west-1-zeta.neon-proxy-scram.yaml

@@ -8,7 +8,10 @@ settings:
   authBackend: "console"
   authEndpoint: "http://console-staging.local/management/api/v2"
   domain: "*.eu-west-1.aws.neon.build"
-  sentryEnvironment: "development"
+  sentryEnvironment: "staging"
+  wssPort: 8443
+  metricCollectionEndpoint: "http://console-staging.local/billing/api/v1/usage_events"
+  metricCollectionInterval: "1min"
 
 # -- Additional labels for neon-proxy pods
 podLabels:
@@ -23,6 +26,7 @@ exposedService:
     service.beta.kubernetes.io/aws-load-balancer-nlb-target-type: ip
     service.beta.kubernetes.io/aws-load-balancer-scheme: internet-facing
     external-dns.alpha.kubernetes.io/hostname: eu-west-1.aws.neon.build
+  httpsPort: 443
 
 #metrics:
 #  enabled: true

.github/helm-values/dev-eu-west-1-zeta.neon-storage-broker.yaml

@@ -49,4 +49,4 @@ extraManifests:
           - "{{ .Release.Namespace }}"
 
 settings:
-  sentryEnvironment: "development"
+  sentryEnvironment: "staging"

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

@@ -8,7 +8,9 @@ settings:
   authBackend: "link"
   authEndpoint: "https://console.stage.neon.tech/authenticate_proxy_request/"
   uri: "https://console.stage.neon.tech/psql_session/"
-  sentryEnvironment: "development"
+  sentryEnvironment: "staging"
+  metricCollectionEndpoint: "http://console-staging.local/billing/api/v1/usage_events"
+  metricCollectionInterval: "1min"
 
 # -- Additional labels for neon-proxy-link pods
 podLabels:

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

@@ -8,7 +8,10 @@ settings:
   authBackend: "console"
   authEndpoint: "http://console-staging.local/management/api/v2"
   domain: "*.cloud.stage.neon.tech"
-  sentryEnvironment: "development"
+  sentryEnvironment: "staging"
+  wssPort: 8443
+  metricCollectionEndpoint: "http://console-staging.local/billing/api/v1/usage_events"
+  metricCollectionInterval: "1min"
 
 # -- Additional labels for neon-proxy pods
 podLabels:
@@ -23,6 +26,7 @@ exposedService:
     service.beta.kubernetes.io/aws-load-balancer-nlb-target-type: ip
     service.beta.kubernetes.io/aws-load-balancer-scheme: internet-facing
     external-dns.alpha.kubernetes.io/hostname: neon-proxy-scram-legacy.beta.us-east-2.aws.neon.build
+  httpsPort: 443
 
 #metrics:
 #  enabled: true

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

@@ -8,7 +8,10 @@ settings:
   authBackend: "console"
   authEndpoint: "http://console-staging.local/management/api/v2"
   domain: "*.us-east-2.aws.neon.build"
-  sentryEnvironment: "development"
+  sentryEnvironment: "staging"
+  wssPort: 8443
+  metricCollectionEndpoint: "http://console-staging.local/billing/api/v1/usage_events"
+  metricCollectionInterval: "1min"
 
 # -- Additional labels for neon-proxy pods
 podLabels:
@@ -23,6 +26,7 @@ exposedService:
     service.beta.kubernetes.io/aws-load-balancer-nlb-target-type: ip
     service.beta.kubernetes.io/aws-load-balancer-scheme: internet-facing
     external-dns.alpha.kubernetes.io/hostname: us-east-2.aws.neon.build
+  httpsPort: 443
 
 #metrics:
 #  enabled: true

.github/helm-values/dev-us-east-2-beta.neon-storage-broker.yaml

@@ -49,4 +49,4 @@ extraManifests:
           - "{{ .Release.Namespace }}"
 
 settings:
-  sentryEnvironment: "development"
+  sentryEnvironment: "staging"

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

@@ -9,6 +9,9 @@ settings:
   authEndpoint: "http://console-release.local/management/api/v2"
   domain: "*.ap-southeast-1.aws.neon.tech"
   sentryEnvironment: "production"
+  wssPort: 8443
+  metricCollectionEndpoint: "http://console-release.local/billing/api/v1/usage_events"
+  metricCollectionInterval: "10min"
 
 # -- Additional labels for neon-proxy pods
 podLabels:
@@ -23,6 +26,7 @@ exposedService:
     service.beta.kubernetes.io/aws-load-balancer-nlb-target-type: ip
     service.beta.kubernetes.io/aws-load-balancer-scheme: internet-facing
     external-dns.alpha.kubernetes.io/hostname: ap-southeast-1.aws.neon.tech
+  httpsPort: 443
 
 #metrics:
 #  enabled: true

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

@@ -9,6 +9,9 @@ settings:
   authEndpoint: "http://console-release.local/management/api/v2"
   domain: "*.eu-central-1.aws.neon.tech"
   sentryEnvironment: "production"
+  wssPort: 8443
+  metricCollectionEndpoint: "http://console-release.local/billing/api/v1/usage_events"
+  metricCollectionInterval: "10min"
 
 # -- Additional labels for neon-proxy pods
 podLabels:
@@ -23,6 +26,7 @@ exposedService:
     service.beta.kubernetes.io/aws-load-balancer-nlb-target-type: ip
     service.beta.kubernetes.io/aws-load-balancer-scheme: internet-facing
     external-dns.alpha.kubernetes.io/hostname: eu-central-1.aws.neon.tech
+  httpsPort: 443
 
 #metrics:
 #  enabled: true

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

@@ -9,6 +9,9 @@ settings:
   authEndpoint: "http://console-release.local/management/api/v2"
   domain: "*.us-east-2.aws.neon.tech"
   sentryEnvironment: "production"
+  wssPort: 8443
+  metricCollectionEndpoint: "http://console-release.local/billing/api/v1/usage_events"
+  metricCollectionInterval: "10min"
 
 # -- Additional labels for neon-proxy pods
 podLabels:
@@ -23,6 +26,7 @@ exposedService:
     service.beta.kubernetes.io/aws-load-balancer-nlb-target-type: ip
     service.beta.kubernetes.io/aws-load-balancer-scheme: internet-facing
     external-dns.alpha.kubernetes.io/hostname: us-east-2.aws.neon.tech
+  httpsPort: 443
 
 #metrics:
 #  enabled: true

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

@@ -9,6 +9,9 @@ settings:
   authEndpoint: "http://console-release.local/management/api/v2"
   domain: "*.us-west-2.aws.neon.tech"
   sentryEnvironment: "production"
+  wssPort: 8443
+  metricCollectionEndpoint: "http://console-release.local/billing/api/v1/usage_events"
+  metricCollectionInterval: "10min"
 
 # -- Additional labels for neon-proxy pods
 podLabels:
@@ -23,6 +26,7 @@ exposedService:
     service.beta.kubernetes.io/aws-load-balancer-nlb-target-type: ip
     service.beta.kubernetes.io/aws-load-balancer-scheme: internet-facing
     external-dns.alpha.kubernetes.io/hostname: us-west-2.aws.neon.tech
+  httpsPort: 443
 
 #metrics:
 #  enabled: true

.github/helm-values/production.proxy-scram.yaml

@@ -3,6 +3,9 @@ settings:
   authEndpoint: "http://console-release.local/management/api/v2"
   domain: "*.cloud.neon.tech"
   sentryEnvironment: "production"
+  wssPort: 8443
+  metricCollectionEndpoint: "http://console-release.local/billing/api/v1/usage_events"
+  metricCollectionInterval: "10min"
 
 podLabels:
   zenith_service: proxy-scram
@@ -16,6 +19,7 @@ exposedService:
     service.beta.kubernetes.io/aws-load-balancer-nlb-target-type: ip
     service.beta.kubernetes.io/aws-load-balancer-scheme: internet-facing
     external-dns.alpha.kubernetes.io/hostname: '*.cloud.neon.tech'
+  httpsPort: 443
 
 metrics:
   enabled: true

.github/workflows/benchmarking.yml

@@ -489,3 +489,108 @@ jobs:
         slack-message: "Periodic TPC-H perf testing ${{ matrix.platform }}: ${{ job.status }}\n${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}"
       env:
         SLACK_BOT_TOKEN: ${{ secrets.SLACK_BOT_TOKEN }}
+
+  user-examples-compare:
+    if: success() || failure()
+    needs: [ tpch-compare ]
+
+    strategy:
+      fail-fast: false
+      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 == 'refs/heads/main' ) }}
+      PLATFORM: ${{ matrix.platform }}
+
+    runs-on: [ self-hosted, us-east-2, x64 ]
+    container:
+      image: 369495373322.dkr.ecr.eu-central-1.amazonaws.com/rust:pinned
+      options: --init
+
+    timeout-minutes: 360 # 6h
+
+    steps:
+    - uses: actions/checkout@v3
+
+    - name: Download Neon artifact
+      uses: ./.github/actions/download
+      with:
+        name: neon-${{ runner.os }}-release-artifact
+        path: /tmp/neon/
+        prefix: latest
+
+    - name: Add Postgres binaries to PATH
+      run: |
+        ${POSTGRES_DISTRIB_DIR}/v${DEFAULT_PG_VERSION}/bin/pgbench --version
+        echo "${POSTGRES_DISTRIB_DIR}/v${DEFAULT_PG_VERSION}/bin" >> $GITHUB_PATH
+
+    - name: Set up Connection String
+      id: set-up-connstr
+      run: |
+        case "${PLATFORM}" in
+          neon-captest-prefetch)
+            CONNSTR=${{ secrets.BENCHMARK_USER_EXAMPLE_CAPTEST_CONNSTR }}
+            ;;
+          rds-aurora)
+            CONNSTR=${{ secrets.BENCHMARK_USER_EXAMPLE_RDS_AURORA_CONNSTR }}
+            ;;
+          rds-postgres)
+            CONNSTR=${{ secrets.BENCHMARK_USER_EXAMPLE_RDS_POSTGRES_CONNSTR }}
+            ;;
+          *)
+            echo 2>&1 "Unknown PLATFORM=${PLATFORM}. Allowed only 'neon-captest-prefetch', 'rds-aurora', or 'rds-postgres'"
+            exit 1
+            ;;
+        esac
+
+        echo "connstr=${CONNSTR}" >> $GITHUB_OUTPUT
+
+        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:
+        build_type: ${{ env.BUILD_TYPE }}
+        test_selection: performance/test_perf_olap.py
+        run_in_parallel: false
+        save_perf_report: ${{ env.SAVE_PERF_REPORT }}
+        extra_params: -m remote_cluster --timeout 21600 -k test_user_examples
+      env:
+        VIP_VAP_ACCESS_TOKEN: "${{ secrets.VIP_VAP_ACCESS_TOKEN }}"
+        PERF_TEST_RESULT_CONNSTR: "${{ secrets.PERF_TEST_RESULT_CONNSTR }}"
+        BENCHMARK_CONNSTR: ${{ steps.set-up-connstr.outputs.connstr }}
+
+    - name: Create Allure report
+      if: success() || failure()
+      uses: ./.github/actions/allure-report
+      with:
+        action: generate
+        build_type: ${{ env.BUILD_TYPE }}
+
+    - name: Post to a Slack channel
+      if: ${{ github.event.schedule && failure() }}
+      uses: slackapi/slack-github-action@v1
+      with:
+        channel-id: "C033QLM5P7D" # dev-staging-stream
+        slack-message: "Periodic TPC-H perf testing ${{ matrix.platform }}: ${{ job.status }}\n${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}"
+      env:
+        SLACK_BOT_TOKEN: ${{ secrets.SLACK_BOT_TOKEN }}

.github/workflows/build_and_test.yml

@@ -111,6 +111,7 @@ jobs:
       # Some of our rust modules use FFI and need those to be checked
       - name: Get postgres headers
         run: make postgres-headers -j$(nproc)
+
       - name: Run cargo clippy
         run: ./run_clippy.sh
 
@@ -126,6 +127,11 @@ jobs:
           cargo hakari generate --diff  # workspace-hack Cargo.toml is up-to-date
           cargo hakari manage-deps --dry-run  # all workspace crates depend on workspace-hack
 
+      # https://github.com/EmbarkStudios/cargo-deny
+      - name: Check rust licenses/bans/advisories/sources
+        if: ${{ !cancelled() }}
+        run: cargo deny check
+
   build-neon:
     runs-on: [ self-hosted, dev, x64 ]
     container:
@@ -177,13 +183,12 @@ jobs:
       # corresponding Cargo.toml files for their descriptions.
       - name: Set env variables
         run: |
+          CARGO_FEATURES="--features testing"
           if [[ $BUILD_TYPE == "debug" ]]; then
             cov_prefix="scripts/coverage --profraw-prefix=$GITHUB_JOB --dir=/tmp/coverage run"
-            CARGO_FEATURES="--features testing"
             CARGO_FLAGS="--locked $CARGO_FEATURES"
           elif [[ $BUILD_TYPE == "release" ]]; then
             cov_prefix=""
-            CARGO_FEATURES="--features testing,profiling"
             CARGO_FLAGS="--locked --release $CARGO_FEATURES"
           fi
           echo "cov_prefix=${cov_prefix}" >> $GITHUB_ENV
@@ -590,6 +595,8 @@ jobs:
     defaults:
       run:
         shell: sh -eu {0}
+    env:
+      VM_INFORMANT_VERSION: 0.1.1
 
     steps:
       - name: Downloading latest vm-builder
@@ -601,9 +608,22 @@ jobs:
         run: |
           docker pull 369495373322.dkr.ecr.eu-central-1.amazonaws.com/compute-node-${{ matrix.version }}:${{needs.tag.outputs.build-tag}}
 
+      - name: Downloading VM informant version ${{ env.VM_INFORMANT_VERSION }}
+        run: |
+          curl -fL https://github.com/neondatabase/autoscaling/releases/download/${{ env.VM_INFORMANT_VERSION }}/vm-informant -o vm-informant
+          chmod +x vm-informant
+
+      - name: Adding VM informant to compute-node image
+        run: |
+          ID=$(docker create 369495373322.dkr.ecr.eu-central-1.amazonaws.com/compute-node-${{ matrix.version }}:${{needs.tag.outputs.build-tag}})
+          docker cp vm-informant $ID:/bin/vm-informant
+          docker commit $ID temp-vm-compute-node
+          docker rm -f $ID
+
       - name: Build vm image
         run: |
-          ./vm-builder -src=369495373322.dkr.ecr.eu-central-1.amazonaws.com/compute-node-${{ matrix.version }}:${{needs.tag.outputs.build-tag}} -dst=369495373322.dkr.ecr.eu-central-1.amazonaws.com/vm-compute-node-${{ matrix.version }}:${{needs.tag.outputs.build-tag}}
+          # note: as of 2023-01-12, vm-builder requires a trailing ":latest" for local images
+          ./vm-builder -src=temp-vm-compute-node:latest -dst=369495373322.dkr.ecr.eu-central-1.amazonaws.com/vm-compute-node-${{ matrix.version }}:${{needs.tag.outputs.build-tag}}
 
       - name: Pushing vm-compute-node image
         run: |
@@ -789,6 +809,8 @@ jobs:
     strategy:
       matrix:
         include: ${{fromJSON(needs.calculate-deploy-targets.outputs.matrix-include)}}
+    environment:
+      name: prod-old
     steps:
       - name: Checkout
         uses: actions/checkout@v3
@@ -834,7 +856,9 @@ jobs:
         shell: bash
     strategy:
       matrix:
-        target_region: [ us-east-2 ]
+        target_region: [ eu-west-1, us-east-2 ]
+    environment:
+      name: dev-${{ matrix.target_region }}
     steps:
       - name: Checkout
         uses: actions/checkout@v3
@@ -906,6 +930,8 @@ jobs:
     strategy:
       matrix:
         target_region: [ us-east-2, us-west-2, eu-central-1, ap-southeast-1 ]
+    environment:
+      name: prod-${{ matrix.target_region }}
     steps:
       - name: Checkout
         uses: actions/checkout@v3
@@ -945,6 +971,8 @@ jobs:
     strategy:
       matrix:
         include: ${{fromJSON(needs.calculate-deploy-targets.outputs.matrix-include)}}
+    environment:
+      name: prod-old
     env:
       KUBECONFIG: .kubeconfig
     steps:
@@ -970,8 +998,8 @@ jobs:
       - name: Re-deploy proxy
         run: |
           DOCKER_TAG=${{needs.tag.outputs.build-tag}}
-          helm upgrade ${{ matrix.proxy_job }}       neondatabase/neon-proxy --namespace neon-proxy --install -f .github/helm-values/${{ matrix.proxy_config }}.yaml       --set image.tag=${DOCKER_TAG} --set settings.sentryUrl=${{ secrets.SENTRY_URL_PROXY }} --wait --timeout 15m0s
-          helm upgrade ${{ matrix.proxy_job }}-scram neondatabase/neon-proxy --namespace neon-proxy --install -f .github/helm-values/${{ matrix.proxy_config }}-scram.yaml --set image.tag=${DOCKER_TAG} --set settings.sentryUrl=${{ secrets.SENTRY_URL_PROXY }} --wait --timeout 15m0s
+          helm upgrade ${{ matrix.proxy_job }}       neondatabase/neon-proxy --namespace neon-proxy --install --atomic -f .github/helm-values/${{ matrix.proxy_config }}.yaml       --set image.tag=${DOCKER_TAG} --set settings.sentryUrl=${{ secrets.SENTRY_URL_PROXY }} --wait --timeout 15m0s
+          helm upgrade ${{ matrix.proxy_job }}-scram neondatabase/neon-proxy --namespace neon-proxy --install --atomic -f .github/helm-values/${{ matrix.proxy_config }}-scram.yaml --set image.tag=${DOCKER_TAG} --set settings.sentryUrl=${{ secrets.SENTRY_URL_PROXY }} --wait --timeout 15m0s
 
   deploy-storage-broker:
     name: deploy storage broker on old staging and old prod
@@ -988,6 +1016,8 @@ jobs:
     strategy:
       matrix:
         include: ${{fromJSON(needs.calculate-deploy-targets.outputs.matrix-include)}}
+    environment:
+      name: prod-old
     env:
       KUBECONFIG: .kubeconfig
     steps:
@@ -1036,6 +1066,8 @@ jobs:
             target_cluster: dev-eu-west-1-zeta
             deploy_link_proxy: false
             deploy_legacy_scram_proxy: false
+    environment:
+      name: dev-${{ matrix.target_region }}
     steps:
       - name: Checkout
         uses: actions/checkout@v3
@@ -1051,19 +1083,19 @@ jobs:
       - name: Re-deploy scram proxy
         run: |
           DOCKER_TAG=${{needs.tag.outputs.build-tag}}
-          helm upgrade neon-proxy-scram neondatabase/neon-proxy --namespace neon-proxy --create-namespace --install -f .github/helm-values/${{ matrix.target_cluster }}.neon-proxy-scram.yaml --set image.tag=${DOCKER_TAG} --set settings.sentryUrl=${{ secrets.SENTRY_URL_PROXY }} --wait --timeout 15m0s
+          helm upgrade neon-proxy-scram neondatabase/neon-proxy --namespace neon-proxy --create-namespace --install --atomic -f .github/helm-values/${{ matrix.target_cluster }}.neon-proxy-scram.yaml --set image.tag=${DOCKER_TAG} --set settings.sentryUrl=${{ secrets.SENTRY_URL_PROXY }} --wait --timeout 15m0s
 
       - name: Re-deploy link proxy
         if: matrix.deploy_link_proxy
         run: |
           DOCKER_TAG=${{needs.tag.outputs.build-tag}}
-          helm upgrade neon-proxy-link neondatabase/neon-proxy --namespace neon-proxy --create-namespace --install -f .github/helm-values/${{ matrix.target_cluster }}.neon-proxy-link.yaml --set image.tag=${DOCKER_TAG} --set settings.sentryUrl=${{ secrets.SENTRY_URL_PROXY }} --wait --timeout 15m0s
+          helm upgrade neon-proxy-link neondatabase/neon-proxy --namespace neon-proxy --create-namespace --install --atomic -f .github/helm-values/${{ matrix.target_cluster }}.neon-proxy-link.yaml --set image.tag=${DOCKER_TAG} --set settings.sentryUrl=${{ secrets.SENTRY_URL_PROXY }} --wait --timeout 15m0s
 
       - name: Re-deploy legacy scram proxy
         if: matrix.deploy_legacy_scram_proxy
         run: |
           DOCKER_TAG=${{needs.tag.outputs.build-tag}}
-          helm upgrade neon-proxy-scram-legacy neondatabase/neon-proxy --namespace neon-proxy --create-namespace --install -f .github/helm-values/${{ matrix.target_cluster }}.neon-proxy-scram-legacy.yaml --set image.tag=${DOCKER_TAG} --set settings.sentryUrl=${{ secrets.SENTRY_URL_PROXY }} --wait --timeout 15m0s
+          helm upgrade neon-proxy-scram-legacy neondatabase/neon-proxy --namespace neon-proxy --create-namespace --install --atomic -f .github/helm-values/${{ matrix.target_cluster }}.neon-proxy-scram-legacy.yaml --set image.tag=${DOCKER_TAG} --set settings.sentryUrl=${{ secrets.SENTRY_URL_PROXY }} --wait --timeout 15m0s
 
   deploy-storage-broker-dev-new:
     runs-on: [ self-hosted, dev, x64 ]
@@ -1083,6 +1115,8 @@ jobs:
             target_cluster: dev-us-east-2-beta
           - target_region:  eu-west-1
             target_cluster: dev-eu-west-1-zeta
+    environment:
+      name: dev-${{ matrix.target_region }}
     steps:
       - name: Checkout
         uses: actions/checkout@v3
@@ -1121,6 +1155,8 @@ jobs:
             target_cluster: prod-eu-central-1-gamma
           - target_region: ap-southeast-1
             target_cluster: prod-ap-southeast-1-epsilon
+    environment:
+      name: prod-${{ matrix.target_region }}
     steps:
       - name: Checkout
         uses: actions/checkout@v3
@@ -1136,7 +1172,7 @@ jobs:
       - name: Re-deploy proxy
         run: |
           DOCKER_TAG=${{needs.tag.outputs.build-tag}}
-          helm upgrade neon-proxy-scram neondatabase/neon-proxy --namespace neon-proxy --create-namespace --install -f .github/helm-values/${{ matrix.target_cluster }}.neon-proxy-scram.yaml --set image.tag=${DOCKER_TAG} --set settings.sentryUrl=${{ secrets.SENTRY_URL_PROXY }} --wait --timeout 15m0s
+          helm upgrade neon-proxy-scram neondatabase/neon-proxy --namespace neon-proxy --create-namespace --install --atomic -f .github/helm-values/${{ matrix.target_cluster }}.neon-proxy-scram.yaml --set image.tag=${DOCKER_TAG} --set settings.sentryUrl=${{ secrets.SENTRY_URL_PROXY }} --wait --timeout 15m0s
 
   deploy-storage-broker-prod-new:
     runs-on: prod
@@ -1160,6 +1196,8 @@ jobs:
             target_cluster: prod-eu-central-1-gamma
           - target_region: ap-southeast-1
             target_cluster: prod-ap-southeast-1-epsilon
+    environment:
+      name: prod-${{ matrix.target_region }}
     steps:
       - name: Checkout
         uses: actions/checkout@v3

Cargo.toml

@@ -1,14 +1,3 @@
-# 'named-profiles' feature was stabilized in cargo 1.57. This line makes the
-# build work with older cargo versions.
-#
-# We have this because as of this writing, the latest cargo Debian package
-# that's available is 1.56. (Confusingly, the Debian package version number
-# is 0.57, whereas 'cargo --version' says 1.56.)
-#
-# See https://tracker.debian.org/pkg/cargo for the current status of the
-# package. When that gets updated, we can remove this.
-cargo-features = ["named-profiles"]
-
 [workspace]
 members = [
     "compute_tools",
@@ -21,6 +10,143 @@ members = [
     "libs/*",
 ]
 
+[workspace.package]
+edition = "2021"
+license = "Apache-2.0"
+
+## All dependency versions, used in the project
+[workspace.dependencies]
+anyhow = { version = "1.0", features = ["backtrace"] }
+async-stream = "0.3"
+async-trait = "0.1"
+atty = "0.2.14"
+aws-config = { version = "0.51.0", default-features = false, features=["rustls"] }
+aws-sdk-s3 = "0.21.0"
+aws-smithy-http = "0.51.0"
+aws-types = "0.51.0"
+base64 = "0.13.0"
+bincode = "1.3"
+bindgen = "0.61"
+bstr = "1.0"
+byteorder = "1.4"
+bytes = "1.0"
+chrono = { version = "0.4", default-features = false, features = ["clock"] }
+clap = "4.0"
+close_fds = "0.3.2"
+comfy-table = "6.1"
+const_format = "0.2"
+crc32c = "0.6"
+crossbeam-utils = "0.8.5"
+fail = "0.5.0"
+fs2 = "0.4.3"
+futures = "0.3"
+futures-core = "0.3"
+futures-util = "0.3"
+git-version = "0.3"
+hashbrown = "0.13"
+hex = "0.4"
+hex-literal = "0.3"
+hmac = "0.12.1"
+hostname = "0.3.1"
+humantime = "2.1"
+humantime-serde = "1.1.1"
+hyper = "0.14"
+hyper-tungstenite = "0.9"
+itertools = "0.10"
+jsonwebtoken = "8"
+libc = "0.2"
+md5 = "0.7.0"
+memoffset = "0.8"
+nix = "0.26"
+notify = "5.0.0"
+num-traits = "0.2.15"
+once_cell = "1.13"
+parking_lot = "0.12"
+pin-project-lite = "0.2"
+prometheus = {version = "0.13", default_features=false, features = ["process"]} # removes protobuf dependency
+prost = "0.11"
+rand = "0.8"
+regex = "1.4"
+reqwest = { version = "0.11", default-features = false, features = ["rustls-tls"] }
+routerify = "3"
+rstar = "0.9.3"
+rustls = "0.20"
+rustls-pemfile = "1"
+rustls-split = "0.3"
+scopeguard = "1.1"
+sentry = { version = "0.29", default-features = false, features = ["backtrace", "contexts", "panic", "rustls", "reqwest" ] }
+serde = { version = "1.0", features = ["derive"] }
+serde_json = "1"
+serde_with = "2.0"
+sha2 = "0.10.2"
+signal-hook = "0.3"
+socket2 = "0.4.4"
+strum = "0.24"
+strum_macros = "0.24"
+svg_fmt = "0.4.1"
+tar = "0.4"
+thiserror = "1.0"
+tls-listener = { version = "0.6", features = ["rustls", "hyper-h1"] }
+tokio = { version = "1.17", features = ["macros"] }
+tokio-postgres-rustls = "0.9.0"
+tokio-rustls = "0.23"
+tokio-stream = "0.1"
+tokio-util = { version = "0.7", features = ["io"] }
+toml = "0.5"
+toml_edit = { version = "0.17", features = ["easy"] }
+tonic = {version = "0.8", features = ["tls", "tls-roots"]}
+tracing = "0.1"
+tracing-subscriber = { version = "0.3", features = ["env-filter"] }
+url = "2.2"
+uuid = { version = "1.2", features = ["v4", "serde"] }
+walkdir = "2.3.2"
+webpki-roots = "0.22.5"
+x509-parser = "0.14"
+
+## TODO replace this with tracing
+env_logger = "0.10"
+log = "0.4"
+
+## TODO switch when the new release is made
+amplify_num = { git = "https://github.com/rust-amplify/rust-amplify.git", tag = "v4.0.0-beta.1" }
+
+## Libraries from neondatabase/ git forks, ideally with changes to be upstreamed
+postgres = { git = "https://github.com/neondatabase/rust-postgres.git", rev="43e6db254a97fdecbce33d8bc0890accfd74495e" }
+postgres-protocol = { git = "https://github.com/neondatabase/rust-postgres.git", rev="43e6db254a97fdecbce33d8bc0890accfd74495e" }
+postgres-types = { git = "https://github.com/neondatabase/rust-postgres.git", rev="43e6db254a97fdecbce33d8bc0890accfd74495e" }
+tokio-postgres = { git = "https://github.com/neondatabase/rust-postgres.git", rev="43e6db254a97fdecbce33d8bc0890accfd74495e" }
+tokio-tar = { git = "https://github.com/neondatabase/tokio-tar.git", rev="404df61437de0feef49ba2ccdbdd94eb8ad6e142" }
+
+## Local libraries
+consumption_metrics = { version = "0.1", path = "./libs/consumption_metrics/" }
+metrics = { version = "0.1", path = "./libs/metrics/" }
+pageserver_api = { version = "0.1", path = "./libs/pageserver_api/" }
+postgres_connection = { version = "0.1", path = "./libs/postgres_connection/" }
+postgres_ffi = { version = "0.1", path = "./libs/postgres_ffi/" }
+pq_proto = { version = "0.1", path = "./libs/pq_proto/" }
+remote_storage = { version = "0.1", path = "./libs/remote_storage/" }
+safekeeper_api = { version = "0.1", path = "./libs/safekeeper_api" }
+storage_broker = { version = "0.1", path = "./storage_broker/" } # Note: main broker code is inside the binary crate, so linking with the library shouldn't be heavy.
+tenant_size_model = { version = "0.1", path = "./libs/tenant_size_model/" }
+utils = { version = "0.1", path = "./libs/utils/" }
+
+## Common library dependency
+workspace_hack = { version = "0.1", path = "./workspace_hack/" }
+
+## Build dependencies
+criterion = "0.4"
+rcgen = "0.10"
+rstest = "0.16"
+tempfile = "3.2"
+tonic-build = "0.8"
+
+# This is only needed for proxy's tests.
+# TODO: we should probably fork `tokio-postgres-rustls` instead.
+[patch.crates-io]
+tokio-postgres = { git = "https://github.com/neondatabase/rust-postgres.git", rev="43e6db254a97fdecbce33d8bc0890accfd74495e" }
+
+################# Binary contents sections
+
 [profile.release]
 # This is useful for profiling and, to some extent, debug.
 # Besides, debug info should not affect the performance.
@@ -81,9 +207,3 @@ inherits = "release"
 debug = false # true = 2 = all symbols, 1 = line only
 opt-level = "z"
 lto = true
-
-
-# This is only needed for proxy's tests.
-# TODO: we should probably fork `tokio-postgres-rustls` instead.
-[patch.crates-io]
-tokio-postgres = { git = "https://github.com/neondatabase/rust-postgres.git", rev="43e6db254a97fdecbce33d8bc0890accfd74495e" }

Dockerfile.compute-node-v14

@@ -29,7 +29,13 @@ RUN cd postgres && \
     make MAKELEVEL=0 -j $(getconf _NPROCESSORS_ONLN) -s -C contrib/ install && \
     # Install headers
     make MAKELEVEL=0 -j $(getconf _NPROCESSORS_ONLN) -s -C src/include install && \
-    make MAKELEVEL=0 -j $(getconf _NPROCESSORS_ONLN) -s -C src/interfaces/libpq install
+    make MAKELEVEL=0 -j $(getconf _NPROCESSORS_ONLN) -s -C src/interfaces/libpq install && \
+    # Enable some of contrib extensions
+    echo 'trusted = true' >> /usr/local/pgsql/share/extension/bloom.control && \
+    echo 'trusted = true' >> /usr/local/pgsql/share/extension/pgrowlocks.control && \
+    echo 'trusted = true' >> /usr/local/pgsql/share/extension/intagg.control && \
+    echo 'trusted = true' >> /usr/local/pgsql/share/extension/pgstattuple.control && \
+    echo 'trusted = true' >> /usr/local/pgsql/share/extension/earthdistance.control
 
 #########################################################################################
 #
@@ -55,7 +61,9 @@ RUN wget https://download.osgeo.org/postgis/source/postgis-3.3.1.tar.gz && \
     echo 'trusted = true' >> /usr/local/pgsql/share/extension/postgis.control && \
     echo 'trusted = true' >> /usr/local/pgsql/share/extension/postgis_raster.control && \
     echo 'trusted = true' >> /usr/local/pgsql/share/extension/postgis_tiger_geocoder.control && \
-    echo 'trusted = true' >> /usr/local/pgsql/share/extension/postgis_topology.control
+    echo 'trusted = true' >> /usr/local/pgsql/share/extension/postgis_topology.control && \
+    echo 'trusted = true' >> /usr/local/pgsql/share/extension/address_standardizer.control && \
+    echo 'trusted = true' >> /usr/local/pgsql/share/extension/address_standardizer_data_us.control
 
 #########################################################################################
 #

Dockerfile.compute-node-v15

@@ -29,7 +29,13 @@ RUN cd postgres && \
     make MAKELEVEL=0 -j $(getconf _NPROCESSORS_ONLN) -s -C contrib/ install && \
     # Install headers
     make MAKELEVEL=0 -j $(getconf _NPROCESSORS_ONLN) -s -C src/include install && \
-    make MAKELEVEL=0 -j $(getconf _NPROCESSORS_ONLN) -s -C src/interfaces/libpq install
+    make MAKELEVEL=0 -j $(getconf _NPROCESSORS_ONLN) -s -C src/interfaces/libpq install && \
+    # Enable some of contrib extensions
+    echo 'trusted = true' >> /usr/local/pgsql/share/extension/bloom.control && \
+    echo 'trusted = true' >> /usr/local/pgsql/share/extension/pgrowlocks.control && \
+    echo 'trusted = true' >> /usr/local/pgsql/share/extension/intagg.control && \
+    echo 'trusted = true' >> /usr/local/pgsql/share/extension/pgstattuple.control && \
+    echo 'trusted = true' >> /usr/local/pgsql/share/extension/earthdistance.control
 
 #########################################################################################
 #
@@ -55,7 +61,9 @@ RUN wget https://download.osgeo.org/postgis/source/postgis-3.3.1.tar.gz && \
     echo 'trusted = true' >> /usr/local/pgsql/share/extension/postgis.control && \
     echo 'trusted = true' >> /usr/local/pgsql/share/extension/postgis_raster.control && \
     echo 'trusted = true' >> /usr/local/pgsql/share/extension/postgis_tiger_geocoder.control && \
-    echo 'trusted = true' >> /usr/local/pgsql/share/extension/postgis_topology.control
+    echo 'trusted = true' >> /usr/local/pgsql/share/extension/postgis_topology.control && \
+    echo 'trusted = true' >> /usr/local/pgsql/share/extension/address_standardizer.control && \
+    echo 'trusted = true' >> /usr/local/pgsql/share/extension/address_standardizer_data_us.control
 
 #########################################################################################
 #

README.md

@@ -31,7 +31,8 @@ libssl-dev clang pkg-config libpq-dev cmake postgresql-client protobuf-compiler
 * On Fedora, these packages are needed:
 ```bash
 dnf install flex bison readline-devel zlib-devel openssl-devel \
-  libseccomp-devel perl clang cmake postgresql postgresql-contrib protobuf-compiler
+  libseccomp-devel perl clang cmake postgresql postgresql-contrib protobuf-compiler \
+  protobuf-devel
 ```
 
 2. [Install Rust](https://www.rust-lang.org/tools/install)
@@ -117,11 +118,8 @@ Python (3.9 or higher), and install python3 packages using `./scripts/pysync` (r
 # Later that would be responsibility of a package install script
 > ./target/debug/neon_local init
 Starting pageserver at '127.0.0.1:64000' in '.neon'.
-pageserver started, pid: 2545906
-Successfully initialized timeline de200bd42b49cc1814412c7e592dd6e9
-Stopped pageserver 1 process with pid 2545906
 
-# start pageserver and safekeeper
+# start pageserver, safekeeper, and broker for their intercommunication
 > ./target/debug/neon_local start
 Starting neon broker at 127.0.0.1:50051
 storage_broker started, pid: 2918372
@@ -130,6 +128,12 @@ pageserver started, pid: 2918386
 Starting safekeeper at '127.0.0.1:5454' in '.neon/safekeepers/sk1'.
 safekeeper 1 started, pid: 2918437
 
+# create initial tenant and use it as a default for every future neon_local invocation
+> ./target/debug/neon_local tenant create --set-default
+tenant 9ef87a5bf0d92544f6fafeeb3239695c successfully created on the pageserver
+Created an initial timeline 'de200bd42b49cc1814412c7e592dd6e9' at Lsn 0/16B5A50 for tenant: 9ef87a5bf0d92544f6fafeeb3239695c
+Setting tenant 9ef87a5bf0d92544f6fafeeb3239695c as a default one
+
 # start postgres compute node
 > ./target/debug/neon_local pg start main
 Starting new postgres (v14) main on timeline de200bd42b49cc1814412c7e592dd6e9 ...

compute_tools/Cargo.toml

@@ -1,23 +1,25 @@
 [package]
 name = "compute_tools"
 version = "0.1.0"
-edition = "2021"
+edition.workspace = true
+license.workspace = true
 
 [dependencies]
-anyhow = "1.0"
-chrono = { version = "0.4", default-features = false, features = ["clock"] }
-clap = "4.0"
-env_logger = "0.9"
-futures = "0.3.13"
-hyper = { version = "0.14", features = ["full"] }
-log = { version = "0.4", features = ["std", "serde"] }
-notify = "5.0.0"
-postgres = { git = "https://github.com/neondatabase/rust-postgres.git", rev="43e6db254a97fdecbce33d8bc0890accfd74495e" }
-regex = "1"
-serde = { version = "1.0", features = ["derive"] }
-serde_json = "1"
-tar = "0.4"
-tokio = { version = "1.17", features = ["macros", "rt", "rt-multi-thread"] }
-tokio-postgres = { git = "https://github.com/neondatabase/rust-postgres.git", rev="43e6db254a97fdecbce33d8bc0890accfd74495e" }
-url = "2.2.2"
-workspace_hack = { version = "0.1", path = "../workspace_hack" }
+anyhow.workspace = true
+chrono.workspace = true
+clap.workspace = true
+futures.workspace = true
+hyper = { workspace = true, features = ["full"] }
+notify.workspace = true
+postgres.workspace = true
+regex.workspace = true
+serde.workspace = true
+serde_json.workspace = true
+tar.workspace = true
+tokio = { workspace = true, features = ["rt", "rt-multi-thread"] }
+tokio-postgres.workspace = true
+tracing.workspace = true
+tracing-subscriber.workspace = true
+url.workspace = true
+
+workspace_hack.workspace = true

compute_tools/README.md

@@ -19,6 +19,10 @@ Also `compute_ctl` spawns two separate service threads:
 - `http-endpoint` runs a Hyper HTTP API server, which serves readiness and the
   last activity requests.
 
+If the `vm-informant` binary is present at `/bin/vm-informant`, it will also be started. For VM
+compute nodes, `vm-informant` communicates with the VM autoscaling system. It coordinates
+downscaling and (eventually) will request immediate upscaling under resource pressure.
+
 Usage example:
 ```sh
 compute_ctl -D /var/db/postgres/compute \

compute_tools/src/bin/compute_ctl.rs

@@ -18,6 +18,10 @@
 //! - `http-endpoint` runs a Hyper HTTP API server, which serves readiness and the
 //!   last activity requests.
 //!
+//! If the `vm-informant` binary is present at `/bin/vm-informant`, it will also be started. For VM
+//! compute nodes, `vm-informant` communicates with the VM autoscaling system. It coordinates
+//! downscaling and (eventually) will request immediate upscaling under resource pressure.
+//!
 //! Usage example:
 //! ```sh
 //! compute_ctl -D /var/db/postgres/compute \
@@ -36,10 +40,11 @@ use std::{thread, time::Duration};
 use anyhow::{Context, Result};
 use chrono::Utc;
 use clap::Arg;
-use log::{error, info};
+use tracing::{error, info};
 
 use compute_tools::compute::{ComputeMetrics, ComputeNode, ComputeState, ComputeStatus};
 use compute_tools::http::api::launch_http_server;
+use compute_tools::informant::spawn_vm_informant_if_present;
 use compute_tools::logger::*;
 use compute_tools::monitor::launch_monitor;
 use compute_tools::params::*;
@@ -48,7 +53,6 @@ use compute_tools::spec::*;
 use url::Url;
 
 fn main() -> Result<()> {
-    // TODO: re-use `utils::logging` later
     init_logger(DEFAULT_LOG_LEVEL)?;
 
     let matches = cli().get_matches();
@@ -114,30 +118,48 @@ fn main() -> Result<()> {
     // requests, while configuration is still in progress.
     let _http_handle = launch_http_server(&compute).expect("cannot launch http endpoint thread");
     let _monitor_handle = launch_monitor(&compute).expect("cannot launch compute monitor thread");
+    // Also spawn the thread responsible for handling the VM informant -- if it's present
+    let _vm_informant_handle = spawn_vm_informant_if_present().expect("cannot launch VM informant");
 
-    // Run compute (Postgres) and hang waiting on it.
-    match compute.prepare_and_run() {
-        Ok(ec) => {
-            let code = ec.code().unwrap_or(1);
-            info!("Postgres exited with code {}, shutting down", code);
-            exit(code)
-        }
-        Err(error) => {
-            error!("could not start the compute node: {:?}", error);
-
+    // Start Postgres
+    let mut delay_exit = false;
+    let mut exit_code = None;
+    let pg = match compute.start_compute() {
+        Ok(pg) => Some(pg),
+        Err(err) => {
+            error!("could not start the compute node: {:?}", err);
             let mut state = compute.state.write().unwrap();
-            state.error = Some(format!("{:?}", error));
+            state.error = Some(format!("{:?}", err));
             state.status = ComputeStatus::Failed;
             drop(state);
-
-            // Keep serving HTTP requests, so the cloud control plane was able to
-            // get the actual error.
-            info!("giving control plane 30s to collect the error before shutdown");
-            thread::sleep(Duration::from_secs(30));
-            info!("shutting down");
-            Err(error)
+            delay_exit = true;
+            None
         }
+    };
+
+    // Wait for the child Postgres process forever. In this state Ctrl+C will
+    // propagate to Postgres and it will be shut down as well.
+    if let Some(mut pg) = pg {
+        let ecode = pg
+            .wait()
+            .expect("failed to start waiting on Postgres process");
+        info!("Postgres exited with code {}, shutting down", ecode);
+        exit_code = ecode.code()
     }
+
+    if let Err(err) = compute.check_for_core_dumps() {
+        error!("error while checking for core dumps: {err:?}");
+    }
+
+    // If launch failed, keep serving HTTP requests for a while, so the cloud
+    // control plane can get the actual error.
+    if delay_exit {
+        info!("giving control plane 30s to collect the error before shutdown");
+        thread::sleep(Duration::from_secs(30));
+        info!("shutting down");
+    }
+
+    exit(exit_code.unwrap_or(1))
 }
 
 fn cli() -> clap::Command {

compute_tools/src/checker.rs

@@ -1,10 +1,11 @@
 use anyhow::{anyhow, Result};
-use log::error;
 use postgres::Client;
 use tokio_postgres::NoTls;
+use tracing::{error, instrument};
 
 use crate::compute::ComputeNode;
 
+#[instrument(skip_all)]
 pub fn create_writability_check_data(client: &mut Client) -> Result<()> {
     let query = "
     CREATE TABLE IF NOT EXISTS health_check (
@@ -21,6 +22,7 @@ pub fn create_writability_check_data(client: &mut Client) -> Result<()> {
     Ok(())
 }
 
+#[instrument(skip_all)]
 pub async fn check_writability(compute: &ComputeNode) -> Result<()> {
     let (client, connection) = tokio_postgres::connect(compute.connstr.as_str(), NoTls).await?;
     if client.is_closed() {

compute_tools/src/compute.rs

@@ -17,15 +17,15 @@
 use std::fs;
 use std::os::unix::fs::PermissionsExt;
 use std::path::Path;
-use std::process::{Command, ExitStatus, Stdio};
+use std::process::{Command, Stdio};
 use std::sync::atomic::{AtomicU64, Ordering};
 use std::sync::RwLock;
 
 use anyhow::{Context, Result};
 use chrono::{DateTime, Utc};
-use log::{info, warn};
 use postgres::{Client, NoTls};
 use serde::{Serialize, Serializer};
+use tracing::{info, instrument, warn};
 
 use crate::checker::create_writability_check_data;
 use crate::config;
@@ -121,6 +121,7 @@ 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))]
     fn get_basebackup(&self, lsn: &str) -> Result<()> {
         let start_time = Utc::now();
 
@@ -154,6 +155,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))]
     fn sync_safekeepers(&self) -> Result<String> {
         let start_time = Utc::now();
 
@@ -196,6 +198,7 @@ impl ComputeNode {
 
     /// Do all the preparations like PGDATA directory creation, configuration,
     /// safekeepers sync, basebackup, etc.
+    #[instrument(skip(self))]
     pub fn prepare_pgdata(&self) -> Result<()> {
         let spec = &self.spec;
         let pgdata_path = Path::new(&self.pgdata);
@@ -229,9 +232,8 @@ impl ComputeNode {
 
     /// Start Postgres as a child process and manage DBs/roles.
     /// After that this will hang waiting on the postmaster process to exit.
-    pub fn run(&self) -> Result<ExitStatus> {
-        let start_time = Utc::now();
-
+    #[instrument(skip(self))]
+    pub fn start_postgres(&self) -> Result<std::process::Child> {
         let pgdata_path = Path::new(&self.pgdata);
 
         // Run postgres as a child process.
@@ -242,10 +244,15 @@ impl ComputeNode {
 
         wait_for_postgres(&mut pg, pgdata_path)?;
 
+        Ok(pg)
+    }
+
+    #[instrument(skip(self))]
+    pub fn apply_config(&self) -> Result<()> {
         // If connection fails,
         // it may be the old node with `zenith_admin` superuser.
         //
-        // In this case we need to connect with old `zenith_admin`name
+        // In this case we need to connect with old `zenith_admin` name
         // and create new user. We cannot simply rename connected user,
         // but we can create a new one and grant it all privileges.
         let mut client = match Client::connect(self.connstr.as_str(), NoTls) {
@@ -271,6 +278,7 @@ impl ComputeNode {
             Ok(client) => client,
         };
 
+        // Proceed with post-startup configuration. Note, that order of operations is important.
         handle_roles(&self.spec, &mut client)?;
         handle_databases(&self.spec, &mut client)?;
         handle_role_deletions(self, &mut client)?;
@@ -279,8 +287,34 @@ impl ComputeNode {
 
         // 'Close' connection
         drop(client);
-        let startup_end_time = Utc::now();
 
+        info!(
+            "finished configuration of compute for project {}",
+            self.spec.cluster.cluster_id
+        );
+
+        Ok(())
+    }
+
+    #[instrument(skip(self))]
+    pub fn start_compute(&self) -> Result<std::process::Child> {
+        info!(
+            "starting compute for project {}, operation {}, tenant {}, timeline {}",
+            self.spec.cluster.cluster_id,
+            self.spec.operation_uuid.as_ref().unwrap(),
+            self.tenant,
+            self.timeline,
+        );
+
+        self.prepare_pgdata()?;
+
+        let start_time = Utc::now();
+
+        let pg = self.start_postgres()?;
+
+        self.apply_config()?;
+
+        let startup_end_time = Utc::now();
         self.metrics.config_ms.store(
             startup_end_time
                 .signed_duration_since(start_time)
@@ -300,34 +334,7 @@ impl ComputeNode {
 
         self.set_status(ComputeStatus::Running);
 
-        info!(
-            "finished configuration of compute for project {}",
-            self.spec.cluster.cluster_id
-        );
-
-        // Wait for child Postgres process basically forever. In this state Ctrl+C
-        // will propagate to Postgres and it will be shut down as well.
-        let ecode = pg
-            .wait()
-            .expect("failed to start waiting on Postgres process");
-
-        self.check_for_core_dumps()
-            .expect("failed to check for core dumps");
-
-        Ok(ecode)
-    }
-
-    pub fn prepare_and_run(&self) -> Result<ExitStatus> {
-        info!(
-            "starting compute for project {}, operation {}, tenant {}, timeline {}",
-            self.spec.cluster.cluster_id,
-            self.spec.operation_uuid.as_ref().unwrap(),
-            self.tenant,
-            self.timeline,
-        );
-
-        self.prepare_pgdata()?;
-        self.run()
+        Ok(pg)
     }
 
     // Look for core dumps and collect backtraces.
@@ -340,7 +347,7 @@ impl ComputeNode {
     //
     // Use that as a default location and pattern, except macos where core dumps are written
     // to /cores/ directory by default.
-    fn check_for_core_dumps(&self) -> Result<()> {
+    pub fn check_for_core_dumps(&self) -> Result<()> {
         let core_dump_dir = match std::env::consts::OS {
             "macos" => Path::new("/cores/"),
             _ => Path::new(&self.pgdata),

compute_tools/src/http/api.rs

@@ -6,32 +6,14 @@ use std::thread;
 use anyhow::Result;
 use hyper::service::{make_service_fn, service_fn};
 use hyper::{Body, Method, Request, Response, Server, StatusCode};
-use log::{error, info};
 use serde_json;
+use tracing::{error, info};
 
-use crate::compute::{ComputeNode, ComputeStatus};
+use crate::compute::ComputeNode;
 
 // Service function to handle all available routes.
 async fn routes(req: Request<Body>, compute: Arc<ComputeNode>) -> Response<Body> {
     match (req.method(), req.uri().path()) {
-        // Timestamp of the last Postgres activity in the plain text.
-        // DEPRECATED in favour of /status
-        (&Method::GET, "/last_activity") => {
-            info!("serving /last_active GET request");
-            let state = compute.state.read().unwrap();
-
-            // Use RFC3339 format for consistency.
-            Response::new(Body::from(state.last_active.to_rfc3339()))
-        }
-
-        // Has compute setup process finished? -> true/false.
-        // DEPRECATED in favour of /status
-        (&Method::GET, "/ready") => {
-            info!("serving /ready GET request");
-            let status = compute.get_status();
-            Response::new(Body::from(format!("{}", status == ComputeStatus::Running)))
-        }
-
         // Serialized compute state.
         (&Method::GET, "/status") => {
             info!("serving /status GET request");
@@ -46,16 +28,6 @@ async fn routes(req: Request<Body>, compute: Arc<ComputeNode>) -> Response<Body>
             Response::new(Body::from(serde_json::to_string(&compute.metrics).unwrap()))
         }
 
-        // DEPRECATED, use POST instead
-        (&Method::GET, "/check_writability") => {
-            info!("serving /check_writability GET request");
-            let res = crate::checker::check_writability(&compute).await;
-            match res {
-                Ok(_) => Response::new(Body::from("true")),
-                Err(e) => Response::new(Body::from(e.to_string())),
-            }
-        }
-
         (&Method::POST, "/check_writability") => {
             info!("serving /check_writability POST request");
             let res = crate::checker::check_writability(&compute).await;

compute_tools/src/http/openapi_spec.yaml

@@ -37,58 +37,7 @@ paths:
               schema:
                 $ref: "#/components/schemas/ComputeMetrics"
 
-  /ready:
-    get:
-      deprecated: true
-      tags:
-      - "info"
-      summary: Check whether compute startup process finished successfully
-      description: ""
-      operationId: computeIsReady
-      responses:
-        "200":
-          description: Compute is ready ('true') or not ('false')
-          content:
-            text/plain:
-              schema:
-                type: string
-                example: "true"
-
-  /last_activity:
-    get:
-      deprecated: true
-      tags:
-      - "info"
-      summary: Get timestamp of the last compute activity
-      description: ""
-      operationId: getLastComputeActivityTS
-      responses:
-        "200":
-          description: Timestamp of the last compute activity
-          content:
-            text/plain:
-              schema:
-                type: string
-                example: "2022-10-12T07:20:50.52Z"
-
   /check_writability:
-    get:
-      deprecated: true
-      tags:
-      - "check"
-      summary: Check that we can write new data on this compute
-      description: ""
-      operationId: checkComputeWritabilityDeprecated
-      responses:
-        "200":
-          description: Check result
-          content:
-            text/plain:
-              schema:
-                type: string
-                description: Error text or 'true' if check passed
-                example: "true"
-
     post:
       tags:
       - "check"

compute_tools/src/informant.rs

@@ -0,0 +1,50 @@
+use std::path::Path;
+use std::process;
+use std::thread;
+use std::time::Duration;
+use tracing::{info, warn};
+
+use anyhow::{Context, Result};
+
+const VM_INFORMANT_PATH: &str = "/bin/vm-informant";
+const RESTART_INFORMANT_AFTER_MILLIS: u64 = 5000;
+
+/// Launch a thread to start the VM informant if it's present (and restart, on failure)
+pub fn spawn_vm_informant_if_present() -> Result<Option<thread::JoinHandle<()>>> {
+    let exists = Path::new(VM_INFORMANT_PATH)
+        .try_exists()
+        .context("could not check if path exists")?;
+
+    if !exists {
+        return Ok(None);
+    }
+
+    Ok(Some(
+        thread::Builder::new()
+            .name("run-vm-informant".into())
+            .spawn(move || run_informant())?,
+    ))
+}
+
+fn run_informant() -> ! {
+    let restart_wait = Duration::from_millis(RESTART_INFORMANT_AFTER_MILLIS);
+
+    info!("starting VM informant");
+
+    loop {
+        let mut cmd = process::Command::new(VM_INFORMANT_PATH);
+        // Block on subprocess:
+        let result = cmd.status();
+
+        match result {
+            Err(e) => warn!("failed to run VM informant at {VM_INFORMANT_PATH:?}: {e}"),
+            Ok(status) if !status.success() => {
+                warn!("{VM_INFORMANT_PATH} exited with code {status:?}, retrying")
+            }
+            Ok(_) => info!("{VM_INFORMANT_PATH} ended gracefully (unexpectedly). Retrying"),
+        }
+
+        // Wait before retrying
+        thread::sleep(restart_wait);
+    }
+}

compute_tools/src/lib.rs

@@ -8,6 +8,7 @@ pub mod http;
 #[macro_use]
 pub mod logger;
 pub mod compute;
+pub mod informant;
 pub mod monitor;
 pub mod params;
 pub mod pg_helpers;

compute_tools/src/logger.rs

@@ -1,42 +1,20 @@
-use std::io::Write;
-
 use anyhow::Result;
-use chrono::Utc;
-use env_logger::{Builder, Env};
-
-macro_rules! info_println {
-    ($($tts:tt)*) => {
-        if log_enabled!(Level::Info) {
-            println!($($tts)*);
-        }
-    }
-}
-
-macro_rules! info_print {
-    ($($tts:tt)*) => {
-        if log_enabled!(Level::Info) {
-            print!($($tts)*);
-        }
-    }
-}
+use tracing_subscriber::layer::SubscriberExt;
+use tracing_subscriber::prelude::*;
 
 /// Initialize `env_logger` using either `default_level` or
 /// `RUST_LOG` environment variable as default log level.
 pub fn init_logger(default_level: &str) -> Result<()> {
-    let env = Env::default().filter_or("RUST_LOG", default_level);
+    let env_filter = tracing_subscriber::EnvFilter::try_from_default_env()
+        .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new(default_level));
 
-    Builder::from_env(env)
-        .format(|buf, record| {
-            let thread_handle = std::thread::current();
-            writeln!(
-                buf,
-                "{} [{}] {}: {}",
-                Utc::now().format("%Y-%m-%d %H:%M:%S%.3f %Z"),
-                thread_handle.name().unwrap_or("main"),
-                record.level(),
-                record.args()
-            )
-        })
+    let fmt_layer = tracing_subscriber::fmt::layer()
+        .with_target(false)
+        .with_writer(std::io::stderr);
+
+    tracing_subscriber::registry()
+        .with(env_filter)
+        .with(fmt_layer)
         .init();
 
     Ok(())

compute_tools/src/monitor.rs

@@ -3,8 +3,8 @@ use std::{thread, time};
 
 use anyhow::Result;
 use chrono::{DateTime, Utc};
-use log::{debug, info};
 use postgres::{Client, NoTls};
+use tracing::{debug, info};
 
 use crate::compute::ComputeNode;
 
@@ -52,10 +52,16 @@ fn watch_compute_activity(compute: &ComputeNode) {
                     let mut idle_backs: Vec<DateTime<Utc>> = vec![];
 
                     for b in backs.into_iter() {
-                        let state: String = b.get("state");
-                        let change: String = b.get("state_change");
+                        let state: String = match b.try_get("state") {
+                            Ok(state) => state,
+                            Err(_) => continue,
+                        };
 
                         if state == "idle" {
+                            let change: String = match b.try_get("state_change") {
+                                Ok(state_change) => state_change,
+                                Err(_) => continue,
+                            };
                             let change = DateTime::parse_from_rfc3339(&change);
                             match change {
                                 Ok(t) => idle_backs.push(t.with_timezone(&Utc)),

compute_tools/src/pg_helpers.rs

@@ -11,6 +11,7 @@ use anyhow::{bail, Result};
 use notify::{RecursiveMode, Watcher};
 use postgres::{Client, Transaction};
 use serde::Deserialize;
+use tracing::{debug, instrument};
 
 const POSTGRES_WAIT_TIMEOUT: Duration = Duration::from_millis(60 * 1000); // milliseconds
 
@@ -129,8 +130,8 @@ impl Role {
     /// Serialize a list of role parameters into a Postgres-acceptable
     /// string of arguments.
     pub fn to_pg_options(&self) -> String {
-        // XXX: consider putting LOGIN as a default option somewhere higher, e.g. in Rails.
-        // For now we do not use generic `options` for roles. Once used, add
+        // XXX: consider putting LOGIN as a default option somewhere higher, e.g. in control-plane.
+        // For now, we do not use generic `options` for roles. Once used, add
         // `self.options.as_pg_options()` somewhere here.
         let mut params: String = "LOGIN".to_string();
 
@@ -229,6 +230,7 @@ pub fn get_existing_dbs(client: &mut Client) -> Result<Vec<Database>> {
 /// Wait for Postgres to become ready to accept connections. It's ready to
 /// accept connections when the state-field in `pgdata/postmaster.pid` says
 /// 'ready'.
+#[instrument(skip(pg))]
 pub fn wait_for_postgres(pg: &mut Child, pgdata: &Path) -> Result<()> {
     let pid_path = pgdata.join("postmaster.pid");
 
@@ -287,18 +289,18 @@ pub fn wait_for_postgres(pg: &mut Child, pgdata: &Path) -> Result<()> {
         }
 
         let res = rx.recv_timeout(Duration::from_millis(100));
-        log::debug!("woken up by notify: {res:?}");
+        debug!("woken up by notify: {res:?}");
         // If there are multiple events in the channel already, we only need to be
         // check once. Swallow the extra events before we go ahead to check the
         // pid file.
         while let Ok(res) = rx.try_recv() {
-            log::debug!("swallowing extra event: {res:?}");
+            debug!("swallowing extra event: {res:?}");
         }
 
         // Check that we can open pid file first.
         if let Ok(file) = File::open(&pid_path) {
             if !postmaster_pid_seen {
-                log::debug!("postmaster.pid appeared");
+                debug!("postmaster.pid appeared");
                 watcher
                     .unwatch(pgdata)
                     .expect("Failed to remove pgdata dir watch");
@@ -314,7 +316,7 @@ pub fn wait_for_postgres(pg: &mut Child, pgdata: &Path) -> Result<()> {
             // Pid file could be there and we could read it, but it could be empty, for example.
             if let Some(Ok(line)) = last_line {
                 let status = line.trim();
-                log::debug!("last line of postmaster.pid: {status:?}");
+                debug!("last line of postmaster.pid: {status:?}");
 
                 // Now Postgres is ready to accept connections
                 if status == "ready" {
@@ -330,7 +332,7 @@ pub fn wait_for_postgres(pg: &mut Child, pgdata: &Path) -> Result<()> {
         }
     }
 
-    log::info!("PostgreSQL is now running, continuing to configure it");
+    tracing::info!("PostgreSQL is now running, continuing to configure it");
 
     Ok(())
 }

compute_tools/src/spec.rs

@@ -2,10 +2,10 @@ use std::path::Path;
 use std::str::FromStr;
 
 use anyhow::Result;
-use log::{info, log_enabled, warn, Level};
 use postgres::config::Config;
 use postgres::{Client, NoTls};
 use serde::Deserialize;
+use tracing::{info, info_span, instrument, span_enabled, warn, Level};
 
 use crate::compute::ComputeNode;
 use crate::config;
@@ -79,23 +79,25 @@ pub fn update_pg_hba(pgdata_path: &Path) -> Result<()> {
 
 /// Given a cluster spec json and open transaction it handles roles creation,
 /// deletion and update.
+#[instrument(skip_all)]
 pub fn handle_roles(spec: &ComputeSpec, client: &mut Client) -> Result<()> {
     let mut xact = client.transaction()?;
     let existing_roles: Vec<Role> = get_existing_roles(&mut xact)?;
 
     // Print a list of existing Postgres roles (only in debug mode)
-    info!("postgres roles:");
-    for r in &existing_roles {
-        info_println!(
-            "{} - {}:{}",
-            " ".repeat(27 + 5),
-            r.name,
-            if r.encrypted_password.is_some() {
-                "[FILTERED]"
-            } else {
-                "(null)"
-            }
-        );
+    if span_enabled!(Level::INFO) {
+        info!("postgres roles:");
+        for r in &existing_roles {
+            info!(
+                "    - {}:{}",
+                r.name,
+                if r.encrypted_password.is_some() {
+                    "[FILTERED]"
+                } else {
+                    "(null)"
+                }
+            );
+        }
     }
 
     // Process delta operations first
@@ -136,58 +138,68 @@ pub fn handle_roles(spec: &ComputeSpec, client: &mut Client) -> Result<()> {
     info!("cluster spec roles:");
     for role in &spec.cluster.roles {
         let name = &role.name;
-
-        info_print!(
-            "{} - {}:{}",
-            " ".repeat(27 + 5),
-            name,
-            if role.encrypted_password.is_some() {
-                "[FILTERED]"
-            } else {
-                "(null)"
-            }
-        );
-
         // XXX: with a limited number of roles it is fine, but consider making it a HashMap
         let pg_role = existing_roles.iter().find(|r| r.name == *name);
 
-        if let Some(r) = pg_role {
-            let mut update_role = false;
-
+        enum RoleAction {
+            None,
+            Update,
+            Create,
+        }
+        let action = if let Some(r) = pg_role {
             if (r.encrypted_password.is_none() && role.encrypted_password.is_some())
                 || (r.encrypted_password.is_some() && role.encrypted_password.is_none())
             {
-                update_role = true;
+                RoleAction::Update
             } else if let Some(pg_pwd) = &r.encrypted_password {
                 // Check whether password changed or not (trim 'md5:' prefix first)
-                update_role = pg_pwd[3..] != *role.encrypted_password.as_ref().unwrap();
+                if pg_pwd[3..] != *role.encrypted_password.as_ref().unwrap() {
+                    RoleAction::Update
+                } else {
+                    RoleAction::None
+                }
+            } else {
+                RoleAction::None
             }
+        } else {
+            RoleAction::Create
+        };
 
-            if update_role {
+        match action {
+            RoleAction::None => {}
+            RoleAction::Update => {
                 let mut query: String = format!("ALTER ROLE {} ", name.pg_quote());
-                info_print!(" -> update");
-
                 query.push_str(&role.to_pg_options());
                 xact.execute(query.as_str(), &[])?;
             }
-        } else {
-            info!("role name: '{}'", &name);
-            let mut query: String = format!("CREATE ROLE {} ", name.pg_quote());
-            info!("role create query: '{}'", &query);
-            info_print!(" -> create");
+            RoleAction::Create => {
+                let mut query: String = format!("CREATE ROLE {} ", name.pg_quote());
+                info!("role create query: '{}'", &query);
+                query.push_str(&role.to_pg_options());
+                xact.execute(query.as_str(), &[])?;
 
-            query.push_str(&role.to_pg_options());
-            xact.execute(query.as_str(), &[])?;
-
-            let grant_query = format!(
-                "GRANT pg_read_all_data, pg_write_all_data TO {}",
-                name.pg_quote()
-            );
-            xact.execute(grant_query.as_str(), &[])?;
-            info!("role grant query: '{}'", &grant_query);
+                let grant_query = format!(
+                    "GRANT pg_read_all_data, pg_write_all_data TO {}",
+                    name.pg_quote()
+                );
+                xact.execute(grant_query.as_str(), &[])?;
+                info!("role grant query: '{}'", &grant_query);
+            }
         }
 
-        info_print!("\n");
+        if span_enabled!(Level::INFO) {
+            let pwd = if role.encrypted_password.is_some() {
+                "[FILTERED]"
+            } else {
+                "(null)"
+            };
+            let action_str = match action {
+                RoleAction::None => "",
+                RoleAction::Create => " -> create",
+                RoleAction::Update => " -> update",
+            };
+            info!("   - {}:{}{}", name, pwd, action_str);
+        }
     }
 
     xact.commit()?;
@@ -196,23 +208,32 @@ pub fn handle_roles(spec: &ComputeSpec, client: &mut Client) -> Result<()> {
 }
 
 /// Reassign all dependent objects and delete requested roles.
+#[instrument(skip_all)]
 pub fn handle_role_deletions(node: &ComputeNode, client: &mut Client) -> Result<()> {
-    let spec = &node.spec;
-
-    // First, reassign all dependent objects to db owners.
-    if let Some(ops) = &spec.delta_operations {
+    if let Some(ops) = &node.spec.delta_operations {
+        // First, reassign all dependent objects to db owners.
         info!("reassigning dependent objects of to-be-deleted roles");
+
+        // Fetch existing roles. We could've exported and used `existing_roles` from
+        // `handle_roles()`, but we only make this list there before creating new roles.
+        // Which is probably fine as we never create to-be-deleted roles, but that'd
+        // just look a bit untidy. Anyway, the entire `pg_roles` should be in shared
+        // buffers already, so this shouldn't be a big deal.
+        let mut xact = client.transaction()?;
+        let existing_roles: Vec<Role> = get_existing_roles(&mut xact)?;
+        xact.commit()?;
+
         for op in ops {
-            if op.action == "delete_role" {
+            // Check that role is still present in Postgres, as this could be a
+            // restart with the same spec after role deletion.
+            if op.action == "delete_role" && existing_roles.iter().any(|r| r.name == op.name) {
                 reassign_owned_objects(node, &op.name)?;
             }
         }
-    }
 
-    // Second, proceed with role deletions.
-    let mut xact = client.transaction()?;
-    if let Some(ops) = &spec.delta_operations {
+        // Second, proceed with role deletions.
         info!("processing role deletions");
+        let mut xact = client.transaction()?;
         for op in ops {
             // We do not check either role exists or not,
             // Postgres will take care of it for us
@@ -223,6 +244,7 @@ pub fn handle_role_deletions(node: &ComputeNode, client: &mut Client) -> Result<
                 xact.execute(query.as_str(), &[])?;
             }
         }
+        xact.commit()?;
     }
 
     Ok(())
@@ -263,13 +285,16 @@ fn reassign_owned_objects(node: &ComputeNode, role_name: &PgIdent) -> Result<()>
 /// like `CREATE DATABASE` and `DROP DATABASE` do not support it. Statement-level
 /// atomicity should be enough here due to the order of operations and various checks,
 /// which together provide us idempotency.
+#[instrument(skip_all)]
 pub fn handle_databases(spec: &ComputeSpec, client: &mut Client) -> Result<()> {
     let existing_dbs: Vec<Database> = get_existing_dbs(client)?;
 
     // Print a list of existing Postgres databases (only in debug mode)
-    info!("postgres databases:");
-    for r in &existing_dbs {
-        info_println!("{} - {}:{}", " ".repeat(27 + 5), r.name, r.owner);
+    if span_enabled!(Level::INFO) {
+        info!("postgres databases:");
+        for r in &existing_dbs {
+            info!("    {}:{}", r.name, r.owner);
+        }
     }
 
     // Process delta operations first
@@ -312,12 +337,15 @@ pub fn handle_databases(spec: &ComputeSpec, client: &mut Client) -> Result<()> {
     for db in &spec.cluster.databases {
         let name = &db.name;
 
-        info_print!("{} - {}:{}", " ".repeat(27 + 5), db.name, db.owner);
-
         // XXX: with a limited number of databases it is fine, but consider making it a HashMap
         let pg_db = existing_dbs.iter().find(|r| r.name == *name);
 
-        if let Some(r) = pg_db {
+        enum DatabaseAction {
+            None,
+            Update,
+            Create,
+        }
+        let action = if let Some(r) = pg_db {
             // XXX: db owner name is returned as quoted string from Postgres,
             // when quoting is needed.
             let new_owner = if r.owner.starts_with('"') {
@@ -327,24 +355,42 @@ pub fn handle_databases(spec: &ComputeSpec, client: &mut Client) -> Result<()> {
             };
 
             if new_owner != r.owner {
+                // Update the owner
+                DatabaseAction::Update
+            } else {
+                DatabaseAction::None
+            }
+        } else {
+            DatabaseAction::Create
+        };
+
+        match action {
+            DatabaseAction::None => {}
+            DatabaseAction::Update => {
                 let query: String = format!(
                     "ALTER DATABASE {} OWNER TO {}",
                     name.pg_quote(),
                     db.owner.pg_quote()
                 );
-                info_print!(" -> update");
-
+                let _ = info_span!("executing", query).entered();
                 client.execute(query.as_str(), &[])?;
             }
-        } else {
-            let mut query: String = format!("CREATE DATABASE {} ", name.pg_quote());
-            info_print!(" -> create");
+            DatabaseAction::Create => {
+                let mut query: String = format!("CREATE DATABASE {} ", name.pg_quote());
+                query.push_str(&db.to_pg_options());
+                let _ = info_span!("executing", query).entered();
+                client.execute(query.as_str(), &[])?;
+            }
+        };
 
-            query.push_str(&db.to_pg_options());
-            client.execute(query.as_str(), &[])?;
+        if span_enabled!(Level::INFO) {
+            let action_str = match action {
+                DatabaseAction::None => "",
+                DatabaseAction::Create => " -> create",
+                DatabaseAction::Update => " -> update",
+            };
+            info!("   - {}:{}{}", db.name, db.owner, action_str);
         }
-
-        info_print!("\n");
     }
 
     Ok(())
@@ -352,6 +398,7 @@ pub fn handle_databases(spec: &ComputeSpec, client: &mut Client) -> Result<()> {
 
 /// Grant CREATE ON DATABASE to the database owner and do some other alters and grants
 /// to allow users creating trusted extensions and re-creating `public` schema, for example.
+#[instrument(skip_all)]
 pub fn handle_grants(node: &ComputeNode, client: &mut Client) -> Result<()> {
     let spec = &node.spec;
 

control_plane/Cargo.toml

@@ -1,31 +1,31 @@
 [package]
 name = "control_plane"
 version = "0.1.0"
-edition = "2021"
+edition.workspace = true
+license.workspace = true
 
 [dependencies]
-anyhow = "1.0"
-clap = "4.0"
-comfy-table = "6.1"
-git-version = "0.3.5"
-nix = "0.25"
-once_cell = "1.13.0"
-postgres = { git = "https://github.com/neondatabase/rust-postgres.git", rev = "43e6db254a97fdecbce33d8bc0890accfd74495e" }
-regex = "1"
-reqwest = { version = "0.11", default-features = false, features = ["blocking", "json", "rustls-tls"] }
-serde = { version = "1.0", features = ["derive"] }
-serde_with = "2.0"
-tar = "0.4.38"
-thiserror = "1"
-toml = "0.5"
-url = "2.2.2"
-
+anyhow.workspace = true
+clap.workspace = true
+comfy-table.workspace = true
+git-version.workspace = true
+nix.workspace = true
+once_cell.workspace = true
+postgres.workspace = true
+regex.workspace = true
+reqwest = { workspace = true, features = ["blocking", "json"] }
+serde.workspace = true
+serde_with.workspace = true
+tar.workspace = true
+thiserror.workspace = true
+toml.workspace = true
+url.workspace = true
 # Note: Do not directly depend on pageserver or safekeeper; use pageserver_api or safekeeper_api
 # instead, so that recompile times are better.
-pageserver_api = { path = "../libs/pageserver_api" }
-postgres_connection = { path = "../libs/postgres_connection" }
-safekeeper_api = { path = "../libs/safekeeper_api" }
-# Note: main broker code is inside the binary crate, so linking with the library shouldn't be heavy.
-storage_broker = { version = "0.1", path = "../storage_broker" }
-utils = { path = "../libs/utils" }
-workspace_hack = { version = "0.1", path = "../workspace_hack" }
+pageserver_api.workspace = true
+safekeeper_api.workspace = true
+postgres_connection.workspace = true
+storage_broker.workspace = true
+utils.workspace = true
+
+workspace_hack.workspace = true

control_plane/src/background_process.rs

@@ -136,22 +136,6 @@ where
     anyhow::bail!("{process_name} did not start in {RETRY_UNTIL_SECS} seconds");
 }
 
-/// Send SIGTERM to child process
-pub fn send_stop_child_process(child: &std::process::Child) -> anyhow::Result<()> {
-    let pid = child.id();
-    match kill(
-        nix::unistd::Pid::from_raw(pid.try_into().unwrap()),
-        Signal::SIGTERM,
-    ) {
-        Ok(()) => Ok(()),
-        Err(Errno::ESRCH) => {
-            println!("child process with pid {pid} does not exist");
-            Ok(())
-        }
-        Err(e) => anyhow::bail!("Failed to send signal to child process with pid {pid}: {e}"),
-    }
-}
-
 /// Stops the process, using the pid file given. Returns Ok also if the process is already not running.
 pub fn stop_process(immediate: bool, process_name: &str, pid_file: &Path) -> anyhow::Result<()> {
     let pid = match pid_file::read(pid_file)

control_plane/src/bin/neon_local.rs

@@ -263,7 +263,7 @@ fn get_tenant_id(sub_match: &ArgMatches, env: &local_env::LocalEnv) -> anyhow::R
     } else if let Some(default_id) = env.default_tenant_id {
         Ok(default_id)
     } else {
-        bail!("No tenant id. Use --tenant-id, or set 'default_tenant_id' in the config file");
+        anyhow::bail!("No tenant id. Use --tenant-id, or set a default tenant");
     }
 }
 
@@ -284,8 +284,6 @@ fn parse_timeline_id(sub_match: &ArgMatches) -> anyhow::Result<Option<TimelineId
 }
 
 fn handle_init(init_match: &ArgMatches) -> anyhow::Result<LocalEnv> {
-    let initial_timeline_id_arg = parse_timeline_id(init_match)?;
-
     // Create config file
     let toml_file: String = if let Some(config_path) = init_match.get_one::<PathBuf>("config") {
         // load and parse the file
@@ -309,30 +307,16 @@ fn handle_init(init_match: &ArgMatches) -> anyhow::Result<LocalEnv> {
         LocalEnv::parse_config(&toml_file).context("Failed to create neon configuration")?;
     env.init(pg_version)
         .context("Failed to initialize neon repository")?;
-    let initial_tenant_id = env
-        .default_tenant_id
-        .expect("default_tenant_id should be generated by the `env.init()` call above");
 
     // Initialize pageserver, create initial tenant and timeline.
     let pageserver = PageServerNode::from_env(&env);
-    let initial_timeline_id = pageserver
-        .initialize(
-            Some(initial_tenant_id),
-            initial_timeline_id_arg,
-            &pageserver_config_overrides(init_match),
-            pg_version,
-        )
+    pageserver
+        .initialize(&pageserver_config_overrides(init_match))
         .unwrap_or_else(|e| {
             eprintln!("pageserver init failed: {e:?}");
             exit(1);
         });
 
-    env.register_branch_mapping(
-        DEFAULT_BRANCH_NAME.to_owned(),
-        initial_tenant_id,
-        initial_timeline_id,
-    )?;
-
     Ok(env)
 }
 
@@ -388,6 +372,17 @@ fn handle_tenant(tenant_match: &ArgMatches, env: &mut local_env::LocalEnv) -> an
             println!(
                 "Created an initial timeline '{new_timeline_id}' at Lsn {last_record_lsn} for tenant: {new_tenant_id}",
             );
+
+            if create_match.get_flag("set-default") {
+                println!("Setting tenant {new_tenant_id} as a default one");
+                env.default_tenant_id = Some(new_tenant_id);
+            }
+        }
+        Some(("set-default", set_default_match)) => {
+            let tenant_id =
+                parse_tenant_id(set_default_match)?.context("No tenant id specified")?;
+            println!("Setting tenant {tenant_id} as a default one");
+            env.default_tenant_id = Some(tenant_id);
         }
         Some(("config", create_match)) => {
             let tenant_id = get_tenant_id(create_match, env)?;
@@ -928,9 +923,8 @@ fn cli() -> Command {
         .version(GIT_VERSION)
         .subcommand(
             Command::new("init")
-                .about("Initialize a new Neon repository")
+                .about("Initialize a new Neon repository, preparing configs for services to start with")
                 .arg(pageserver_config_args.clone())
-                .arg(timeline_id_arg.clone().help("Use a specific timeline id when creating a tenant and its initial timeline"))
                 .arg(
                     Arg::new("config")
                         .long("config")
@@ -992,11 +986,14 @@ fn cli() -> Command {
                 .arg(timeline_id_arg.clone().help("Use a specific timeline id when creating a tenant and its initial timeline"))
                 .arg(Arg::new("config").short('c').num_args(1).action(ArgAction::Append).required(false))
                 .arg(pg_version_arg.clone())
+                .arg(Arg::new("set-default").long("set-default").action(ArgAction::SetTrue).required(false)
+                    .help("Use this tenant in future CLI commands where tenant_id is needed, but not specified"))
                 )
+            .subcommand(Command::new("set-default").arg(tenant_id_arg.clone().required(true))
+                .about("Set a particular tenant as default in future CLI commands where tenant_id is needed, but not specified"))
             .subcommand(Command::new("config")
                 .arg(tenant_id_arg.clone())
-                .arg(Arg::new("config").short('c').num_args(1).action(ArgAction::Append).required(false))
-                )
+                .arg(Arg::new("config").short('c').num_args(1).action(ArgAction::Append).required(false)))
         )
         .subcommand(
             Command::new("pageserver")

control_plane/src/compute.rs

@@ -201,7 +201,7 @@ impl PostgresNode {
             .stderr(Stdio::piped());
 
         if let Some(token) = auth_token {
-            cmd.env("ZENITH_AUTH_TOKEN", token);
+            cmd.env("NEON_AUTH_TOKEN", token);
         }
 
         let sync_handle = cmd
@@ -304,17 +304,17 @@ impl PostgresNode {
 
             // Set up authentication
             //
-            // $ZENITH_AUTH_TOKEN will be replaced with value from environment
+            // $NEON_AUTH_TOKEN will be replaced with value from environment
             // variable during compute pg startup. It is done this way because
             // otherwise user will be able to retrieve the value using SHOW
             // command or pg_settings
             let password = if let AuthType::NeonJWT = auth_type {
-                "$ZENITH_AUTH_TOKEN"
+                "$NEON_AUTH_TOKEN"
             } else {
                 ""
             };
             // NOTE avoiding spaces in connection string, because it is less error prone if we forward it somewhere.
-            // Also note that not all parameters are supported here. Because in compute we substitute $ZENITH_AUTH_TOKEN
+            // Also note that not all parameters are supported here. Because in compute we substitute $NEON_AUTH_TOKEN
             // We parse this string and build it back with token from env var, and for simplicity rebuild
             // uses only needed variables namely host, port, user, password.
             format!("postgresql://no_user:{password}@{host}:{port}")
@@ -323,7 +323,7 @@ impl PostgresNode {
         conf.append_line("");
         conf.append("neon.pageserver_connstring", &pageserver_connstr);
         if let AuthType::NeonJWT = auth_type {
-            conf.append("neon.safekeeper_token_env", "$ZENITH_AUTH_TOKEN");
+            conf.append("neon.safekeeper_token_env", "$NEON_AUTH_TOKEN");
         }
         conf.append("neon.tenant_id", &self.tenant_id.to_string());
         conf.append("neon.timeline_id", &self.timeline_id.to_string());
@@ -448,7 +448,7 @@ impl PostgresNode {
             self.env.pg_lib_dir(self.pg_version)?.to_str().unwrap(),
         );
         if let Some(token) = auth_token {
-            cmd.env("ZENITH_AUTH_TOKEN", token);
+            cmd.env("NEON_AUTH_TOKEN", token);
         }
 
         let pg_ctl = cmd.output().context("pg_ctl failed")?;

control_plane/src/local_env.rs

@@ -296,11 +296,6 @@ impl LocalEnv {
             env.neon_distrib_dir = env::current_exe()?.parent().unwrap().to_owned();
         }
 
-        // If no initial tenant ID was given, generate it.
-        if env.default_tenant_id.is_none() {
-            env.default_tenant_id = Some(TenantId::generate());
-        }
-
         env.base_data_dir = base_path();
 
         Ok(env)

control_plane/src/pageserver.rs

@@ -7,7 +7,7 @@ use std::path::PathBuf;
 use std::process::{Child, Command};
 use std::{io, result};
 
-use anyhow::{bail, ensure, Context};
+use anyhow::{bail, Context};
 use pageserver_api::models::{
     TenantConfigRequest, TenantCreateRequest, TenantInfo, TimelineCreateRequest, TimelineInfo,
 };
@@ -130,83 +130,15 @@ impl PageServerNode {
         overrides
     }
 
-    /// Initializes a pageserver node by creating its config with the overrides provided,
-    /// and creating an initial tenant and timeline afterwards.
-    pub fn initialize(
-        &self,
-        create_tenant: Option<TenantId>,
-        initial_timeline_id: Option<TimelineId>,
-        config_overrides: &[&str],
-        pg_version: u32,
-    ) -> anyhow::Result<TimelineId> {
+    /// Initializes a pageserver node by creating its config with the overrides provided.
+    pub fn initialize(&self, config_overrides: &[&str]) -> anyhow::Result<()> {
         // First, run `pageserver --init` and wait for it to write a config into FS and exit.
         self.pageserver_init(config_overrides).with_context(|| {
             format!(
                 "Failed to run init for pageserver node {}",
                 self.env.pageserver.id,
             )
-        })?;
-
-        // Then, briefly start it fully to run HTTP commands on it,
-        // to create initial tenant and timeline.
-        // We disable the remote storage, since we stop pageserver right after the timeline creation,
-        // hence most of the uploads will either aborted or not started: no point to start them at all.
-        let disabled_remote_storage_override = "remote_storage={}";
-        let mut pageserver_process = self
-            .start_node(
-                &[disabled_remote_storage_override],
-                // Previous overrides will be taken from the config created before, don't overwrite them.
-                false,
-            )
-            .with_context(|| {
-                format!(
-                    "Failed to start a process for pageserver node {}",
-                    self.env.pageserver.id,
-                )
-            })?;
-
-        let init_result = self
-            .try_init_timeline(create_tenant, initial_timeline_id, pg_version)
-            .context("Failed to create initial tenant and timeline for pageserver");
-        match &init_result {
-            Ok(initial_timeline_id) => {
-                println!("Successfully initialized timeline {initial_timeline_id}")
-            }
-            Err(e) => eprintln!("{e:#}"),
-        }
-        background_process::send_stop_child_process(&pageserver_process)?;
-
-        let exit_code = pageserver_process.wait()?;
-        ensure!(
-            exit_code.success(),
-            format!(
-                "pageserver init failed with exit code {:?}",
-                exit_code.code()
-            )
-        );
-        println!(
-            "Stopped pageserver {} process with pid {}",
-            self.env.pageserver.id,
-            pageserver_process.id(),
-        );
-        init_result
-    }
-
-    fn try_init_timeline(
-        &self,
-        new_tenant_id: Option<TenantId>,
-        new_timeline_id: Option<TimelineId>,
-        pg_version: u32,
-    ) -> anyhow::Result<TimelineId> {
-        let initial_tenant_id = self.tenant_create(new_tenant_id, HashMap::new())?;
-        let initial_timeline_info = self.timeline_create(
-            initial_tenant_id,
-            new_timeline_id,
-            None,
-            None,
-            Some(pg_version),
-        )?;
-        Ok(initial_timeline_info.timeline_id)
+        })
     }
 
     pub fn repo_path(&self) -> PathBuf {
@@ -320,7 +252,7 @@ impl PageServerNode {
             let token = self
                 .env
                 .generate_auth_token(&Claims::new(None, Scope::SafekeeperData))?;
-            vec![("ZENITH_AUTH_TOKEN".to_owned(), token)]
+            vec![("NEON_AUTH_TOKEN".to_owned(), token)]
         } else {
             Vec::new()
         })

deny.toml

@@ -0,0 +1,90 @@
+# This file was auto-generated using `cargo deny init`.
+# cargo-deny is a cargo plugin that lets you lint your project's
+# dependency graph to ensure all your dependencies conform
+# to your expectations and requirements.
+
+# Root options
+targets = []
+all-features = false
+no-default-features = false
+feature-depth = 1
+
+# This section is considered when running `cargo deny check advisories`
+# More documentation for the advisories section can be found here:
+# https://embarkstudios.github.io/cargo-deny/checks/advisories/cfg.html
+[advisories]
+db-urls = ["https://github.com/rustsec/advisory-db"]
+vulnerability = "deny"
+unmaintained = "warn"
+yanked = "warn"
+notice = "warn"
+ignore = []
+
+# This section is considered when running `cargo deny check licenses`
+# More documentation for the licenses section can be found here:
+# https://embarkstudios.github.io/cargo-deny/checks/licenses/cfg.html
+[licenses]
+unlicensed = "deny"
+allow = [
+    "Apache-2.0",
+    "Artistic-2.0",
+    "BSD-2-Clause",
+    "BSD-3-Clause",
+    "ISC",
+    "MIT",
+    "MPL-2.0",
+    "OpenSSL",
+    "Unicode-DFS-2016",
+]
+deny = []
+copyleft = "warn"
+allow-osi-fsf-free = "neither"
+default = "deny"
+confidence-threshold = 0.8
+exceptions = [
+    # Zlib license has some restrictions if we decide to change sth
+    { allow = ["Zlib"], name = "const_format_proc_macros", version = "*" },
+    { allow = ["Zlib"], name = "const_format", version = "*" },
+]
+
+[[licenses.clarify]]
+name = "ring"
+version = "*"
+expression = "MIT AND ISC AND OpenSSL"
+license-files = [
+    { path = "LICENSE", hash = 0xbd0eed23 }
+]
+
+[licenses.private]
+ignore = true
+registries = []
+
+# This section is considered when running `cargo deny check bans`.
+# More documentation about the 'bans' section can be found here:
+# https://embarkstudios.github.io/cargo-deny/checks/bans/cfg.html
+[bans]
+multiple-versions = "warn"
+wildcards = "allow"
+highlight = "all"
+workspace-default-features = "allow"
+external-default-features = "allow"
+allow = []
+deny = []
+skip = []
+skip-tree = []
+
+# This section is considered when running `cargo deny check sources`.
+# More documentation about the 'sources' section can be found here:
+# https://embarkstudios.github.io/cargo-deny/checks/sources/cfg.html
+[sources]
+unknown-registry = "warn"
+unknown-git = "warn"
+allow-registry = ["https://github.com/rust-lang/crates.io-index"]
+allow-git = []
+
+[sources.allow-org]
+github = [
+    "neondatabase",
+]
+gitlab = []
+bitbucket = []

docs/authentication.md

@@ -65,7 +65,7 @@ There is no administrative API except those provided by PostgreSQL.
 
 #### Outgoing connections
 Compute connects to Pageserver for getting pages.
-The connection string is configured by the `neon.pageserver_connstring` PostgreSQL GUC, e.g. `postgresql://no_user:$ZENITH_AUTH_TOKEN@localhost:15028`.
+The connection string is configured by the `neon.pageserver_connstring` PostgreSQL GUC, e.g. `postgresql://no_user:$NEON_AUTH_TOKEN@localhost:15028`.
 The environment variable inside the connection string is substituted with
 the JWT token.
 
@@ -77,7 +77,7 @@ If the GUC is unset, no token is passed.
 
 Note that both tokens can be (and typically are) the same;
 the scope is the tenant and the token is usually passed through the
-`$ZENITH_AUTH_TOKEN` environment variable.
+`$NEON_AUTH_TOKEN` environment variable.
 
 ### Pageserver
 #### Overview
@@ -114,7 +114,7 @@ either of three values:
 Pageserver makes a connection to a Safekeeper for each active timeline.
 As Pageserver may want to access any timeline it has on the disk,
 it is given a blanket JWT token to access any data on any Safekeeper.
-This token is passed through an environment variable called `ZENITH_AUTH_TOKEN`
+This token is passed through an environment variable called `NEON_AUTH_TOKEN`
 (non-configurable as of writing this text).
 
 A better way _may be_ to store JWT token for each timeline next to it,

docs/consumption_metrics.md

@@ -0,0 +1,115 @@
+### Overview
+Pageserver and proxy periodically collect consumption metrics and push them to a HTTP endpoint.
+
+This doc describes current implementation details.
+For design details see [the RFC](./rfcs/021-metering.md) and [the discussion on Github](https://github.com/neondatabase/neon/pull/2884).
+
+- The metrics are collected in a separate thread, and the collection interval and endpoint are configurable.
+
+- Metrics are cached, so that we don't send unchanged metrics on every iteration.
+
+- Metrics are sent in batches of 1000 (see CHUNK_SIZE const) metrics max with no particular grouping guarantees.
+
+batch format is
+```json
+
+{ "events" : [metric1, metric2, ...]]}
+
+```
+See metric format examples below.
+
+- All metrics values are in bytes, unless otherwise specified.
+
+- Currently no retries are implemented.
+
+### Pageserver metrics
+
+#### Configuration
+The endpoint and the collection interval are specified in the pageserver config file (or can be passed as command line arguments):
+`metric_collection_endpoint` defaults to None, which means that metric collection is disabled by default.
+`metric_collection_interval` defaults to 10min
+
+#### Metrics
+
+Currently, the following metrics are collected:
+
+- `written_size`
+
+Amount of WAL produced , by a timeline, i.e. last_record_lsn
+This is an absolute, per-timeline metric.
+
+- `resident_size`
+
+Size of all the layer files in the tenant's directory on disk on the pageserver.
+This is an absolute, per-tenant metric.
+
+- `remote_storage_size`
+
+Size of the remote storage (S3) directory.
+This is an absolute, per-tenant metric.
+
+- `timeline_logical_size`
+Logical size of the data in the timeline
+This is an absolute, per-timeline metric.
+
+- `synthetic_storage_size`
+Size of all tenant's branches including WAL
+This is the same metric that `tenant/{tenant_id}/size` endpoint returns.
+This is an absolute, per-tenant metric.
+
+Synthetic storage size is calculated in a separate thread, so it might be slightly outdated.
+
+#### Format example
+
+```json
+{
+"metric": "remote_storage_size",
+"type": "absolute",
+"time": "2022-12-28T11:07:19.317310284Z",
+"idempotency_key": "2022-12-28 11:07:19.317310324 UTC-1-4019",
+"value": 12345454,
+"tenant_id": "5d07d9ce9237c4cd845ea7918c0afa7d",
+"timeline_id": "a03ebb4f5922a1c56ff7485cc8854143",
+}
+```
+
+`idempotency_key` is a unique key for each metric, so that we can deduplicate metrics.
+It is a combination of the time, node_id and a random number.
+
+### Proxy consumption metrics
+
+#### Configuration
+The endpoint and the collection interval can be passed as command line arguments for proxy:
+`metric_collection_endpoint` no default, which means that metric collection is disabled by default.
+`metric_collection_interval` no default
+
+#### Metrics
+
+Currently, only one proxy metric is collected:
+
+- `proxy_io_bytes_per_client`
+Outbound traffic per client.
+This is an incremental, per-endpoint metric.
+
+#### Format example
+
+```json
+{
+"metric": "proxy_io_bytes_per_client",
+"type": "incremental",
+"start_time": "2022-12-28T11:07:19.317310284Z",
+"stop_time": "2022-12-28T11:07:19.317310284Z",
+"idempotency_key": "2022-12-28 11:07:19.317310324 UTC-1-4019",
+"value": 12345454,
+"endpoint_id": "5d07d9ce9237c4cd845ea7918c0afa7d",
+}
+```
+
+The metric is incremental, so the value is the difference between the current and the previous value.
+If there is no previous value, the value, the value is the current value and the `start_time` equals `stop_time`.
+
+### TODO
+
+- [ ] Handle errors better: currently if one tenant fails to gather metrics, the whole iteration fails and metrics are not sent for any tenant.
+- [ ] Add retries
+- [ ] Tune the interval

docs/rfcs/021-metering.md

@@ -0,0 +1,186 @@
+# Consumption tracking
+
+
+# Goals
+
+This proposal is made with two mostly but not entirely overlapping goals:
+
+* Collect info that is needed for consumption-based billing
+* Cross-check AWS bills
+
+
+# Metrics
+
+There are six metrics to collect:
+
+* CPU time. Wall clock seconds * the current number of cores. We have a fixed ratio of memory to cores, so the current memory size is the function of the number of cores. Measured per each `endpoint`.
+
+* Traffic. In/out traffic on the proxy. Measured per each `endpoint`.
+
+* Written size. Amount of data we write. That is different from both traffic and storage size, as only during the writing we
+
+  a) occupy some disk bandwidth on safekeepers
+
+  b) necessarily cross AZ boundaries delivering WAL to all safekeepers
+
+  Each timeline/branch has at most one writer, so the data is collected per branch.
+
+* Synthetic storage size. That is what is exposed now with pageserver's `/v1/tenant/{}/size`. Looks like now it is per-tenant. (Side note: can we make it per branch to show as branch physical size in UI?)
+
+* Real storage size. That is the size of the tenant directory on the pageservers disk. Per-tenant.
+
+* S3 storage size. That is the size of the tenant data on S3. Per-tenant.
+
+That info should be enough to build an internal model that predicts AWS price (hence tracking `written data` and `real storage size`). As for the billing model we probably can get away with mentioning only `CPU time`, `synthetic storage size`, and `traffic` consumption.
+
+# Services participating in metrics collection
+
+## Proxy
+
+For actual implementation details check `/docs/consumption_metrics.md`
+
+Proxy is the only place that knows about traffic flow, so it tracks it and reports it with quite a small interval, let's say 1 minute. A small interval is needed here since the proxy is stateless, and any restart will reset accumulated consumption. Also proxy should report deltas since the last report, not an absolute value of the counter. Such kind of events is easier to integrate over a period of time to get the amount of traffic during some time interval.
+
+Example event:
+
+```json
+{
+"metric": "proxy_io_bytes_per_client",
+"type": "incremental",
+"start_time": "2022-12-28T11:07:19.317310284Z",
+"stop_time": "2022-12-28T11:07:19.317310284Z",
+"idempotency_key": "2022-12-28 11:07:19.317310324 UTC-1-4019",
+"value": 12345454,
+"endpoint_id": "5d07d9ce9237c4cd845ea7918c0afa7d",
+}
+```
+
+Since we report deltas over some period of time, it makes sense to include `event_start_time`/`event_stop_time` where `event_start_time` is the time of the previous report. That will allow us to identify metering gaps better (e.g., failed send/delivery).
+
+When there is no active connection proxy can avoid reporting anything. Also, deltas are additive, so several console instances serving the same user and endpoint can report traffic without coordination.
+
+## Console
+
+The console knows about start/stop events, so it knows the amount of CPU time allocated to each endpoint. It also knows about operation successes and failures and can avoid billing clients after unsuccessful 'suspend' events. The console doesn't know the current compute size within the allowed limits on the endpoint. So with CPU time, we do the following:
+
+* While we don't yet have the autoscaling console can report `cpu time` as the number of seconds since the last `start_compute` event.
+
+* When we have autoscaling, `autoscaler-agent` can report `cpu time`*`compute_units_count` in the same increments as the proxy reports traffic.
+
+Example event:
+
+```json
+{
+    "metric": "effective_compute_seconds",
+    "type": "increment",
+    "endpoint_id": "blazing-warrior-34",
+    "event_start_time": ...,
+    "event_stop_time": ...,
+    "value": 12345454,
+}
+```
+
+I'd also suggest reporting one value, `cpu time`*`compute_units_count`, instead of two separate fields as it makes event schema simpler (it is possible to treat it the same way as traffic) and preserves additivity.
+
+## Pageserver
+
+For actual implementation details check `/docs/consumption_metrics.md`
+
+Pageserver knows / has access to / can calculate the rest of the metrics:
+
+* Written size -- that is basically `last_received_lsn`,
+* Synthetic storage size -- there is a way to calculate it, albeit a costly one,
+* Real storage size -- there is a way to calculate it using a layer map or filesystem,
+* S3 storage size -- can calculate it by S3 API calls
+
+Some of those metrics are expensive to calculate, so the reporting period here is driven mainly by implementation details. We can set it to, for example, once per hour. Not a big deal since the pageserver is stateful, and all metrics can be reported as an absolute value, not increments. At the same time, a smaller reporting period improves UX, so it would be good to have something more real-time.
+
+`written size` is primarily a safekeeper-related metric, but since it is available on both pageserver and safekeeper, we can avoid reporting anything from the safekeeper.
+
+Example event:
+
+```json
+{
+"metric": "remote_storage_size",
+"type": "absolute",
+"time": "2022-12-28T11:07:19.317310284Z",
+"idempotency_key": "2022-12-28 11:07:19.317310324 UTC-1-4019",
+"value": 12345454,
+"tenant_id": "5d07d9ce9237c4cd845ea7918c0afa7d",
+"timeline_id": "a03ebb4f5922a1c56ff7485cc8854143",
+}
+```
+
+# Data collection
+
+## Push vs. pull
+
+We already have pull-based Prometheus metrics, so it is tempting to use them here too. However, in our setup, it is hard to tell when some metric changes. For example, garbage collection will constantly free some disk space over a week, even if the project is down for that week. We could also iterate through all existing tenants/branches/endpoints, but that means some amount of code to do that properly and most likely we will end up with some per-metric hacks in the collector to cut out some of the tenants that are surely not changing that metric.
+
+With the push model, it is easier to publish data only about actively changing metrics -- pageserver knows when it performs s3 offloads, garbage collection and starts/stops consuming data from the safekeeper; proxy knows about connected clients; console / autoscaler-agent knows about active cpu time.
+
+Hence, let's go with a push-based model.
+
+## Common bus vs. proxying through the console
+
+We can implement such push systems in a few ways:
+
+a. Each component pushes its metrics to the "common bus", namely segment, Kafka, or something similar. That approach scales well, but it would be harder to test it locally, will introduce new dependencies, we will have to distribute secrets for that connection to all of the components, etc. We would also have to loop back some of the events and their aggregates to the console, as we want to show some that metrics to the user in real-time.
+
+b. Each component can call HTTP `POST` with its events to the console, and the console can forward it to the segment for later integration with metronome / orb / onebill / etc. With that approach, only the console has to speak with segment. Also since that data passes through the console, the console can save the latest metrics values, so there is no need for constant feedback of that events back from the segment.
+
+# Implementation
+
+Each (proxy|pageserver|autoscaler-agent) sends consumption events to the single endpoint in the console:
+
+```json
+POST /usage_events HTTP/1.1
+Content-Type: application/json
+
+[
+{
+"metric": "remote_storage_size",
+"type": "absolute",
+"time": "2022-12-28T11:07:19.317310284Z",
+"idempotency_key": "2022-12-28 11:07:19.317310324 UTC-1-4019",
+"value": 12345454,
+"tenant_id": "5d07d9ce9237c4cd845ea7918c0afa7d",
+"timeline_id": "a03ebb4f5922a1c56ff7485cc8854143",
+},
+...
+]
+```
+
+![data flow](./images/metering.jpg)
+
+Events could be either:
+* `incremental` -- change in consumption since the previous event or service restart. That is `effective_cpu_seconds`, `traffic_in_bytes`, and `traffic_out_bytes`.
+* `absolute` -- that is the current value of a metric. All of the size-related metrics are absolute.
+
+Each service can post events at its own pace and bundle together data from different tenants/endpoints.
+
+The console algorithm upon receive of events could be the following:
+
+1. Create and send a segment event with the same content (possibly enriching it with tenant/timeline data for endpoint-based events).
+2. Update the latest state of per-tenant and per-endpoint metrics in the database.
+3. Check whether any of that metrics is above the allowed threshold and stop the project if necessary.
+
+Since all the data comes in batches, we can do the batch update to reduce the number of queries in the database. Proxy traffic is probably the most frequent metric, so with batching, we will have extra `number_of_proxies` requests to the database each minute. This is most likely fine for now but will generate many dead tuples in the console database. If that is the case, we can change step 2 to the following:
+
+2.1. Check if there $tenant_$metric / $endpoint_$metric key in Redis
+
+2.2. If no stored value is found and the metric is incremental, then fetch the current value from DWH (which keeps aggregated value for all the events) and publish it.
+
+2.3. Publish a new value (absolute metric) or add an increment to the stored value (incremental metric)
+
+## Consumption watchdog
+
+Since all the data goes through the console, we don't have to run any background thread/coroutines to check whether consumption is within the allowed limits. We only change consumption with `POST /usage_events`, so limit checks could be applied in the same handler.
+
+## Extensibility
+
+If we need to add a new metric (e.g. s3 traffic or something else), the console code should, by default, process it and publish segment event, even if the metric name is unknown to the console.
+
+## Naming & schema
+
+Each metric name should end up with units -- now `_seconds` and `_bytes`, and segment event should always have `tenant_id` and `timeline_id`/`endpoint_id` where applicable.

docs/sourcetree.md

@@ -18,10 +18,6 @@ Intended to be used in integration tests and in CLI tools for local installation
 Documentation of the Neon features and concepts.
 Now it is mostly dev documentation.
 
-`/monitoring`:
-
-TODO
-
 `/pageserver`:
 
 Neon storage service.
@@ -98,6 +94,13 @@ cargo hakari manage-deps
 
 If you don't have hakari installed (`error: no such subcommand: hakari`), install it by running `cargo install cargo-hakari`.
 
+### Checking Rust 3rd-parties
+[Cargo deny](https://embarkstudios.github.io/cargo-deny/index.html) is a cargo plugin that lets us lint project's dependency graph to ensure all dependencies conform to requirements. It detects security issues, matches licenses, and ensures crates only come from trusted sources.
+
+```bash
+cargo deny check
+```
+
 ## Using Python
 Note that Debian/Ubuntu Python packages are stale, as it commonly happens,
 so manual installation of dependencies is not recommended.

libs/consumption_metrics/Cargo.toml

@@ -0,0 +1,16 @@
+[package]
+name = "consumption_metrics"
+version = "0.1.0"
+edition = "2021"
+license = "Apache-2.0"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
+anyhow = "1.0.68"
+chrono = { version = "0.4", default-features = false, features = ["clock", "serde"] }
+rand = "0.8.3"
+serde = "1.0.152"
+serde_with = "2.1.0"
+utils = { version = "0.1.0", path = "../utils" }
+workspace_hack = { version = "0.1.0", path = "../../workspace_hack" }

libs/consumption_metrics/src/lib.rs

@@ -0,0 +1,50 @@
+//!
+//! Shared code for consumption metics collection
+//!
+use chrono::{DateTime, Utc};
+use rand::Rng;
+use serde::Serialize;
+
+#[derive(Serialize, Debug, Clone, Eq, PartialEq, Ord, PartialOrd)]
+#[serde(tag = "type")]
+pub enum EventType {
+    #[serde(rename = "absolute")]
+    Absolute { time: DateTime<Utc> },
+    #[serde(rename = "incremental")]
+    Incremental {
+        start_time: DateTime<Utc>,
+        stop_time: DateTime<Utc>,
+    },
+}
+
+#[derive(Serialize, Debug, Clone, Eq, PartialEq, Ord, PartialOrd)]
+pub struct Event<Extra> {
+    #[serde(flatten)]
+    #[serde(rename = "type")]
+    pub kind: EventType,
+
+    pub metric: &'static str,
+    pub idempotency_key: String,
+    pub value: u64,
+
+    #[serde(flatten)]
+    pub extra: Extra,
+}
+
+pub fn idempotency_key(node_id: String) -> String {
+    format!(
+        "{}-{}-{:04}",
+        Utc::now(),
+        node_id,
+        rand::thread_rng().gen_range(0..=9999)
+    )
+}
+
+pub const CHUNK_SIZE: usize = 1000;
+
+// Just a wrapper around a slice of events
+// to serialize it as `{"events" : [ ] }
+#[derive(serde::Serialize)]
+pub struct EventChunk<'a, T> {
+    pub events: &'a [T],
+}

libs/metrics/Cargo.toml

@@ -1,10 +1,12 @@
 [package]
 name = "metrics"
 version = "0.1.0"
-edition = "2021"
+edition.workspace = true
+license.workspace = true
 
 [dependencies]
-prometheus = {version = "0.13", default_features=false, features = ["process"]} # removes protobuf dependency
-libc = "0.2"
-once_cell = "1.13.0"
-workspace_hack = { version = "0.1", path = "../../workspace_hack" }
+prometheus.workspace = true
+libc.workspace = true
+once_cell.workspace = true
+
+workspace_hack.workspace = true

libs/pageserver_api/Cargo.toml

@@ -1,16 +1,17 @@
 [package]
 name = "pageserver_api"
 version = "0.1.0"
-edition = "2021"
+edition.workspace = true
+license.workspace = true
 
 [dependencies]
-serde = { version = "1.0", features = ["derive"] }
-serde_with = "2.0"
-const_format = "0.2.21"
-anyhow = { version = "1.0", features = ["backtrace"] }
-bytes = "1.0.1"
-byteorder = "1.4.3"
+serde.workspace = true
+serde_with.workspace = true
+const_format.workspace = true
+anyhow.workspace = true
+bytes.workspace = true
+byteorder.workspace = true
+utils.workspace = true
+postgres_ffi.workspace = true
 
-utils = { path = "../utils" }
-postgres_ffi = { path = "../postgres_ffi" }
-workspace_hack = { version = "0.1", path = "../../workspace_hack" }
+workspace_hack.workspace = true

libs/pageserver_api/src/models.rs

@@ -1,4 +1,4 @@
-use std::num::NonZeroU64;
+use std::num::{NonZeroU64, NonZeroUsize};
 
 use byteorder::{BigEndian, ReadBytesExt};
 use serde::{Deserialize, Serialize};
@@ -210,6 +210,11 @@ pub struct TimelineInfo {
     pub state: TimelineState,
 }
 
+#[derive(Debug, Serialize, Deserialize)]
+pub struct DownloadRemoteLayersTaskSpawnRequest {
+    pub max_concurrent_downloads: NonZeroUsize,
+}
+
 #[derive(Debug, Serialize, Deserialize, Clone)]
 pub struct DownloadRemoteLayersTaskInfo {
     pub task_id: String,

libs/postgres_connection/Cargo.toml

@@ -1,17 +1,17 @@
 [package]
 name = "postgres_connection"
 version = "0.1.0"
-edition = "2021"
-
-# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+edition.workspace = true
+license.workspace = true
 
 [dependencies]
-anyhow = "1.0"
-itertools = "0.10.3"
-postgres = { git = "https://github.com/neondatabase/rust-postgres.git", rev = "43e6db254a97fdecbce33d8bc0890accfd74495e" }
-tokio-postgres = { git = "https://github.com/neondatabase/rust-postgres.git", rev="43e6db254a97fdecbce33d8bc0890accfd74495e" }
-url = "2.2.2"
-workspace_hack = { version = "0.1", path = "../../workspace_hack" }
+anyhow.workspace = true
+itertools.workspace = true
+postgres.workspace = true
+tokio-postgres.workspace = true
+url.workspace = true
+
+workspace_hack.workspace = true
 
 [dev-dependencies]
-once_cell = "1.13.0"
+once_cell.workspace = true

libs/postgres_ffi/Cargo.toml

@@ -1,29 +1,31 @@
 [package]
 name = "postgres_ffi"
 version = "0.1.0"
-edition = "2021"
+edition.workspace = true
+license.workspace = true
 
 [dependencies]
-rand = "0.8.3"
-regex = "1.4.5"
-bytes = "1.0.1"
-byteorder = "1.4.3"
-anyhow = "1.0"
-crc32c = "0.6.0"
-hex = "0.4.3"
-once_cell = "1.13.0"
-log = "0.4.14"
-memoffset = "0.7"
-thiserror = "1.0"
-serde = { version = "1.0", features = ["derive"] }
-utils = { path = "../utils" }
-workspace_hack = { version = "0.1", path = "../../workspace_hack" }
+rand.workspace = true
+regex.workspace = true
+bytes.workspace = true
+byteorder.workspace = true
+anyhow.workspace = true
+crc32c.workspace = true
+hex.workspace = true
+once_cell.workspace = true
+log.workspace = true
+memoffset.workspace = true
+thiserror.workspace = true
+serde.workspace = true
+utils.workspace = true
+
+workspace_hack.workspace = true
 
 [dev-dependencies]
-env_logger = "0.9"
-postgres = { git = "https://github.com/neondatabase/rust-postgres.git", rev="43e6db254a97fdecbce33d8bc0890accfd74495e" }
+env_logger.workspace = true
+postgres.workspace = true
 wal_craft = { path = "wal_craft" }
 
 [build-dependencies]
-anyhow = "1.0"
-bindgen = "0.61"
+anyhow.workspace = true
+bindgen.workspace = true

libs/postgres_ffi/wal_craft/Cargo.toml

@@ -1,17 +1,17 @@
 [package]
 name = "wal_craft"
 version = "0.1.0"
-edition = "2021"
-
-# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+edition.workspace = true
+license.workspace = true
 
 [dependencies]
-anyhow = "1.0"
-clap = "4.0"
-env_logger = "0.9"
-log = "0.4"
-once_cell = "1.13.0"
-postgres = { git = "https://github.com/neondatabase/rust-postgres.git", rev="43e6db254a97fdecbce33d8bc0890accfd74495e" }
-postgres_ffi = { path = "../" }
-tempfile = "3.2"
-workspace_hack = { version = "0.1", path = "../../../workspace_hack" }
+anyhow.workspace = true
+clap.workspace = true
+env_logger.workspace = true
+log.workspace = true
+once_cell.workspace = true
+postgres.workspace = true
+postgres_ffi.workspace = true
+tempfile.workspace = true
+
+workspace_hack.workspace = true

libs/pq_proto/Cargo.toml

@@ -1,16 +1,18 @@
 [package]
 name = "pq_proto"
 version = "0.1.0"
-edition = "2021"
+edition.workspace = true
+license.workspace = true
 
 [dependencies]
-anyhow = "1.0"
-bytes = "1.0.1"
-pin-project-lite = "0.2.7"
-postgres-protocol = { git = "https://github.com/neondatabase/rust-postgres.git", rev="43e6db254a97fdecbce33d8bc0890accfd74495e" }
-rand = "0.8.3"
-serde = { version = "1.0", features = ["derive"] }
-tokio = { version = "1.17", features = ["macros"] }
-tracing = "0.1"
+anyhow.workspace = true
+bytes.workspace = true
+pin-project-lite.workspace = true
+postgres-protocol.workspace = true
+rand.workspace = true
+serde.workspace = true
+tokio.workspace = true
+tracing.workspace = true
+thiserror.workspace = true
 
-workspace_hack = { version = "0.1", path = "../../workspace_hack" }
+workspace_hack.workspace = true

libs/pq_proto/src/lib.rs

@@ -5,7 +5,7 @@
 // Tools for calling certain async methods in sync contexts.
 pub mod sync;
 
-use anyhow::{bail, ensure, Context, Result};
+use anyhow::{ensure, Context, Result};
 use bytes::{Buf, BufMut, Bytes, BytesMut};
 use postgres_protocol::PG_EPOCH;
 use serde::{Deserialize, Serialize};
@@ -194,6 +194,35 @@ macro_rules! retry_read {
     };
 }
 
+/// An error occured during connection being open.
+#[derive(thiserror::Error, Debug)]
+pub enum ConnectionError {
+    /// IO error during writing to or reading from the connection socket.
+    #[error("Socket IO error: {0}")]
+    Socket(std::io::Error),
+    /// Invalid packet was received from client
+    #[error("Protocol error: {0}")]
+    Protocol(String),
+    /// Failed to parse a protocol mesage
+    #[error("Message parse error: {0}")]
+    MessageParse(anyhow::Error),
+}
+
+impl From<anyhow::Error> for ConnectionError {
+    fn from(e: anyhow::Error) -> Self {
+        Self::MessageParse(e)
+    }
+}
+
+impl ConnectionError {
+    pub fn into_io_error(self) -> io::Error {
+        match self {
+            ConnectionError::Socket(io) => io,
+            other => io::Error::new(io::ErrorKind::Other, other.to_string()),
+        }
+    }
+}
+
 impl FeMessage {
     /// Read one message from the stream.
     /// This function returns `Ok(None)` in case of EOF.
@@ -216,7 +245,9 @@ impl FeMessage {
     /// }
     /// ```
     #[inline(never)]
-    pub fn read(stream: &mut (impl io::Read + Unpin)) -> anyhow::Result<Option<FeMessage>> {
+    pub fn read(
+        stream: &mut (impl io::Read + Unpin),
+    ) -> Result<Option<FeMessage>, ConnectionError> {
         Self::read_fut(&mut AsyncishRead(stream)).wait()
     }
 
@@ -224,7 +255,7 @@ impl FeMessage {
     /// See documentation for `Self::read`.
     pub fn read_fut<Reader>(
         stream: &mut Reader,
-    ) -> SyncFuture<Reader, impl Future<Output = anyhow::Result<Option<FeMessage>>> + '_>
+    ) -> SyncFuture<Reader, impl Future<Output = Result<Option<FeMessage>, ConnectionError>> + '_>
     where
         Reader: tokio::io::AsyncRead + Unpin,
     {
@@ -238,17 +269,21 @@ impl FeMessage {
             let tag = match retry_read!(stream.read_u8().await) {
                 Ok(b) => b,
                 Err(e) if e.kind() == io::ErrorKind::UnexpectedEof => return Ok(None),
-                Err(e) => return Err(e.into()),
+                Err(e) => return Err(ConnectionError::Socket(e)),
             };
 
             // The message length includes itself, so it better be at least 4.
-            let len = retry_read!(stream.read_u32().await)?
+            let len = retry_read!(stream.read_u32().await)
+                .map_err(ConnectionError::Socket)?
                 .checked_sub(4)
-                .context("invalid message length")?;
+                .ok_or_else(|| ConnectionError::Protocol("invalid message length".to_string()))?;
 
             let body = {
                 let mut buffer = vec![0u8; len as usize];
-                stream.read_exact(&mut buffer).await?;
+                stream
+                    .read_exact(&mut buffer)
+                    .await
+                    .map_err(ConnectionError::Socket)?;
                 Bytes::from(buffer)
             };
 
@@ -265,7 +300,11 @@ impl FeMessage {
                 b'c' => Ok(Some(FeMessage::CopyDone)),
                 b'f' => Ok(Some(FeMessage::CopyFail)),
                 b'p' => Ok(Some(FeMessage::PasswordMessage(body))),
-                tag => bail!("unknown message tag: {},'{:?}'", tag, body),
+                tag => {
+                    return Err(ConnectionError::Protocol(format!(
+                        "unknown message tag: {tag},'{body:?}'"
+                    )))
+                }
             }
         })
     }
@@ -275,7 +314,9 @@ impl FeStartupPacket {
     /// Read startup message from the stream.
     // XXX: It's tempting yet undesirable to accept `stream` by value,
     // since such a change will cause user-supplied &mut references to be consumed
-    pub fn read(stream: &mut (impl io::Read + Unpin)) -> anyhow::Result<Option<FeMessage>> {
+    pub fn read(
+        stream: &mut (impl io::Read + Unpin),
+    ) -> Result<Option<FeMessage>, ConnectionError> {
         Self::read_fut(&mut AsyncishRead(stream)).wait()
     }
 
@@ -284,7 +325,7 @@ impl FeStartupPacket {
     // since such a change will cause user-supplied &mut references to be consumed
     pub fn read_fut<Reader>(
         stream: &mut Reader,
-    ) -> SyncFuture<Reader, impl Future<Output = anyhow::Result<Option<FeMessage>>> + '_>
+    ) -> SyncFuture<Reader, impl Future<Output = Result<Option<FeMessage>, ConnectionError>> + '_>
     where
         Reader: tokio::io::AsyncRead + Unpin,
     {
@@ -302,31 +343,41 @@ impl FeStartupPacket {
             let len = match retry_read!(stream.read_u32().await) {
                 Ok(len) => len as usize,
                 Err(e) if e.kind() == io::ErrorKind::UnexpectedEof => return Ok(None),
-                Err(e) => return Err(e.into()),
+                Err(e) => return Err(ConnectionError::Socket(e)),
             };
 
             #[allow(clippy::manual_range_contains)]
             if len < 4 || len > MAX_STARTUP_PACKET_LENGTH {
-                bail!("invalid message length");
+                return Err(ConnectionError::Protocol(format!(
+                    "invalid message length {len}"
+                )));
             }
 
-            let request_code = retry_read!(stream.read_u32().await)?;
+            let request_code =
+                retry_read!(stream.read_u32().await).map_err(ConnectionError::Socket)?;
 
             // the rest of startup packet are params
             let params_len = len - 8;
             let mut params_bytes = vec![0u8; params_len];
-            stream.read_exact(params_bytes.as_mut()).await?;
+            stream
+                .read_exact(params_bytes.as_mut())
+                .await
+                .map_err(ConnectionError::Socket)?;
 
             // Parse params depending on request code
             let req_hi = request_code >> 16;
             let req_lo = request_code & ((1 << 16) - 1);
             let message = match (req_hi, req_lo) {
                 (RESERVED_INVALID_MAJOR_VERSION, CANCEL_REQUEST_CODE) => {
-                    ensure!(params_len == 8, "expected 8 bytes for CancelRequest params");
+                    if params_len != 8 {
+                        return Err(ConnectionError::Protocol(
+                            "expected 8 bytes for CancelRequest params".to_string(),
+                        ));
+                    }
                     let mut cursor = Cursor::new(params_bytes);
                     FeStartupPacket::CancelRequest(CancelKeyData {
-                        backend_pid: cursor.read_i32().await?,
-                        cancel_key: cursor.read_i32().await?,
+                        backend_pid: cursor.read_i32().await.map_err(ConnectionError::Socket)?,
+                        cancel_key: cursor.read_i32().await.map_err(ConnectionError::Socket)?,
                     })
                 }
                 (RESERVED_INVALID_MAJOR_VERSION, NEGOTIATE_SSL_CODE) => {
@@ -338,7 +389,9 @@ impl FeStartupPacket {
                     FeStartupPacket::GssEncRequest
                 }
                 (RESERVED_INVALID_MAJOR_VERSION, unrecognized_code) => {
-                    bail!("Unrecognized request code {}", unrecognized_code)
+                    return Err(ConnectionError::Protocol(format!(
+                        "Unrecognized request code {unrecognized_code}"
+                    )));
                 }
                 // TODO bail if protocol major_version is not 3?
                 (major_version, minor_version) => {
@@ -346,15 +399,21 @@ impl FeStartupPacket {
                     // See `postgres: ProcessStartupPacket, build_startup_packet`.
                     let mut tokens = str::from_utf8(&params_bytes)
                         .context("StartupMessage params: invalid utf-8")?
-                        .strip_suffix('\0') // drop packet's own null terminator
-                        .context("StartupMessage params: missing null terminator")?
+                        .strip_suffix('\0') // drop packet's own null
+                        .ok_or_else(|| {
+                            ConnectionError::Protocol(
+                                "StartupMessage params: missing null terminator".to_string(),
+                            )
+                        })?
                         .split_terminator('\0');
 
                     let mut params = HashMap::new();
                     while let Some(name) = tokens.next() {
-                        let value = tokens
-                            .next()
-                            .context("StartupMessage params: key without value")?;
+                        let value = tokens.next().ok_or_else(|| {
+                            ConnectionError::Protocol(
+                                "StartupMessage params: key without value".to_string(),
+                            )
+                        })?;
 
                         params.insert(name.to_owned(), value.to_owned());
                     }
@@ -458,7 +517,7 @@ pub enum BeMessage<'a> {
     CloseComplete,
     // None means column is NULL
     DataRow(&'a [Option<&'a [u8]>]),
-    ErrorResponse(&'a str),
+    ErrorResponse(&'a str, Option<&'a [u8; 5]>),
     /// Single byte - used in response to SSLRequest/GSSENCRequest.
     EncryptionResponse(bool),
     NoData,
@@ -606,7 +665,7 @@ fn write_body<R>(buf: &mut BytesMut, f: impl FnOnce(&mut BytesMut) -> R) -> R {
 }
 
 /// Safe write of s into buf as cstring (String in the protocol).
-fn write_cstr(s: impl AsRef<[u8]>, buf: &mut BytesMut) -> Result<(), io::Error> {
+fn write_cstr(s: impl AsRef<[u8]>, buf: &mut BytesMut) -> io::Result<()> {
     let bytes = s.as_ref();
     if bytes.contains(&0) {
         return Err(io::Error::new(
@@ -626,6 +685,8 @@ fn read_cstr(buf: &mut Bytes) -> anyhow::Result<Bytes> {
     Ok(result)
 }
 
+pub const SQLSTATE_INTERNAL_ERROR: &[u8; 5] = b"XX000";
+
 impl<'a> BeMessage<'a> {
     /// Write message to the given buf.
     // Unlike the reading side, we use BytesMut
@@ -765,10 +826,7 @@ impl<'a> BeMessage<'a> {
             // First byte of each field represents type of this field. Set just enough fields
             // to satisfy rust-postgres client: 'S' -- severity, 'C' -- error, 'M' -- error
             // message text.
-            BeMessage::ErrorResponse(error_msg) => {
-                // For all the errors set Severity to Error and error code to
-                // 'internal error'.
-
+            BeMessage::ErrorResponse(error_msg, pg_error_code) => {
                 // 'E' signalizes ErrorResponse messages
                 buf.put_u8(b'E');
                 write_body(buf, |buf| {
@@ -776,7 +834,9 @@ impl<'a> BeMessage<'a> {
                     buf.put_slice(b"ERROR\0");
 
                     buf.put_u8(b'C'); // SQLSTATE error code
-                    buf.put_slice(b"CXX000\0");
+                    buf.put_slice(&terminate_code(
+                        pg_error_code.unwrap_or(SQLSTATE_INTERNAL_ERROR),
+                    ));
 
                     buf.put_u8(b'M'); // the message
                     write_cstr(error_msg, buf)?;
@@ -799,7 +859,7 @@ impl<'a> BeMessage<'a> {
                     buf.put_slice(b"NOTICE\0");
 
                     buf.put_u8(b'C'); // SQLSTATE error code
-                    buf.put_slice(b"CXX000\0");
+                    buf.put_slice(&terminate_code(SQLSTATE_INTERNAL_ERROR));
 
                     buf.put_u8(b'M'); // the message
                     write_cstr(error_msg.as_bytes(), buf)?;
@@ -1087,3 +1147,12 @@ mod tests {
         let _ = FeStartupPacket::read_fut(stream).await;
     }
 }
+
+fn terminate_code(code: &[u8; 5]) -> [u8; 6] {
+    let mut terminated = [0; 6];
+    for (i, &elem) in code.iter().enumerate() {
+        terminated[i] = elem;
+    }
+
+    terminated
+}

libs/remote_storage/Cargo.toml

@@ -1,27 +1,28 @@
 [package]
 name = "remote_storage"
 version = "0.1.0"
-edition = "2021"
+edition.workspace = true
+license.workspace = true
 
 [dependencies]
-anyhow = { version = "1.0", features = ["backtrace"] }
-async-trait = "0.1"
-metrics = { version = "0.1", path = "../metrics" }
-utils = { version = "0.1", path = "../utils" }
-once_cell = "1.13.0"
-aws-smithy-http = "0.51.0"
-aws-types = "0.51.0"
-aws-config = { version = "0.51.0", default-features = false, features=["rustls"] }
-aws-sdk-s3 = "0.21.0"
-hyper = { version = "0.14", features = ["stream"] }
-serde = { version = "1.0", features = ["derive"] }
-serde_json = "1"
-tokio = { version = "1.17", features = ["sync", "macros", "fs", "io-util"] }
-tokio-util = { version = "0.7", features = ["io"] }
-toml_edit = { version = "0.14", features = ["easy"] }
-tracing = "0.1.27"
+anyhow.workspace = true
+async-trait.workspace = true
+once_cell.workspace = true
+aws-smithy-http.workspace = true
+aws-types.workspace = true
+aws-config.workspace = true
+aws-sdk-s3.workspace = true
+hyper = { workspace = true, features = ["stream"] }
+serde.workspace = true
+serde_json.workspace = true
+tokio = { workspace = true, features = ["sync", "fs", "io-util"] }
+tokio-util.workspace = true
+toml_edit.workspace = true
+tracing.workspace = true
+metrics.workspace = true
+utils.workspace = true
 
-workspace_hack = { version = "0.1", path = "../../workspace_hack" }
+workspace_hack.workspace = true
 
 [dev-dependencies]
-tempfile = "3.2"
+tempfile.workspace = true

libs/safekeeper_api/Cargo.toml

@@ -1,12 +1,13 @@
 [package]
 name = "safekeeper_api"
 version = "0.1.0"
-edition = "2021"
+edition.workspace = true
+license.workspace = true
 
 [dependencies]
-serde = { version = "1.0", features = ["derive"] }
-serde_with = "2.0"
-const_format = "0.2.21"
+serde.workspace = true
+serde_with.workspace = true
+const_format.workspace = true
+utils.workspace = true
 
-utils = { path = "../utils" }
-workspace_hack = { version = "0.1", path = "../../workspace_hack" }
+workspace_hack.workspace = true

libs/tenant_size_model/Cargo.toml

@@ -1,8 +1,11 @@
 [package]
 name = "tenant_size_model"
 version = "0.1.0"
-edition = "2021"
+edition.workspace = true
 publish = false
+license.workspace = true
 
 [dependencies]
-workspace_hack = { version = "0.1", path = "../../workspace_hack" }
+anyhow.workspace = true
+
+workspace_hack.workspace = true

libs/tenant_size_model/src/lib.rs

@@ -1,6 +1,8 @@
 use std::borrow::Cow;
 use std::collections::HashMap;
 
+use anyhow::Context;
+
 /// Pricing model or history size builder.
 ///
 /// Maintains knowledge of the branches and their modifications. Generic over the branch name key
@@ -132,22 +134,25 @@ impl<K: std::hash::Hash + Eq + 'static> Storage<K> {
         op: Cow<'static, str>,
         lsn: u64,
         size: Option<u64>,
-    ) where
+    ) -> anyhow::Result<()>
+    where
         K: std::borrow::Borrow<Q>,
-        Q: std::hash::Hash + Eq,
+        Q: std::hash::Hash + Eq + std::fmt::Debug,
     {
-        let lastseg_id = *self.branches.get(branch).unwrap();
+        let Some(lastseg_id) = self.branches.get(branch).copied() else { anyhow::bail!("branch not found: {branch:?}") };
         let newseg_id = self.segments.len();
         let lastseg = &mut self.segments[lastseg_id];
 
         assert!(lsn > lastseg.end_lsn);
 
+        let Some(start_size) = lastseg.end_size else { anyhow::bail!("no end_size on latest segment for {branch:?}") };
+
         let newseg = Segment {
             op,
             parent: Some(lastseg_id),
             start_lsn: lastseg.end_lsn,
             end_lsn: lsn,
-            start_size: lastseg.end_size.unwrap(),
+            start_size,
             end_size: size,
             children_after: Vec::new(),
             needed: false,
@@ -156,6 +161,8 @@ impl<K: std::hash::Hash + Eq + 'static> Storage<K> {
 
         self.segments.push(newseg);
         *self.branches.get_mut(branch).expect("read already") = newseg_id;
+
+        Ok(())
     }
 
     /// Advances the branch with the named operation, by the relative LSN and logical size bytes.
@@ -165,21 +172,24 @@ impl<K: std::hash::Hash + Eq + 'static> Storage<K> {
         op: Cow<'static, str>,
         lsn_bytes: u64,
         size_bytes: i64,
-    ) where
+    ) -> anyhow::Result<()>
+    where
         K: std::borrow::Borrow<Q>,
-        Q: std::hash::Hash + Eq,
+        Q: std::hash::Hash + Eq + std::fmt::Debug,
     {
-        let lastseg_id = *self.branches.get(branch).unwrap();
+        let Some(lastseg_id) = self.branches.get(branch).copied() else { anyhow::bail!("branch not found: {branch:?}") };
         let newseg_id = self.segments.len();
         let lastseg = &mut self.segments[lastseg_id];
 
+        let Some(last_end_size) = lastseg.end_size else { anyhow::bail!("no end_size on latest segment for {branch:?}") };
+
         let newseg = Segment {
             op,
             parent: Some(lastseg_id),
             start_lsn: lastseg.end_lsn,
             end_lsn: lastseg.end_lsn + lsn_bytes,
-            start_size: lastseg.end_size.unwrap(),
-            end_size: Some((lastseg.end_size.unwrap() as i64 + size_bytes) as u64),
+            start_size: last_end_size,
+            end_size: Some((last_end_size as i64 + size_bytes) as u64),
             children_after: Vec::new(),
             needed: false,
         };
@@ -187,50 +197,54 @@ impl<K: std::hash::Hash + Eq + 'static> Storage<K> {
 
         self.segments.push(newseg);
         *self.branches.get_mut(branch).expect("read already") = newseg_id;
+        Ok(())
     }
 
-    pub fn insert<Q: ?Sized>(&mut self, branch: &Q, bytes: u64)
+    pub fn insert<Q: ?Sized>(&mut self, branch: &Q, bytes: u64) -> anyhow::Result<()>
     where
         K: std::borrow::Borrow<Q>,
-        Q: std::hash::Hash + Eq,
+        Q: std::hash::Hash + Eq + std::fmt::Debug,
     {
-        self.modify_branch(branch, "insert".into(), bytes, bytes as i64);
+        self.modify_branch(branch, "insert".into(), bytes, bytes as i64)
     }
 
-    pub fn update<Q: ?Sized>(&mut self, branch: &Q, bytes: u64)
+    pub fn update<Q: ?Sized>(&mut self, branch: &Q, bytes: u64) -> anyhow::Result<()>
     where
         K: std::borrow::Borrow<Q>,
-        Q: std::hash::Hash + Eq,
+        Q: std::hash::Hash + Eq + std::fmt::Debug,
     {
-        self.modify_branch(branch, "update".into(), bytes, 0i64);
+        self.modify_branch(branch, "update".into(), bytes, 0i64)
     }
 
-    pub fn delete<Q: ?Sized>(&mut self, branch: &Q, bytes: u64)
+    pub fn delete<Q: ?Sized>(&mut self, branch: &Q, bytes: u64) -> anyhow::Result<()>
     where
         K: std::borrow::Borrow<Q>,
-        Q: std::hash::Hash + Eq,
+        Q: std::hash::Hash + Eq + std::fmt::Debug,
     {
-        self.modify_branch(branch, "delete".into(), bytes, -(bytes as i64));
+        self.modify_branch(branch, "delete".into(), bytes, -(bytes as i64))
     }
 
-    /// Panics if the parent branch cannot be found.
-    pub fn branch<Q: ?Sized>(&mut self, parent: &Q, name: K)
+    pub fn branch<Q: ?Sized>(&mut self, parent: &Q, name: K) -> anyhow::Result<()>
     where
-        K: std::borrow::Borrow<Q>,
-        Q: std::hash::Hash + Eq,
+        K: std::borrow::Borrow<Q> + std::fmt::Debug,
+        Q: std::hash::Hash + Eq + std::fmt::Debug,
     {
         // Find the right segment
-        let branchseg_id = *self
-            .branches
-            .get(parent)
-            .expect("should had found the parent by key");
+        let branchseg_id = *self.branches.get(parent).with_context(|| {
+            format!(
+                "should had found the parent {:?} by key. in branches {:?}",
+                parent, self.branches
+            )
+        })?;
+
         let _branchseg = &mut self.segments[branchseg_id];
 
         // Create branch name for it
         self.branches.insert(name, branchseg_id);
+        Ok(())
     }
 
-    pub fn calculate(&mut self, retention_period: u64) -> SegmentSize {
+    pub fn calculate(&mut self, retention_period: u64) -> anyhow::Result<SegmentSize> {
         // Phase 1: Mark all the segments that need to be retained
         for (_branch, &last_seg_id) in self.branches.iter() {
             let last_seg = &self.segments[last_seg_id];
@@ -255,7 +269,7 @@ impl<K: std::hash::Hash + Eq + 'static> Storage<K> {
         self.size_from_snapshot_later(0)
     }
 
-    fn size_from_wal(&self, seg_id: usize) -> SegmentSize {
+    fn size_from_wal(&self, seg_id: usize) -> anyhow::Result<SegmentSize> {
         let seg = &self.segments[seg_id];
 
         let this_size = seg.end_lsn - seg.start_lsn;
@@ -266,10 +280,10 @@ impl<K: std::hash::Hash + Eq + 'static> Storage<K> {
         for &child_id in seg.children_after.iter() {
             // try each child both ways
             let child = &self.segments[child_id];
-            let p1 = self.size_from_wal(child_id);
+            let p1 = self.size_from_wal(child_id)?;
 
             let p = if !child.needed {
-                let p2 = self.size_from_snapshot_later(child_id);
+                let p2 = self.size_from_snapshot_later(child_id)?;
                 if p1.total() < p2.total() {
                     p1
                 } else {
@@ -280,15 +294,15 @@ impl<K: std::hash::Hash + Eq + 'static> Storage<K> {
             };
             children.push(p);
         }
-        SegmentSize {
+        Ok(SegmentSize {
             seg_id,
             method: if seg.needed { WalNeeded } else { Wal },
             this_size,
             children,
-        }
+        })
     }
 
-    fn size_from_snapshot_later(&self, seg_id: usize) -> SegmentSize {
+    fn size_from_snapshot_later(&self, seg_id: usize) -> anyhow::Result<SegmentSize> {
         // If this is needed, then it's time to do the snapshot and continue
         // with wal method.
         let seg = &self.segments[seg_id];
@@ -299,10 +313,10 @@ impl<K: std::hash::Hash + Eq + 'static> Storage<K> {
             for &child_id in seg.children_after.iter() {
                 // try each child both ways
                 let child = &self.segments[child_id];
-                let p1 = self.size_from_wal(child_id);
+                let p1 = self.size_from_wal(child_id)?;
 
                 let p = if !child.needed {
-                    let p2 = self.size_from_snapshot_later(child_id);
+                    let p2 = self.size_from_snapshot_later(child_id)?;
                     if p1.total() < p2.total() {
                         p1
                     } else {
@@ -313,12 +327,12 @@ impl<K: std::hash::Hash + Eq + 'static> Storage<K> {
                 };
                 children.push(p);
             }
-            SegmentSize {
+            Ok(SegmentSize {
                 seg_id,
                 method: WalNeeded,
                 this_size: seg.start_size,
                 children,
-            }
+            })
         } else {
             // If any of the direct children are "needed", need to be able to reconstruct here
             let mut children_needed = false;
@@ -333,7 +347,7 @@ impl<K: std::hash::Hash + Eq + 'static> Storage<K> {
             let method1 = if !children_needed {
                 let mut children = Vec::new();
                 for child in seg.children_after.iter() {
-                    children.push(self.size_from_snapshot_later(*child));
+                    children.push(self.size_from_snapshot_later(*child)?);
                 }
                 Some(SegmentSize {
                     seg_id,
@@ -349,20 +363,25 @@ impl<K: std::hash::Hash + Eq + 'static> Storage<K> {
             let method2 = if children_needed || seg.children_after.len() >= 2 {
                 let mut children = Vec::new();
                 for child in seg.children_after.iter() {
-                    children.push(self.size_from_wal(*child));
+                    children.push(self.size_from_wal(*child)?);
                 }
+                let Some(this_size) = seg.end_size else { anyhow::bail!("no end_size at junction {seg_id}") };
                 Some(SegmentSize {
                     seg_id,
                     method: SnapshotAfter,
-                    this_size: seg.end_size.unwrap(),
+                    this_size,
                     children,
                 })
             } else {
                 None
             };
 
-            match (method1, method2) {
-                (None, None) => panic!(),
+            Ok(match (method1, method2) {
+                (None, None) => anyhow::bail!(
+                    "neither method was applicable: children_after={}, children_needed={}",
+                    seg.children_after.len(),
+                    children_needed
+                ),
                 (Some(method), None) => method,
                 (None, Some(method)) => method,
                 (Some(method1), Some(method2)) => {
@@ -372,7 +391,7 @@ impl<K: std::hash::Hash + Eq + 'static> Storage<K> {
                         method2
                     }
                 }
-            }
+            })
         }
     }
 

libs/tenant_size_model/src/main.rs

@@ -7,118 +7,118 @@
 use tenant_size_model::{Segment, SegmentSize, Storage};
 
 // Main branch only. Some updates on it.
-fn scenario_1() -> (Vec<Segment>, SegmentSize) {
+fn scenario_1() -> anyhow::Result<(Vec<Segment>, SegmentSize)> {
     // Create main branch
     let mut storage = Storage::new("main");
 
     // Bulk load 5 GB of data to it
-    storage.insert("main", 5_000);
+    storage.insert("main", 5_000)?;
 
     // Stream of updates
     for _ in 0..5 {
-        storage.update("main", 1_000);
+        storage.update("main", 1_000)?;
     }
 
-    let size = storage.calculate(1000);
+    let size = storage.calculate(1000)?;
 
-    (storage.into_segments(), size)
+    Ok((storage.into_segments(), size))
 }
 
 // Main branch only. Some updates on it.
-fn scenario_2() -> (Vec<Segment>, SegmentSize) {
+fn scenario_2() -> anyhow::Result<(Vec<Segment>, SegmentSize)> {
     // Create main branch
     let mut storage = Storage::new("main");
 
     // Bulk load 5 GB of data to it
-    storage.insert("main", 5_000);
+    storage.insert("main", 5_000)?;
 
     // Stream of updates
     for _ in 0..5 {
-        storage.update("main", 1_000);
+        storage.update("main", 1_000)?;
     }
 
     // Branch
-    storage.branch("main", "child");
-    storage.update("child", 1_000);
+    storage.branch("main", "child")?;
+    storage.update("child", 1_000)?;
 
     // More updates on parent
-    storage.update("main", 1_000);
+    storage.update("main", 1_000)?;
 
-    let size = storage.calculate(1000);
+    let size = storage.calculate(1000)?;
 
-    (storage.into_segments(), size)
+    Ok((storage.into_segments(), size))
 }
 
 // Like 2, but more updates on main
-fn scenario_3() -> (Vec<Segment>, SegmentSize) {
+fn scenario_3() -> anyhow::Result<(Vec<Segment>, SegmentSize)> {
     // Create main branch
     let mut storage = Storage::new("main");
 
     // Bulk load 5 GB of data to it
-    storage.insert("main", 5_000);
+    storage.insert("main", 5_000)?;
 
     // Stream of updates
     for _ in 0..5 {
-        storage.update("main", 1_000);
+        storage.update("main", 1_000)?;
     }
 
     // Branch
-    storage.branch("main", "child");
-    storage.update("child", 1_000);
+    storage.branch("main", "child")?;
+    storage.update("child", 1_000)?;
 
     // More updates on parent
     for _ in 0..5 {
-        storage.update("main", 1_000);
+        storage.update("main", 1_000)?;
     }
 
-    let size = storage.calculate(1000);
+    let size = storage.calculate(1000)?;
 
-    (storage.into_segments(), size)
+    Ok((storage.into_segments(), size))
 }
 
 // Diverged branches
-fn scenario_4() -> (Vec<Segment>, SegmentSize) {
+fn scenario_4() -> anyhow::Result<(Vec<Segment>, SegmentSize)> {
     // Create main branch
     let mut storage = Storage::new("main");
 
     // Bulk load 5 GB of data to it
-    storage.insert("main", 5_000);
+    storage.insert("main", 5_000)?;
 
     // Stream of updates
     for _ in 0..5 {
-        storage.update("main", 1_000);
+        storage.update("main", 1_000)?;
     }
 
     // Branch
-    storage.branch("main", "child");
-    storage.update("child", 1_000);
+    storage.branch("main", "child")?;
+    storage.update("child", 1_000)?;
 
     // More updates on parent
     for _ in 0..8 {
-        storage.update("main", 1_000);
+        storage.update("main", 1_000)?;
     }
 
-    let size = storage.calculate(1000);
+    let size = storage.calculate(1000)?;
 
-    (storage.into_segments(), size)
+    Ok((storage.into_segments(), size))
 }
 
-fn scenario_5() -> (Vec<Segment>, SegmentSize) {
+fn scenario_5() -> anyhow::Result<(Vec<Segment>, SegmentSize)> {
     let mut storage = Storage::new("a");
-    storage.insert("a", 5000);
-    storage.branch("a", "b");
-    storage.update("b", 4000);
-    storage.update("a", 2000);
-    storage.branch("a", "c");
-    storage.insert("c", 4000);
-    storage.insert("a", 2000);
+    storage.insert("a", 5000)?;
+    storage.branch("a", "b")?;
+    storage.update("b", 4000)?;
+    storage.update("a", 2000)?;
+    storage.branch("a", "c")?;
+    storage.insert("c", 4000)?;
+    storage.insert("a", 2000)?;
 
-    let size = storage.calculate(5000);
+    let size = storage.calculate(5000)?;
 
-    (storage.into_segments(), size)
+    Ok((storage.into_segments(), size))
 }
 
-fn scenario_6() -> (Vec<Segment>, SegmentSize) {
+fn scenario_6() -> anyhow::Result<(Vec<Segment>, SegmentSize)> {
     use std::borrow::Cow;
 
     const NO_OP: Cow<'static, str> = Cow::Borrowed("");
@@ -133,18 +133,18 @@ fn scenario_6() -> (Vec<Segment>, SegmentSize) {
 
     let mut storage = Storage::new(None);
 
-    storage.branch(&None, branches[0]); // at 0
-    storage.modify_branch(&branches[0], NO_OP, 108951064, 43696128); // at 108951064
-    storage.branch(&branches[0], branches[1]); // at 108951064
-    storage.modify_branch(&branches[1], NO_OP, 15560408, -1851392); // at 124511472
-    storage.modify_branch(&branches[0], NO_OP, 174464360, -1531904); // at 283415424
-    storage.branch(&branches[0], branches[2]); // at 283415424
-    storage.modify_branch(&branches[2], NO_OP, 15906192, 8192); // at 299321616
-    storage.modify_branch(&branches[0], NO_OP, 18909976, 32768); // at 302325400
+    storage.branch(&None, branches[0])?; // at 0
+    storage.modify_branch(&branches[0], NO_OP, 108951064, 43696128)?; // at 108951064
+    storage.branch(&branches[0], branches[1])?; // at 108951064
+    storage.modify_branch(&branches[1], NO_OP, 15560408, -1851392)?; // at 124511472
+    storage.modify_branch(&branches[0], NO_OP, 174464360, -1531904)?; // at 283415424
+    storage.branch(&branches[0], branches[2])?; // at 283415424
+    storage.modify_branch(&branches[2], NO_OP, 15906192, 8192)?; // at 299321616
+    storage.modify_branch(&branches[0], NO_OP, 18909976, 32768)?; // at 302325400
 
-    let size = storage.calculate(100_000);
+    let size = storage.calculate(100_000)?;
 
-    (storage.into_segments(), size)
+    Ok((storage.into_segments(), size))
 }
 
 fn main() {
@@ -163,7 +163,8 @@ fn main() {
             eprintln!("invalid scenario {}", other);
             std::process::exit(1);
         }
-    };
+    }
+    .unwrap();
 
     graphviz_tree(&segments, &size);
 }
@@ -251,7 +252,7 @@ fn graphviz_tree(segments: &[Segment], tree: &SegmentSize) {
 
 #[test]
 fn scenarios_return_same_size() {
-    type ScenarioFn = fn() -> (Vec<Segment>, SegmentSize);
+    type ScenarioFn = fn() -> anyhow::Result<(Vec<Segment>, SegmentSize)>;
     let truths: &[(u32, ScenarioFn, _)] = &[
         (line!(), scenario_1, 8000),
         (line!(), scenario_2, 9000),
@@ -262,7 +263,7 @@ fn scenarios_return_same_size() {
     ];
 
     for (line, scenario, expected) in truths {
-        let (_, size) = scenario();
+        let (_, size) = scenario().unwrap();
         assert_eq!(*expected, size.total_children(), "scenario on line {line}");
     }
 }

libs/utils/Cargo.toml

@@ -1,47 +1,49 @@
 [package]
 name = "utils"
 version = "0.1.0"
-edition = "2021"
+edition.workspace = true
+license.workspace = true
 
 [dependencies]
-sentry = "0.29.0"
-async-trait = "0.1"
-anyhow = "1.0"
-bincode = "1.3"
-bytes = "1.0.1"
-hyper = { version = "0.14.7", features = ["full"] }
-routerify = "3"
-serde = { version = "1.0", features = ["derive"] }
-serde_json = "1"
-thiserror = "1.0"
-tokio = { version = "1.17", features = ["macros"]}
-tokio-rustls = "0.23"
-tracing = "0.1"
-tracing-subscriber = { version = "0.3", features = ["env-filter", "json"] }
-nix = "0.25"
-signal-hook = "0.3.10"
-rand = "0.8.3"
-jsonwebtoken = "8"
-hex = { version = "0.4.3", features = ["serde"] }
-rustls = "0.20.2"
-rustls-split = "0.3.0"
-git-version = "0.3.5"
-serde_with = "2.0"
-once_cell = "1.13.0"
-strum = "0.24"
-strum_macros = "0.24"
+sentry.workspace = true
+async-trait.workspace = true
+anyhow.workspace = true
+bincode.workspace = true
+bytes.workspace = true
+hyper = { workspace = true, features = ["full"] }
+routerify.workspace = true
+serde.workspace = true
+serde_json.workspace = true
+thiserror.workspace = true
+tokio.workspace = true
+tokio-rustls.workspace = true
+tracing.workspace = true
+tracing-subscriber = { workspace = true, features = ["json"] }
+nix.workspace = true
+signal-hook.workspace = true
+rand.workspace = true
+jsonwebtoken.workspace = true
+hex = { workspace = true, features = ["serde"] }
+rustls.workspace = true
+rustls-split.workspace = true
+git-version.workspace = true
+serde_with.workspace = true
+once_cell.workspace = true
+strum.workspace = true
+strum_macros.workspace = true
 
-metrics = { path = "../metrics" }
-pq_proto = { path = "../pq_proto" }
-workspace_hack = { version = "0.1", path = "../../workspace_hack" }
+metrics.workspace = true
+pq_proto.workspace = true
+
+workspace_hack.workspace = true
 
 [dev-dependencies]
-byteorder = "1.4.3"
-bytes = "1.0.1"
-hex-literal = "0.3"
-tempfile = "3.2"
-criterion = "0.4"
-rustls-pemfile = "1"
+byteorder.workspace = true
+bytes.workspace = true
+hex-literal.workspace = true
+tempfile.workspace = true
+criterion.workspace = true
+rustls-pemfile.workspace = true
 
 [[bench]]
 name = "benchmarks"

libs/utils/src/logging.rs

@@ -8,6 +8,7 @@ use strum_macros::{EnumString, EnumVariantNames};
 pub enum LogFormat {
     Plain,
     Json,
+    Test,
 }
 
 impl LogFormat {
@@ -39,6 +40,7 @@ pub fn init(log_format: LogFormat) -> anyhow::Result<()> {
     match log_format {
         LogFormat::Json => base_logger.json().init(),
         LogFormat::Plain => base_logger.init(),
+        LogFormat::Test => base_logger.with_test_writer().init(),
     }
 
     Ok(())

libs/utils/src/postgres_backend.rs

@@ -3,11 +3,11 @@
 //! implementation determining how to process the queries. Currently its API
 //! is rather narrow, but we can extend it once required.
 
+use crate::postgres_backend_async::{log_query_error, short_error, QueryError};
 use crate::sock_split::{BidiStream, ReadStream, WriteStream};
-use anyhow::{bail, ensure, Context, Result};
+use anyhow::Context;
 use bytes::{Bytes, BytesMut};
 use pq_proto::{BeMessage, FeMessage, FeStartupPacket};
-use rand::Rng;
 use serde::{Deserialize, Serialize};
 use std::fmt;
 use std::io::{self, Write};
@@ -22,25 +22,32 @@ pub trait Handler {
     /// postgres_backend will issue ReadyForQuery after calling this (this
     /// might be not what we want after CopyData streaming, but currently we don't
     /// care).
-    fn process_query(&mut self, pgb: &mut PostgresBackend, query_string: &str) -> Result<()>;
+    fn process_query(
+        &mut self,
+        pgb: &mut PostgresBackend,
+        query_string: &str,
+    ) -> Result<(), QueryError>;
 
     /// Called on startup packet receival, allows to process params.
     ///
     /// If Ok(false) is returned postgres_backend will skip auth -- that is needed for new users
     /// creation is the proxy code. That is quite hacky and ad-hoc solution, may be we could allow
     /// to override whole init logic in implementations.
-    fn startup(&mut self, _pgb: &mut PostgresBackend, _sm: &FeStartupPacket) -> Result<()> {
+    fn startup(
+        &mut self,
+        _pgb: &mut PostgresBackend,
+        _sm: &FeStartupPacket,
+    ) -> Result<(), QueryError> {
         Ok(())
     }
 
-    /// Check auth md5
-    fn check_auth_md5(&mut self, _pgb: &mut PostgresBackend, _md5_response: &[u8]) -> Result<()> {
-        bail!("MD5 auth failed")
-    }
-
     /// Check auth jwt
-    fn check_auth_jwt(&mut self, _pgb: &mut PostgresBackend, _jwt_response: &[u8]) -> Result<()> {
-        bail!("JWT auth failed")
+    fn check_auth_jwt(
+        &mut self,
+        _pgb: &mut PostgresBackend,
+        _jwt_response: &[u8],
+    ) -> Result<(), QueryError> {
+        Err(QueryError::Other(anyhow::anyhow!("JWT auth failed")))
     }
 
     fn is_shutdown_requested(&self) -> bool {
@@ -61,7 +68,6 @@ pub enum ProtoState {
 #[derive(Debug, PartialEq, Eq, Clone, Copy, Serialize, Deserialize)]
 pub enum AuthType {
     Trust,
-    MD5,
     // This mimics postgres's AuthenticationCleartextPassword but instead of password expects JWT
     NeonJWT,
 }
@@ -72,9 +78,8 @@ impl FromStr for AuthType {
     fn from_str(s: &str) -> Result<Self, Self::Err> {
         match s {
             "Trust" => Ok(Self::Trust),
-            "MD5" => Ok(Self::MD5),
             "NeonJWT" => Ok(Self::NeonJWT),
-            _ => bail!("invalid value \"{s}\" for auth type"),
+            _ => anyhow::bail!("invalid value \"{s}\" for auth type"),
         }
     }
 }
@@ -83,7 +88,6 @@ impl fmt::Display for AuthType {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         f.write_str(match self {
             AuthType::Trust => "Trust",
-            AuthType::MD5 => "MD5",
             AuthType::NeonJWT => "NeonJWT",
         })
     }
@@ -134,7 +138,6 @@ pub struct PostgresBackend {
 
     pub state: ProtoState,
 
-    md5_salt: [u8; 4],
     auth_type: AuthType,
 
     peer_addr: SocketAddr,
@@ -164,7 +167,7 @@ pub fn is_socket_read_timed_out(error: &anyhow::Error) -> bool {
 }
 
 // Cast a byte slice to a string slice, dropping null terminator if there's one.
-fn cstr_to_str(bytes: &[u8]) -> Result<&str> {
+fn cstr_to_str(bytes: &[u8]) -> anyhow::Result<&str> {
     let without_null = bytes.strip_suffix(&[0]).unwrap_or(bytes);
     std::str::from_utf8(without_null).map_err(|e| e.into())
 }
@@ -187,7 +190,6 @@ impl PostgresBackend {
             stream: Some(Stream::Bidirectional(BidiStream::from_tcp(socket))),
             buf_out: BytesMut::with_capacity(10 * 1024),
             state: ProtoState::Initialization,
-            md5_salt: [0u8; 4],
             auth_type,
             tls_config,
             peer_addr,
@@ -199,10 +201,10 @@ impl PostgresBackend {
     }
 
     /// Get direct reference (into the Option) to the read stream.
-    fn get_stream_in(&mut self) -> Result<&mut BidiStream> {
+    fn get_stream_in(&mut self) -> anyhow::Result<&mut BidiStream> {
         match &mut self.stream {
             Some(Stream::Bidirectional(stream)) => Ok(stream),
-            _ => bail!("reader taken"),
+            _ => anyhow::bail!("reader taken"),
         }
     }
 
@@ -226,7 +228,7 @@ impl PostgresBackend {
     }
 
     /// Read full message or return None if connection is closed.
-    pub fn read_message(&mut self) -> Result<Option<FeMessage>> {
+    pub fn read_message(&mut self) -> Result<Option<FeMessage>, QueryError> {
         let (state, stream) = (self.state, self.get_stream_in()?);
 
         use ProtoState::*;
@@ -234,6 +236,7 @@ impl PostgresBackend {
             Initialization | Encrypted => FeStartupPacket::read(stream),
             Authentication | Established => FeMessage::read(stream),
         }
+        .map_err(QueryError::from)
     }
 
     /// Write message into internal output buffer.
@@ -257,7 +260,7 @@ impl PostgresBackend {
     }
 
     // Wrapper for run_message_loop() that shuts down socket when we are done
-    pub fn run(mut self, handler: &mut impl Handler) -> Result<()> {
+    pub fn run(mut self, handler: &mut impl Handler) -> Result<(), QueryError> {
         let ret = self.run_message_loop(handler);
         if let Some(stream) = self.stream.as_mut() {
             let _ = stream.shutdown(Shutdown::Both);
@@ -265,7 +268,7 @@ impl PostgresBackend {
         ret
     }
 
-    fn run_message_loop(&mut self, handler: &mut impl Handler) -> Result<()> {
+    fn run_message_loop(&mut self, handler: &mut impl Handler) -> Result<(), QueryError> {
         trace!("postgres backend to {:?} started", self.peer_addr);
 
         let mut unnamed_query_string = Bytes::new();
@@ -274,7 +277,7 @@ impl PostgresBackend {
             match self.read_message() {
                 Ok(message) => {
                     if let Some(msg) = message {
-                        trace!("got message {:?}", msg);
+                        trace!("got message {msg:?}");
 
                         match self.process_message(handler, msg, &mut unnamed_query_string)? {
                             ProcessMsgResult::Continue => continue,
@@ -285,10 +288,12 @@ impl PostgresBackend {
                     }
                 }
                 Err(e) => {
-                    // If it is a timeout error, continue the loop
-                    if !is_socket_read_timed_out(&e) {
-                        return Err(e);
+                    if let QueryError::Other(e) = &e {
+                        if is_socket_read_timed_out(e) {
+                            continue;
+                        }
                     }
+                    return Err(e);
                 }
             }
         }
@@ -306,7 +311,7 @@ impl PostgresBackend {
             }
             stream => {
                 self.stream = stream;
-                bail!("can't start TLs without bidi stream");
+                anyhow::bail!("can't start TLs without bidi stream");
             }
         }
     }
@@ -316,17 +321,16 @@ impl PostgresBackend {
         handler: &mut impl Handler,
         msg: FeMessage,
         unnamed_query_string: &mut Bytes,
-    ) -> Result<ProcessMsgResult> {
+    ) -> Result<ProcessMsgResult, QueryError> {
         // Allow only startup and password messages during auth. Otherwise client would be able to bypass auth
         // TODO: change that to proper top-level match of protocol state with separate message handling for each state
-        if self.state < ProtoState::Established {
-            ensure!(
-                matches!(
-                    msg,
-                    FeMessage::PasswordMessage(_) | FeMessage::StartupPacket(_)
-                ),
-                "protocol violation"
-            );
+        if self.state < ProtoState::Established
+            && !matches!(
+                msg,
+                FeMessage::PasswordMessage(_) | FeMessage::StartupPacket(_)
+            )
+        {
+            return Err(QueryError::Other(anyhow::anyhow!("protocol violation")));
         }
 
         let have_tls = self.tls_config.is_some();
@@ -350,8 +354,13 @@ impl PostgresBackend {
                     }
                     FeStartupPacket::StartupMessage { .. } => {
                         if have_tls && !matches!(self.state, ProtoState::Encrypted) {
-                            self.write_message(&BeMessage::ErrorResponse("must connect with TLS"))?;
-                            bail!("client did not connect with TLS");
+                            self.write_message(&BeMessage::ErrorResponse(
+                                "must connect with TLS",
+                                None,
+                            ))?;
+                            return Err(QueryError::Other(anyhow::anyhow!(
+                                "client did not connect with TLS"
+                            )));
                         }
 
                         // NB: startup() may change self.auth_type -- we are using that in proxy code
@@ -367,13 +376,6 @@ impl PostgresBackend {
                                     .write_message(&BeMessage::ReadyForQuery)?;
                                 self.state = ProtoState::Established;
                             }
-                            AuthType::MD5 => {
-                                rand::thread_rng().fill(&mut self.md5_salt);
-                                self.write_message(&BeMessage::AuthenticationMD5Password(
-                                    self.md5_salt,
-                                ))?;
-                                self.state = ProtoState::Authentication;
-                            }
                             AuthType::NeonJWT => {
                                 self.write_message(&BeMessage::AuthenticationCleartextPassword)?;
                                 self.state = ProtoState::Authentication;
@@ -393,20 +395,15 @@ impl PostgresBackend {
 
                 match self.auth_type {
                     AuthType::Trust => unreachable!(),
-                    AuthType::MD5 => {
-                        let (_, md5_response) = m.split_last().context("protocol violation")?;
-
-                        if let Err(e) = handler.check_auth_md5(self, md5_response) {
-                            self.write_message(&BeMessage::ErrorResponse(&e.to_string()))?;
-                            bail!("auth failed: {}", e);
-                        }
-                    }
                     AuthType::NeonJWT => {
                         let (_, jwt_response) = m.split_last().context("protocol violation")?;
 
                         if let Err(e) = handler.check_auth_jwt(self, jwt_response) {
-                            self.write_message(&BeMessage::ErrorResponse(&e.to_string()))?;
-                            bail!("auth failed: {}", e);
+                            self.write_message(&BeMessage::ErrorResponse(
+                                &e.to_string(),
+                                Some(e.pg_error_code()),
+                            ))?;
+                            return Err(e);
                         }
                     }
                 }
@@ -420,33 +417,14 @@ impl PostgresBackend {
                 // remove null terminator
                 let query_string = cstr_to_str(&body)?;
 
-                trace!("got query {:?}", query_string);
-                // xxx distinguish fatal and recoverable errors?
+                trace!("got query {query_string:?}");
                 if let Err(e) = handler.process_query(self, query_string) {
-                    // ":?" uses the alternate formatting style, which makes anyhow display the
-                    // full cause of the error, not just the top-level context + its trace.
-                    // We don't want to send that in the ErrorResponse though,
-                    // because it's not relevant to the compute node logs.
-                    //
-                    // We also don't want to log full stacktrace when the error is primitive,
-                    // such as usual connection closed.
-                    let short_error = format!("{:#}", e);
-                    let root_cause = e.root_cause().to_string();
-                    if root_cause.contains("connection closed unexpectedly")
-                        || root_cause.contains("Broken pipe (os error 32)")
-                    {
-                        error!(
-                            "query handler for '{}' failed: {}",
-                            query_string, short_error
-                        );
-                    } else {
-                        error!("query handler for '{}' failed: {:?}", query_string, e);
-                    }
-                    self.write_message_noflush(&BeMessage::ErrorResponse(&short_error))?;
-                    // TODO: untangle convoluted control flow
-                    if e.to_string().contains("failed to run") {
-                        return Ok(ProcessMsgResult::Break);
-                    }
+                    log_query_error(query_string, &e);
+                    let short_error = short_error(&e);
+                    self.write_message_noflush(&BeMessage::ErrorResponse(
+                        &short_error,
+                        Some(e.pg_error_code()),
+                    ))?;
                 }
                 self.write_message(&BeMessage::ReadyForQuery)?;
             }
@@ -471,11 +449,13 @@ impl PostgresBackend {
 
             FeMessage::Execute(_) => {
                 let query_string = cstr_to_str(unnamed_query_string)?;
-                trace!("got execute {:?}", query_string);
-                // xxx distinguish fatal and recoverable errors?
+                trace!("got execute {query_string:?}");
                 if let Err(e) = handler.process_query(self, query_string) {
-                    error!("query handler for '{}' failed: {:?}", query_string, e);
-                    self.write_message(&BeMessage::ErrorResponse(&e.to_string()))?;
+                    log_query_error(query_string, &e);
+                    self.write_message(&BeMessage::ErrorResponse(
+                        &e.to_string(),
+                        Some(e.pg_error_code()),
+                    ))?;
                 }
                 // NOTE there is no ReadyForQuery message. This handler is used
                 // for basebackup and it uses CopyOut which doesn't require
@@ -494,7 +474,9 @@ impl PostgresBackend {
             // We prefer explicit pattern matching to wildcards, because
             // this helps us spot the places where new variants are missing
             FeMessage::CopyData(_) | FeMessage::CopyDone | FeMessage::CopyFail => {
-                bail!("unexpected message type: {:?}", msg);
+                return Err(QueryError::Other(anyhow::anyhow!(
+                    "unexpected message type: {msg:?}"
+                )));
             }
         }
 

libs/utils/src/postgres_backend_async.rs

@@ -4,45 +4,87 @@
 //! is rather narrow, but we can extend it once required.
 
 use crate::postgres_backend::AuthType;
-use anyhow::{bail, Context, Result};
-use bytes::{Bytes, BytesMut};
-use pq_proto::{BeMessage, FeMessage, FeStartupPacket};
-use rand::Rng;
-use std::future::Future;
+use anyhow::Context;
+use bytes::{Buf, Bytes, BytesMut};
+use pq_proto::{BeMessage, ConnectionError, FeMessage, FeStartupPacket, SQLSTATE_INTERNAL_ERROR};
+use std::io;
 use std::net::SocketAddr;
 use std::pin::Pin;
 use std::sync::Arc;
 use std::task::Poll;
-use tracing::{debug, error, trace};
+use std::{future::Future, task::ready};
+use tracing::{debug, error, info, trace};
 
 use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt, BufReader};
 use tokio_rustls::TlsAcceptor;
 
+pub fn is_expected_io_error(e: &io::Error) -> bool {
+    use io::ErrorKind::*;
+    matches!(
+        e.kind(),
+        ConnectionRefused | ConnectionAborted | ConnectionReset
+    )
+}
+
+/// An error, occurred during query processing:
+/// either during the connection ([`ConnectionError`]) or before/after it.
+#[derive(thiserror::Error, Debug)]
+pub enum QueryError {
+    /// The connection was lost while processing the query.
+    #[error(transparent)]
+    Disconnected(#[from] ConnectionError),
+    /// Some other error
+    #[error(transparent)]
+    Other(#[from] anyhow::Error),
+}
+
+impl From<io::Error> for QueryError {
+    fn from(e: io::Error) -> Self {
+        Self::Disconnected(ConnectionError::Socket(e))
+    }
+}
+
+impl QueryError {
+    pub fn pg_error_code(&self) -> &'static [u8; 5] {
+        match self {
+            Self::Disconnected(_) => b"08006",         // connection failure
+            Self::Other(_) => SQLSTATE_INTERNAL_ERROR, // internal error
+        }
+    }
+}
+
 #[async_trait::async_trait]
 pub trait Handler {
     /// Handle single query.
     /// postgres_backend will issue ReadyForQuery after calling this (this
     /// might be not what we want after CopyData streaming, but currently we don't
     /// care).
-    async fn process_query(&mut self, pgb: &mut PostgresBackend, query_string: &str) -> Result<()>;
+    async fn process_query(
+        &mut self,
+        pgb: &mut PostgresBackend,
+        query_string: &str,
+    ) -> Result<(), QueryError>;
 
     /// Called on startup packet receival, allows to process params.
     ///
     /// If Ok(false) is returned postgres_backend will skip auth -- that is needed for new users
     /// creation is the proxy code. That is quite hacky and ad-hoc solution, may be we could allow
     /// to override whole init logic in implementations.
-    fn startup(&mut self, _pgb: &mut PostgresBackend, _sm: &FeStartupPacket) -> Result<()> {
+    fn startup(
+        &mut self,
+        _pgb: &mut PostgresBackend,
+        _sm: &FeStartupPacket,
+    ) -> Result<(), QueryError> {
         Ok(())
     }
 
-    /// Check auth md5
-    fn check_auth_md5(&mut self, _pgb: &mut PostgresBackend, _md5_response: &[u8]) -> Result<()> {
-        bail!("MD5 auth failed")
-    }
-
     /// Check auth jwt
-    fn check_auth_jwt(&mut self, _pgb: &mut PostgresBackend, _jwt_response: &[u8]) -> Result<()> {
-        bail!("JWT auth failed")
+    fn check_auth_jwt(
+        &mut self,
+        _pgb: &mut PostgresBackend,
+        _jwt_response: &[u8],
+    ) -> Result<(), QueryError> {
+        Err(QueryError::Other(anyhow::anyhow!("JWT auth failed")))
     }
 }
 
@@ -76,17 +118,14 @@ impl AsyncWrite for Stream {
         self: Pin<&mut Self>,
         cx: &mut std::task::Context<'_>,
         buf: &[u8],
-    ) -> Poll<Result<usize, std::io::Error>> {
+    ) -> Poll<io::Result<usize>> {
         match self.get_mut() {
             Self::Unencrypted(stream) => Pin::new(stream).poll_write(cx, buf),
             Self::Tls(stream) => Pin::new(stream).poll_write(cx, buf),
             Self::Broken => unreachable!(),
         }
     }
-    fn poll_flush(
-        self: Pin<&mut Self>,
-        cx: &mut std::task::Context<'_>,
-    ) -> Poll<Result<(), std::io::Error>> {
+    fn poll_flush(self: Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> Poll<io::Result<()>> {
         match self.get_mut() {
             Self::Unencrypted(stream) => Pin::new(stream).poll_flush(cx),
             Self::Tls(stream) => Pin::new(stream).poll_flush(cx),
@@ -96,7 +135,7 @@ impl AsyncWrite for Stream {
     fn poll_shutdown(
         self: Pin<&mut Self>,
         cx: &mut std::task::Context<'_>,
-    ) -> Poll<Result<(), std::io::Error>> {
+    ) -> Poll<io::Result<()>> {
         match self.get_mut() {
             Self::Unencrypted(stream) => Pin::new(stream).poll_shutdown(cx),
             Self::Tls(stream) => Pin::new(stream).poll_shutdown(cx),
@@ -109,7 +148,7 @@ impl AsyncRead for Stream {
         self: Pin<&mut Self>,
         cx: &mut std::task::Context<'_>,
         buf: &mut tokio::io::ReadBuf<'_>,
-    ) -> Poll<Result<(), std::io::Error>> {
+    ) -> Poll<io::Result<()>> {
         match self.get_mut() {
             Self::Unencrypted(stream) => Pin::new(stream).poll_read(cx, buf),
             Self::Tls(stream) => Pin::new(stream).poll_read(cx, buf),
@@ -120,12 +159,14 @@ impl AsyncRead for Stream {
 
 pub struct PostgresBackend {
     stream: Stream,
+
     // Output buffer. c.f. BeMessage::write why we are using BytesMut here.
+    // The data between 0 and "current position" as tracked by the bytes::Buf
+    // implementation of BytesMut, have already been written.
     buf_out: BytesMut,
 
     pub state: ProtoState,
 
-    md5_salt: [u8; 4],
     auth_type: AuthType,
 
     peer_addr: SocketAddr,
@@ -143,7 +184,7 @@ pub fn query_from_cstring(query_string: Bytes) -> Vec<u8> {
 }
 
 // Cast a byte slice to a string slice, dropping null terminator if there's one.
-fn cstr_to_str(bytes: &[u8]) -> Result<&str> {
+fn cstr_to_str(bytes: &[u8]) -> anyhow::Result<&str> {
     let without_null = bytes.strip_suffix(&[0]).unwrap_or(bytes);
     std::str::from_utf8(without_null).map_err(|e| e.into())
 }
@@ -153,14 +194,13 @@ impl PostgresBackend {
         socket: tokio::net::TcpStream,
         auth_type: AuthType,
         tls_config: Option<Arc<rustls::ServerConfig>>,
-    ) -> std::io::Result<Self> {
+    ) -> io::Result<Self> {
         let peer_addr = socket.peer_addr()?;
 
         Ok(Self {
             stream: Stream::Unencrypted(BufReader::new(socket)),
             buf_out: BytesMut::with_capacity(10 * 1024),
             state: ProtoState::Initialization,
-            md5_salt: [0u8; 4],
             auth_type,
             tls_config,
             peer_addr,
@@ -172,30 +212,65 @@ impl PostgresBackend {
     }
 
     /// Read full message or return None if connection is closed.
-    pub async fn read_message(&mut self) -> Result<Option<FeMessage>> {
+    pub async fn read_message(&mut self) -> Result<Option<FeMessage>, QueryError> {
         use ProtoState::*;
         match self.state {
             Initialization | Encrypted => FeStartupPacket::read_fut(&mut self.stream).await,
             Authentication | Established => FeMessage::read_fut(&mut self.stream).await,
             Closed => Ok(None),
         }
+        .map_err(QueryError::from)
     }
 
     /// Flush output buffer into the socket.
-    pub async fn flush(&mut self) -> std::io::Result<&mut Self> {
-        self.stream.write_all(&self.buf_out).await?;
+    pub async fn flush(&mut self) -> io::Result<()> {
+        while self.buf_out.has_remaining() {
+            let bytes_written = self.stream.write(self.buf_out.chunk()).await?;
+            self.buf_out.advance(bytes_written);
+        }
         self.buf_out.clear();
-        Ok(self)
+        Ok(())
     }
 
     /// Write message into internal output buffer.
-    pub fn write_message(&mut self, message: &BeMessage<'_>) -> Result<&mut Self, std::io::Error> {
+    pub fn write_message(&mut self, message: &BeMessage<'_>) -> io::Result<&mut Self> {
         BeMessage::write(&mut self.buf_out, message)?;
         Ok(self)
     }
 
+    /// Returns an AsyncWrite implementation that wraps all the data written
+    /// to it in CopyData messages, and writes them to the connection
+    ///
+    /// The caller is responsible for sending CopyOutResponse and CopyDone messages.
+    pub fn copyout_writer(&mut self) -> CopyDataWriter {
+        CopyDataWriter { pgb: self }
+    }
+
+    /// A polling function that tries to write all the data from 'buf_out' to the
+    /// underlying stream.
+    fn poll_write_buf(
+        &mut self,
+        cx: &mut std::task::Context<'_>,
+    ) -> Poll<Result<(), std::io::Error>> {
+        while self.buf_out.has_remaining() {
+            match ready!(Pin::new(&mut self.stream).poll_write(cx, self.buf_out.chunk())) {
+                Ok(bytes_written) => self.buf_out.advance(bytes_written),
+                Err(err) => return Poll::Ready(Err(err)),
+            }
+        }
+        Poll::Ready(Ok(()))
+    }
+
+    fn poll_flush(&mut self, cx: &mut std::task::Context<'_>) -> Poll<Result<(), std::io::Error>> {
+        Pin::new(&mut self.stream).poll_flush(cx)
+    }
+
     // Wrapper for run_message_loop() that shuts down socket when we are done
-    pub async fn run<F, S>(mut self, handler: &mut impl Handler, shutdown_watcher: F) -> Result<()>
+    pub async fn run<F, S>(
+        mut self,
+        handler: &mut impl Handler,
+        shutdown_watcher: F,
+    ) -> Result<(), QueryError>
     where
         F: Fn() -> S,
         S: Future,
@@ -209,7 +284,7 @@ impl PostgresBackend {
         &mut self,
         handler: &mut impl Handler,
         shutdown_watcher: F,
-    ) -> Result<()>
+    ) -> Result<(), QueryError>
     where
         F: Fn() -> S,
         S: Future,
@@ -245,7 +320,7 @@ impl PostgresBackend {
                         return Ok(());
                     }
                 }
-                Ok::<(), anyhow::Error>(())
+                Ok::<(), QueryError>(())
             } => {
                 // Handshake complete.
                 result?;
@@ -290,14 +365,14 @@ impl PostgresBackend {
             self.stream = Stream::Tls(Box::new(tls_stream));
             return Ok(());
         };
-        bail!("TLS already started");
+        anyhow::bail!("TLS already started");
     }
 
     async fn process_handshake_message(
         &mut self,
         handler: &mut impl Handler,
         msg: FeMessage,
-    ) -> Result<ProcessMsgResult> {
+    ) -> Result<ProcessMsgResult, QueryError> {
         assert!(self.state < ProtoState::Established);
         let have_tls = self.tls_config.is_some();
         match msg {
@@ -320,8 +395,13 @@ impl PostgresBackend {
                     }
                     FeStartupPacket::StartupMessage { .. } => {
                         if have_tls && !matches!(self.state, ProtoState::Encrypted) {
-                            self.write_message(&BeMessage::ErrorResponse("must connect with TLS"))?;
-                            bail!("client did not connect with TLS");
+                            self.write_message(&BeMessage::ErrorResponse(
+                                "must connect with TLS",
+                                None,
+                            ))?;
+                            return Err(QueryError::Other(anyhow::anyhow!(
+                                "client did not connect with TLS"
+                            )));
                         }
 
                         // NB: startup() may change self.auth_type -- we are using that in proxy code
@@ -337,13 +417,6 @@ impl PostgresBackend {
                                     .write_message(&BeMessage::ReadyForQuery)?;
                                 self.state = ProtoState::Established;
                             }
-                            AuthType::MD5 => {
-                                rand::thread_rng().fill(&mut self.md5_salt);
-                                self.write_message(&BeMessage::AuthenticationMD5Password(
-                                    self.md5_salt,
-                                ))?;
-                                self.state = ProtoState::Authentication;
-                            }
                             AuthType::NeonJWT => {
                                 self.write_message(&BeMessage::AuthenticationCleartextPassword)?;
                                 self.state = ProtoState::Authentication;
@@ -364,20 +437,15 @@ impl PostgresBackend {
 
                 match self.auth_type {
                     AuthType::Trust => unreachable!(),
-                    AuthType::MD5 => {
-                        let (_, md5_response) = m.split_last().context("protocol violation")?;
-
-                        if let Err(e) = handler.check_auth_md5(self, md5_response) {
-                            self.write_message(&BeMessage::ErrorResponse(&e.to_string()))?;
-                            bail!("auth failed: {}", e);
-                        }
-                    }
                     AuthType::NeonJWT => {
                         let (_, jwt_response) = m.split_last().context("protocol violation")?;
 
                         if let Err(e) = handler.check_auth_jwt(self, jwt_response) {
-                            self.write_message(&BeMessage::ErrorResponse(&e.to_string()))?;
-                            bail!("auth failed: {}", e);
+                            self.write_message(&BeMessage::ErrorResponse(
+                                &e.to_string(),
+                                Some(e.pg_error_code()),
+                            ))?;
+                            return Err(e);
                         }
                     }
                 }
@@ -400,33 +468,28 @@ impl PostgresBackend {
         handler: &mut impl Handler,
         msg: FeMessage,
         unnamed_query_string: &mut Bytes,
-    ) -> Result<ProcessMsgResult> {
+    ) -> Result<ProcessMsgResult, QueryError> {
         // Allow only startup and password messages during auth. Otherwise client would be able to bypass auth
         // TODO: change that to proper top-level match of protocol state with separate message handling for each state
         assert!(self.state == ProtoState::Established);
 
         match msg {
             FeMessage::StartupPacket(_) | FeMessage::PasswordMessage(_) => {
-                bail!("protocol violation");
+                return Err(QueryError::Other(anyhow::anyhow!("protocol violation")));
             }
 
             FeMessage::Query(body) => {
                 // remove null terminator
                 let query_string = cstr_to_str(&body)?;
 
-                trace!("got query {:?}", query_string);
-                // xxx distinguish fatal and recoverable errors?
+                trace!("got query {query_string:?}");
                 if let Err(e) = handler.process_query(self, query_string).await {
-                    // ":?" uses the alternate formatting style, which makes anyhow display the
-                    // full cause of the error, not just the top-level context + its trace.
-                    // We don't want to send that in the ErrorResponse though,
-                    // because it's not relevant to the compute node logs.
-                    error!("query handler for '{}' failed: {:?}", query_string, e);
-                    self.write_message(&BeMessage::ErrorResponse(&e.to_string()))?;
-                    // TODO: untangle convoluted control flow
-                    if e.to_string().contains("failed to run") {
-                        return Ok(ProcessMsgResult::Break);
-                    }
+                    log_query_error(query_string, &e);
+                    let short_error = short_error(&e);
+                    self.write_message(&BeMessage::ErrorResponse(
+                        &short_error,
+                        Some(e.pg_error_code()),
+                    ))?;
                 }
                 self.write_message(&BeMessage::ReadyForQuery)?;
             }
@@ -451,11 +514,13 @@ impl PostgresBackend {
 
             FeMessage::Execute(_) => {
                 let query_string = cstr_to_str(unnamed_query_string)?;
-                trace!("got execute {:?}", query_string);
-                // xxx distinguish fatal and recoverable errors?
+                trace!("got execute {query_string:?}");
                 if let Err(e) = handler.process_query(self, query_string).await {
-                    error!("query handler for '{}' failed: {:?}", query_string, e);
-                    self.write_message(&BeMessage::ErrorResponse(&e.to_string()))?;
+                    log_query_error(query_string, &e);
+                    self.write_message(&BeMessage::ErrorResponse(
+                        &e.to_string(),
+                        Some(e.pg_error_code()),
+                    ))?;
                 }
                 // NOTE there is no ReadyForQuery message. This handler is used
                 // for basebackup and it uses CopyOut which doesn't require
@@ -474,10 +539,96 @@ impl PostgresBackend {
             // We prefer explicit pattern matching to wildcards, because
             // this helps us spot the places where new variants are missing
             FeMessage::CopyData(_) | FeMessage::CopyDone | FeMessage::CopyFail => {
-                bail!("unexpected message type: {:?}", msg);
+                return Err(QueryError::Other(anyhow::anyhow!(
+                    "unexpected message type: {:?}",
+                    msg
+                )));
             }
         }
 
         Ok(ProcessMsgResult::Continue)
     }
 }
+
+///
+/// A futures::AsyncWrite implementation that wraps all data written to it in CopyData
+/// messages.
+///
+
+pub struct CopyDataWriter<'a> {
+    pgb: &'a mut PostgresBackend,
+}
+
+impl<'a> AsyncWrite for CopyDataWriter<'a> {
+    fn poll_write(
+        self: Pin<&mut Self>,
+        cx: &mut std::task::Context<'_>,
+        buf: &[u8],
+    ) -> Poll<Result<usize, std::io::Error>> {
+        let this = self.get_mut();
+
+        // It's not strictly required to flush between each message, but makes it easier
+        // to view in wireshark, and usually the messages that the callers write are
+        // decently-sized anyway.
+        match ready!(this.pgb.poll_write_buf(cx)) {
+            Ok(()) => {}
+            Err(err) => return Poll::Ready(Err(err)),
+        }
+
+        // CopyData
+        // XXX: if the input is large, we should split it into multiple messages.
+        // Not sure what the threshold should be, but the ultimate hard limit is that
+        // the length cannot exceed u32.
+        this.pgb.write_message(&BeMessage::CopyData(buf))?;
+
+        Poll::Ready(Ok(buf.len()))
+    }
+
+    fn poll_flush(
+        self: Pin<&mut Self>,
+        cx: &mut std::task::Context<'_>,
+    ) -> Poll<Result<(), std::io::Error>> {
+        let this = self.get_mut();
+        match ready!(this.pgb.poll_write_buf(cx)) {
+            Ok(()) => {}
+            Err(err) => return Poll::Ready(Err(err)),
+        }
+        this.pgb.poll_flush(cx)
+    }
+    fn poll_shutdown(
+        self: Pin<&mut Self>,
+        cx: &mut std::task::Context<'_>,
+    ) -> Poll<Result<(), std::io::Error>> {
+        let this = self.get_mut();
+        match ready!(this.pgb.poll_write_buf(cx)) {
+            Ok(()) => {}
+            Err(err) => return Poll::Ready(Err(err)),
+        }
+        this.pgb.poll_flush(cx)
+    }
+}
+
+pub fn short_error(e: &QueryError) -> String {
+    match e {
+        QueryError::Disconnected(connection_error) => connection_error.to_string(),
+        QueryError::Other(e) => format!("{e:#}"),
+    }
+}
+
+pub(super) fn log_query_error(query: &str, e: &QueryError) {
+    match e {
+        QueryError::Disconnected(ConnectionError::Socket(io_error)) => {
+            if is_expected_io_error(io_error) {
+                info!("query handler for '{query}' failed with expected io error: {io_error}");
+            } else {
+                error!("query handler for '{query}' failed with io error: {io_error}");
+            }
+        }
+        QueryError::Disconnected(other_connection_error) => {
+            error!("query handler for '{query}' failed with connection error: {other_connection_error:?}")
+        }
+        QueryError::Other(e) => {
+            error!("query handler for '{query}' failed: {e:?}");
+        }
+    }
+}

libs/utils/tests/ssl_test.rs

@@ -9,7 +9,10 @@ use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt};
 use bytes::{Buf, BufMut, Bytes, BytesMut};
 use once_cell::sync::Lazy;
 
-use utils::postgres_backend::{AuthType, Handler, PostgresBackend};
+use utils::{
+    postgres_backend::{AuthType, Handler, PostgresBackend},
+    postgres_backend_async::QueryError,
+};
 
 fn make_tcp_pair() -> (TcpStream, TcpStream) {
     let listener = TcpListener::bind("127.0.0.1:0").unwrap();
@@ -105,7 +108,7 @@ fn ssl() {
             &mut self,
             _pgb: &mut PostgresBackend,
             query_string: &str,
-        ) -> anyhow::Result<()> {
+        ) -> Result<(), QueryError> {
             self.got_query = query_string == QUERY;
             Ok(())
         }
@@ -152,7 +155,7 @@ fn no_ssl() {
             &mut self,
             _pgb: &mut PostgresBackend,
             _query_string: &str,
-        ) -> anyhow::Result<()> {
+        ) -> Result<(), QueryError> {
             panic!()
         }
     }
@@ -212,7 +215,7 @@ fn server_forces_ssl() {
             &mut self,
             _pgb: &mut PostgresBackend,
             _query_string: &str,
-        ) -> anyhow::Result<()> {
+        ) -> Result<(), QueryError> {
             panic!()
         }
     }

pageserver/Cargo.toml

@@ -1,7 +1,8 @@
 [package]
 name = "pageserver"
 version = "0.1.0"
-edition = "2021"
+edition.workspace = true
+license.workspace = true
 
 [features]
 default = []
@@ -9,72 +10,69 @@ default = []
 # which adds some runtime cost to run tests on outage conditions
 testing = ["fail/failpoints"]
 
-profiling = ["pprof"]
-
 [dependencies]
-amplify_num = { git = "https://github.com/hlinnaka/rust-amplify.git", branch = "unsigned-int-perf" }
-anyhow = { version = "1.0", features = ["backtrace"] }
-async-stream = "0.3"
-async-trait = "0.1"
-byteorder = "1.4.3"
-bytes = "1.0.1"
-chrono = { version = "0.4.23", default-features = false, features = ["clock", "serde"] }
-clap = { version = "4.0", features = ["string"] }
-close_fds = "0.3.2"
-const_format = "0.2.21"
-crc32c = "0.6.0"
-crossbeam-utils = "0.8.5"
-fail = "0.5.0"
-futures = "0.3.13"
-git-version = "0.3.5"
-hex = "0.4.3"
-humantime = "2.1.0"
-humantime-serde = "1.1.1"
-hyper = "0.14"
-itertools = "0.10.3"
-nix = "0.25"
-num-traits = "0.2.15"
-once_cell = "1.13.0"
-pin-project-lite = "0.2.7"
-postgres = { git = "https://github.com/neondatabase/rust-postgres.git", rev="43e6db254a97fdecbce33d8bc0890accfd74495e" }
-postgres-protocol = { git = "https://github.com/neondatabase/rust-postgres.git", rev="43e6db254a97fdecbce33d8bc0890accfd74495e" }
-postgres-types = { git = "https://github.com/neondatabase/rust-postgres.git", rev="43e6db254a97fdecbce33d8bc0890accfd74495e" }
-pprof = { git = "https://github.com/neondatabase/pprof-rs.git", branch = "wallclock-profiling", features = ["flamegraph"], optional = true }
-rand = "0.8.3"
-regex = "1.4.5"
-rstar = "0.9.3"
-scopeguard = "1.1.0"
-serde = { version = "1.0", features = ["derive"] }
-serde_json = { version = "1.0", features = ["raw_value"] }
-serde_with = "2.0"
-signal-hook = "0.3.10"
-svg_fmt = "0.4.1"
-tar = "0.4.33"
-thiserror = "1.0"
-tokio = { version = "1.17", features = ["process", "sync", "macros", "fs", "rt", "io-util", "time"] }
-tokio-postgres = { git = "https://github.com/neondatabase/rust-postgres.git", rev="43e6db254a97fdecbce33d8bc0890accfd74495e" }
-tokio-util = { version = "0.7.3", features = ["io", "io-util"] }
-toml_edit = { version = "0.14", features = ["easy"] }
-tracing = "0.1.36"
-url = "2"
-walkdir = "2.3.2"
-
-metrics = { path = "../libs/metrics" }
-pageserver_api = { path = "../libs/pageserver_api" }
-postgres_connection = { path = "../libs/postgres_connection" }
-postgres_ffi = { path = "../libs/postgres_ffi" }
-pq_proto = { path = "../libs/pq_proto" }
-remote_storage = { path = "../libs/remote_storage" }
-storage_broker = { version = "0.1", path = "../storage_broker" }
-tenant_size_model = { path = "../libs/tenant_size_model" }
-utils = { path = "../libs/utils" }
-workspace_hack = { version = "0.1", path = "../workspace_hack" }
-reqwest = "0.11.13"
+amplify_num.workspace = true
+anyhow.workspace = true
+async-stream.workspace = true
+async-trait.workspace = true
+byteorder.workspace = true
+bytes.workspace = true
+chrono = { workspace = true, features = ["serde"] }
+clap = { workspace = true, features = ["string"] }
+close_fds.workspace = true
+const_format.workspace = true
+consumption_metrics.workspace = true
+crc32c.workspace = true
+crossbeam-utils.workspace = true
+fail.workspace = true
+futures.workspace = true
+git-version.workspace = true
+hex.workspace = true
+humantime.workspace = true
+humantime-serde.workspace = true
+hyper.workspace = true
+itertools.workspace = true
+nix.workspace = true
+num-traits.workspace = true
+once_cell.workspace = true
+pin-project-lite.workspace = true
+postgres.workspace = true
+postgres-protocol.workspace = true
+postgres-types.workspace = true
+rand.workspace = true
+regex.workspace = true
+rstar.workspace = true
+scopeguard.workspace = true
+serde.workspace = true
+serde_json = { workspace = true, features = ["raw_value"] }
+serde_with.workspace = true
+signal-hook.workspace = true
+svg_fmt.workspace = true
+tokio-tar.workspace = true
+thiserror.workspace = true
+tokio = { workspace = true, features = ["process", "sync", "fs", "rt", "io-util", "time"] }
+tokio-postgres.workspace = true
+tokio-util.workspace = true
+toml_edit.workspace = true
+tracing.workspace = true
+url.workspace = true
+walkdir.workspace = true
+metrics.workspace = true
+pageserver_api.workspace = true
+postgres_connection.workspace = true
+postgres_ffi.workspace = true
+pq_proto.workspace = true
+remote_storage.workspace = true
+storage_broker.workspace = true
+tenant_size_model.workspace = true
+utils.workspace = true
+workspace_hack.workspace = true
+reqwest.workspace = true
 
 [dev-dependencies]
-criterion = "0.4"
-hex-literal = "0.3"
-tempfile = "3.2"
+criterion.workspace = true
+hex-literal.workspace = true
+tempfile.workspace = true
 
 [[bench]]
 name = "bench_layer_map"

pageserver/benches/bench_walredo.rs

@@ -30,33 +30,44 @@ fn redo_scenarios(c: &mut Criterion) {
     let conf = PageServerConf::dummy_conf(repo_dir.path().to_path_buf());
     let conf = Box::leak(Box::new(conf));
     let tenant_id = TenantId::generate();
-    // std::fs::create_dir_all(conf.tenant_path(&tenant_id)).unwrap();
-    let mut manager = PostgresRedoManager::new(conf, tenant_id);
-    manager.launch_process(14).unwrap();
+
+    let manager = PostgresRedoManager::new(conf, tenant_id);
 
     let manager = Arc::new(manager);
 
+    tracing::info!("executing first");
+    short().execute(&manager).unwrap();
+    tracing::info!("first executed");
+
     let thread_counts = [1, 2, 4, 8, 16];
 
-    for thread_count in thread_counts {
-        c.bench_with_input(
-            BenchmarkId::new("short-50record", thread_count),
-            &thread_count,
-            |b, thread_count| {
-                add_multithreaded_walredo_requesters(b, *thread_count, &manager, short, 50);
-            },
-        );
-    }
+    let mut group = c.benchmark_group("short");
+    group.sampling_mode(criterion::SamplingMode::Flat);
 
     for thread_count in thread_counts {
-        c.bench_with_input(
-            BenchmarkId::new("medium-10record", thread_count),
+        group.bench_with_input(
+            BenchmarkId::new("short", thread_count),
             &thread_count,
             |b, thread_count| {
-                add_multithreaded_walredo_requesters(b, *thread_count, &manager, medium, 10);
+                add_multithreaded_walredo_requesters(b, *thread_count, &manager, short);
             },
         );
     }
+    drop(group);
+
+    let mut group = c.benchmark_group("medium");
+    group.sampling_mode(criterion::SamplingMode::Flat);
+
+    for thread_count in thread_counts {
+        group.bench_with_input(
+            BenchmarkId::new("medium", thread_count),
+            &thread_count,
+            |b, thread_count| {
+                add_multithreaded_walredo_requesters(b, *thread_count, &manager, medium);
+            },
+        );
+    }
+    drop(group);
 }
 
 /// Sets up `threads` number of requesters to `request_redo`, with the given input.
@@ -65,46 +76,66 @@ fn add_multithreaded_walredo_requesters(
     threads: u32,
     manager: &Arc<PostgresRedoManager>,
     input_factory: fn() -> Request,
-    request_repeats: usize,
 ) {
-    b.iter_batched_ref(
-        || {
-            // barrier for all of the threads, and the benchmarked thread
-            let barrier = Arc::new(Barrier::new(threads as usize + 1));
+    assert_ne!(threads, 0);
 
-            let jhs = (0..threads)
-                .map(|_| {
-                    std::thread::spawn({
-                        let manager = manager.clone();
-                        let barrier = barrier.clone();
-                        move || {
-                            let input = std::iter::repeat(input_factory())
-                                .take(request_repeats)
-                                .collect::<Vec<_>>();
+    if threads == 1 {
+        b.iter_batched_ref(
+            || Some(input_factory()),
+            |input| execute_all(input.take(), manager),
+            criterion::BatchSize::PerIteration,
+        );
+    } else {
+        let (work_tx, work_rx) = std::sync::mpsc::sync_channel(threads as usize);
 
-                            barrier.wait();
+        let work_rx = std::sync::Arc::new(std::sync::Mutex::new(work_rx));
 
-                            execute_all(input, &manager).unwrap();
+        let barrier = Arc::new(Barrier::new(threads as usize + 1));
 
-                            barrier.wait();
+        let jhs = (0..threads)
+            .map(|_| {
+                std::thread::spawn({
+                    let manager = manager.clone();
+                    let barrier = barrier.clone();
+                    let work_rx = work_rx.clone();
+                    move || loop {
+                        // queue up and wait if we want to go another round
+                        if work_rx.lock().unwrap().recv().is_err() {
+                            break;
                         }
-                    })
+
+                        let input = Some(input_factory());
+
+                        barrier.wait();
+
+                        execute_all(input, &manager).unwrap();
+
+                        barrier.wait();
+                    }
                 })
-                .collect::<Vec<_>>();
+            })
+            .collect::<Vec<_>>();
 
-            (barrier, JoinOnDrop(jhs))
-        },
-        |input| {
-            let barrier = &input.0;
+        let _jhs = JoinOnDrop(jhs);
 
-            // start the work
-            barrier.wait();
+        b.iter_batched(
+            || {
+                for _ in 0..threads {
+                    work_tx.send(()).unwrap()
+                }
+            },
+            |()| {
+                // start the work
+                barrier.wait();
 
-            // wait for work to complete
-            barrier.wait();
-        },
-        criterion::BatchSize::PerIteration,
-    );
+                // wait for work to complete
+                barrier.wait();
+            },
+            criterion::BatchSize::PerIteration,
+        );
+
+        drop(work_tx);
+    }
 }
 
 struct JoinOnDrop(Vec<std::thread::JoinHandle<()>>);
@@ -121,7 +152,10 @@ impl Drop for JoinOnDrop {
     }
 }
 
-fn execute_all(input: Vec<Request>, manager: &PostgresRedoManager) -> Result<(), WalRedoError> {
+fn execute_all<I>(input: I, manager: &PostgresRedoManager) -> Result<(), WalRedoError>
+where
+    I: IntoIterator<Item = Request>,
+{
     // just fire all requests as fast as possible
     input.into_iter().try_for_each(|req| {
         let page = req.execute(manager)?;
@@ -143,6 +177,7 @@ macro_rules! lsn {
     }};
 }
 
+/// Short payload, 1132 bytes.
 // pg_records are copypasted from log, where they are put with Debug impl of Bytes, which uses \0
 // for null bytes.
 #[allow(clippy::octal_escapes)]
@@ -172,6 +207,7 @@ fn short() -> Request {
     }
 }
 
+/// Medium sized payload, serializes as 26393 bytes.
 // see [`short`]
 #[allow(clippy::octal_escapes)]
 fn medium() -> Request {

pageserver/src/basebackup.rs

@@ -10,20 +10,24 @@
 //! This module is responsible for creation of such tarball
 //! from data stored in object storage.
 //!
-use anyhow::{anyhow, bail, ensure, Context, Result};
+use anyhow::{anyhow, bail, ensure, Context};
 use bytes::{BufMut, BytesMut};
 use fail::fail_point;
-use itertools::Itertools;
 use std::fmt::Write as FmtWrite;
-use std::io;
-use std::io::Write;
-use std::sync::Arc;
 use std::time::SystemTime;
-use tar::{Builder, EntryType, Header};
+use tokio::io;
+use tokio::io::AsyncWrite;
 use tracing::*;
 
-use crate::task_mgr;
-use crate::tenant::{with_ondemand_download, PageReconstructResult, Timeline};
+/// NB: This relies on a modified version of tokio_tar that does *not* write the
+/// end-of-archive marker (1024 zero bytes), when the Builder struct is dropped
+/// without explicitly calling 'finish' or 'into_inner'!
+///
+/// See https://github.com/neondatabase/tokio-tar/pull/1
+///
+use tokio_tar::{Builder, EntryType, Header};
+
+use crate::tenant::Timeline;
 use pageserver_api::reltag::{RelTag, SlruKind};
 
 use postgres_ffi::pg_constants::{DEFAULTTABLESPACE_OID, GLOBALTABLESPACE_OID};
@@ -34,116 +38,130 @@ use postgres_ffi::PG_TLI;
 use postgres_ffi::{BLCKSZ, RELSEG_SIZE, WAL_SEGMENT_SIZE};
 use utils::lsn::Lsn;
 
+/// Create basebackup with non-rel data in it.
+/// Only include relational data if 'full_backup' is true.
+///
+/// Currently we use empty 'req_lsn' in two cases:
+///  * During the basebackup right after timeline creation
+///  * When working without safekeepers. In this situation it is important to match the lsn
+///    we are taking basebackup on with the lsn that is used in pageserver's walreceiver
+///    to start the replication.
+pub async fn send_basebackup_tarball<'a, W>(
+    write: &'a mut W,
+    timeline: &'a Timeline,
+    req_lsn: Option<Lsn>,
+    prev_lsn: Option<Lsn>,
+    full_backup: bool,
+) -> anyhow::Result<()>
+where
+    W: AsyncWrite + Send + Sync + Unpin,
+{
+    // Compute postgres doesn't have any previous WAL files, but the first
+    // record that it's going to write needs to include the LSN of the
+    // previous record (xl_prev). We include prev_record_lsn in the
+    // "zenith.signal" file, so that postgres can read it during startup.
+    //
+    // We don't keep full history of record boundaries in the page server,
+    // however, only the predecessor of the latest record on each
+    // timeline. So we can only provide prev_record_lsn when you take a
+    // base backup at the end of the timeline, i.e. at last_record_lsn.
+    // Even at the end of the timeline, we sometimes don't have a valid
+    // prev_lsn value; that happens if the timeline was just branched from
+    // an old LSN and it doesn't have any WAL of its own yet. We will set
+    // prev_lsn to Lsn(0) if we cannot provide the correct value.
+    let (backup_prev, backup_lsn) = if let Some(req_lsn) = req_lsn {
+        // Backup was requested at a particular LSN. The caller should've
+        // already checked that it's a valid LSN.
+
+        // If the requested point is the end of the timeline, we can
+        // provide prev_lsn. (get_last_record_rlsn() might return it as
+        // zero, though, if no WAL has been generated on this timeline
+        // yet.)
+        let end_of_timeline = timeline.get_last_record_rlsn();
+        if req_lsn == end_of_timeline.last {
+            (end_of_timeline.prev, req_lsn)
+        } else {
+            (Lsn(0), req_lsn)
+        }
+    } else {
+        // Backup was requested at end of the timeline.
+        let end_of_timeline = timeline.get_last_record_rlsn();
+        (end_of_timeline.prev, end_of_timeline.last)
+    };
+
+    // Consolidate the derived and the provided prev_lsn values
+    let prev_lsn = if let Some(provided_prev_lsn) = prev_lsn {
+        if backup_prev != Lsn(0) {
+            ensure!(backup_prev == provided_prev_lsn);
+        }
+        provided_prev_lsn
+    } else {
+        backup_prev
+    };
+
+    info!(
+        "taking basebackup lsn={}, prev_lsn={} (full_backup={})",
+        backup_lsn, prev_lsn, full_backup
+    );
+
+    let basebackup = Basebackup {
+        ar: Builder::new_non_terminated(write),
+        timeline,
+        lsn: backup_lsn,
+        prev_record_lsn: prev_lsn,
+        full_backup,
+    };
+    basebackup
+        .send_tarball()
+        .instrument(info_span!("send_tarball", backup_lsn=%backup_lsn))
+        .await
+}
+
 /// This is short-living object only for the time of tarball creation,
 /// created mostly to avoid passing a lot of parameters between various functions
 /// used for constructing tarball.
-pub struct Basebackup<'a, W>
+struct Basebackup<'a, W>
 where
-    W: Write,
+    W: AsyncWrite + Send + Sync + Unpin,
 {
-    ar: Builder<AbortableWrite<W>>,
-    timeline: &'a Arc<Timeline>,
-    pub lsn: Lsn,
+    ar: Builder<&'a mut W>,
+    timeline: &'a Timeline,
+    lsn: Lsn,
     prev_record_lsn: Lsn,
     full_backup: bool,
-    finished: bool,
 }
 
-// Create basebackup with non-rel data in it.
-// Only include relational data if 'full_backup' is true.
-//
-// Currently we use empty lsn in two cases:
-//  * During the basebackup right after timeline creation
-//  * When working without safekeepers. In this situation it is important to match the lsn
-//    we are taking basebackup on with the lsn that is used in pageserver's walreceiver
-//    to start the replication.
 impl<'a, W> Basebackup<'a, W>
 where
-    W: Write,
+    W: AsyncWrite + Send + Sync + Unpin,
 {
-    pub fn new(
-        write: W,
-        timeline: &'a Arc<Timeline>,
-        req_lsn: Option<Lsn>,
-        prev_lsn: Option<Lsn>,
-        full_backup: bool,
-    ) -> Result<Basebackup<'a, W>> {
-        // Compute postgres doesn't have any previous WAL files, but the first
-        // record that it's going to write needs to include the LSN of the
-        // previous record (xl_prev). We include prev_record_lsn in the
-        // "zenith.signal" file, so that postgres can read it during startup.
-        //
-        // We don't keep full history of record boundaries in the page server,
-        // however, only the predecessor of the latest record on each
-        // timeline. So we can only provide prev_record_lsn when you take a
-        // base backup at the end of the timeline, i.e. at last_record_lsn.
-        // Even at the end of the timeline, we sometimes don't have a valid
-        // prev_lsn value; that happens if the timeline was just branched from
-        // an old LSN and it doesn't have any WAL of its own yet. We will set
-        // prev_lsn to Lsn(0) if we cannot provide the correct value.
-        let (backup_prev, backup_lsn) = if let Some(req_lsn) = req_lsn {
-            // Backup was requested at a particular LSN. The caller should've
-            // already checked that it's a valid LSN.
-
-            // If the requested point is the end of the timeline, we can
-            // provide prev_lsn. (get_last_record_rlsn() might return it as
-            // zero, though, if no WAL has been generated on this timeline
-            // yet.)
-            let end_of_timeline = timeline.get_last_record_rlsn();
-            if req_lsn == end_of_timeline.last {
-                (end_of_timeline.prev, req_lsn)
-            } else {
-                (Lsn(0), req_lsn)
-            }
-        } else {
-            // Backup was requested at end of the timeline.
-            let end_of_timeline = timeline.get_last_record_rlsn();
-            (end_of_timeline.prev, end_of_timeline.last)
-        };
-
-        // Consolidate the derived and the provided prev_lsn values
-        let prev_lsn = if let Some(provided_prev_lsn) = prev_lsn {
-            if backup_prev != Lsn(0) {
-                ensure!(backup_prev == provided_prev_lsn)
-            }
-            provided_prev_lsn
-        } else {
-            backup_prev
-        };
-
-        info!(
-            "taking basebackup lsn={}, prev_lsn={} (full_backup={})",
-            backup_lsn, prev_lsn, full_backup
-        );
-
-        Ok(Basebackup {
-            ar: Builder::new(AbortableWrite::new(write)),
-            timeline,
-            lsn: backup_lsn,
-            prev_record_lsn: prev_lsn,
-            full_backup,
-            finished: false,
-        })
-    }
-
-    pub fn send_tarball(mut self) -> anyhow::Result<()> {
+    async fn send_tarball(mut self) -> anyhow::Result<()> {
         // TODO include checksum
 
         // Create pgdata subdirs structure
         for dir in PGDATA_SUBDIRS.iter() {
             let header = new_tar_header_dir(dir)?;
-            self.ar.append(&header, &mut io::empty())?;
+            self.ar
+                .append(&header, &mut io::empty())
+                .await
+                .context("could not add directory to basebackup tarball")?;
         }
 
-        // Send empty config files.
+        // Send config files.
         for filepath in PGDATA_SPECIAL_FILES.iter() {
             if *filepath == "pg_hba.conf" {
                 let data = PG_HBA.as_bytes();
                 let header = new_tar_header(filepath, data.len() as u64)?;
-                self.ar.append(&header, data)?;
+                self.ar
+                    .append(&header, data)
+                    .await
+                    .context("could not add config file to basebackup tarball")?;
             } else {
                 let header = new_tar_header(filepath, 0)?;
-                self.ar.append(&header, &mut io::empty())?;
+                self.ar
+                    .append(&header, &mut io::empty())
+                    .await
+                    .context("could not add config file to basebackup tarball")?;
             }
         }
 
@@ -153,30 +171,24 @@ where
             SlruKind::MultiXactOffsets,
             SlruKind::MultiXactMembers,
         ] {
-            for segno in
-                with_ondemand_download_sync(|| self.timeline.list_slru_segments(kind, self.lsn))?
-            {
-                self.add_slru_segment(kind, segno)?;
+            for segno in self.timeline.list_slru_segments(kind, self.lsn).await? {
+                self.add_slru_segment(kind, segno).await?;
             }
         }
 
         // Create tablespace directories
-        for ((spcnode, dbnode), has_relmap_file) in
-            with_ondemand_download_sync(|| self.timeline.list_dbdirs(self.lsn))?
-        {
-            self.add_dbdir(spcnode, dbnode, has_relmap_file)?;
+        for ((spcnode, dbnode), has_relmap_file) in self.timeline.list_dbdirs(self.lsn).await? {
+            self.add_dbdir(spcnode, dbnode, has_relmap_file).await?;
 
             // Gather and send relational files in each database if full backup is requested.
             if self.full_backup {
-                for rel in with_ondemand_download_sync(|| {
-                    self.timeline.list_rels(spcnode, dbnode, self.lsn)
-                })? {
-                    self.add_rel(rel)?;
+                for rel in self.timeline.list_rels(spcnode, dbnode, self.lsn).await? {
+                    self.add_rel(rel).await?;
                 }
             }
         }
-        for xid in with_ondemand_download_sync(|| self.timeline.list_twophase_files(self.lsn))? {
-            self.add_twophase_file(xid)?;
+        for xid in self.timeline.list_twophase_files(self.lsn).await? {
+            self.add_twophase_file(xid).await?;
         }
 
         fail_point!("basebackup-before-control-file", |_| {
@@ -184,44 +196,44 @@ where
         });
 
         // Generate pg_control and bootstrap WAL segment.
-        self.add_pgcontrol_file()?;
-        self.ar.finish()?;
-        self.finished = true;
+        self.add_pgcontrol_file().await?;
+        self.ar.finish().await?;
         debug!("all tarred up!");
         Ok(())
     }
 
-    fn add_rel(&mut self, tag: RelTag) -> anyhow::Result<()> {
-        let nblocks =
-            with_ondemand_download_sync(|| self.timeline.get_rel_size(tag, self.lsn, false))?;
-
-        // Function that adds relation segment data to archive
-        let mut add_file = |segment_index, data: &Vec<u8>| -> anyhow::Result<()> {
-            let file_name = tag.to_segfile_name(segment_index as u32);
-            let header = new_tar_header(&file_name, data.len() as u64)?;
-            self.ar.append(&header, data.as_slice())?;
-            Ok(())
-        };
+    async fn add_rel(&mut self, tag: RelTag) -> anyhow::Result<()> {
+        let nblocks = self.timeline.get_rel_size(tag, self.lsn, false).await?;
 
         // If the relation is empty, create an empty file
         if nblocks == 0 {
-            add_file(0, &vec![])?;
+            let file_name = tag.to_segfile_name(0);
+            let header = new_tar_header(&file_name, 0)?;
+            self.ar.append(&header, &mut io::empty()).await?;
             return Ok(());
         }
 
         // Add a file for each chunk of blocks (aka segment)
-        let chunks = (0..nblocks).chunks(RELSEG_SIZE as usize);
-        for (seg, blocks) in chunks.into_iter().enumerate() {
+        let mut startblk = 0;
+        let mut seg = 0;
+        while startblk < nblocks {
+            let endblk = std::cmp::min(startblk + RELSEG_SIZE, nblocks);
+
             let mut segment_data: Vec<u8> = vec![];
-            for blknum in blocks {
+            for blknum in startblk..endblk {
                 let img = self
                     .timeline
                     .get_rel_page_at_lsn(tag, blknum, self.lsn, false)
-                    .no_ondemand_download()?;
+                    .await?;
                 segment_data.extend_from_slice(&img[..]);
             }
 
-            add_file(seg, &segment_data)?;
+            let file_name = tag.to_segfile_name(seg as u32);
+            let header = new_tar_header(&file_name, segment_data.len() as u64)?;
+            self.ar.append(&header, segment_data.as_slice()).await?;
+
+            seg += 1;
+            startblk = endblk;
         }
 
         Ok(())
@@ -230,17 +242,18 @@ where
     //
     // Generate SLRU segment files from repository.
     //
-    fn add_slru_segment(&mut self, slru: SlruKind, segno: u32) -> anyhow::Result<()> {
-        let nblocks = with_ondemand_download_sync(|| {
-            self.timeline.get_slru_segment_size(slru, segno, self.lsn)
-        })?;
+    async fn add_slru_segment(&mut self, slru: SlruKind, segno: u32) -> anyhow::Result<()> {
+        let nblocks = self
+            .timeline
+            .get_slru_segment_size(slru, segno, self.lsn)
+            .await?;
 
         let mut slru_buf: Vec<u8> = Vec::with_capacity(nblocks as usize * BLCKSZ as usize);
         for blknum in 0..nblocks {
-            let img = with_ondemand_download_sync(|| {
-                self.timeline
-                    .get_slru_page_at_lsn(slru, segno, blknum, self.lsn)
-            })?;
+            let img = self
+                .timeline
+                .get_slru_page_at_lsn(slru, segno, blknum, self.lsn)
+                .await?;
 
             if slru == SlruKind::Clog {
                 ensure!(img.len() == BLCKSZ as usize || img.len() == BLCKSZ as usize + 8);
@@ -253,7 +266,7 @@ where
 
         let segname = format!("{}/{:>04X}", slru.to_str(), segno);
         let header = new_tar_header(&segname, slru_buf.len() as u64)?;
-        self.ar.append(&header, slru_buf.as_slice())?;
+        self.ar.append(&header, slru_buf.as_slice()).await?;
 
         trace!("Added to basebackup slru {} relsize {}", segname, nblocks);
         Ok(())
@@ -265,16 +278,17 @@ where
     // Each directory contains a PG_VERSION file, and the default database
     // directories also contain pg_filenode.map files.
     //
-    fn add_dbdir(
+    async fn add_dbdir(
         &mut self,
         spcnode: u32,
         dbnode: u32,
         has_relmap_file: bool,
     ) -> anyhow::Result<()> {
         let relmap_img = if has_relmap_file {
-            let img = with_ondemand_download_sync(|| {
-                self.timeline.get_relmap_file(spcnode, dbnode, self.lsn)
-            })?;
+            let img = self
+                .timeline
+                .get_relmap_file(spcnode, dbnode, self.lsn)
+                .await?;
             ensure!(img.len() == 512);
             Some(img)
         } else {
@@ -284,14 +298,14 @@ where
         if spcnode == GLOBALTABLESPACE_OID {
             let pg_version_str = self.timeline.pg_version.to_string();
             let header = new_tar_header("PG_VERSION", pg_version_str.len() as u64)?;
-            self.ar.append(&header, pg_version_str.as_bytes())?;
+            self.ar.append(&header, pg_version_str.as_bytes()).await?;
 
             info!("timeline.pg_version {}", self.timeline.pg_version);
 
             if let Some(img) = relmap_img {
                 // filenode map for global tablespace
                 let header = new_tar_header("global/pg_filenode.map", img.len() as u64)?;
-                self.ar.append(&header, &img[..])?;
+                self.ar.append(&header, &img[..]).await?;
             } else {
                 warn!("global/pg_filenode.map is missing");
             }
@@ -310,7 +324,7 @@ where
                 && self
                     .timeline
                     .list_rels(spcnode, dbnode, self.lsn)
-                    .no_ondemand_download()?
+                    .await?
                     .is_empty()
             {
                 return Ok(());
@@ -321,18 +335,18 @@ where
             // Append dir path for each database
             let path = format!("base/{}", dbnode);
             let header = new_tar_header_dir(&path)?;
-            self.ar.append(&header, &mut io::empty())?;
+            self.ar.append(&header, &mut io::empty()).await?;
 
             if let Some(img) = relmap_img {
                 let dst_path = format!("base/{}/PG_VERSION", dbnode);
 
                 let pg_version_str = self.timeline.pg_version.to_string();
                 let header = new_tar_header(&dst_path, pg_version_str.len() as u64)?;
-                self.ar.append(&header, pg_version_str.as_bytes())?;
+                self.ar.append(&header, pg_version_str.as_bytes()).await?;
 
                 let relmap_path = format!("base/{}/pg_filenode.map", dbnode);
                 let header = new_tar_header(&relmap_path, img.len() as u64)?;
-                self.ar.append(&header, &img[..])?;
+                self.ar.append(&header, &img[..]).await?;
             }
         };
         Ok(())
@@ -341,8 +355,8 @@ where
     //
     // Extract twophase state files
     //
-    fn add_twophase_file(&mut self, xid: TransactionId) -> anyhow::Result<()> {
-        let img = with_ondemand_download_sync(|| self.timeline.get_twophase_file(xid, self.lsn))?;
+    async fn add_twophase_file(&mut self, xid: TransactionId) -> anyhow::Result<()> {
+        let img = self.timeline.get_twophase_file(xid, self.lsn).await?;
 
         let mut buf = BytesMut::new();
         buf.extend_from_slice(&img[..]);
@@ -350,7 +364,7 @@ where
         buf.put_u32_le(crc);
         let path = format!("pg_twophase/{:>08X}", xid);
         let header = new_tar_header(&path, buf.len() as u64)?;
-        self.ar.append(&header, &buf[..])?;
+        self.ar.append(&header, &buf[..]).await?;
 
         Ok(())
     }
@@ -359,7 +373,7 @@ where
     // Add generated pg_control file and bootstrap WAL segment.
     // Also send zenith.signal file with extra bootstrap data.
     //
-    fn add_pgcontrol_file(&mut self) -> anyhow::Result<()> {
+    async fn add_pgcontrol_file(&mut self) -> anyhow::Result<()> {
         // add zenith.signal file
         let mut zenith_signal = String::new();
         if self.prev_record_lsn == Lsn(0) {
@@ -371,17 +385,23 @@ where
         } else {
             write!(zenith_signal, "PREV LSN: {}", self.prev_record_lsn)?;
         }
-        self.ar.append(
-            &new_tar_header("zenith.signal", zenith_signal.len() as u64)?,
-            zenith_signal.as_bytes(),
-        )?;
+        self.ar
+            .append(
+                &new_tar_header("zenith.signal", zenith_signal.len() as u64)?,
+                zenith_signal.as_bytes(),
+            )
+            .await?;
 
-        let checkpoint_bytes =
-            with_ondemand_download_sync(|| self.timeline.get_checkpoint(self.lsn))
-                .context("failed to get checkpoint bytes")?;
-        let pg_control_bytes =
-            with_ondemand_download_sync(|| self.timeline.get_control_file(self.lsn))
-                .context("failed get control bytes")?;
+        let checkpoint_bytes = self
+            .timeline
+            .get_checkpoint(self.lsn)
+            .await
+            .context("failed to get checkpoint bytes")?;
+        let pg_control_bytes = self
+            .timeline
+            .get_control_file(self.lsn)
+            .await
+            .context("failed get control bytes")?;
 
         let (pg_control_bytes, system_identifier) = postgres_ffi::generate_pg_control(
             &pg_control_bytes,
@@ -392,7 +412,7 @@ where
 
         //send pg_control
         let header = new_tar_header("global/pg_control", pg_control_bytes.len() as u64)?;
-        self.ar.append(&header, &pg_control_bytes[..])?;
+        self.ar.append(&header, &pg_control_bytes[..]).await?;
 
         //send wal segment
         let segno = self.lsn.segment_number(WAL_SEGMENT_SIZE);
@@ -404,24 +424,11 @@ where
             postgres_ffi::generate_wal_segment(segno, system_identifier, self.timeline.pg_version)
                 .map_err(|e| anyhow!(e).context("Failed generating wal segment"))?;
         ensure!(wal_seg.len() == WAL_SEGMENT_SIZE);
-        self.ar.append(&header, &wal_seg[..])?;
+        self.ar.append(&header, &wal_seg[..]).await?;
         Ok(())
     }
 }
 
-impl<'a, W> Drop for Basebackup<'a, W>
-where
-    W: Write,
-{
-    /// If the basebackup was not finished, prevent the Archive::drop() from
-    /// writing the end-of-archive marker.
-    fn drop(&mut self) {
-        if !self.finished {
-            self.ar.get_mut().abort();
-        }
-    }
-}
-
 //
 // Create new tarball entry header
 //
@@ -457,57 +464,3 @@ fn new_tar_header_dir(path: &str) -> anyhow::Result<Header> {
     header.set_cksum();
     Ok(header)
 }
-
-/// A wrapper that passes through all data to the underlying Write,
-/// until abort() is called.
-///
-/// tar::Builder has an annoying habit of finishing the archive with
-/// a valid tar end-of-archive marker (two 512-byte sectors of zeros),
-/// even if an error occurs and we don't finish building the archive.
-/// We'd rather abort writing the tarball immediately than construct
-/// a seemingly valid but incomplete archive. This wrapper allows us
-/// to swallow the end-of-archive marker that Builder::drop() emits,
-/// without writing it to the underlying sink.
-///
-struct AbortableWrite<W> {
-    w: W,
-    aborted: bool,
-}
-
-impl<W> AbortableWrite<W> {
-    pub fn new(w: W) -> Self {
-        AbortableWrite { w, aborted: false }
-    }
-
-    pub fn abort(&mut self) {
-        self.aborted = true;
-    }
-}
-
-impl<W> Write for AbortableWrite<W>
-where
-    W: Write,
-{
-    fn write(&mut self, data: &[u8]) -> io::Result<usize> {
-        if self.aborted {
-            Ok(data.len())
-        } else {
-            self.w.write(data)
-        }
-    }
-    fn flush(&mut self) -> io::Result<()> {
-        if self.aborted {
-            Ok(())
-        } else {
-            self.w.flush()
-        }
-    }
-}
-
-fn with_ondemand_download_sync<F, T>(f: F) -> anyhow::Result<T>
-where
-    F: Send + Fn() -> PageReconstructResult<T>,
-    T: Send,
-{
-    task_mgr::COMPUTE_REQUEST_RUNTIME.block_on(with_ondemand_download(f))
-}

pageserver/src/billing_metrics.rs

@@ -1,283 +0,0 @@
-//!
-//! Periodically collect consumption metrics for all active tenants
-//! and push them to a HTTP endpoint.
-//! Cache metrics to send only the updated ones.
-//!
-
-use anyhow;
-use tracing::*;
-use utils::id::TimelineId;
-
-use crate::task_mgr;
-use crate::tenant::mgr;
-use pageserver_api::models::TenantState;
-use utils::id::TenantId;
-
-use serde::{Deserialize, Serialize};
-use std::collections::HashMap;
-use std::fmt;
-use std::str::FromStr;
-use std::time::Duration;
-
-use chrono::{DateTime, Utc};
-use reqwest::Url;
-
-/// BillingMetric struct that defines the format for one metric entry
-/// i.e.
-///
-/// ```json
-/// {
-/// "metric": "remote_storage_size",
-/// "type": "absolute",
-/// "tenant_id": "5d07d9ce9237c4cd845ea7918c0afa7d",
-/// "timeline_id": "00000000000000000000000000000000",
-/// "time": ...,
-/// "value": 12345454,
-/// }
-/// ```
-#[derive(Serialize, Deserialize, Debug, Clone, Eq, PartialEq, Ord, PartialOrd)]
-pub struct BillingMetric {
-    pub metric: BillingMetricKind,
-    pub metric_type: &'static str,
-    pub tenant_id: TenantId,
-    pub timeline_id: Option<TimelineId>,
-    pub time: DateTime<Utc>,
-    pub value: u64,
-}
-
-impl BillingMetric {
-    pub fn new_absolute(
-        metric: BillingMetricKind,
-        tenant_id: TenantId,
-        timeline_id: Option<TimelineId>,
-        value: u64,
-    ) -> Self {
-        Self {
-            metric,
-            metric_type: "absolute",
-            tenant_id,
-            timeline_id,
-            time: Utc::now(),
-            value,
-        }
-    }
-}
-
-#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Ord, PartialOrd, Serialize, Deserialize)]
-#[serde(rename_all = "snake_case")]
-pub enum BillingMetricKind {
-    /// Amount of WAL produced , by a timeline, i.e. last_record_lsn
-    /// This is an absolute, per-timeline metric.
-    WrittenSize,
-    /// Size of all tenant branches including WAL
-    /// This is an absolute, per-tenant metric.
-    /// This is the same metric that tenant/tenant_id/size endpoint returns.
-    SyntheticStorageSize,
-    /// Size of all the layer files in the tenant's directory on disk on the pageserver.
-    /// This is an absolute, per-tenant metric.
-    /// See also prometheus metric RESIDENT_PHYSICAL_SIZE.
-    ResidentSize,
-    /// Size of the remote storage (S3) directory.
-    /// This is an absolute, per-tenant metric.
-    RemoteStorageSize,
-}
-
-impl FromStr for BillingMetricKind {
-    type Err = anyhow::Error;
-
-    fn from_str(s: &str) -> Result<Self, Self::Err> {
-        match s {
-            "written_size" => Ok(Self::WrittenSize),
-            "synthetic_storage_size" => Ok(Self::SyntheticStorageSize),
-            "resident_size" => Ok(Self::ResidentSize),
-            "remote_storage_size" => Ok(Self::RemoteStorageSize),
-            _ => anyhow::bail!("invalid value \"{s}\" for metric type"),
-        }
-    }
-}
-
-impl fmt::Display for BillingMetricKind {
-    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        f.write_str(match self {
-            BillingMetricKind::WrittenSize => "written_size",
-            BillingMetricKind::SyntheticStorageSize => "synthetic_storage_size",
-            BillingMetricKind::ResidentSize => "resident_size",
-            BillingMetricKind::RemoteStorageSize => "remote_storage_size",
-        })
-    }
-}
-
-#[derive(Debug, Clone, PartialEq, Eq, Hash)]
-pub struct BillingMetricsKey {
-    tenant_id: TenantId,
-    timeline_id: Option<TimelineId>,
-    metric: BillingMetricKind,
-}
-
-#[derive(serde::Serialize)]
-struct EventChunk<'a> {
-    events: &'a [BillingMetric],
-}
-
-/// Main thread that serves metrics collection
-pub async fn collect_metrics(
-    metric_collection_endpoint: &Url,
-    metric_collection_interval: Duration,
-) -> anyhow::Result<()> {
-    let mut ticker = tokio::time::interval(metric_collection_interval);
-
-    info!("starting collect_metrics");
-
-    // define client here to reuse it for all requests
-    let client = reqwest::Client::new();
-    let mut cached_metrics: HashMap<BillingMetricsKey, u64> = HashMap::new();
-
-    loop {
-        tokio::select! {
-            _ = task_mgr::shutdown_watcher() => {
-                info!("collect_metrics received cancellation request");
-                return Ok(());
-            },
-            _ = ticker.tick() => {
-                collect_metrics_task(&client, &mut cached_metrics, metric_collection_endpoint).await?;
-            }
-        }
-    }
-}
-
-/// One iteration of metrics collection
-///
-/// Gather per-tenant and per-timeline metrics and send them to the `metric_collection_endpoint`.
-/// Cache metrics to avoid sending the same metrics multiple times.
-pub async fn collect_metrics_task(
-    client: &reqwest::Client,
-    cached_metrics: &mut HashMap<BillingMetricsKey, u64>,
-    metric_collection_endpoint: &reqwest::Url,
-) -> anyhow::Result<()> {
-    let mut current_metrics: Vec<(BillingMetricsKey, u64)> = Vec::new();
-    trace!(
-        "starting collect_metrics_task. metric_collection_endpoint: {}",
-        metric_collection_endpoint
-    );
-
-    // get list of tenants
-    let tenants = mgr::list_tenants().await;
-
-    // iterate through list of Active tenants and collect metrics
-    for (tenant_id, tenant_state) in tenants {
-        if tenant_state != TenantState::Active {
-            continue;
-        }
-
-        let tenant = mgr::get_tenant(tenant_id, true).await?;
-
-        let mut tenant_resident_size = 0;
-
-        // iterate through list of timelines in tenant
-        for timeline in tenant.list_timelines().iter() {
-            let timeline_written_size = u64::from(timeline.get_last_record_lsn());
-
-            current_metrics.push((
-                BillingMetricsKey {
-                    tenant_id,
-                    timeline_id: Some(timeline.timeline_id),
-                    metric: BillingMetricKind::WrittenSize,
-                },
-                timeline_written_size,
-            ));
-
-            let timeline_resident_size = timeline.get_resident_physical_size();
-            tenant_resident_size += timeline_resident_size;
-
-            debug!(
-                "per-timeline current metrics for tenant: {}: timeline {} resident_size={} last_record_lsn {} (as bytes)",
-                tenant_id, timeline.timeline_id, timeline_resident_size, timeline_written_size)
-        }
-
-        let tenant_remote_size = tenant.get_remote_size().await?;
-        debug!(
-            "collected current metrics for tenant: {}: state={:?} resident_size={} remote_size={}",
-            tenant_id, tenant_state, tenant_resident_size, tenant_remote_size
-        );
-
-        current_metrics.push((
-            BillingMetricsKey {
-                tenant_id,
-                timeline_id: None,
-                metric: BillingMetricKind::ResidentSize,
-            },
-            tenant_resident_size,
-        ));
-
-        current_metrics.push((
-            BillingMetricsKey {
-                tenant_id,
-                timeline_id: None,
-                metric: BillingMetricKind::RemoteStorageSize,
-            },
-            tenant_remote_size,
-        ));
-
-        // TODO add SyntheticStorageSize metric
-    }
-
-    // Filter metrics
-    current_metrics.retain(|(curr_key, curr_val)| match cached_metrics.get(curr_key) {
-        Some(val) => val != curr_val,
-        None => true,
-    });
-
-    if current_metrics.is_empty() {
-        trace!("no new metrics to send");
-        return Ok(());
-    }
-
-    // Send metrics.
-    // Split into chunks of 1000 metrics to avoid exceeding the max request size
-    const CHUNK_SIZE: usize = 1000;
-    let chunks = current_metrics.chunks(CHUNK_SIZE);
-
-    let mut chunk_to_send: Vec<BillingMetric> = Vec::with_capacity(1000);
-
-    for chunk in chunks {
-        chunk_to_send.clear();
-        // enrich metrics with timestamp and metric_kind before sending
-        chunk_to_send.extend(chunk.iter().map(|(curr_key, curr_val)| {
-            BillingMetric::new_absolute(
-                curr_key.metric,
-                curr_key.tenant_id,
-                curr_key.timeline_id,
-                *curr_val,
-            )
-        }));
-
-        let chunk_json = serde_json::value::to_raw_value(&EventChunk {
-            events: &chunk_to_send,
-        })
-        .expect("BillingMetric should not fail serialization");
-
-        let res = client
-            .post(metric_collection_endpoint.clone())
-            .json(&chunk_json)
-            .send()
-            .await;
-
-        match res {
-            Ok(res) => {
-                if res.status().is_success() {
-                    // update cached metrics after they were sent successfully
-                    for (curr_key, curr_val) in chunk.iter() {
-                        cached_metrics.insert(curr_key.clone(), *curr_val);
-                    }
-                } else {
-                    error!("metrics endpoint refused the sent metrics: {:?}", res);
-                }
-            }
-            Err(err) => {
-                error!("failed to send metrics: {:?}", err);
-            }
-        }
-    }
-
-    Ok(())
-}

pageserver/src/bin/pageserver.rs

@@ -13,7 +13,7 @@ use tracing::*;
 use metrics::set_build_info_metric;
 use pageserver::{
     config::{defaults::*, PageServerConf},
-    http, page_cache, page_service, profiling, task_mgr,
+    http, page_cache, page_service, task_mgr,
     task_mgr::TaskKind,
     task_mgr::{
         BACKGROUND_RUNTIME, COMPUTE_REQUEST_RUNTIME, MGMT_REQUEST_RUNTIME, WALRECEIVER_RUNTIME,
@@ -40,8 +40,6 @@ const FEATURES: &[&str] = &[
     "testing",
     #[cfg(feature = "fail/failpoints")]
     "fail/failpoints",
-    #[cfg(feature = "profiling")]
-    "profiling",
 ];
 
 fn version() -> String {
@@ -247,15 +245,12 @@ fn start_pageserver(conf: &'static PageServerConf) -> anyhow::Result<()> {
     // Install signal handlers
     let signals = signals::install_shutdown_handlers()?;
 
-    // Start profiler (if enabled)
-    let profiler_guard = profiling::init_profiler(conf);
-
     // Launch broker client
     WALRECEIVER_RUNTIME.block_on(pageserver::walreceiver::init_broker_client(conf))?;
 
     // Initialize authentication for incoming connections
     let auth = match &conf.auth_type {
-        AuthType::Trust | AuthType::MD5 => None,
+        AuthType::Trust => None,
         AuthType::NeonJWT => {
             // unwrap is ok because check is performed when creating config, so path is set and file exists
             let key_path = conf.auth_validation_public_key_path.as_ref().unwrap();
@@ -264,19 +259,35 @@ fn start_pageserver(conf: &'static PageServerConf) -> anyhow::Result<()> {
     };
     info!("Using auth: {:#?}", conf.auth_type);
 
-    match var("ZENITH_AUTH_TOKEN") {
-        Ok(v) => {
+    // TODO: remove ZENITH_AUTH_TOKEN once it's not used anywhere in development/staging/prod configuration.
+    match (var("ZENITH_AUTH_TOKEN"), var("NEON_AUTH_TOKEN")) {
+        (old, Ok(v)) => {
             info!("Loaded JWT token for authentication with Safekeeper");
+            if let Ok(v_old) = old {
+                warn!(
+                    "JWT token for Safekeeper is specified twice, ZENITH_AUTH_TOKEN is deprecated"
+                );
+                if v_old != v {
+                    warn!("JWT token for Safekeeper has two different values, choosing NEON_AUTH_TOKEN");
+                }
+            }
             pageserver::config::SAFEKEEPER_AUTH_TOKEN
                 .set(Arc::new(v))
                 .map_err(|_| anyhow!("Could not initialize SAFEKEEPER_AUTH_TOKEN"))?;
         }
-        Err(VarError::NotPresent) => {
+        (Ok(v), _) => {
+            info!("Loaded JWT token for authentication with Safekeeper");
+            warn!("Please update pageserver configuration: the JWT token should be NEON_AUTH_TOKEN, not ZENITH_AUTH_TOKEN");
+            pageserver::config::SAFEKEEPER_AUTH_TOKEN
+                .set(Arc::new(v))
+                .map_err(|_| anyhow!("Could not initialize SAFEKEEPER_AUTH_TOKEN"))?;
+        }
+        (_, Err(VarError::NotPresent)) => {
             info!("No JWT token for authentication with Safekeeper detected");
         }
-        Err(e) => {
+        (_, Err(e)) => {
             return Err(e).with_context(|| {
-                "Failed to either load to detect non-present ZENITH_AUTH_TOKEN environment variable"
+                "Failed to either load to detect non-present NEON_AUTH_TOKEN environment variable"
             })
         }
     };
@@ -322,9 +333,11 @@ fn start_pageserver(conf: &'static PageServerConf) -> anyhow::Result<()> {
                 "consumption metrics collection",
                 true,
                 async move {
-                    pageserver::billing_metrics::collect_metrics(
+                    pageserver::consumption_metrics::collect_metrics(
                         metric_collection_endpoint,
                         conf.metric_collection_interval,
+                        conf.synthetic_size_calculation_interval,
+                        conf.id,
                     )
                     .instrument(info_span!("metrics_collection"))
                     .await?;
@@ -355,7 +368,6 @@ fn start_pageserver(conf: &'static PageServerConf) -> anyhow::Result<()> {
                 "Got {}. Terminating in immediate shutdown mode",
                 signal.name()
             );
-            profiling::exit_profiler(conf, &profiler_guard);
             std::process::exit(111);
         }
 
@@ -364,7 +376,6 @@ fn start_pageserver(conf: &'static PageServerConf) -> anyhow::Result<()> {
                 "Got {}. Terminating gracefully in fast shutdown mode",
                 signal.name()
             );
-            profiling::exit_profiler(conf, &profiler_guard);
             BACKGROUND_RUNTIME.block_on(pageserver::shutdown_pageserver(0));
             unreachable!()
         }

pageserver/src/config.rs

@@ -59,6 +59,8 @@ pub mod defaults {
 
     pub const DEFAULT_METRIC_COLLECTION_INTERVAL: &str = "10 min";
     pub const DEFAULT_METRIC_COLLECTION_ENDPOINT: Option<reqwest::Url> = None;
+    pub const DEFAULT_SYNTHETIC_SIZE_CALCULATION_INTERVAL: &str = "10 min";
+
     ///
     /// Default built-in configuration file.
     ///
@@ -83,6 +85,7 @@ pub mod defaults {
 #concurrent_tenant_size_logical_size_queries = '{DEFAULT_CONCURRENT_TENANT_SIZE_LOGICAL_SIZE_QUERIES}'
 
 #metric_collection_interval = '{DEFAULT_METRIC_COLLECTION_INTERVAL}'
+#synthetic_size_calculation_interval = '{DEFAULT_SYNTHETIC_SIZE_CALCULATION_INTERVAL}'
 
 # [tenant_config]
 #checkpoint_distance = {DEFAULT_CHECKPOINT_DISTANCE} # in bytes
@@ -138,7 +141,6 @@ pub struct PageServerConf {
     pub auth_validation_public_key_path: Option<PathBuf>,
     pub remote_storage_config: Option<RemoteStorageConfig>,
 
-    pub profiling: ProfilingConfig,
     pub default_tenant_conf: TenantConf,
 
     /// Storage broker endpoints to connect to.
@@ -153,6 +155,7 @@ pub struct PageServerConf {
     // How often to collect metrics and send them to the metrics endpoint.
     pub metric_collection_interval: Duration,
     pub metric_collection_endpoint: Option<Url>,
+    pub synthetic_size_calculation_interval: Duration,
 
     pub test_remote_failures: u64,
 }
@@ -165,25 +168,6 @@ pub struct PageServerConf {
 /// startup code to the connection code through a dozen layers.
 pub static SAFEKEEPER_AUTH_TOKEN: OnceCell<Arc<String>> = OnceCell::new();
 
-#[derive(Debug, Clone, PartialEq, Eq)]
-pub enum ProfilingConfig {
-    Disabled,
-    PageRequests,
-}
-
-impl FromStr for ProfilingConfig {
-    type Err = anyhow::Error;
-
-    fn from_str(s: &str) -> Result<ProfilingConfig, Self::Err> {
-        let result = match s {
-            "disabled"  => ProfilingConfig::Disabled,
-            "page_requests"  => ProfilingConfig::PageRequests,
-            _ => bail!("invalid value \"{s}\" for profiling option, valid values are \"disabled\" and \"page_requests\""),
-        };
-        Ok(result)
-    }
-}
-
 // use dedicated enum for builder to better indicate the intention
 // and avoid possible confusion with nested options
 pub enum BuilderValue<T> {
@@ -226,7 +210,6 @@ struct PageServerConfigBuilder {
 
     id: BuilderValue<NodeId>,
 
-    profiling: BuilderValue<ProfilingConfig>,
     broker_endpoint: BuilderValue<Uri>,
     broker_keepalive_interval: BuilderValue<Duration>,
 
@@ -236,6 +219,7 @@ struct PageServerConfigBuilder {
 
     metric_collection_interval: BuilderValue<Duration>,
     metric_collection_endpoint: BuilderValue<Option<Url>>,
+    synthetic_size_calculation_interval: BuilderValue<Duration>,
 
     test_remote_failures: BuilderValue<u64>,
 }
@@ -262,7 +246,6 @@ impl Default for PageServerConfigBuilder {
             auth_validation_public_key_path: Set(None),
             remote_storage_config: Set(None),
             id: NotSet,
-            profiling: Set(ProfilingConfig::Disabled),
             broker_endpoint: Set(storage_broker::DEFAULT_ENDPOINT
                 .parse()
                 .expect("failed to parse default broker endpoint")),
@@ -277,6 +260,10 @@ impl Default for PageServerConfigBuilder {
                 DEFAULT_METRIC_COLLECTION_INTERVAL,
             )
             .expect("cannot parse default metric collection interval")),
+            synthetic_size_calculation_interval: Set(humantime::parse_duration(
+                DEFAULT_SYNTHETIC_SIZE_CALCULATION_INTERVAL,
+            )
+            .expect("cannot parse default synthetic size calculation interval")),
             metric_collection_endpoint: Set(DEFAULT_METRIC_COLLECTION_ENDPOINT),
 
             test_remote_failures: Set(0),
@@ -348,10 +335,6 @@ impl PageServerConfigBuilder {
         self.id = BuilderValue::Set(node_id)
     }
 
-    pub fn profiling(&mut self, profiling: ProfilingConfig) {
-        self.profiling = BuilderValue::Set(profiling)
-    }
-
     pub fn log_format(&mut self, log_format: LogFormat) {
         self.log_format = BuilderValue::Set(log_format)
     }
@@ -368,6 +351,14 @@ impl PageServerConfigBuilder {
         self.metric_collection_endpoint = BuilderValue::Set(metric_collection_endpoint)
     }
 
+    pub fn synthetic_size_calculation_interval(
+        &mut self,
+        synthetic_size_calculation_interval: Duration,
+    ) {
+        self.synthetic_size_calculation_interval =
+            BuilderValue::Set(synthetic_size_calculation_interval)
+    }
+
     pub fn test_remote_failures(&mut self, fail_first: u64) {
         self.test_remote_failures = BuilderValue::Set(fail_first);
     }
@@ -405,7 +396,6 @@ impl PageServerConfigBuilder {
                 .remote_storage_config
                 .ok_or(anyhow!("missing remote_storage_config"))?,
             id: self.id.ok_or(anyhow!("missing id"))?,
-            profiling: self.profiling.ok_or(anyhow!("missing profiling"))?,
             // TenantConf is handled separately
             default_tenant_conf: TenantConf::default(),
             broker_endpoint: self
@@ -426,6 +416,9 @@ impl PageServerConfigBuilder {
             metric_collection_endpoint: self
                 .metric_collection_endpoint
                 .ok_or(anyhow!("missing metric_collection_endpoint"))?,
+            synthetic_size_calculation_interval: self
+                .synthetic_size_calculation_interval
+                .ok_or(anyhow!("missing synthetic_size_calculation_interval"))?,
             test_remote_failures: self
                 .test_remote_failures
                 .ok_or(anyhow!("missing test_remote_failuers"))?,
@@ -588,7 +581,6 @@ impl PageServerConf {
                     t_conf = Self::parse_toml_tenant_conf(item)?;
                 }
                 "id" => builder.id(NodeId(parse_toml_u64(key, item)?)),
-                "profiling" => builder.profiling(parse_toml_from_str(key, item)?),
                 "broker_endpoint" => builder.broker_endpoint(parse_toml_string(key, item)?.parse().context("failed to parse broker endpoint")?),
                 "broker_keepalive_interval" => builder.broker_keepalive_interval(parse_toml_duration(key, item)?),
                 "log_format" => builder.log_format(
@@ -605,7 +597,8 @@ impl PageServerConf {
                     let endpoint = parse_toml_string(key, item)?.parse().context("failed to parse metric_collection_endpoint")?;
                     builder.metric_collection_endpoint(Some(endpoint));
                 },
-
+                "synthetic_size_calculation_interval" =>
+                    builder.synthetic_size_calculation_interval(parse_toml_duration(key, item)?),
                 "test_remote_failures" => builder.test_remote_failures(parse_toml_u64(key, item)?),
                 _ => bail!("unrecognized pageserver option '{key}'"),
             }
@@ -722,7 +715,6 @@ impl PageServerConf {
             auth_type: AuthType::Trust,
             auth_validation_public_key_path: None,
             remote_storage_config: None,
-            profiling: ProfilingConfig::Disabled,
             default_tenant_conf: TenantConf::default(),
             broker_endpoint: storage_broker::DEFAULT_ENDPOINT.parse().unwrap(),
             broker_keepalive_interval: Duration::from_secs(5000),
@@ -730,6 +722,7 @@ impl PageServerConf {
             concurrent_tenant_size_logical_size_queries: ConfigurableSemaphore::default(),
             metric_collection_interval: Duration::from_secs(60),
             metric_collection_endpoint: defaults::DEFAULT_METRIC_COLLECTION_ENDPOINT,
+            synthetic_size_calculation_interval: Duration::from_secs(60),
             test_remote_failures: 0,
         }
     }
@@ -863,6 +856,7 @@ id = 10
 
 metric_collection_interval = '222 s'
 metric_collection_endpoint = 'http://localhost:80/metrics'
+synthetic_size_calculation_interval = '333 s'
 log_format = 'json'
 
 "#;
@@ -898,7 +892,6 @@ log_format = 'json'
                 auth_type: AuthType::Trust,
                 auth_validation_public_key_path: None,
                 remote_storage_config: None,
-                profiling: ProfilingConfig::Disabled,
                 default_tenant_conf: TenantConf::default(),
                 broker_endpoint: storage_broker::DEFAULT_ENDPOINT.parse().unwrap(),
                 broker_keepalive_interval: humantime::parse_duration(
@@ -910,6 +903,9 @@ log_format = 'json'
                     defaults::DEFAULT_METRIC_COLLECTION_INTERVAL
                 )?,
                 metric_collection_endpoint: defaults::DEFAULT_METRIC_COLLECTION_ENDPOINT,
+                synthetic_size_calculation_interval: humantime::parse_duration(
+                    defaults::DEFAULT_SYNTHETIC_SIZE_CALCULATION_INTERVAL
+                )?,
                 test_remote_failures: 0,
             },
             "Correct defaults should be used when no config values are provided"
@@ -949,7 +945,6 @@ log_format = 'json'
                 auth_type: AuthType::Trust,
                 auth_validation_public_key_path: None,
                 remote_storage_config: None,
-                profiling: ProfilingConfig::Disabled,
                 default_tenant_conf: TenantConf::default(),
                 broker_endpoint: storage_broker::DEFAULT_ENDPOINT.parse().unwrap(),
                 broker_keepalive_interval: Duration::from_secs(5),
@@ -957,6 +952,7 @@ log_format = 'json'
                 concurrent_tenant_size_logical_size_queries: ConfigurableSemaphore::default(),
                 metric_collection_interval: Duration::from_secs(222),
                 metric_collection_endpoint: Some(Url::parse("http://localhost:80/metrics")?),
+                synthetic_size_calculation_interval: Duration::from_secs(333),
                 test_remote_failures: 0,
             },
             "Should be able to parse all basic config values correctly"

pageserver/src/consumption_metrics.rs

@@ -0,0 +1,292 @@
+//!
+//! Periodically collect consumption metrics for all active tenants
+//! and push them to a HTTP endpoint.
+//! Cache metrics to send only the updated ones.
+//!
+use crate::task_mgr::{self, TaskKind, BACKGROUND_RUNTIME};
+use crate::tenant::mgr;
+use anyhow;
+use chrono::Utc;
+use consumption_metrics::{idempotency_key, Event, EventChunk, EventType, CHUNK_SIZE};
+use pageserver_api::models::TenantState;
+use reqwest::Url;
+use serde::Serialize;
+use serde_with::{serde_as, DisplayFromStr};
+use std::collections::HashMap;
+use std::time::Duration;
+use tracing::*;
+use utils::id::{NodeId, TenantId, TimelineId};
+
+const WRITTEN_SIZE: &str = "written_size";
+const SYNTHETIC_STORAGE_SIZE: &str = "synthetic_storage_size";
+const RESIDENT_SIZE: &str = "resident_size";
+const REMOTE_STORAGE_SIZE: &str = "remote_storage_size";
+const TIMELINE_LOGICAL_SIZE: &str = "timeline_logical_size";
+
+#[serde_as]
+#[derive(Serialize)]
+struct Ids {
+    #[serde_as(as = "DisplayFromStr")]
+    tenant_id: TenantId,
+    #[serde_as(as = "Option<DisplayFromStr>")]
+    #[serde(skip_serializing_if = "Option::is_none")]
+    timeline_id: Option<TimelineId>,
+}
+
+/// Key that uniquely identifies the object, this metric describes.
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+pub struct PageserverConsumptionMetricsKey {
+    pub tenant_id: TenantId,
+    pub timeline_id: Option<TimelineId>,
+    pub metric: &'static str,
+}
+
+/// Main thread that serves metrics collection
+pub async fn collect_metrics(
+    metric_collection_endpoint: &Url,
+    metric_collection_interval: Duration,
+    synthetic_size_calculation_interval: Duration,
+    node_id: NodeId,
+) -> anyhow::Result<()> {
+    let mut ticker = tokio::time::interval(metric_collection_interval);
+
+    info!("starting collect_metrics");
+
+    // spin up background worker that caclulates tenant sizes
+    task_mgr::spawn(
+        BACKGROUND_RUNTIME.handle(),
+        TaskKind::CalculateSyntheticSize,
+        None,
+        None,
+        "synthetic size calculation",
+        false,
+        async move {
+            calculate_synthetic_size_worker(synthetic_size_calculation_interval)
+                .instrument(info_span!("synthetic_size_worker"))
+                .await?;
+            Ok(())
+        },
+    );
+
+    // define client here to reuse it for all requests
+    let client = reqwest::Client::new();
+    let mut cached_metrics: HashMap<PageserverConsumptionMetricsKey, u64> = HashMap::new();
+
+    loop {
+        tokio::select! {
+            _ = task_mgr::shutdown_watcher() => {
+                info!("collect_metrics received cancellation request");
+                return Ok(());
+            },
+            _ = ticker.tick() => {
+                if let Err(err) = collect_metrics_iteration(&client, &mut cached_metrics, metric_collection_endpoint, node_id).await
+                {
+                    error!("metrics collection failed: {err:?}");
+                }
+            }
+        }
+    }
+}
+
+/// One iteration of metrics collection
+///
+/// Gather per-tenant and per-timeline metrics and send them to the `metric_collection_endpoint`.
+/// Cache metrics to avoid sending the same metrics multiple times.
+///
+/// TODO
+/// - refactor this function (chunking+sending part) to reuse it in proxy module;
+/// - improve error handling. Now if one tenant fails to collect metrics,
+/// the whole iteration fails and metrics for other tenants are not collected.
+pub async fn collect_metrics_iteration(
+    client: &reqwest::Client,
+    cached_metrics: &mut HashMap<PageserverConsumptionMetricsKey, u64>,
+    metric_collection_endpoint: &reqwest::Url,
+    node_id: NodeId,
+) -> anyhow::Result<()> {
+    let mut current_metrics: Vec<(PageserverConsumptionMetricsKey, u64)> = Vec::new();
+    trace!(
+        "starting collect_metrics_iteration. metric_collection_endpoint: {}",
+        metric_collection_endpoint
+    );
+
+    // get list of tenants
+    let tenants = mgr::list_tenants().await;
+
+    // iterate through list of Active tenants and collect metrics
+    for (tenant_id, tenant_state) in tenants {
+        if tenant_state != TenantState::Active {
+            continue;
+        }
+
+        let tenant = mgr::get_tenant(tenant_id, true).await?;
+
+        let mut tenant_resident_size = 0;
+
+        // iterate through list of timelines in tenant
+        for timeline in tenant.list_timelines().iter() {
+            // collect per-timeline metrics only for active timelines
+            if timeline.is_active() {
+                let timeline_written_size = u64::from(timeline.get_last_record_lsn());
+
+                current_metrics.push((
+                    PageserverConsumptionMetricsKey {
+                        tenant_id,
+                        timeline_id: Some(timeline.timeline_id),
+                        metric: WRITTEN_SIZE,
+                    },
+                    timeline_written_size,
+                ));
+
+                let (timeline_logical_size, is_exact) = timeline.get_current_logical_size()?;
+                // Only send timeline logical size when it is fully calculated.
+                if is_exact {
+                    current_metrics.push((
+                        PageserverConsumptionMetricsKey {
+                            tenant_id,
+                            timeline_id: Some(timeline.timeline_id),
+                            metric: TIMELINE_LOGICAL_SIZE,
+                        },
+                        timeline_logical_size,
+                    ));
+                }
+            }
+
+            let timeline_resident_size = timeline.get_resident_physical_size();
+            tenant_resident_size += timeline_resident_size;
+        }
+
+        let tenant_remote_size = tenant.get_remote_size().await?;
+        debug!(
+            "collected current metrics for tenant: {}: state={:?} resident_size={} remote_size={}",
+            tenant_id, tenant_state, tenant_resident_size, tenant_remote_size
+        );
+
+        current_metrics.push((
+            PageserverConsumptionMetricsKey {
+                tenant_id,
+                timeline_id: None,
+                metric: RESIDENT_SIZE,
+            },
+            tenant_resident_size,
+        ));
+
+        current_metrics.push((
+            PageserverConsumptionMetricsKey {
+                tenant_id,
+                timeline_id: None,
+                metric: REMOTE_STORAGE_SIZE,
+            },
+            tenant_remote_size,
+        ));
+
+        // Note that this metric is calculated in a separate bgworker
+        // Here we only use cached value, which may lag behind the real latest one
+        let tenant_synthetic_size = tenant.get_cached_synthetic_size();
+        current_metrics.push((
+            PageserverConsumptionMetricsKey {
+                tenant_id,
+                timeline_id: None,
+                metric: SYNTHETIC_STORAGE_SIZE,
+            },
+            tenant_synthetic_size,
+        ));
+    }
+
+    // Filter metrics
+    current_metrics.retain(|(curr_key, curr_val)| match cached_metrics.get(curr_key) {
+        Some(val) => val != curr_val,
+        None => true,
+    });
+
+    if current_metrics.is_empty() {
+        trace!("no new metrics to send");
+        return Ok(());
+    }
+
+    // Send metrics.
+    // Split into chunks of 1000 metrics to avoid exceeding the max request size
+    let chunks = current_metrics.chunks(CHUNK_SIZE);
+
+    let mut chunk_to_send: Vec<Event<Ids>> = Vec::with_capacity(CHUNK_SIZE);
+
+    for chunk in chunks {
+        chunk_to_send.clear();
+
+        // enrich metrics with type,timestamp and idempotency key before sending
+        chunk_to_send.extend(chunk.iter().map(|(curr_key, curr_val)| Event {
+            kind: EventType::Absolute { time: Utc::now() },
+            metric: curr_key.metric,
+            idempotency_key: idempotency_key(node_id.to_string()),
+            value: *curr_val,
+            extra: Ids {
+                tenant_id: curr_key.tenant_id,
+                timeline_id: curr_key.timeline_id,
+            },
+        }));
+
+        let chunk_json = serde_json::value::to_raw_value(&EventChunk {
+            events: &chunk_to_send,
+        })
+        .expect("PageserverConsumptionMetric should not fail serialization");
+
+        let res = client
+            .post(metric_collection_endpoint.clone())
+            .json(&chunk_json)
+            .send()
+            .await;
+
+        match res {
+            Ok(res) => {
+                if res.status().is_success() {
+                    // update cached metrics after they were sent successfully
+                    for (curr_key, curr_val) in chunk.iter() {
+                        cached_metrics.insert(curr_key.clone(), *curr_val);
+                    }
+                } else {
+                    error!("metrics endpoint refused the sent metrics: {:?}", res);
+                }
+            }
+            Err(err) => {
+                error!("failed to send metrics: {:?}", err);
+            }
+        }
+    }
+
+    Ok(())
+}
+
+/// Caclculate synthetic size for each active tenant
+pub async fn calculate_synthetic_size_worker(
+    synthetic_size_calculation_interval: Duration,
+) -> anyhow::Result<()> {
+    info!("starting calculate_synthetic_size_worker");
+
+    let mut ticker = tokio::time::interval(synthetic_size_calculation_interval);
+
+    loop {
+        tokio::select! {
+            _ = task_mgr::shutdown_watcher() => {
+                return Ok(());
+            },
+        _ = ticker.tick() => {
+
+                let tenants = mgr::list_tenants().await;
+                // iterate through list of Active tenants and collect metrics
+                for (tenant_id, tenant_state) in tenants {
+
+                    if tenant_state != TenantState::Active {
+                        continue;
+                    }
+
+                    if let Ok(tenant) = mgr::get_tenant(tenant_id, true).await
+                    {
+                        if let Err(e) = tenant.calculate_synthetic_size().await {
+                            error!("failed to calculate synthetic size for tenant {}: {}", tenant_id, e);
+                        }
+                    }
+
+                }
+            }
+        }
+    }
+}

pageserver/src/http/routes.rs

@@ -3,6 +3,7 @@ use std::sync::Arc;
 use anyhow::{anyhow, Context, Result};
 use hyper::StatusCode;
 use hyper::{Body, Request, Response, Uri};
+use pageserver_api::models::DownloadRemoteLayersTaskSpawnRequest;
 use remote_storage::GenericRemoteStorage;
 use tokio_util::sync::CancellationToken;
 use tracing::*;
@@ -13,7 +14,7 @@ use super::models::{
 };
 use crate::pgdatadir_mapping::LsnForTimestamp;
 use crate::tenant::config::TenantConfOpt;
-use crate::tenant::{with_ondemand_download, Timeline};
+use crate::tenant::{PageReconstructError, Timeline};
 use crate::{config::PageServerConf, tenant::mgr};
 use utils::{
     auth::JwtAuth,
@@ -77,6 +78,15 @@ fn check_permission(request: &Request<Body>, tenant_id: Option<TenantId>) -> Res
     })
 }
 
+fn apierror_from_prerror(err: PageReconstructError) -> ApiError {
+    match err {
+        PageReconstructError::Other(err) => ApiError::InternalServerError(err),
+        PageReconstructError::WalRedo(err) => {
+            ApiError::InternalServerError(anyhow::Error::new(err))
+        }
+    }
+}
+
 // Helper function to construct a TimelineInfo struct for a timeline
 async fn build_timeline_info(
     timeline: &Arc<Timeline>,
@@ -120,7 +130,7 @@ fn build_timeline_info_common(timeline: &Arc<Timeline>) -> anyhow::Result<Timeli
         lsn @ Lsn(_) => Some(lsn),
     };
     let current_logical_size = match timeline.get_current_logical_size() {
-        Ok(size) => Some(size),
+        Ok((size, _)) => Some(size),
         Err(err) => {
             error!("Timeline info creation failed to get current logical size: {err:?}");
             None
@@ -298,9 +308,10 @@ async fn get_lsn_by_timestamp_handler(request: Request<Body>) -> Result<Response
         .await
         .and_then(|tenant| tenant.get_timeline(timeline_id, true))
         .map_err(ApiError::NotFound)?;
-    let result = with_ondemand_download(|| timeline.find_lsn_for_timestamp(timestamp_pg))
+    let result = timeline
+        .find_lsn_for_timestamp(timestamp_pg)
         .await
-        .map_err(ApiError::InternalServerError)?;
+        .map_err(apierror_from_prerror)?;
 
     let result = match result {
         LsnForTimestamp::Present(lsn) => format!("{lsn}"),
@@ -585,7 +596,7 @@ async fn tenant_create_handler(mut request: Request<Body>) -> Result<Response<Bo
             // is Active when this function returns.
             if let res @ Err(_) = tenant.wait_to_become_active().await {
                 // This shouldn't happen because we just created the tenant directory
-                // in tenant_mgr::create_tenant, and there aren't any remote timelines
+                // in tenant::mgr::create_tenant, and there aren't any remote timelines
                 // to load, so, nothing can really fail during load.
                 // Don't do cleanup because we don't know how we got here.
                 // The tenant will likely be in `Broken` state and subsequent
@@ -738,17 +749,17 @@ async fn timeline_compact_handler(request: Request<Body>) -> Result<Response<Bod
     let timeline_id: TimelineId = parse_request_param(&request, "timeline_id")?;
     check_permission(&request, Some(tenant_id))?;
 
-    let tenant = mgr::get_tenant(tenant_id, true)
-        .await
-        .map_err(ApiError::NotFound)?;
-    let timeline = tenant
-        .get_timeline(timeline_id, true)
-        .map_err(ApiError::NotFound)?;
-    timeline
-        .compact()
+    let result_receiver = mgr::immediate_compact(tenant_id, timeline_id)
         .await
+        .context("spawn compaction task")
         .map_err(ApiError::InternalServerError)?;
 
+    let result: anyhow::Result<()> = result_receiver
+        .await
+        .context("receive compaction result")
+        .map_err(ApiError::InternalServerError)?;
+    result.map_err(ApiError::InternalServerError)?;
+
     json_response(StatusCode::OK, ())
 }
 
@@ -778,10 +789,11 @@ async fn timeline_checkpoint_handler(request: Request<Body>) -> Result<Response<
 }
 
 async fn timeline_download_remote_layers_handler_post(
-    request: Request<Body>,
+    mut request: Request<Body>,
 ) -> Result<Response<Body>, ApiError> {
     let tenant_id: TenantId = parse_request_param(&request, "tenant_id")?;
     let timeline_id: TimelineId = parse_request_param(&request, "timeline_id")?;
+    let body: DownloadRemoteLayersTaskSpawnRequest = json_request(&mut request).await?;
     check_permission(&request, Some(tenant_id))?;
 
     let tenant = mgr::get_tenant(tenant_id, true)
@@ -790,7 +802,7 @@ async fn timeline_download_remote_layers_handler_post(
     let timeline = tenant
         .get_timeline(timeline_id, true)
         .map_err(ApiError::NotFound)?;
-    match timeline.spawn_download_all_remote_layers().await {
+    match timeline.spawn_download_all_remote_layers(body).await {
         Ok(st) => json_response(StatusCode::ACCEPTED, st),
         Err(st) => json_response(StatusCode::CONFLICT, st),
     }
@@ -904,6 +916,7 @@ pub fn make_router(
             "/v1/tenant/:tenant_id/timeline/:timeline_id/download_remote_layers",
             timeline_download_remote_layers_handler_post,
         )
+        .post("/add_forced_now", handle_add_forced_now)
         .get(
             "/v1/tenant/:tenant_id/timeline/:timeline_id/download_remote_layers",
             timeline_download_remote_layers_handler_get,
@@ -914,3 +927,14 @@ pub fn make_router(
         )
         .any(handler_404))
 }
+
+async fn handle_add_forced_now(req: Request<Body>) -> Result<Response<Body>, ApiError> {
+    let now = get_query_param(&req, "now")?;
+
+    let now = chrono::DateTime::parse_from_rfc3339(&now).unwrap();
+    let now = now.with_timezone(&chrono::Utc);
+
+    crate::tenant::timeline::Timeline::force_next_now(now.into());
+
+    json_response(StatusCode::OK, ())
+}

pageserver/src/import_datadir.rs

@@ -2,12 +2,13 @@
 //! Import data and WAL from a PostgreSQL data directory and WAL segments into
 //! a neon Timeline.
 //!
-use std::fs::File;
-use std::io::{Read, Seek, SeekFrom};
 use std::path::{Path, PathBuf};
 
 use anyhow::{bail, ensure, Context, Result};
 use bytes::Bytes;
+use futures::StreamExt;
+use tokio::io::{AsyncRead, AsyncReadExt};
+use tokio_tar::Archive;
 use tracing::*;
 use walkdir::WalkDir;
 
@@ -42,7 +43,7 @@ pub fn get_lsn_from_controlfile(path: &Path) -> Result<Lsn> {
 /// This is currently only used to import a cluster freshly created by initdb.
 /// The code that deals with the checkpoint would not work right if the
 /// cluster was not shut down cleanly.
-pub fn import_timeline_from_postgres_datadir(
+pub async fn import_timeline_from_postgres_datadir(
     tline: &Timeline,
     pgdata_path: &Path,
     pgdata_lsn: Lsn,
@@ -65,9 +66,11 @@ pub fn import_timeline_from_postgres_datadir(
             let absolute_path = entry.path();
             let relative_path = absolute_path.strip_prefix(pgdata_path)?;
 
-            let file = File::open(absolute_path)?;
+            let mut file = tokio::fs::File::open(absolute_path).await?;
             let len = metadata.len() as usize;
-            if let Some(control_file) = import_file(&mut modification, relative_path, file, len)? {
+            if let Some(control_file) =
+                import_file(&mut modification, relative_path, &mut file, len).await?
+            {
                 pg_control = Some(control_file);
             }
             modification.flush()?;
@@ -96,18 +99,19 @@ pub fn import_timeline_from_postgres_datadir(
         tline,
         Lsn(pg_control.checkPointCopy.redo),
         pgdata_lsn,
-    )?;
+    )
+    .await?;
 
     Ok(())
 }
 
 // subroutine of import_timeline_from_postgres_datadir(), to load one relation file.
-fn import_rel<Reader: Read>(
-    modification: &mut DatadirModification,
+async fn import_rel(
+    modification: &mut DatadirModification<'_>,
     path: &Path,
     spcoid: Oid,
     dboid: Oid,
-    mut reader: Reader,
+    reader: &mut (impl AsyncRead + Send + Sync + Unpin),
     len: usize,
 ) -> anyhow::Result<()> {
     // Does it look like a relation file?
@@ -139,7 +143,11 @@ fn import_rel<Reader: Read>(
     // Call put_rel_creation for every segment of the relation,
     // because there is no guarantee about the order in which we are processing segments.
     // ignore "relation already exists" error
-    if let Err(e) = modification.put_rel_creation(rel, nblocks as u32) {
+    //
+    // FIXME: use proper error type for this, instead of parsing the error message.
+    // Or better yet, keep track of which relations we've already created
+    // https://github.com/neondatabase/neon/issues/3309
+    if let Err(e) = modification.put_rel_creation(rel, nblocks as u32).await {
         if e.to_string().contains("already exists") {
             debug!("relation {} already exists. we must be extending it", rel);
         } else {
@@ -148,7 +156,7 @@ fn import_rel<Reader: Read>(
     }
 
     loop {
-        let r = reader.read_exact(&mut buf);
+        let r = reader.read_exact(&mut buf).await;
         match r {
             Ok(_) => {
                 modification.put_rel_page_image(rel, blknum, Bytes::copy_from_slice(&buf))?;
@@ -174,18 +182,18 @@ fn import_rel<Reader: Read>(
     //
     // If we process rel segments out of order,
     // put_rel_extend will skip the update.
-    modification.put_rel_extend(rel, blknum)?;
+    modification.put_rel_extend(rel, blknum).await?;
 
     Ok(())
 }
 
 /// Import an SLRU segment file
 ///
-fn import_slru<Reader: Read>(
-    modification: &mut DatadirModification,
+async fn import_slru(
+    modification: &mut DatadirModification<'_>,
     slru: SlruKind,
     path: &Path,
-    mut reader: Reader,
+    reader: &mut (impl AsyncRead + Send + Sync + Unpin),
     len: usize,
 ) -> anyhow::Result<()> {
     info!("importing slru file {path:?}");
@@ -202,11 +210,13 @@ fn import_slru<Reader: Read>(
 
     ensure!(nblocks <= pg_constants::SLRU_PAGES_PER_SEGMENT as usize);
 
-    modification.put_slru_segment_creation(slru, segno, nblocks as u32)?;
+    modification
+        .put_slru_segment_creation(slru, segno, nblocks as u32)
+        .await?;
 
     let mut rpageno = 0;
     loop {
-        let r = reader.read_exact(&mut buf);
+        let r = reader.read_exact(&mut buf).await;
         match r {
             Ok(_) => {
                 modification.put_slru_page_image(
@@ -237,19 +247,20 @@ fn import_slru<Reader: Read>(
 
 /// Scan PostgreSQL WAL files in given directory and load all records between
 /// 'startpoint' and 'endpoint' into the repository.
-fn import_wal(
+async fn import_wal(
     walpath: &Path,
     tline: &Timeline,
     startpoint: Lsn,
     endpoint: Lsn,
 ) -> anyhow::Result<()> {
+    use std::io::Read;
     let mut waldecoder = WalStreamDecoder::new(startpoint, tline.pg_version);
 
     let mut segno = startpoint.segment_number(WAL_SEGMENT_SIZE);
     let mut offset = startpoint.segment_offset(WAL_SEGMENT_SIZE);
     let mut last_lsn = startpoint;
 
-    let mut walingest = WalIngest::new(tline, startpoint).no_ondemand_download()?;
+    let mut walingest = WalIngest::new(tline, startpoint).await?;
 
     while last_lsn <= endpoint {
         // FIXME: assume postgresql tli 1 for now
@@ -265,10 +276,11 @@ fn import_wal(
         }
 
         // Slurp the WAL file
-        let mut file = File::open(&path)?;
+        let mut file = std::fs::File::open(&path)?;
 
         if offset > 0 {
-            file.seek(SeekFrom::Start(offset as u64))?;
+            use std::io::Seek;
+            file.seek(std::io::SeekFrom::Start(offset as u64))?;
         }
 
         let nread = file.read_to_end(&mut buf)?;
@@ -286,7 +298,7 @@ fn import_wal(
             if let Some((lsn, recdata)) = waldecoder.poll_decode()? {
                 walingest
                     .ingest_record(recdata, lsn, &mut modification, &mut decoded)
-                    .no_ondemand_download()?;
+                    .await?;
                 last_lsn = lsn;
 
                 nrecords += 1;
@@ -310,9 +322,9 @@ fn import_wal(
     Ok(())
 }
 
-pub fn import_basebackup_from_tar<Reader: Read>(
+pub async fn import_basebackup_from_tar(
     tline: &Timeline,
-    reader: Reader,
+    reader: &mut (impl AsyncRead + Send + Sync + Unpin),
     base_lsn: Lsn,
 ) -> Result<()> {
     info!("importing base at {base_lsn}");
@@ -322,21 +334,24 @@ pub fn import_basebackup_from_tar<Reader: Read>(
     let mut pg_control: Option<ControlFileData> = None;
 
     // Import base
-    for base_tar_entry in tar::Archive::new(reader).entries()? {
-        let entry = base_tar_entry?;
+    let mut entries = Archive::new(reader).entries()?;
+    while let Some(base_tar_entry) = entries.next().await {
+        let mut entry = base_tar_entry?;
         let header = entry.header();
         let len = header.entry_size()? as usize;
         let file_path = header.path()?.into_owned();
 
         match header.entry_type() {
-            tar::EntryType::Regular => {
-                if let Some(res) = import_file(&mut modification, file_path.as_ref(), entry, len)? {
+            tokio_tar::EntryType::Regular => {
+                if let Some(res) =
+                    import_file(&mut modification, file_path.as_ref(), &mut entry, len).await?
+                {
                     // We found the pg_control file.
                     pg_control = Some(res);
                 }
                 modification.flush()?;
             }
-            tar::EntryType::Directory => {
+            tokio_tar::EntryType::Directory => {
                 debug!("directory {:?}", file_path);
             }
             _ => {
@@ -356,9 +371,9 @@ pub fn import_basebackup_from_tar<Reader: Read>(
     Ok(())
 }
 
-pub fn import_wal_from_tar<Reader: Read>(
+pub async fn import_wal_from_tar(
     tline: &Timeline,
-    reader: Reader,
+    reader: &mut (impl AsyncRead + Send + Sync + Unpin),
     start_lsn: Lsn,
     end_lsn: Lsn,
 ) -> Result<()> {
@@ -367,20 +382,23 @@ pub fn import_wal_from_tar<Reader: Read>(
     let mut segno = start_lsn.segment_number(WAL_SEGMENT_SIZE);
     let mut offset = start_lsn.segment_offset(WAL_SEGMENT_SIZE);
     let mut last_lsn = start_lsn;
-    let mut walingest = WalIngest::new(tline, start_lsn).no_ondemand_download()?;
+    let mut walingest = WalIngest::new(tline, start_lsn).await?;
 
     // Ingest wal until end_lsn
     info!("importing wal until {}", end_lsn);
-    let mut pg_wal_tar = tar::Archive::new(reader);
-    let mut pg_wal_entries_iter = pg_wal_tar.entries()?;
+    let mut pg_wal_tar = Archive::new(reader);
+    let mut pg_wal_entries = pg_wal_tar.entries()?;
     while last_lsn <= end_lsn {
         let bytes = {
-            let entry = pg_wal_entries_iter.next().expect("expected more wal")?;
+            let mut entry = pg_wal_entries
+                .next()
+                .await
+                .ok_or_else(|| anyhow::anyhow!("expected more wal"))??;
             let header = entry.header();
             let file_path = header.path()?.into_owned();
 
             match header.entry_type() {
-                tar::EntryType::Regular => {
+                tokio_tar::EntryType::Regular => {
                     // FIXME: assume postgresql tli 1 for now
                     let expected_filename = XLogFileName(1, segno, WAL_SEGMENT_SIZE);
                     let file_name = file_path
@@ -390,9 +408,9 @@ pub fn import_wal_from_tar<Reader: Read>(
                     ensure!(expected_filename == file_name);
 
                     debug!("processing wal file {:?}", file_path);
-                    read_all_bytes(entry)?
+                    read_all_bytes(&mut entry).await?
                 }
-                tar::EntryType::Directory => {
+                tokio_tar::EntryType::Directory => {
                     debug!("directory {:?}", file_path);
                     continue;
                 }
@@ -414,7 +432,7 @@ pub fn import_wal_from_tar<Reader: Read>(
             if let Some((lsn, recdata)) = waldecoder.poll_decode()? {
                 walingest
                     .ingest_record(recdata, lsn, &mut modification, &mut decoded)
-                    .no_ondemand_download()?;
+                    .await?;
                 last_lsn = lsn;
 
                 debug!("imported record at {} (end {})", lsn, end_lsn);
@@ -433,7 +451,7 @@ pub fn import_wal_from_tar<Reader: Read>(
     }
 
     // Log any extra unused files
-    for e in &mut pg_wal_entries_iter {
+    while let Some(e) = pg_wal_entries.next().await {
         let entry = e?;
         let header = entry.header();
         let file_path = header.path()?.into_owned();
@@ -443,10 +461,10 @@ pub fn import_wal_from_tar<Reader: Read>(
     Ok(())
 }
 
-fn import_file<Reader: Read>(
-    modification: &mut DatadirModification,
+async fn import_file(
+    modification: &mut DatadirModification<'_>,
     file_path: &Path,
-    reader: Reader,
+    reader: &mut (impl AsyncRead + Send + Sync + Unpin),
     len: usize,
 ) -> Result<Option<ControlFileData>> {
     let file_name = match file_path.file_name() {
@@ -466,7 +484,7 @@ fn import_file<Reader: Read>(
 
         match file_name.as_ref() {
             "pg_control" => {
-                let bytes = read_all_bytes(reader)?;
+                let bytes = read_all_bytes(reader).await?;
 
                 // Extract the checkpoint record and import it separately.
                 let pg_control = ControlFileData::decode(&bytes[..])?;
@@ -479,15 +497,15 @@ fn import_file<Reader: Read>(
                 return Ok(Some(pg_control));
             }
             "pg_filenode.map" => {
-                let bytes = read_all_bytes(reader)?;
-                modification.put_relmap_file(spcnode, dbnode, bytes)?;
+                let bytes = read_all_bytes(reader).await?;
+                modification.put_relmap_file(spcnode, dbnode, bytes).await?;
                 debug!("imported relmap file")
             }
             "PG_VERSION" => {
                 debug!("ignored PG_VERSION file");
             }
             _ => {
-                import_rel(modification, file_path, spcnode, dbnode, reader, len)?;
+                import_rel(modification, file_path, spcnode, dbnode, reader, len).await?;
                 debug!("imported rel creation");
             }
         }
@@ -502,44 +520,46 @@ fn import_file<Reader: Read>(
 
         match file_name.as_ref() {
             "pg_filenode.map" => {
-                let bytes = read_all_bytes(reader)?;
-                modification.put_relmap_file(spcnode, dbnode, bytes)?;
+                let bytes = read_all_bytes(reader).await?;
+                modification.put_relmap_file(spcnode, dbnode, bytes).await?;
                 debug!("imported relmap file")
             }
             "PG_VERSION" => {
                 debug!("ignored PG_VERSION file");
             }
             _ => {
-                import_rel(modification, file_path, spcnode, dbnode, reader, len)?;
+                import_rel(modification, file_path, spcnode, dbnode, reader, len).await?;
                 debug!("imported rel creation");
             }
         }
     } else if file_path.starts_with("pg_xact") {
         let slru = SlruKind::Clog;
 
-        import_slru(modification, slru, file_path, reader, len)?;
+        import_slru(modification, slru, file_path, reader, len).await?;
         debug!("imported clog slru");
     } else if file_path.starts_with("pg_multixact/offsets") {
         let slru = SlruKind::MultiXactOffsets;
 
-        import_slru(modification, slru, file_path, reader, len)?;
+        import_slru(modification, slru, file_path, reader, len).await?;
         debug!("imported multixact offsets slru");
     } else if file_path.starts_with("pg_multixact/members") {
         let slru = SlruKind::MultiXactMembers;
 
-        import_slru(modification, slru, file_path, reader, len)?;
+        import_slru(modification, slru, file_path, reader, len).await?;
         debug!("imported multixact members slru");
     } else if file_path.starts_with("pg_twophase") {
         let xid = u32::from_str_radix(file_name.as_ref(), 16)?;
 
-        let bytes = read_all_bytes(reader)?;
-        modification.put_twophase_file(xid, Bytes::copy_from_slice(&bytes[..]))?;
+        let bytes = read_all_bytes(reader).await?;
+        modification
+            .put_twophase_file(xid, Bytes::copy_from_slice(&bytes[..]))
+            .await?;
         debug!("imported twophase file");
     } else if file_path.starts_with("pg_wal") {
         debug!("found wal file in base section. ignore it");
     } else if file_path.starts_with("zenith.signal") {
         // Parse zenith signal file to set correct previous LSN
-        let bytes = read_all_bytes(reader)?;
+        let bytes = read_all_bytes(reader).await?;
         // zenith.signal format is "PREV LSN: prev_lsn"
         // TODO write serialization and deserialization in the same place.
         let zenith_signal = std::str::from_utf8(&bytes)?.trim();
@@ -576,8 +596,8 @@ fn import_file<Reader: Read>(
     Ok(None)
 }
 
-fn read_all_bytes<Reader: Read>(mut reader: Reader) -> Result<Bytes> {
+async fn read_all_bytes(reader: &mut (impl AsyncRead + Send + Sync + Unpin)) -> Result<Bytes> {
     let mut buf: Vec<u8> = vec![];
-    reader.read_to_end(&mut buf)?;
+    reader.read_to_end(&mut buf).await?;
     Ok(Bytes::copy_from_slice(&buf[..]))
 }

pageserver/src/lib.rs

@@ -1,7 +1,7 @@
 mod auth;
 pub mod basebackup;
-pub mod billing_metrics;
 pub mod config;
+pub mod consumption_metrics;
 pub mod http;
 pub mod import_datadir;
 pub mod keyspace;
@@ -9,7 +9,6 @@ pub(crate) mod metrics;
 pub mod page_cache;
 pub mod page_service;
 pub mod pgdatadir_mapping;
-pub mod profiling;
 pub mod repository;
 pub mod task_mgr;
 pub mod tenant;

pageserver/src/metrics.rs

@@ -209,15 +209,34 @@ pub static NUM_ONDISK_LAYERS: Lazy<IntGauge> = Lazy::new(|| {
 
 // remote storage metrics
 
-static REMOTE_UPLOAD_QUEUE_UNFINISHED_TASKS: Lazy<IntGaugeVec> = Lazy::new(|| {
+/// NB: increment _after_ recording the current value into [`REMOTE_TIMELINE_CLIENT_CALLS_STARTED_HIST`].
+static REMOTE_TIMELINE_CLIENT_CALLS_UNFINISHED_GAUGE: Lazy<IntGaugeVec> = Lazy::new(|| {
     register_int_gauge_vec!(
-        "pageserver_remote_upload_queue_unfinished_tasks",
-        "Number of tasks in the upload queue that are not finished yet.",
+        "pageserver_remote_timeline_client_calls_unfinished",
+        "Number of ongoing calls to remote timeline client. \
+         Used to populate pageserver_remote_timeline_client_calls_started. \
+         This metric is not useful for sampling from Prometheus, but useful in tests.",
         &["tenant_id", "timeline_id", "file_kind", "op_kind"],
     )
     .expect("failed to define a metric")
 });
 
+static REMOTE_TIMELINE_CLIENT_CALLS_STARTED_HIST: Lazy<HistogramVec> = Lazy::new(|| {
+    register_histogram_vec!(
+        "pageserver_remote_timeline_client_calls_started",
+        "When calling a remote timeline client method, we record the current value \
+         of the calls_unfinished gauge in this histogram. Plot the histogram \
+         over time in a heatmap to visualize how many operations were ongoing \
+         at a given instant. It gives you a better idea of the queue depth \
+         than plotting the gauge directly, since operations may complete faster \
+         than the sampling interval.",
+        &["tenant_id", "timeline_id", "file_kind", "op_kind"],
+        // The calls_unfinished gauge is an integer gauge, hence we have integer buckets.
+        vec![0.0, 1.0, 2.0, 4.0, 6.0, 8.0, 10.0, 15.0, 20.0, 40.0, 60.0, 80.0, 100.0, 500.0],
+    )
+    .expect("failed to define a metric")
+});
+
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
 pub enum RemoteOpKind {
     Upload,
@@ -248,15 +267,12 @@ impl RemoteOpFileKind {
     }
 }
 
-pub static REMOTE_OPERATION_KINDS: &[&str] = &["upload", "download", "delete"];
-pub static REMOTE_OPERATION_FILE_KINDS: &[&str] = &["layer", "index"];
-pub static REMOTE_OPERATION_STATUSES: &[&str] = &["success", "failure"];
-
 pub static REMOTE_OPERATION_TIME: Lazy<HistogramVec> = Lazy::new(|| {
     register_histogram_vec!(
         "pageserver_remote_operation_seconds",
         "Time spent on remote storage operations. \
-        Grouped by tenant, timeline, operation_kind and status",
+        Grouped by tenant, timeline, operation_kind and status. \
+        Does not account for time spent waiting in remote timeline client's queues.",
         &["tenant_id", "timeline_id", "file_kind", "op_kind", "status"]
     )
     .expect("failed to define a metric")
@@ -475,21 +491,6 @@ impl Drop for TimelineMetrics {
         for op in SMGR_QUERY_TIME_OPERATIONS {
             let _ = SMGR_QUERY_TIME.remove_label_values(&[op, tenant_id, timeline_id]);
         }
-
-        let _ = REMOTE_UPLOAD_QUEUE_UNFINISHED_TASKS.remove_label_values(&[tenant_id, timeline_id]);
-        for file_kind in REMOTE_OPERATION_FILE_KINDS {
-            for op in REMOTE_OPERATION_KINDS {
-                for status in REMOTE_OPERATION_STATUSES {
-                    let _ = REMOTE_OPERATION_TIME.remove_label_values(&[
-                        tenant_id,
-                        timeline_id,
-                        file_kind,
-                        op,
-                        status,
-                    ]);
-                }
-            }
-        }
     }
 }
 
@@ -510,7 +511,8 @@ pub struct RemoteTimelineClientMetrics {
     timeline_id: String,
     remote_physical_size_gauge: Mutex<Option<UIntGauge>>,
     remote_operation_time: Mutex<HashMap<(&'static str, &'static str, &'static str), Histogram>>,
-    unfinished_tasks: Mutex<HashMap<(&'static str, &'static str), IntGauge>>,
+    calls_unfinished_gauge: Mutex<HashMap<(&'static str, &'static str), IntGauge>>,
+    calls_started_hist: Mutex<HashMap<(&'static str, &'static str), Histogram>>,
 }
 
 impl RemoteTimelineClientMetrics {
@@ -519,7 +521,8 @@ impl RemoteTimelineClientMetrics {
             tenant_id: tenant_id.to_string(),
             timeline_id: timeline_id.to_string(),
             remote_operation_time: Mutex::new(HashMap::default()),
-            unfinished_tasks: Mutex::new(HashMap::default()),
+            calls_unfinished_gauge: Mutex::new(HashMap::default()),
+            calls_started_hist: Mutex::new(HashMap::default()),
             remote_physical_size_gauge: Mutex::new(None),
         }
     }
@@ -558,16 +561,37 @@ impl RemoteTimelineClientMetrics {
         });
         metric.clone()
     }
-    pub fn unfinished_tasks(
+    fn calls_unfinished_gauge(
         &self,
         file_kind: &RemoteOpFileKind,
         op_kind: &RemoteOpKind,
     ) -> IntGauge {
         // XXX would be nice to have an upgradable RwLock
-        let mut guard = self.unfinished_tasks.lock().unwrap();
+        let mut guard = self.calls_unfinished_gauge.lock().unwrap();
         let key = (file_kind.as_str(), op_kind.as_str());
         let metric = guard.entry(key).or_insert_with(move || {
-            REMOTE_UPLOAD_QUEUE_UNFINISHED_TASKS
+            REMOTE_TIMELINE_CLIENT_CALLS_UNFINISHED_GAUGE
+                .get_metric_with_label_values(&[
+                    &self.tenant_id.to_string(),
+                    &self.timeline_id.to_string(),
+                    key.0,
+                    key.1,
+                ])
+                .unwrap()
+        });
+        metric.clone()
+    }
+
+    fn calls_started_hist(
+        &self,
+        file_kind: &RemoteOpFileKind,
+        op_kind: &RemoteOpKind,
+    ) -> Histogram {
+        // XXX would be nice to have an upgradable RwLock
+        let mut guard = self.calls_started_hist.lock().unwrap();
+        let key = (file_kind.as_str(), op_kind.as_str());
+        let metric = guard.entry(key).or_insert_with(move || {
+            REMOTE_TIMELINE_CLIENT_CALLS_STARTED_HIST
                 .get_metric_with_label_values(&[
                     &self.tenant_id.to_string(),
                     &self.timeline_id.to_string(),
@@ -580,6 +604,58 @@ impl RemoteTimelineClientMetrics {
     }
 }
 
+/// See [`RemoteTimelineClientMetrics::call_begin`].
+#[must_use]
+pub(crate) struct RemoteTimelineClientCallMetricGuard(Option<IntGauge>);
+
+impl RemoteTimelineClientCallMetricGuard {
+    /// Consume this guard object without decrementing the metric.
+    /// The caller vouches to do this manually, so that the prior increment of the gauge will cancel out.
+    pub fn will_decrement_manually(mut self) {
+        self.0 = None; // prevent drop() from decrementing
+    }
+}
+
+impl Drop for RemoteTimelineClientCallMetricGuard {
+    fn drop(&mut self) {
+        if let RemoteTimelineClientCallMetricGuard(Some(guard)) = self {
+            guard.dec();
+        }
+    }
+}
+
+impl RemoteTimelineClientMetrics {
+    /// Increment the metrics that track ongoing calls to the remote timeline client instance.
+    ///
+    /// Drop the returned guard object once the operation is finished to decrement the values.
+    /// Or, use [`RemoteTimelineClientCallMetricGuard::will_decrement_manually`] and [`call_end`] if that
+    /// is more suitable.
+    /// Never do both.
+    pub(crate) fn call_begin(
+        &self,
+        file_kind: &RemoteOpFileKind,
+        op_kind: &RemoteOpKind,
+    ) -> RemoteTimelineClientCallMetricGuard {
+        let unfinished_metric = self.calls_unfinished_gauge(file_kind, op_kind);
+        self.calls_started_hist(file_kind, op_kind)
+            .observe(unfinished_metric.get() as f64);
+        unfinished_metric.inc();
+        RemoteTimelineClientCallMetricGuard(Some(unfinished_metric))
+    }
+
+    /// Manually decrement the metric instead of using the guard object.
+    /// Using the guard object is generally preferable.
+    /// See [`call_begin`] for more context.
+    pub(crate) fn call_end(&self, file_kind: &RemoteOpFileKind, op_kind: &RemoteOpKind) {
+        let unfinished_metric = self.calls_unfinished_gauge(file_kind, op_kind);
+        debug_assert!(
+            unfinished_metric.get() > 0,
+            "begin and end should cancel out"
+        );
+        unfinished_metric.dec();
+    }
+}
+
 impl Drop for RemoteTimelineClientMetrics {
     fn drop(&mut self) {
         let RemoteTimelineClientMetrics {
@@ -587,13 +663,22 @@ impl Drop for RemoteTimelineClientMetrics {
             timeline_id,
             remote_physical_size_gauge,
             remote_operation_time,
-            unfinished_tasks,
+            calls_unfinished_gauge,
+            calls_started_hist,
         } = self;
         for ((a, b, c), _) in remote_operation_time.get_mut().unwrap().drain() {
             let _ = REMOTE_OPERATION_TIME.remove_label_values(&[tenant_id, timeline_id, a, b, c]);
         }
-        for ((a, b), _) in unfinished_tasks.get_mut().unwrap().drain() {
-            let _ = REMOTE_UPLOAD_QUEUE_UNFINISHED_TASKS.remove_label_values(&[
+        for ((a, b), _) in calls_unfinished_gauge.get_mut().unwrap().drain() {
+            let _ = REMOTE_TIMELINE_CLIENT_CALLS_UNFINISHED_GAUGE.remove_label_values(&[
+                tenant_id,
+                timeline_id,
+                a,
+                b,
+            ]);
+        }
+        for ((a, b), _) in calls_started_hist.get_mut().unwrap().drain() {
+            let _ = REMOTE_TIMELINE_CLIENT_CALLS_STARTED_HIST.remove_label_values(&[
                 tenant_id,
                 timeline_id,
                 a,

pageserver/src/page_service.rs

@@ -9,7 +9,7 @@
 //  custom protocol.
 //
 
-use anyhow::{bail, ensure, Context, Result};
+use anyhow::Context;
 use bytes::Buf;
 use bytes::Bytes;
 use futures::{Stream, StreamExt};
@@ -19,6 +19,8 @@ use pageserver_api::models::{
     PagestreamFeMessage, PagestreamGetPageRequest, PagestreamGetPageResponse,
     PagestreamNblocksRequest, PagestreamNblocksResponse,
 };
+use pq_proto::ConnectionError;
+use pq_proto::FeStartupPacket;
 use pq_proto::{BeMessage, FeMessage, RowDescriptor};
 use std::io;
 use std::net::TcpListener;
@@ -26,11 +28,9 @@ use std::str;
 use std::str::FromStr;
 use std::sync::Arc;
 use std::time::Duration;
-use tokio::pin;
-use tokio_util::io::StreamReader;
-use tokio_util::io::SyncIoBridge;
 use tracing::*;
 use utils::id::ConnectionId;
+use utils::postgres_backend_async::QueryError;
 use utils::{
     auth::{Claims, JwtAuth, Scope},
     id::{TenantId, TimelineId},
@@ -42,10 +42,9 @@ use utils::{
 
 use crate::auth::check_permission;
 use crate::basebackup;
-use crate::config::{PageServerConf, ProfilingConfig};
+use crate::config::PageServerConf;
 use crate::import_datadir::import_wal_from_tar;
 use crate::metrics::{LIVE_CONNECTIONS_COUNT, SMGR_QUERY_TIME};
-use crate::profiling::profpoint_start;
 use crate::task_mgr;
 use crate::task_mgr::TaskKind;
 use crate::tenant::mgr;
@@ -64,8 +63,8 @@ fn copyin_stream(pgb: &mut PostgresBackend) -> impl Stream<Item = io::Result<Byt
                 _ = task_mgr::shutdown_watcher() => {
                     // We were requested to shut down.
                     let msg = format!("pageserver is shutting down");
-                    let _ = pgb.write_message(&BeMessage::ErrorResponse(&msg));
-                    Err(anyhow::anyhow!(msg))
+                    let _ = pgb.write_message(&BeMessage::ErrorResponse(&msg, None));
+                    Err(QueryError::Other(anyhow::anyhow!(msg)))
                 }
 
                 msg = pgb.read_message() => { msg }
@@ -78,14 +77,15 @@ fn copyin_stream(pgb: &mut PostgresBackend) -> impl Stream<Item = io::Result<Byt
                         FeMessage::CopyDone => { break },
                         FeMessage::Sync => continue,
                         FeMessage::Terminate => {
-                            let msg = format!("client terminated connection with Terminate message during COPY");
-                            pgb.write_message(&BeMessage::ErrorResponse(&msg))?;
+                            let msg = "client terminated connection with Terminate message during COPY";
+                            let query_error_error = QueryError::Disconnected(ConnectionError::Socket(io::Error::new(io::ErrorKind::ConnectionReset, msg)));
+                            pgb.write_message(&BeMessage::ErrorResponse(msg, Some(query_error_error.pg_error_code())))?;
                             Err(io::Error::new(io::ErrorKind::ConnectionReset, msg))?;
                             break;
                         }
                         m => {
-                            let msg = format!("unexpected message {:?}", m);
-                            pgb.write_message(&BeMessage::ErrorResponse(&msg))?;
+                            let msg = format!("unexpected message {m:?}");
+                            pgb.write_message(&BeMessage::ErrorResponse(&msg, None))?;
                             Err(io::Error::new(io::ErrorKind::Other, msg))?;
                             break;
                         }
@@ -95,12 +95,16 @@ fn copyin_stream(pgb: &mut PostgresBackend) -> impl Stream<Item = io::Result<Byt
                 }
                 Ok(None) => {
                     let msg = "client closed connection during COPY";
-                    pgb.write_message(&BeMessage::ErrorResponse(msg))?;
+                    let query_error_error = QueryError::Disconnected(ConnectionError::Socket(io::Error::new(io::ErrorKind::ConnectionReset, msg)));
+                    pgb.write_message(&BeMessage::ErrorResponse(msg, Some(query_error_error.pg_error_code())))?;
                     pgb.flush().await?;
                     Err(io::Error::new(io::ErrorKind::ConnectionReset, msg))?;
                 }
-                Err(e) => {
-                    Err(io::Error::new(io::ErrorKind::Other, e))?;
+                Err(QueryError::Disconnected(ConnectionError::Socket(io_error))) => {
+                    Err(io_error)?;
+                }
+                Err(other) => {
+                    Err(io::Error::new(io::ErrorKind::Other, other))?;
                 }
             };
         }
@@ -198,23 +202,19 @@ async fn page_service_conn_main(
             // we've been requested to shut down
             Ok(())
         }
-        Err(err) => {
-            let root_cause_io_err_kind = err
-                .root_cause()
-                .downcast_ref::<io::Error>()
-                .map(|e| e.kind());
-
+        Err(QueryError::Disconnected(ConnectionError::Socket(io_error))) => {
             // `ConnectionReset` error happens when the Postgres client closes the connection.
             // As this disconnection happens quite often and is expected,
             // we decided to downgrade the logging level to `INFO`.
             // See: https://github.com/neondatabase/neon/issues/1683.
-            if root_cause_io_err_kind == Some(io::ErrorKind::ConnectionReset) {
+            if io_error.kind() == io::ErrorKind::ConnectionReset {
                 info!("Postgres client disconnected");
                 Ok(())
             } else {
-                Err(err)
+                Err(io_error).context("Postgres connection error")
             }
         }
+        other => other.context("Postgres query error"),
     }
 }
 
@@ -251,9 +251,8 @@ impl PageRequestMetrics {
     }
 }
 
-#[derive(Debug)]
 struct PageServerHandler {
-    conf: &'static PageServerConf,
+    _conf: &'static PageServerConf,
     auth: Option<Arc<JwtAuth>>,
     claims: Option<Claims>,
 }
@@ -261,7 +260,7 @@ struct PageServerHandler {
 impl PageServerHandler {
     pub fn new(conf: &'static PageServerConf, auth: Option<Arc<JwtAuth>>) -> Self {
         PageServerHandler {
-            conf,
+            _conf: conf,
             auth,
             claims: None,
         }
@@ -316,7 +315,7 @@ impl PageServerHandler {
                 Some(FeMessage::CopyData(bytes)) => bytes,
                 Some(FeMessage::Terminate) => break,
                 Some(m) => {
-                    bail!("unexpected message: {m:?} during COPY");
+                    anyhow::bail!("unexpected message: {m:?} during COPY");
                 }
                 None => break, // client disconnected
             };
@@ -373,7 +372,7 @@ impl PageServerHandler {
         base_lsn: Lsn,
         _end_lsn: Lsn,
         pg_version: u32,
-    ) -> anyhow::Result<()> {
+    ) -> Result<(), QueryError> {
         task_mgr::associate_with(Some(tenant_id), Some(timeline_id));
         // Create empty timeline
         info!("creating new timeline");
@@ -395,9 +394,7 @@ impl PageServerHandler {
         pgb.write_message(&BeMessage::CopyInResponse)?;
         pgb.flush().await?;
 
-        let copyin_stream = copyin_stream(pgb);
-        pin!(copyin_stream);
-
+        let mut copyin_stream = Box::pin(copyin_stream(pgb));
         timeline
             .import_basebackup_from_tar(&mut copyin_stream, base_lsn)
             .await?;
@@ -429,11 +426,16 @@ impl PageServerHandler {
         timeline_id: TimelineId,
         start_lsn: Lsn,
         end_lsn: Lsn,
-    ) -> anyhow::Result<()> {
+    ) -> Result<(), QueryError> {
         task_mgr::associate_with(Some(tenant_id), Some(timeline_id));
 
         let timeline = get_active_timeline_with_timeout(tenant_id, timeline_id).await?;
-        ensure!(timeline.get_last_record_lsn() == start_lsn);
+        let last_record_lsn = timeline.get_last_record_lsn();
+        if last_record_lsn != start_lsn {
+            return Err(QueryError::Other(
+                anyhow::anyhow!("Cannot import WAL from Lsn {start_lsn} because timeline does not start from the same lsn: {last_record_lsn}"))
+            );
+        }
 
         // TODO leave clean state on error. For now you can use detach to clean
         // up broken state from a failed import.
@@ -443,8 +445,8 @@ impl PageServerHandler {
         pgb.write_message(&BeMessage::CopyInResponse)?;
         pgb.flush().await?;
         let mut copyin_stream = Box::pin(copyin_stream(pgb));
-        let reader = SyncIoBridge::new(StreamReader::new(&mut copyin_stream));
-        tokio::task::block_in_place(|| import_wal_from_tar(&timeline, reader, start_lsn, end_lsn))?;
+        let mut reader = tokio_util::io::StreamReader::new(&mut copyin_stream);
+        import_wal_from_tar(&timeline, &mut reader, start_lsn, end_lsn).await?;
         info!("wal import complete");
 
         // Drain the rest of the Copy data
@@ -457,7 +459,11 @@ impl PageServerHandler {
         }
 
         // TODO Does it make sense to overshoot?
-        ensure!(timeline.get_last_record_lsn() >= end_lsn);
+        if timeline.get_last_record_lsn() < end_lsn {
+            return Err(QueryError::Other(
+                anyhow::anyhow!("Cannot import WAL from Lsn {start_lsn} because timeline does not start from the same lsn: {last_record_lsn}"))
+            );
+        }
 
         // Flush data to disk, then upload to s3. No need for a forced checkpoint.
         // We only want to persist the data, and it doesn't matter if it's in the
@@ -486,7 +492,7 @@ impl PageServerHandler {
         mut lsn: Lsn,
         latest: bool,
         latest_gc_cutoff_lsn: &RcuReadGuard<Lsn>,
-    ) -> Result<Lsn> {
+    ) -> anyhow::Result<Lsn> {
         if latest {
             // Latest page version was requested. If LSN is given, it is a hint
             // to the page server that there have been no modifications to the
@@ -517,11 +523,11 @@ impl PageServerHandler {
             }
         } else {
             if lsn == Lsn(0) {
-                bail!("invalid LSN(0) in request");
+                anyhow::bail!("invalid LSN(0) in request");
             }
             timeline.wait_lsn(lsn).await?;
         }
-        ensure!(
+        anyhow::ensure!(
             lsn >= **latest_gc_cutoff_lsn,
             "tried to request a page version that was garbage collected. requested at {} gc cutoff {}",
             lsn, **latest_gc_cutoff_lsn
@@ -534,15 +540,12 @@ impl PageServerHandler {
         &self,
         timeline: &Timeline,
         req: &PagestreamExistsRequest,
-    ) -> Result<PagestreamBeMessage> {
+    ) -> anyhow::Result<PagestreamBeMessage> {
         let latest_gc_cutoff_lsn = timeline.get_latest_gc_cutoff_lsn();
         let lsn = Self::wait_or_get_last_lsn(timeline, req.lsn, req.latest, &latest_gc_cutoff_lsn)
             .await?;
 
-        let exists = crate::tenant::with_ondemand_download(|| {
-            timeline.get_rel_exists(req.rel, lsn, req.latest)
-        })
-        .await?;
+        let exists = timeline.get_rel_exists(req.rel, lsn, req.latest).await?;
 
         Ok(PagestreamBeMessage::Exists(PagestreamExistsResponse {
             exists,
@@ -554,15 +557,12 @@ impl PageServerHandler {
         &self,
         timeline: &Timeline,
         req: &PagestreamNblocksRequest,
-    ) -> Result<PagestreamBeMessage> {
+    ) -> anyhow::Result<PagestreamBeMessage> {
         let latest_gc_cutoff_lsn = timeline.get_latest_gc_cutoff_lsn();
         let lsn = Self::wait_or_get_last_lsn(timeline, req.lsn, req.latest, &latest_gc_cutoff_lsn)
             .await?;
 
-        let n_blocks = crate::tenant::with_ondemand_download(|| {
-            timeline.get_rel_size(req.rel, lsn, req.latest)
-        })
-        .await?;
+        let n_blocks = timeline.get_rel_size(req.rel, lsn, req.latest).await?;
 
         Ok(PagestreamBeMessage::Nblocks(PagestreamNblocksResponse {
             n_blocks,
@@ -574,15 +574,14 @@ impl PageServerHandler {
         &self,
         timeline: &Timeline,
         req: &PagestreamDbSizeRequest,
-    ) -> Result<PagestreamBeMessage> {
+    ) -> anyhow::Result<PagestreamBeMessage> {
         let latest_gc_cutoff_lsn = timeline.get_latest_gc_cutoff_lsn();
         let lsn = Self::wait_or_get_last_lsn(timeline, req.lsn, req.latest, &latest_gc_cutoff_lsn)
             .await?;
 
-        let total_blocks = crate::tenant::with_ondemand_download(|| {
-            timeline.get_db_size(DEFAULTTABLESPACE_OID, req.dbnode, lsn, req.latest)
-        })
-        .await?;
+        let total_blocks = timeline
+            .get_db_size(DEFAULTTABLESPACE_OID, req.dbnode, lsn, req.latest)
+            .await?;
         let db_size = total_blocks as i64 * BLCKSZ as i64;
 
         Ok(PagestreamBeMessage::DbSize(PagestreamDbSizeResponse {
@@ -595,7 +594,7 @@ impl PageServerHandler {
         &self,
         timeline: &Timeline,
         req: &PagestreamGetPageRequest,
-    ) -> Result<PagestreamBeMessage> {
+    ) -> anyhow::Result<PagestreamBeMessage> {
         let latest_gc_cutoff_lsn = timeline.get_latest_gc_cutoff_lsn();
         let lsn = Self::wait_or_get_last_lsn(timeline, req.lsn, req.latest, &latest_gc_cutoff_lsn)
             .await?;
@@ -608,14 +607,9 @@ impl PageServerHandler {
         }
         */
 
-        let page = crate::tenant::with_ondemand_download(|| {
-            // FIXME: this profiling now happens at different place than it used to. The
-            // current profiling is based on a thread-local variable, so it doesn't work
-            // across awaits
-            let _profiling_guard = profpoint_start(self.conf, ProfilingConfig::PageRequests);
-            timeline.get_rel_page_at_lsn(req.rel, req.blkno, lsn, req.latest)
-        })
-        .await?;
+        let page = timeline
+            .get_rel_page_at_lsn(req.rel, req.blkno, lsn, req.latest)
+            .await?;
 
         Ok(PagestreamBeMessage::GetPage(PagestreamGetPageResponse {
             page,
@@ -648,17 +642,13 @@ impl PageServerHandler {
         pgb.write_message(&BeMessage::CopyOutResponse)?;
         pgb.flush().await?;
 
-        /* Send a tarball of the latest layer on the timeline */
-        let mut writer = CopyDataSink {
-            pgb,
-            rt: tokio::runtime::Handle::current(),
-        };
-        tokio::task::block_in_place(|| {
-            let basebackup =
-                basebackup::Basebackup::new(&mut writer, &timeline, lsn, prev_lsn, full_backup)?;
-            tracing::Span::current().record("lsn", basebackup.lsn.to_string().as_str());
-            basebackup.send_tarball()
-        })?;
+        // Send a tarball of the latest layer on the timeline
+        {
+            let mut writer = pgb.copyout_writer();
+            basebackup::send_basebackup_tarball(&mut writer, &timeline, lsn, prev_lsn, full_backup)
+                .await?;
+        }
+
         pgb.write_message(&BeMessage::CopyDone)?;
         pgb.flush().await?;
         info!("basebackup complete");
@@ -668,7 +658,7 @@ impl PageServerHandler {
 
     // when accessing management api supply None as an argument
     // when using to authorize tenant pass corresponding tenant id
-    fn check_permission(&self, tenant_id: Option<TenantId>) -> Result<()> {
+    fn check_permission(&self, tenant_id: Option<TenantId>) -> anyhow::Result<()> {
         if self.auth.is_none() {
             // auth is set to Trust, nothing to check so just return ok
             return Ok(());
@@ -690,20 +680,19 @@ impl postgres_backend_async::Handler for PageServerHandler {
         &mut self,
         _pgb: &mut PostgresBackend,
         jwt_response: &[u8],
-    ) -> anyhow::Result<()> {
+    ) -> Result<(), QueryError> {
         // this unwrap is never triggered, because check_auth_jwt only called when auth_type is NeonJWT
         // which requires auth to be present
         let data = self
             .auth
             .as_ref()
             .unwrap()
-            .decode(str::from_utf8(jwt_response)?)?;
+            .decode(str::from_utf8(jwt_response).context("jwt response is not UTF-8")?)?;
 
-        if matches!(data.claims.scope, Scope::Tenant) {
-            ensure!(
-                data.claims.tenant_id.is_some(),
+        if matches!(data.claims.scope, Scope::Tenant) && data.claims.tenant_id.is_none() {
+            return Err(QueryError::Other(anyhow::anyhow!(
                 "jwt token scope is Tenant, but tenant id is missing"
-            )
+            )));
         }
 
         info!(
@@ -715,22 +704,33 @@ impl postgres_backend_async::Handler for PageServerHandler {
         Ok(())
     }
 
+    fn startup(
+        &mut self,
+        _pgb: &mut PostgresBackend,
+        _sm: &FeStartupPacket,
+    ) -> Result<(), QueryError> {
+        Ok(())
+    }
+
     async fn process_query(
         &mut self,
         pgb: &mut PostgresBackend,
         query_string: &str,
-    ) -> anyhow::Result<()> {
-        debug!("process query {:?}", query_string);
+    ) -> Result<(), QueryError> {
+        debug!("process query {query_string:?}");
 
         if query_string.starts_with("pagestream ") {
             let (_, params_raw) = query_string.split_at("pagestream ".len());
             let params = params_raw.split(' ').collect::<Vec<_>>();
-            ensure!(
-                params.len() == 2,
-                "invalid param number for pagestream command"
-            );
-            let tenant_id = TenantId::from_str(params[0])?;
-            let timeline_id = TimelineId::from_str(params[1])?;
+            if params.len() != 2 {
+                return Err(QueryError::Other(anyhow::anyhow!(
+                    "invalid param number for pagestream command"
+                )));
+            }
+            let tenant_id = TenantId::from_str(params[0])
+                .with_context(|| format!("Failed to parse tenant id from {}", params[0]))?;
+            let timeline_id = TimelineId::from_str(params[1])
+                .with_context(|| format!("Failed to parse timeline id from {}", params[1]))?;
 
             self.check_permission(Some(tenant_id))?;
 
@@ -740,18 +740,24 @@ impl postgres_backend_async::Handler for PageServerHandler {
             let (_, params_raw) = query_string.split_at("basebackup ".len());
             let params = params_raw.split_whitespace().collect::<Vec<_>>();
 
-            ensure!(
-                params.len() >= 2,
-                "invalid param number for basebackup command"
-            );
+            if params.len() < 2 {
+                return Err(QueryError::Other(anyhow::anyhow!(
+                    "invalid param number for basebackup command"
+                )));
+            }
 
-            let tenant_id = TenantId::from_str(params[0])?;
-            let timeline_id = TimelineId::from_str(params[1])?;
+            let tenant_id = TenantId::from_str(params[0])
+                .with_context(|| format!("Failed to parse tenant id from {}", params[0]))?;
+            let timeline_id = TimelineId::from_str(params[1])
+                .with_context(|| format!("Failed to parse timeline id from {}", params[1]))?;
 
             self.check_permission(Some(tenant_id))?;
 
             let lsn = if params.len() == 3 {
-                Some(Lsn::from_str(params[2])?)
+                Some(
+                    Lsn::from_str(params[2])
+                        .with_context(|| format!("Failed to parse Lsn from {}", params[2]))?,
+                )
             } else {
                 None
             };
@@ -766,13 +772,16 @@ impl postgres_backend_async::Handler for PageServerHandler {
             let (_, params_raw) = query_string.split_at("get_last_record_rlsn ".len());
             let params = params_raw.split_whitespace().collect::<Vec<_>>();
 
-            ensure!(
-                params.len() == 2,
-                "invalid param number for get_last_record_rlsn command"
-            );
+            if params.len() != 2 {
+                return Err(QueryError::Other(anyhow::anyhow!(
+                    "invalid param number for get_last_record_rlsn command"
+                )));
+            }
 
-            let tenant_id = TenantId::from_str(params[0])?;
-            let timeline_id = TimelineId::from_str(params[1])?;
+            let tenant_id = TenantId::from_str(params[0])
+                .with_context(|| format!("Failed to parse tenant id from {}", params[0]))?;
+            let timeline_id = TimelineId::from_str(params[1])
+                .with_context(|| format!("Failed to parse timeline id from {}", params[1]))?;
 
             self.check_permission(Some(tenant_id))?;
             let timeline = get_active_timeline_with_timeout(tenant_id, timeline_id).await?;
@@ -794,22 +803,31 @@ impl postgres_backend_async::Handler for PageServerHandler {
             let (_, params_raw) = query_string.split_at("fullbackup ".len());
             let params = params_raw.split_whitespace().collect::<Vec<_>>();
 
-            ensure!(
-                params.len() >= 2,
-                "invalid param number for fullbackup command"
-            );
+            if params.len() < 2 {
+                return Err(QueryError::Other(anyhow::anyhow!(
+                    "invalid param number for fullbackup command"
+                )));
+            }
 
-            let tenant_id = TenantId::from_str(params[0])?;
-            let timeline_id = TimelineId::from_str(params[1])?;
+            let tenant_id = TenantId::from_str(params[0])
+                .with_context(|| format!("Failed to parse tenant id from {}", params[0]))?;
+            let timeline_id = TimelineId::from_str(params[1])
+                .with_context(|| format!("Failed to parse timeline id from {}", params[1]))?;
 
             // The caller is responsible for providing correct lsn and prev_lsn.
             let lsn = if params.len() > 2 {
-                Some(Lsn::from_str(params[2])?)
+                Some(
+                    Lsn::from_str(params[2])
+                        .with_context(|| format!("Failed to parse Lsn from {}", params[2]))?,
+                )
             } else {
                 None
             };
             let prev_lsn = if params.len() > 3 {
-                Some(Lsn::from_str(params[3])?)
+                Some(
+                    Lsn::from_str(params[3])
+                        .with_context(|| format!("Failed to parse Lsn from {}", params[3]))?,
+                )
             } else {
                 None
             };
@@ -834,12 +852,21 @@ impl postgres_backend_async::Handler for PageServerHandler {
             //     -c "import basebackup $TENANT $TIMELINE $START_LSN $END_LSN $PG_VERSION"
             let (_, params_raw) = query_string.split_at("import basebackup ".len());
             let params = params_raw.split_whitespace().collect::<Vec<_>>();
-            ensure!(params.len() == 5);
-            let tenant_id = TenantId::from_str(params[0])?;
-            let timeline_id = TimelineId::from_str(params[1])?;
-            let base_lsn = Lsn::from_str(params[2])?;
-            let end_lsn = Lsn::from_str(params[3])?;
-            let pg_version = u32::from_str(params[4])?;
+            if params.len() != 5 {
+                return Err(QueryError::Other(anyhow::anyhow!(
+                    "invalid param number for import basebackup command"
+                )));
+            }
+            let tenant_id = TenantId::from_str(params[0])
+                .with_context(|| format!("Failed to parse tenant id from {}", params[0]))?;
+            let timeline_id = TimelineId::from_str(params[1])
+                .with_context(|| format!("Failed to parse timeline id from {}", params[1]))?;
+            let base_lsn = Lsn::from_str(params[2])
+                .with_context(|| format!("Failed to parse Lsn from {}", params[2]))?;
+            let end_lsn = Lsn::from_str(params[3])
+                .with_context(|| format!("Failed to parse Lsn from {}", params[3]))?;
+            let pg_version = u32::from_str(params[4])
+                .with_context(|| format!("Failed to parse pg_version from {}", params[4]))?;
 
             self.check_permission(Some(tenant_id))?;
 
@@ -857,7 +884,10 @@ impl postgres_backend_async::Handler for PageServerHandler {
                 Ok(()) => pgb.write_message(&BeMessage::CommandComplete(b"SELECT 1"))?,
                 Err(e) => {
                     error!("error importing base backup between {base_lsn} and {end_lsn}: {e:?}");
-                    pgb.write_message(&BeMessage::ErrorResponse(&e.to_string()))?
+                    pgb.write_message(&BeMessage::ErrorResponse(
+                        &e.to_string(),
+                        Some(e.pg_error_code()),
+                    ))?
                 }
             };
         } else if query_string.starts_with("import wal ") {
@@ -867,11 +897,19 @@ impl postgres_backend_async::Handler for PageServerHandler {
             // caller should poll the http api to check when that is done.
             let (_, params_raw) = query_string.split_at("import wal ".len());
             let params = params_raw.split_whitespace().collect::<Vec<_>>();
-            ensure!(params.len() == 4);
-            let tenant_id = TenantId::from_str(params[0])?;
-            let timeline_id = TimelineId::from_str(params[1])?;
-            let start_lsn = Lsn::from_str(params[2])?;
-            let end_lsn = Lsn::from_str(params[3])?;
+            if params.len() != 4 {
+                return Err(QueryError::Other(anyhow::anyhow!(
+                    "invalid param number for import wal command"
+                )));
+            }
+            let tenant_id = TenantId::from_str(params[0])
+                .with_context(|| format!("Failed to parse tenant id from {}", params[0]))?;
+            let timeline_id = TimelineId::from_str(params[1])
+                .with_context(|| format!("Failed to parse timeline id from {}", params[1]))?;
+            let start_lsn = Lsn::from_str(params[2])
+                .with_context(|| format!("Failed to parse Lsn from {}", params[2]))?;
+            let end_lsn = Lsn::from_str(params[3])
+                .with_context(|| format!("Failed to parse Lsn from {}", params[3]))?;
 
             self.check_permission(Some(tenant_id))?;
 
@@ -882,7 +920,10 @@ impl postgres_backend_async::Handler for PageServerHandler {
                 Ok(()) => pgb.write_message(&BeMessage::CommandComplete(b"SELECT 1"))?,
                 Err(e) => {
                     error!("error importing WAL between {start_lsn} and {end_lsn}: {e:?}");
-                    pgb.write_message(&BeMessage::ErrorResponse(&e.to_string()))?
+                    pgb.write_message(&BeMessage::ErrorResponse(
+                        &e.to_string(),
+                        Some(e.pg_error_code()),
+                    ))?
                 }
             };
         } else if query_string.to_ascii_lowercase().starts_with("set ") {
@@ -893,8 +934,13 @@ impl postgres_backend_async::Handler for PageServerHandler {
             // show <tenant_id>
             let (_, params_raw) = query_string.split_at("show ".len());
             let params = params_raw.split(' ').collect::<Vec<_>>();
-            ensure!(params.len() == 1, "invalid param number for config command");
-            let tenant_id = TenantId::from_str(params[0])?;
+            if params.len() != 1 {
+                return Err(QueryError::Other(anyhow::anyhow!(
+                    "invalid param number for config command"
+                )));
+            }
+            let tenant_id = TenantId::from_str(params[0])
+                .with_context(|| format!("Failed to parse tenant id from {}", params[0]))?;
 
             self.check_permission(Some(tenant_id))?;
 
@@ -935,7 +981,9 @@ impl postgres_backend_async::Handler for PageServerHandler {
             ]))?
             .write_message(&BeMessage::CommandComplete(b"SELECT 1"))?;
         } else {
-            bail!("unknown command");
+            return Err(QueryError::Other(anyhow::anyhow!(
+                "unknown command {query_string}"
+            )));
         }
 
         Ok(())
@@ -947,7 +995,7 @@ impl postgres_backend_async::Handler for PageServerHandler {
 /// If the tenant is Loading, waits for it to become Active, for up to 30 s. That
 /// ensures that queries don't fail immediately after pageserver startup, because
 /// all tenants are still loading.
-async fn get_active_tenant_with_timeout(tenant_id: TenantId) -> Result<Arc<Tenant>> {
+async fn get_active_tenant_with_timeout(tenant_id: TenantId) -> anyhow::Result<Arc<Tenant>> {
     let tenant = mgr::get_tenant(tenant_id, false).await?;
     match tokio::time::timeout(Duration::from_secs(30), tenant.wait_to_become_active()).await {
         Ok(wait_result) => wait_result
@@ -961,37 +1009,8 @@ async fn get_active_tenant_with_timeout(tenant_id: TenantId) -> Result<Arc<Tenan
 async fn get_active_timeline_with_timeout(
     tenant_id: TenantId,
     timeline_id: TimelineId,
-) -> Result<Arc<Timeline>> {
+) -> anyhow::Result<Arc<Timeline>> {
     get_active_tenant_with_timeout(tenant_id)
         .await
         .and_then(|tenant| tenant.get_timeline(timeline_id, true))
 }
-
-///
-/// A std::io::Write implementation that wraps all data written to it in CopyData
-/// messages.
-///
-struct CopyDataSink<'a> {
-    pgb: &'a mut PostgresBackend,
-    rt: tokio::runtime::Handle,
-}
-
-impl<'a> io::Write for CopyDataSink<'a> {
-    fn write(&mut self, data: &[u8]) -> io::Result<usize> {
-        // CopyData
-        // FIXME: if the input is large, we should split it into multiple messages.
-        // Not sure what the threshold should be, but the ultimate hard limit is that
-        // the length cannot exceed u32.
-        // FIXME: flush isn't really required, but makes it easier
-        // to view in wireshark
-        self.pgb.write_message(&BeMessage::CopyData(data))?;
-        self.rt.block_on(self.pgb.flush())?;
-        trace!("CopyData sent for {} bytes!", data.len());
-
-        Ok(data.len())
-    }
-    fn flush(&mut self) -> io::Result<()> {
-        // no-op
-        Ok(())
-    }
-}

pageserver/src/pgdatadir_mapping.rs

@@ -6,11 +6,10 @@
 //! walingest.rs handles a few things like implicit relation creation and extension.
 //! Clarify that)
 //!
-use super::tenant::PageReconstructResult;
+use super::tenant::{PageReconstructError, Timeline};
 use crate::keyspace::{KeySpace, KeySpaceAccum};
-use crate::tenant::{with_ondemand_download, Timeline};
+use crate::repository::*;
 use crate::walrecord::NeonWalRecord;
-use crate::{repository::*, try_no_ondemand_download};
 use anyhow::Context;
 use bytes::{Buf, Bytes};
 use pageserver_api::reltag::{RelTag, SlruKind};
@@ -92,76 +91,80 @@ impl Timeline {
     //------------------------------------------------------------------------------
 
     /// Look up given page version.
-    pub fn get_rel_page_at_lsn(
+    pub async fn get_rel_page_at_lsn(
         &self,
         tag: RelTag,
         blknum: BlockNumber,
         lsn: Lsn,
         latest: bool,
-    ) -> PageReconstructResult<Bytes> {
+    ) -> Result<Bytes, PageReconstructError> {
         if tag.relnode == 0 {
-            return PageReconstructResult::from(anyhow::anyhow!("invalid relnode"));
+            return Err(PageReconstructError::Other(anyhow::anyhow!(
+                "invalid relnode"
+            )));
         }
 
-        let nblocks = try_no_ondemand_download!(self.get_rel_size(tag, lsn, latest));
+        let nblocks = self.get_rel_size(tag, lsn, latest).await?;
         if blknum >= nblocks {
             debug!(
                 "read beyond EOF at {} blk {} at {}, size is {}: returning all-zeros page",
                 tag, blknum, lsn, nblocks
             );
-            return PageReconstructResult::Success(ZERO_PAGE.clone());
+            return Ok(ZERO_PAGE.clone());
         }
 
         let key = rel_block_to_key(tag, blknum);
-        self.get(key, lsn)
+        self.get(key, lsn).await
     }
 
     // Get size of a database in blocks
-    pub fn get_db_size(
+    pub async fn get_db_size(
         &self,
         spcnode: Oid,
         dbnode: Oid,
         lsn: Lsn,
         latest: bool,
-    ) -> PageReconstructResult<usize> {
+    ) -> Result<usize, PageReconstructError> {
         let mut total_blocks = 0;
 
-        let rels = try_no_ondemand_download!(self.list_rels(spcnode, dbnode, lsn));
+        let rels = self.list_rels(spcnode, dbnode, lsn).await?;
 
         for rel in rels {
-            let n_blocks = try_no_ondemand_download!(self.get_rel_size(rel, lsn, latest));
+            let n_blocks = self.get_rel_size(rel, lsn, latest).await?;
             total_blocks += n_blocks as usize;
         }
-        PageReconstructResult::Success(total_blocks)
+        Ok(total_blocks)
     }
 
     /// Get size of a relation file
-    pub fn get_rel_size(
+    pub async fn get_rel_size(
         &self,
         tag: RelTag,
         lsn: Lsn,
         latest: bool,
-    ) -> PageReconstructResult<BlockNumber> {
+    ) -> Result<BlockNumber, PageReconstructError> {
         if tag.relnode == 0 {
-            return PageReconstructResult::from(anyhow::anyhow!("invalid relnode"));
+            return Err(PageReconstructError::Other(anyhow::anyhow!(
+                "invalid relnode"
+            )));
         }
 
         if let Some(nblocks) = self.get_cached_rel_size(&tag, lsn) {
-            return PageReconstructResult::Success(nblocks);
+            return Ok(nblocks);
         }
 
         if (tag.forknum == FSM_FORKNUM || tag.forknum == VISIBILITYMAP_FORKNUM)
-            && !try_no_ondemand_download!(self.get_rel_exists(tag, lsn, latest))
+            && !self.get_rel_exists(tag, lsn, latest).await?
         {
             // FIXME: Postgres sometimes calls smgrcreate() to create
             // FSM, and smgrnblocks() on it immediately afterwards,
             // without extending it.  Tolerate that by claiming that
             // any non-existent FSM fork has size 0.
-            return PageReconstructResult::Success(0);
+            return Ok(0);
         }
 
         let key = rel_size_to_key(tag);
-        let mut buf = try_no_ondemand_download!(self.get(key, lsn));
+        let mut buf = self.get(key, lsn).await?;
         let nblocks = buf.get_u32_le();
 
         if latest {
@@ -174,47 +177,49 @@ impl Timeline {
             // associated with most recent value of LSN.
             self.update_cached_rel_size(tag, lsn, nblocks);
         }
-        PageReconstructResult::Success(nblocks)
+        Ok(nblocks)
     }
 
     /// Does relation exist?
-    pub fn get_rel_exists(
+    pub async fn get_rel_exists(
         &self,
         tag: RelTag,
         lsn: Lsn,
         _latest: bool,
-    ) -> PageReconstructResult<bool> {
+    ) -> Result<bool, PageReconstructError> {
         if tag.relnode == 0 {
-            return PageReconstructResult::from(anyhow::anyhow!("invalid relnode"));
+            return Err(PageReconstructError::Other(anyhow::anyhow!(
+                "invalid relnode"
+            )));
         }
 
         // first try to lookup relation in cache
         if let Some(_nblocks) = self.get_cached_rel_size(&tag, lsn) {
-            return PageReconstructResult::Success(true);
+            return Ok(true);
         }
         // fetch directory listing
         let key = rel_dir_to_key(tag.spcnode, tag.dbnode);
-        let buf = try_no_ondemand_download!(self.get(key, lsn));
+        let buf = self.get(key, lsn).await?;
 
         match RelDirectory::des(&buf).context("deserialization failure") {
             Ok(dir) => {
                 let exists = dir.rels.get(&(tag.relnode, tag.forknum)).is_some();
-                PageReconstructResult::Success(exists)
+                Ok(exists)
             }
-            Err(e) => PageReconstructResult::from(e),
+            Err(e) => Err(PageReconstructError::from(e)),
         }
     }
 
     /// Get a list of all existing relations in given tablespace and database.
-    pub fn list_rels(
+    pub async fn list_rels(
         &self,
         spcnode: Oid,
         dbnode: Oid,
         lsn: Lsn,
-    ) -> PageReconstructResult<HashSet<RelTag>> {
+    ) -> Result<HashSet<RelTag>, PageReconstructError> {
         // fetch directory listing
         let key = rel_dir_to_key(spcnode, dbnode);
-        let buf = try_no_ondemand_download!(self.get(key, lsn));
+        let buf = self.get(key, lsn).await?;
 
         match RelDirectory::des(&buf).context("deserialization failure") {
             Ok(dir) => {
@@ -226,53 +231,53 @@ impl Timeline {
                         forknum: *forknum,
                     }));
 
-                PageReconstructResult::Success(rels)
+                Ok(rels)
             }
-            Err(e) => PageReconstructResult::from(e),
+            Err(e) => Err(PageReconstructError::from(e)),
         }
     }
 
     /// Look up given SLRU page version.
-    pub fn get_slru_page_at_lsn(
+    pub async fn get_slru_page_at_lsn(
         &self,
         kind: SlruKind,
         segno: u32,
         blknum: BlockNumber,
         lsn: Lsn,
-    ) -> PageReconstructResult<Bytes> {
+    ) -> Result<Bytes, PageReconstructError> {
         let key = slru_block_to_key(kind, segno, blknum);
-        self.get(key, lsn)
+        self.get(key, lsn).await
     }
 
     /// Get size of an SLRU segment
-    pub fn get_slru_segment_size(
+    pub async fn get_slru_segment_size(
         &self,
         kind: SlruKind,
         segno: u32,
         lsn: Lsn,
-    ) -> PageReconstructResult<BlockNumber> {
+    ) -> Result<BlockNumber, PageReconstructError> {
         let key = slru_segment_size_to_key(kind, segno);
-        let mut buf = try_no_ondemand_download!(self.get(key, lsn));
-        PageReconstructResult::Success(buf.get_u32_le())
+        let mut buf = self.get(key, lsn).await?;
+        Ok(buf.get_u32_le())
     }
 
     /// Get size of an SLRU segment
-    pub fn get_slru_segment_exists(
+    pub async fn get_slru_segment_exists(
         &self,
         kind: SlruKind,
         segno: u32,
         lsn: Lsn,
-    ) -> PageReconstructResult<bool> {
+    ) -> Result<bool, PageReconstructError> {
         // fetch directory listing
         let key = slru_dir_to_key(kind);
-        let buf = try_no_ondemand_download!(self.get(key, lsn));
+        let buf = self.get(key, lsn).await?;
 
         match SlruSegmentDirectory::des(&buf).context("deserialization failure") {
             Ok(dir) => {
                 let exists = dir.segments.get(&segno).is_some();
-                PageReconstructResult::Success(exists)
+                Ok(exists)
             }
-            Err(e) => PageReconstructResult::from(e),
+            Err(e) => Err(PageReconstructError::from(e)),
         }
     }
 
@@ -283,10 +288,10 @@ impl Timeline {
     /// so it's not well defined which LSN you get if there were multiple commits
     /// "in flight" at that point in time.
     ///
-    pub fn find_lsn_for_timestamp(
+    pub async fn find_lsn_for_timestamp(
         &self,
         search_timestamp: TimestampTz,
-    ) -> PageReconstructResult<LsnForTimestamp> {
+    ) -> Result<LsnForTimestamp, PageReconstructError> {
         let gc_cutoff_lsn_guard = self.get_latest_gc_cutoff_lsn();
         let min_lsn = *gc_cutoff_lsn_guard;
         let max_lsn = self.get_last_record_lsn();
@@ -302,12 +307,14 @@ impl Timeline {
             // cannot overflow, high and low are both smaller than u64::MAX / 2
             let mid = (high + low) / 2;
 
-            let cmp = try_no_ondemand_download!(self.is_latest_commit_timestamp_ge_than(
-                search_timestamp,
-                Lsn(mid * 8),
-                &mut found_smaller,
-                &mut found_larger,
-            ));
+            let cmp = self
+                .is_latest_commit_timestamp_ge_than(
+                    search_timestamp,
+                    Lsn(mid * 8),
+                    &mut found_smaller,
+                    &mut found_larger,
+                )
+                .await?;
 
             if cmp {
                 high = mid;
@@ -319,15 +326,15 @@ impl Timeline {
             (false, false) => {
                 // This can happen if no commit records have been processed yet, e.g.
                 // just after importing a cluster.
-                PageReconstructResult::Success(LsnForTimestamp::NoData(max_lsn))
+                Ok(LsnForTimestamp::NoData(max_lsn))
             }
             (true, false) => {
                 // Didn't find any commit timestamps larger than the request
-                PageReconstructResult::Success(LsnForTimestamp::Future(max_lsn))
+                Ok(LsnForTimestamp::Future(max_lsn))
             }
             (false, true) => {
                 // Didn't find any commit timestamps smaller than the request
-                PageReconstructResult::Success(LsnForTimestamp::Past(max_lsn))
+                Ok(LsnForTimestamp::Past(max_lsn))
             }
             (true, true) => {
                 // low is the LSN of the first commit record *after* the search_timestamp,
@@ -337,7 +344,7 @@ impl Timeline {
                 // Otherwise, if you restore to the returned LSN, the database will
                 // include physical changes from later commits that will be marked
                 // as aborted, and will need to be vacuumed away.
-                PageReconstructResult::Success(LsnForTimestamp::Present(Lsn((low - 1) * 8)))
+                Ok(LsnForTimestamp::Present(Lsn((low - 1) * 8)))
             }
         }
     }
@@ -349,26 +356,21 @@ impl Timeline {
     /// Additionally, sets 'found_smaller'/'found_Larger, if encounters any commits
     /// with a smaller/larger timestamp.
     ///
-    pub fn is_latest_commit_timestamp_ge_than(
+    pub async fn is_latest_commit_timestamp_ge_than(
         &self,
         search_timestamp: TimestampTz,
         probe_lsn: Lsn,
         found_smaller: &mut bool,
         found_larger: &mut bool,
-    ) -> PageReconstructResult<bool> {
-        for segno in try_no_ondemand_download!(self.list_slru_segments(SlruKind::Clog, probe_lsn)) {
-            let nblocks = try_no_ondemand_download!(self.get_slru_segment_size(
-                SlruKind::Clog,
-                segno,
-                probe_lsn
-            ));
+    ) -> Result<bool, PageReconstructError> {
+        for segno in self.list_slru_segments(SlruKind::Clog, probe_lsn).await? {
+            let nblocks = self
+                .get_slru_segment_size(SlruKind::Clog, segno, probe_lsn)
+                .await?;
             for blknum in (0..nblocks).rev() {
-                let clog_page = try_no_ondemand_download!(self.get_slru_page_at_lsn(
-                    SlruKind::Clog,
-                    segno,
-                    blknum,
-                    probe_lsn
-                ));
+                let clog_page = self
+                    .get_slru_page_at_lsn(SlruKind::Clog, segno, blknum, probe_lsn)
+                    .await?;
 
                 if clog_page.len() == BLCKSZ as usize + 8 {
                     let mut timestamp_bytes = [0u8; 8];
@@ -377,76 +379,85 @@ impl Timeline {
 
                     if timestamp >= search_timestamp {
                         *found_larger = true;
-                        return PageReconstructResult::Success(true);
+                        return Ok(true);
                     } else {
                         *found_smaller = true;
                     }
                 }
             }
         }
-        PageReconstructResult::Success(false)
+        Ok(false)
     }
 
     /// Get a list of SLRU segments
-    pub fn list_slru_segments(
+    pub async fn list_slru_segments(
         &self,
         kind: SlruKind,
         lsn: Lsn,
-    ) -> PageReconstructResult<HashSet<u32>> {
+    ) -> Result<HashSet<u32>, PageReconstructError> {
         // fetch directory entry
         let key = slru_dir_to_key(kind);
 
-        let buf = try_no_ondemand_download!(self.get(key, lsn));
+        let buf = self.get(key, lsn).await?;
         match SlruSegmentDirectory::des(&buf).context("deserialization failure") {
-            Ok(dir) => PageReconstructResult::Success(dir.segments),
-            Err(e) => PageReconstructResult::from(e),
+            Ok(dir) => Ok(dir.segments),
+            Err(e) => Err(PageReconstructError::from(e)),
         }
     }
 
-    pub fn get_relmap_file(
+    pub async fn get_relmap_file(
         &self,
         spcnode: Oid,
         dbnode: Oid,
         lsn: Lsn,
-    ) -> PageReconstructResult<Bytes> {
+    ) -> Result<Bytes, PageReconstructError> {
         let key = relmap_file_key(spcnode, dbnode);
 
-        let buf = try_no_ondemand_download!(self.get(key, lsn));
-        PageReconstructResult::Success(buf)
+        self.get(key, lsn).await
     }
 
-    pub fn list_dbdirs(&self, lsn: Lsn) -> PageReconstructResult<HashMap<(Oid, Oid), bool>> {
+    pub async fn list_dbdirs(
+        &self,
+        lsn: Lsn,
+    ) -> Result<HashMap<(Oid, Oid), bool>, PageReconstructError> {
         // fetch directory entry
-        let buf = try_no_ondemand_download!(self.get(DBDIR_KEY, lsn));
+        let buf = self.get(DBDIR_KEY, lsn).await?;
 
         match DbDirectory::des(&buf).context("deserialization failure") {
-            Ok(dir) => PageReconstructResult::Success(dir.dbdirs),
-            Err(e) => PageReconstructResult::from(e),
+            Ok(dir) => Ok(dir.dbdirs),
+            Err(e) => Err(PageReconstructError::from(e)),
         }
     }
 
-    pub fn get_twophase_file(&self, xid: TransactionId, lsn: Lsn) -> PageReconstructResult<Bytes> {
+    pub async fn get_twophase_file(
+        &self,
+        xid: TransactionId,
+        lsn: Lsn,
+    ) -> Result<Bytes, PageReconstructError> {
         let key = twophase_file_key(xid);
-        let buf = try_no_ondemand_download!(self.get(key, lsn));
-        PageReconstructResult::Success(buf)
+        let buf = self.get(key, lsn).await?;
+        Ok(buf)
     }
 
-    pub fn list_twophase_files(&self, lsn: Lsn) -> PageReconstructResult<HashSet<TransactionId>> {
+    pub async fn list_twophase_files(
+        &self,
+        lsn: Lsn,
+    ) -> Result<HashSet<TransactionId>, PageReconstructError> {
         // fetch directory entry
-        let buf = try_no_ondemand_download!(self.get(TWOPHASEDIR_KEY, lsn));
+        let buf = self.get(TWOPHASEDIR_KEY, lsn).await?;
 
         match TwoPhaseDirectory::des(&buf).context("deserialization failure") {
-            Ok(dir) => PageReconstructResult::Success(dir.xids),
-            Err(e) => PageReconstructResult::from(e),
+            Ok(dir) => Ok(dir.xids),
+            Err(e) => Err(PageReconstructError::from(e)),
         }
     }
 
-    pub fn get_control_file(&self, lsn: Lsn) -> PageReconstructResult<Bytes> {
-        self.get(CONTROLFILE_KEY, lsn)
+    pub async fn get_control_file(&self, lsn: Lsn) -> Result<Bytes, PageReconstructError> {
+        self.get(CONTROLFILE_KEY, lsn).await
     }
 
-    pub fn get_checkpoint(&self, lsn: Lsn) -> PageReconstructResult<Bytes> {
-        self.get(CHECKPOINT_KEY, lsn)
+    pub async fn get_checkpoint(&self, lsn: Lsn) -> Result<Bytes, PageReconstructError> {
+        self.get(CHECKPOINT_KEY, lsn).await
     }
 
     /// Does the same as get_current_logical_size but counted on demand.
@@ -460,20 +471,24 @@ impl Timeline {
         cancel: CancellationToken,
     ) -> Result<u64, CalculateLogicalSizeError> {
         // Fetch list of database dirs and iterate them
-        let buf = self.get_download(DBDIR_KEY, lsn).await?;
+        let buf = self.get(DBDIR_KEY, lsn).await.context("read dbdir")?;
         let dbdir = DbDirectory::des(&buf).context("deserialize db directory")?;
 
         let mut total_size: u64 = 0;
         for (spcnode, dbnode) in dbdir.dbdirs.keys() {
-            for rel in
-                crate::tenant::with_ondemand_download(|| self.list_rels(*spcnode, *dbnode, lsn))
-                    .await?
+            for rel in self
+                .list_rels(*spcnode, *dbnode, lsn)
+                .await
+                .context("list rels")?
             {
                 if cancel.is_cancelled() {
                     return Err(CalculateLogicalSizeError::Cancelled);
                 }
                 let relsize_key = rel_size_to_key(rel);
-                let mut buf = self.get_download(relsize_key, lsn).await?;
+                let mut buf = self
+                    .get(relsize_key, lsn)
+                    .await
+                    .context("read relation size of {rel:?}")?;
                 let relsize = buf.get_u32_le();
 
                 total_size += relsize as u64;
@@ -494,7 +509,7 @@ impl Timeline {
         result.add_key(DBDIR_KEY);
 
         // Fetch list of database dirs and iterate them
-        let buf = self.get_download(DBDIR_KEY, lsn).await?;
+        let buf = self.get(DBDIR_KEY, lsn).await?;
         let dbdir = DbDirectory::des(&buf).context("deserialization failure")?;
 
         let mut dbs: Vec<(Oid, Oid)> = dbdir.dbdirs.keys().cloned().collect();
@@ -503,15 +518,15 @@ impl Timeline {
             result.add_key(relmap_file_key(spcnode, dbnode));
             result.add_key(rel_dir_to_key(spcnode, dbnode));
 
-            let mut rels: Vec<RelTag> =
-                with_ondemand_download(|| self.list_rels(spcnode, dbnode, lsn))
-                    .await?
-                    .into_iter()
-                    .collect();
+            let mut rels: Vec<RelTag> = self
+                .list_rels(spcnode, dbnode, lsn)
+                .await?
+                .into_iter()
+                .collect();
             rels.sort_unstable();
             for rel in rels {
                 let relsize_key = rel_size_to_key(rel);
-                let mut buf = self.get_download(relsize_key, lsn).await?;
+                let mut buf = self.get(relsize_key, lsn).await?;
                 let relsize = buf.get_u32_le();
 
                 result.add_range(rel_block_to_key(rel, 0)..rel_block_to_key(rel, relsize));
@@ -527,13 +542,13 @@ impl Timeline {
         ] {
             let slrudir_key = slru_dir_to_key(kind);
             result.add_key(slrudir_key);
-            let buf = self.get_download(slrudir_key, lsn).await?;
+            let buf = self.get(slrudir_key, lsn).await?;
             let dir = SlruSegmentDirectory::des(&buf).context("deserialization failure")?;
             let mut segments: Vec<u32> = dir.segments.iter().cloned().collect();
             segments.sort_unstable();
             for segno in segments {
                 let segsize_key = slru_segment_size_to_key(kind, segno);
-                let mut buf = self.get_download(segsize_key, lsn).await?;
+                let mut buf = self.get(segsize_key, lsn).await?;
                 let segsize = buf.get_u32_le();
 
                 result.add_range(
@@ -545,7 +560,7 @@ impl Timeline {
 
         // Then pg_twophase
         result.add_key(TWOPHASEDIR_KEY);
-        let buf = self.get_download(TWOPHASEDIR_KEY, lsn).await?;
+        let buf = self.get(TWOPHASEDIR_KEY, lsn).await?;
         let twophase_dir = TwoPhaseDirectory::des(&buf).context("deserialization failure")?;
         let mut xids: Vec<TransactionId> = twophase_dir.xids.iter().cloned().collect();
         xids.sort_unstable();
@@ -703,9 +718,14 @@ impl<'a> DatadirModification<'a> {
     }
 
     /// Store a relmapper file (pg_filenode.map) in the repository
-    pub fn put_relmap_file(&mut self, spcnode: Oid, dbnode: Oid, img: Bytes) -> anyhow::Result<()> {
+    pub async fn put_relmap_file(
+        &mut self,
+        spcnode: Oid,
+        dbnode: Oid,
+        img: Bytes,
+    ) -> anyhow::Result<()> {
         // Add it to the directory (if it doesn't exist already)
-        let buf = self.get(DBDIR_KEY).no_ondemand_download()?;
+        let buf = self.get(DBDIR_KEY).await?;
         let mut dbdir = DbDirectory::des(&buf)?;
 
         let r = dbdir.dbdirs.insert((spcnode, dbnode), true);
@@ -731,9 +751,13 @@ impl<'a> DatadirModification<'a> {
         Ok(())
     }
 
-    pub fn put_twophase_file(&mut self, xid: TransactionId, img: Bytes) -> anyhow::Result<()> {
+    pub async fn put_twophase_file(
+        &mut self,
+        xid: TransactionId,
+        img: Bytes,
+    ) -> anyhow::Result<()> {
         // Add it to the directory entry
-        let buf = self.get(TWOPHASEDIR_KEY).no_ondemand_download()?;
+        let buf = self.get(TWOPHASEDIR_KEY).await?;
         let mut dir = TwoPhaseDirectory::des(&buf)?;
         if !dir.xids.insert(xid) {
             anyhow::bail!("twophase file for xid {} already exists", xid);
@@ -757,16 +781,16 @@ impl<'a> DatadirModification<'a> {
         Ok(())
     }
 
-    pub fn drop_dbdir(&mut self, spcnode: Oid, dbnode: Oid) -> anyhow::Result<()> {
+    pub async fn drop_dbdir(&mut self, spcnode: Oid, dbnode: Oid) -> anyhow::Result<()> {
         let req_lsn = self.tline.get_last_record_lsn();
 
         let total_blocks = self
             .tline
             .get_db_size(spcnode, dbnode, req_lsn, true)
-            .no_ondemand_download()?;
+            .await?;
 
         // Remove entry from dbdir
-        let buf = self.get(DBDIR_KEY).no_ondemand_download()?;
+        let buf = self.get(DBDIR_KEY).await?;
         let mut dir = DbDirectory::des(&buf)?;
         if dir.dbdirs.remove(&(spcnode, dbnode)).is_some() {
             let buf = DbDirectory::ser(&dir)?;
@@ -789,11 +813,15 @@ impl<'a> DatadirModification<'a> {
     /// Create a relation fork.
     ///
     /// 'nblocks' is the initial size.
-    pub fn put_rel_creation(&mut self, rel: RelTag, nblocks: BlockNumber) -> anyhow::Result<()> {
+    pub async fn put_rel_creation(
+        &mut self,
+        rel: RelTag,
+        nblocks: BlockNumber,
+    ) -> anyhow::Result<()> {
         anyhow::ensure!(rel.relnode != 0, "invalid relnode");
         // It's possible that this is the first rel for this db in this
         // tablespace.  Create the reldir entry for it if so.
-        let mut dbdir = DbDirectory::des(&self.get(DBDIR_KEY).no_ondemand_download()?)?;
+        let mut dbdir = DbDirectory::des(&self.get(DBDIR_KEY).await?)?;
         let rel_dir_key = rel_dir_to_key(rel.spcnode, rel.dbnode);
         let mut rel_dir = if dbdir.dbdirs.get(&(rel.spcnode, rel.dbnode)).is_none() {
             // Didn't exist. Update dbdir
@@ -805,7 +833,7 @@ impl<'a> DatadirModification<'a> {
             RelDirectory::default()
         } else {
             // reldir already exists, fetch it
-            RelDirectory::des(&self.get(rel_dir_key).no_ondemand_download()?)?
+            RelDirectory::des(&self.get(rel_dir_key).await?)?
         };
 
         // Add the new relation to the rel directory entry, and write it back
@@ -833,17 +861,17 @@ impl<'a> DatadirModification<'a> {
     }
 
     /// Truncate relation
-    pub fn put_rel_truncation(&mut self, rel: RelTag, nblocks: BlockNumber) -> anyhow::Result<()> {
+    pub async fn put_rel_truncation(
+        &mut self,
+        rel: RelTag,
+        nblocks: BlockNumber,
+    ) -> anyhow::Result<()> {
         anyhow::ensure!(rel.relnode != 0, "invalid relnode");
         let last_lsn = self.tline.get_last_record_lsn();
-        if self
-            .tline
-            .get_rel_exists(rel, last_lsn, true)
-            .no_ondemand_download()?
-        {
+        if self.tline.get_rel_exists(rel, last_lsn, true).await? {
             let size_key = rel_size_to_key(rel);
             // Fetch the old size first
-            let old_size = self.get(size_key).no_ondemand_download()?.get_u32_le();
+            let old_size = self.get(size_key).await?.get_u32_le();
 
             // Update the entry with the new size.
             let buf = nblocks.to_le_bytes();
@@ -863,12 +891,16 @@ impl<'a> DatadirModification<'a> {
 
     /// Extend relation
     /// If new size is smaller, do nothing.
-    pub fn put_rel_extend(&mut self, rel: RelTag, nblocks: BlockNumber) -> anyhow::Result<()> {
+    pub async fn put_rel_extend(
+        &mut self,
+        rel: RelTag,
+        nblocks: BlockNumber,
+    ) -> anyhow::Result<()> {
         anyhow::ensure!(rel.relnode != 0, "invalid relnode");
 
         // Put size
         let size_key = rel_size_to_key(rel);
-        let old_size = self.get(size_key).no_ondemand_download()?.get_u32_le();
+        let old_size = self.get(size_key).await?.get_u32_le();
 
         // only extend relation here. never decrease the size
         if nblocks > old_size {
@@ -884,12 +916,12 @@ impl<'a> DatadirModification<'a> {
     }
 
     /// Drop a relation.
-    pub fn put_rel_drop(&mut self, rel: RelTag) -> anyhow::Result<()> {
+    pub async fn put_rel_drop(&mut self, rel: RelTag) -> anyhow::Result<()> {
         anyhow::ensure!(rel.relnode != 0, "invalid relnode");
 
         // Remove it from the directory entry
         let dir_key = rel_dir_to_key(rel.spcnode, rel.dbnode);
-        let buf = self.get(dir_key).no_ondemand_download()?;
+        let buf = self.get(dir_key).await?;
         let mut dir = RelDirectory::des(&buf)?;
 
         if dir.rels.remove(&(rel.relnode, rel.forknum)) {
@@ -900,7 +932,7 @@ impl<'a> DatadirModification<'a> {
 
         // update logical size
         let size_key = rel_size_to_key(rel);
-        let old_size = self.get(size_key).no_ondemand_download()?.get_u32_le();
+        let old_size = self.get(size_key).await?.get_u32_le();
         self.pending_nblocks -= old_size as i64;
 
         // Remove enty from relation size cache
@@ -912,7 +944,7 @@ impl<'a> DatadirModification<'a> {
         Ok(())
     }
 
-    pub fn put_slru_segment_creation(
+    pub async fn put_slru_segment_creation(
         &mut self,
         kind: SlruKind,
         segno: u32,
@@ -920,7 +952,7 @@ impl<'a> DatadirModification<'a> {
     ) -> anyhow::Result<()> {
         // Add it to the directory entry
         let dir_key = slru_dir_to_key(kind);
-        let buf = self.get(dir_key).no_ondemand_download()?;
+        let buf = self.get(dir_key).await?;
         let mut dir = SlruSegmentDirectory::des(&buf)?;
 
         if !dir.segments.insert(segno) {
@@ -956,10 +988,10 @@ impl<'a> DatadirModification<'a> {
     }
 
     /// This method is used for marking truncated SLRU files
-    pub fn drop_slru_segment(&mut self, kind: SlruKind, segno: u32) -> anyhow::Result<()> {
+    pub async fn drop_slru_segment(&mut self, kind: SlruKind, segno: u32) -> anyhow::Result<()> {
         // Remove it from the directory entry
         let dir_key = slru_dir_to_key(kind);
-        let buf = self.get(dir_key).no_ondemand_download()?;
+        let buf = self.get(dir_key).await?;
         let mut dir = SlruSegmentDirectory::des(&buf)?;
 
         if !dir.segments.remove(&segno) {
@@ -983,9 +1015,9 @@ impl<'a> DatadirModification<'a> {
     }
 
     /// This method is used for marking truncated SLRU files
-    pub fn drop_twophase_file(&mut self, xid: TransactionId) -> anyhow::Result<()> {
+    pub async fn drop_twophase_file(&mut self, xid: TransactionId) -> anyhow::Result<()> {
         // Remove it from the directory entry
-        let buf = self.get(TWOPHASEDIR_KEY).no_ondemand_download()?;
+        let buf = self.get(TWOPHASEDIR_KEY).await?;
         let mut dir = TwoPhaseDirectory::des(&buf)?;
 
         if !dir.xids.remove(&xid) {
@@ -1079,7 +1111,7 @@ impl<'a> DatadirModification<'a> {
 
     // Internal helper functions to batch the modifications
 
-    fn get(&self, key: Key) -> PageReconstructResult<Bytes> {
+    async fn get(&self, key: Key) -> Result<Bytes, PageReconstructError> {
         // Have we already updated the same key? Read the pending updated
         // version in that case.
         //
@@ -1087,18 +1119,20 @@ impl<'a> DatadirModification<'a> {
         // value that has been removed, deletion only avoids leaking storage.
         if let Some(value) = self.pending_updates.get(&key) {
             if let Value::Image(img) = value {
-                PageReconstructResult::Success(img.clone())
+                Ok(img.clone())
             } else {
                 // Currently, we never need to read back a WAL record that we
                 // inserted in the same "transaction". All the metadata updates
                 // work directly with Images, and we never need to read actual
                 // data pages. We could handle this if we had to, by calling
                 // the walredo manager, but let's keep it simple for now.
-                PageReconstructResult::from(anyhow::anyhow!("unexpected pending WAL record"))
+                Err(PageReconstructError::from(anyhow::anyhow!(
+                    "unexpected pending WAL record"
+                )))
             }
         } else {
             let lsn = Lsn::max(self.tline.get_last_record_lsn(), self.lsn);
-            self.tline.get(key, lsn)
+            self.tline.get(key, lsn).await
         }
     }
 
@@ -1371,15 +1405,15 @@ fn slru_segment_key_range(kind: SlruKind, segno: u32) -> Range<Key> {
     Key {
         field1: 0x01,
         field2,
-        field3: segno,
-        field4: 0,
+        field3: 1,
+        field4: segno,
         field5: 0,
         field6: 0,
     }..Key {
         field1: 0x01,
         field2,
-        field3: segno,
-        field4: 0,
+        field3: 1,
+        field4: segno,
         field5: 1,
         field6: 0,
     }

pageserver/src/profiling.rs

@@ -1,107 +0,0 @@
-//!
-//! Support for profiling
-//!
-//! This relies on a modified version of the 'pprof-rs' crate. That's not very
-//! nice, so to avoid a hard dependency on that, this is an optional feature.
-//!
-use crate::config::{PageServerConf, ProfilingConfig};
-
-/// The actual implementation is in the `profiling_impl` submodule. If the profiling
-/// feature is not enabled, it's just a dummy implementation that panics if you
-/// try to enabled profiling in the configuration.
-pub use profiling_impl::*;
-
-#[cfg(feature = "profiling")]
-mod profiling_impl {
-    use super::*;
-    use pprof;
-    use std::marker::PhantomData;
-
-    /// Start profiling the current thread. Returns a guard object;
-    /// the profiling continues until the guard is dropped.
-    ///
-    /// Note: profiling is not re-entrant. If you call 'profpoint_start' while
-    /// profiling is already started, nothing happens, and the profiling will be
-    /// stopped when either guard object is dropped.
-    #[inline]
-    pub fn profpoint_start(
-        conf: &crate::config::PageServerConf,
-        point: ProfilingConfig,
-    ) -> Option<ProfilingGuard> {
-        if conf.profiling == point {
-            pprof::start_profiling();
-            Some(ProfilingGuard(PhantomData))
-        } else {
-            None
-        }
-    }
-
-    /// A hack to remove Send and Sync from the ProfilingGuard. Because the
-    /// profiling is attached to current thread.
-    ////
-    /// See comments in https://github.com/rust-lang/rust/issues/68318
-    type PhantomUnsend = std::marker::PhantomData<*mut u8>;
-
-    pub struct ProfilingGuard(PhantomUnsend);
-
-    impl Drop for ProfilingGuard {
-        fn drop(&mut self) {
-            pprof::stop_profiling();
-        }
-    }
-
-    /// Initialize the profiler. This must be called before any 'profpoint_start' calls.
-    pub fn init_profiler(conf: &PageServerConf) -> Option<pprof::ProfilerGuard> {
-        if conf.profiling != ProfilingConfig::Disabled {
-            Some(pprof::ProfilerGuardBuilder::default().build().unwrap())
-        } else {
-            None
-        }
-    }
-
-    /// Exit the profiler. Writes the flamegraph to current workdir.
-    pub fn exit_profiler(_conf: &PageServerConf, profiler_guard: &Option<pprof::ProfilerGuard>) {
-        // Write out the flamegraph
-        if let Some(profiler_guard) = profiler_guard {
-            if let Ok(report) = profiler_guard.report().build() {
-                // this gets written under the workdir
-                let file = std::fs::File::create("flamegraph.svg").unwrap();
-                let mut options = pprof::flamegraph::Options::default();
-                options.image_width = Some(2500);
-                report.flamegraph_with_options(file, &mut options).unwrap();
-            }
-        }
-    }
-}
-
-/// Dummy implementation when compiling without profiling feature or for non-linux OSes.
-#[cfg(not(feature = "profiling"))]
-mod profiling_impl {
-    use super::*;
-
-    pub struct DummyProfilerGuard;
-
-    impl Drop for DummyProfilerGuard {
-        fn drop(&mut self) {
-            // do nothing, this exists to calm Clippy down
-        }
-    }
-
-    pub fn profpoint_start(
-        _conf: &PageServerConf,
-        _point: ProfilingConfig,
-    ) -> Option<DummyProfilerGuard> {
-        None
-    }
-
-    pub fn init_profiler(conf: &PageServerConf) -> Option<DummyProfilerGuard> {
-        if conf.profiling != ProfilingConfig::Disabled {
-            // shouldn't happen, we don't allow profiling in the config if the support
-            // for it is disabled.
-            panic!("profiling enabled but the binary was compiled without profiling support");
-        }
-        None
-    }
-
-    pub fn exit_profiler(_conf: &PageServerConf, _guard: &Option<DummyProfilerGuard>) {}
-}

pageserver/src/task_mgr.rs

@@ -220,6 +220,8 @@ pub enum TaskKind {
 
     // task that drives downloading layers
     DownloadAllRemoteLayers,
+    // Task that calculates synthetis size for all active tenants
+    CalculateSyntheticSize,
 }
 
 #[derive(Default)]

pageserver/src/tenant.rs

@@ -13,13 +13,13 @@
 
 use anyhow::{bail, Context};
 use bytes::Bytes;
+use futures::FutureExt;
 use futures::Stream;
 use pageserver_api::models::TimelineState;
 use remote_storage::DownloadError;
 use remote_storage::GenericRemoteStorage;
 use tokio::sync::watch;
-use tokio_util::io::StreamReader;
-use tokio_util::io::SyncIoBridge;
+use tokio::task::JoinSet;
 use tracing::*;
 use utils::crashsafe::path_with_suffix_extension;
 
@@ -36,9 +36,10 @@ use std::io::Write;
 use std::ops::Bound::Included;
 use std::path::Path;
 use std::path::PathBuf;
-use std::pin::Pin;
 use std::process::Command;
 use std::process::Stdio;
+use std::sync::atomic::AtomicU64;
+use std::sync::atomic::Ordering;
 use std::sync::Arc;
 use std::sync::MutexGuard;
 use std::sync::{Mutex, RwLock};
@@ -89,11 +90,11 @@ pub mod mgr;
 pub mod tasks;
 pub mod upload_queue;
 
-mod timeline;
+pub mod timeline;
 
 pub mod size;
 
-pub use timeline::{with_ondemand_download, PageReconstructError, PageReconstructResult, Timeline};
+pub use timeline::{PageReconstructError, Timeline};
 
 // re-export this function so that page_cache.rs can use it.
 pub use crate::tenant::ephemeral_file::writeback as writeback_ephemeral_file;
@@ -140,6 +141,7 @@ pub struct Tenant {
 
     /// Cached logical sizes updated updated on each [`Tenant::gather_size_inputs`].
     cached_logical_sizes: tokio::sync::Mutex<HashMap<(TimelineId, Lsn), u64>>,
+    cached_synthetic_tenant_size: Arc<AtomicU64>,
 }
 
 /// A timeline with some of its files on disk, being initialized.
@@ -236,21 +238,15 @@ impl UninitializedTimeline<'_> {
     /// Prepares timeline data by loading it from the basebackup archive.
     pub async fn import_basebackup_from_tar(
         self,
-        mut copyin_stream: &mut Pin<&mut impl Stream<Item = io::Result<Bytes>>>,
+        copyin_stream: &mut (impl Stream<Item = io::Result<Bytes>> + Sync + Send + Unpin),
         base_lsn: Lsn,
     ) -> anyhow::Result<Arc<Timeline>> {
         let raw_timeline = self.raw_timeline()?;
 
-        // import_basebackup_from_tar() is not async, mainly because the Tar crate
-        // it uses is not async. So we need to jump through some hoops:
-        // - convert the input from client connection to a synchronous Read
-        // - use block_in_place()
-        let reader = SyncIoBridge::new(StreamReader::new(&mut copyin_stream));
-
-        tokio::task::block_in_place(|| {
-            import_datadir::import_basebackup_from_tar(raw_timeline, reader, base_lsn)
-                .context("Failed to import basebackup")
-        })?;
+        let mut reader = tokio_util::io::StreamReader::new(copyin_stream);
+        import_datadir::import_basebackup_from_tar(raw_timeline, &mut reader, base_lsn)
+            .await
+            .context("Failed to import basebackup")?;
 
         // Flush loop needs to be spawned in order to be able to flush.
         // We want to run proper checkpoint before we mark timeline as available to outside world
@@ -445,8 +441,16 @@ struct RemoteStartupData {
 
 impl Tenant {
     /// Yet another helper for timeline initialization.
-    /// Contains common part for `load_local_timeline` and `load_remote_timeline`
-    async fn setup_timeline(
+    /// Contains the common part of `load_local_timeline` and `load_remote_timeline`.
+    ///
+    /// - Initializes the Timeline struct and inserts it into the tenant's hash map
+    /// - Scans the local timeline directory for layer files and builds the layer map
+    /// - Downloads remote index file and adds remote files to the layer map
+    /// - Schedules remote upload tasks for any files that are present locally but missing from remote storage.
+    ///
+    /// If the operation fails, the timeline is left in the tenant's hash map in Broken state. On success,
+    /// it is marked as Active.
+    async fn timeline_init_and_sync(
         &self,
         timeline_id: TimelineId,
         remote_client: Option<RemoteTimelineClient>,
@@ -489,10 +493,7 @@ impl Tenant {
             // 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(&mut timelines_accessor, true, false) {
-                Ok(initialized_timeline) => {
-                    timelines_accessor.insert(timeline_id, initialized_timeline.clone());
-                    Ok(initialized_timeline)
-                }
+                Ok(new_timeline) => new_timeline,
                 Err(e) => {
                     error!("Failed to initialize timeline {tenant_id}/{timeline_id}: {e:?}");
                     // FIXME using None is a hack, it wont hurt, just ugly.
@@ -508,16 +509,14 @@ impl Tenant {
                             None,
                         )
                         .with_context(|| {
-                            format!(
-                            "Failed to crate broken timeline data for {tenant_id}/{timeline_id}"
-                        )
+                            format!("creating broken timeline data for {tenant_id}/{timeline_id}")
                         })?;
                     broken_timeline.set_state(TimelineState::Broken);
                     timelines_accessor.insert(timeline_id, broken_timeline);
-                    Err(e)
+                    return Err(e);
                 }
             }
-        }?;
+        };
 
         if self.remote_storage.is_some() {
             // Reconcile local state with remote storage, downloading anything that's
@@ -603,7 +602,7 @@ impl Tenant {
                 match tenant_clone.attach().await {
                     Ok(_) => {}
                     Err(e) => {
-                        tenant_clone.set_broken();
+                        tenant_clone.set_broken(&e.to_string());
                         error!("error attaching tenant: {:?}", e);
                     }
                 }
@@ -619,7 +618,7 @@ impl Tenant {
     #[instrument(skip(self), fields(tenant_id=%self.tenant_id))]
     async fn attach(self: &Arc<Tenant>) -> anyhow::Result<()> {
         // Create directory with marker file to indicate attaching state.
-        // The load_local_tenants() function in tenant_mgr relies on the marker file
+        // The load_local_tenants() function in tenant::mgr relies on the marker file
         // to determine whether a tenant has finished attaching.
         let tenant_dir = self.conf.tenant_path(&self.tenant_id);
         let marker_file = self.conf.tenant_attaching_mark_file_path(&self.tenant_id);
@@ -648,26 +647,62 @@ impl Tenant {
             .as_ref()
             .ok_or_else(|| anyhow::anyhow!("cannot attach without remote storage"))?;
 
-        let remote_timelines = remote_timeline_client::list_remote_timelines(
+        let remote_timeline_ids = remote_timeline_client::list_remote_timelines(
             remote_storage,
             self.conf,
             self.tenant_id,
         )
         .await?;
 
-        info!("found {} timelines", remote_timelines.len());
+        info!("found {} timelines", remote_timeline_ids.len());
 
-        let mut timeline_ancestors: HashMap<TimelineId, TimelineMetadata> = HashMap::new();
-        let mut index_parts: HashMap<TimelineId, IndexPart> = HashMap::new();
-        for (timeline_id, index_part) in remote_timelines {
-            let remote_metadata = index_part.parse_metadata().with_context(|| {
-                format!(
-                    "Failed to parse metadata file from remote storage for tenant {} timeline {}",
-                    self.tenant_id, timeline_id
-                )
-            })?;
+        // Download & parse index parts
+        let mut part_downloads = JoinSet::new();
+        for timeline_id in remote_timeline_ids {
+            let client = RemoteTimelineClient::new(
+                remote_storage.clone(),
+                self.conf,
+                self.tenant_id,
+                timeline_id,
+            );
+            part_downloads.spawn(
+                async move {
+                    debug!("starting index part download");
+
+                    let index_part = client
+                        .download_index_file()
+                        .await
+                        .context("download index file")?;
+
+                    let remote_metadata = index_part.parse_metadata().context("parse metadata")?;
+
+                    debug!("finished index part download");
+
+                    Result::<_, anyhow::Error>::Ok((
+                        timeline_id,
+                        client,
+                        index_part,
+                        remote_metadata,
+                    ))
+                }
+                .map(move |res| {
+                    res.with_context(|| format!("download index part for timeline {timeline_id}"))
+                })
+                .instrument(info_span!("download_index_part", timeline=%timeline_id)),
+            );
+        }
+        // Wait for all the download tasks to complete & collect results.
+        let mut remote_clients = HashMap::new();
+        let mut index_parts = HashMap::new();
+        let mut timeline_ancestors = HashMap::new();
+        while let Some(result) = part_downloads.join_next().await {
+            // NB: we already added timeline_id as context to the error
+            let result: Result<_, anyhow::Error> = result.context("joinset task join")?;
+            let (timeline_id, client, index_part, remote_metadata) = result?;
+            debug!("successfully downloaded index part for timeline {timeline_id}");
             timeline_ancestors.insert(timeline_id, remote_metadata);
             index_parts.insert(timeline_id, index_part);
+            remote_clients.insert(timeline_id, client);
         }
 
         // For every timeline, download the metadata file, scan the local directory,
@@ -680,7 +715,7 @@ impl Tenant {
                 timeline_id,
                 index_parts.remove(&timeline_id).unwrap(),
                 remote_metadata,
-                remote_storage.clone(),
+                remote_clients.remove(&timeline_id).unwrap(),
             )
             .await
             .with_context(|| {
@@ -723,22 +758,19 @@ impl Tenant {
         Ok(size)
     }
 
-    #[instrument(skip(self, index_part, remote_metadata, remote_storage), fields(timeline_id=%timeline_id))]
+    #[instrument(skip_all, fields(timeline_id=%timeline_id))]
     async fn load_remote_timeline(
         &self,
         timeline_id: TimelineId,
         index_part: IndexPart,
         remote_metadata: TimelineMetadata,
-        remote_storage: GenericRemoteStorage,
+        remote_client: RemoteTimelineClient,
     ) -> anyhow::Result<()> {
         info!("downloading index file for timeline {}", timeline_id);
         tokio::fs::create_dir_all(self.conf.timeline_path(&timeline_id, &self.tenant_id))
             .await
             .context("Failed to create new timeline directory")?;
 
-        let remote_client =
-            RemoteTimelineClient::new(remote_storage, self.conf, self.tenant_id, timeline_id)?;
-
         let ancestor = if let Some(ancestor_id) = remote_metadata.ancestor_timeline() {
             let timelines = self.timelines.lock().unwrap();
             Some(Arc::clone(timelines.get(&ancestor_id).ok_or_else(
@@ -757,7 +789,7 @@ impl Tenant {
         // cannot be older than the local one
         let local_metadata = None;
 
-        self.setup_timeline(
+        self.timeline_init_and_sync(
             timeline_id,
             Some(remote_client),
             Some(RemoteStartupData {
@@ -834,7 +866,7 @@ impl Tenant {
                 match tenant_clone.load().await {
                     Ok(()) => {}
                     Err(err) => {
-                        tenant_clone.set_broken();
+                        tenant_clone.set_broken(&err.to_string());
                         error!("could not load tenant {tenant_id}: {err:?}");
                     }
                 }
@@ -995,18 +1027,14 @@ impl Tenant {
             None
         };
 
-        let remote_client = self
-            .remote_storage
-            .as_ref()
-            .map(|remote_storage| {
-                RemoteTimelineClient::new(
-                    remote_storage.clone(),
-                    self.conf,
-                    self.tenant_id,
-                    timeline_id,
-                )
-            })
-            .transpose()?;
+        let remote_client = self.remote_storage.as_ref().map(|remote_storage| {
+            RemoteTimelineClient::new(
+                remote_storage.clone(),
+                self.conf,
+                self.tenant_id,
+                timeline_id,
+            )
+        });
 
         let remote_startup_data = match &remote_client {
             Some(remote_client) => match remote_client.download_index_file().await {
@@ -1026,7 +1054,7 @@ impl Tenant {
             None => None,
         };
 
-        self.setup_timeline(
+        self.timeline_init_and_sync(
             timeline_id,
             remote_client,
             remote_startup_data,
@@ -1474,7 +1502,7 @@ impl Tenant {
         });
     }
 
-    pub fn set_broken(&self) {
+    pub fn set_broken(&self, reason: &str) {
         self.state.send_modify(|current_state| {
             match *current_state {
                 TenantState::Active => {
@@ -1483,18 +1511,22 @@ impl Tenant {
                     // 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");
+                    warn!("Changing Active tenant to Broken state, reason: {}", reason);
                 }
                 TenantState::Broken => {
                     // This shouldn't happen either
-                    warn!("Tenant is already broken");
+                    warn!("Tenant is already in Broken state");
                 }
                 TenantState::Stopping => {
                     // This shouldn't happen either
                     *current_state = TenantState::Broken;
-                    warn!("Marking Stopping tenant as Broken");
+                    warn!(
+                        "Marking Stopping tenant as Broken state, reason: {}",
+                        reason
+                    );
                 }
                 TenantState::Loading | TenantState::Attaching => {
+                    info!("Setting tenant as Broken state, reason: {}", reason);
                     *current_state = TenantState::Broken;
                 }
             }
@@ -1696,6 +1728,7 @@ impl Tenant {
             remote_storage,
             state,
             cached_logical_sizes: tokio::sync::Mutex::new(HashMap::new()),
+            cached_synthetic_tenant_size: Arc::new(AtomicU64::new(0)),
         }
     }
 
@@ -1764,9 +1797,12 @@ impl Tenant {
         let mut target_config_file = VirtualFile::open_with_options(
             target_config_path,
             OpenOptions::new()
-                .truncate(true) // This needed for overwriting with small config files
+                // This needed for overwriting with small config files
+                .truncate(true)
                 .write(true)
-                .create_new(first_save),
+                .create_new(first_save)
+                // this will be ignored if create_new(true)
+                .create(true),
         )?;
 
         target_config_file
@@ -1848,7 +1884,12 @@ impl Tenant {
 
         utils::failpoint_sleep_millis_async!("gc_iteration_internal_after_getting_gc_timelines");
 
-        info!("starting on {} timelines", gc_timelines.len());
+        // If there is nothing to GC, we don't want any messages in the INFO log.
+        if !gc_timelines.is_empty() {
+            info!("{} timelines need GC", gc_timelines.len());
+        } else {
+            debug!("{} timelines need GC", gc_timelines.len());
+        }
 
         // Perform GC for each timeline.
         //
@@ -2139,13 +2180,12 @@ impl Tenant {
         let tenant_id = raw_timeline.owning_tenant.tenant_id;
         let unfinished_timeline = raw_timeline.raw_timeline()?;
 
-        tokio::task::block_in_place(|| {
-            import_datadir::import_timeline_from_postgres_datadir(
-                unfinished_timeline,
-                pgdata_path,
-                pgdata_lsn,
-            )
-        })
+        import_datadir::import_timeline_from_postgres_datadir(
+            unfinished_timeline,
+            pgdata_path,
+            pgdata_lsn,
+        )
+        .await
         .with_context(|| {
             format!("Failed to import pgdatadir for timeline {tenant_id}/{timeline_id}")
         })?;
@@ -2201,7 +2241,7 @@ impl Tenant {
                 self.conf,
                 tenant_id,
                 new_timeline_id,
-            )?;
+            );
             remote_client.init_upload_queue_for_empty_remote(&new_metadata)?;
             Some(remote_client)
         } else {
@@ -2329,6 +2369,24 @@ impl Tenant {
 
         size::gather_inputs(self, logical_sizes_at_once, &mut shared_cache).await
     }
+
+    /// Calculate synthetic tenant size
+    /// This is periodically called by background worker.
+    /// result is cached in tenant struct
+    #[instrument(skip_all, fields(tenant_id=%self.tenant_id))]
+    pub async fn calculate_synthetic_size(&self) -> anyhow::Result<u64> {
+        let inputs = self.gather_size_inputs().await?;
+
+        let size = inputs.calculate()?;
+
+        self.cached_synthetic_tenant_size
+            .store(size, Ordering::Relaxed);
+
+        Ok(size)
+    }
+    pub fn get_cached_synthetic_size(&self) -> u64 {
+        self.cached_synthetic_tenant_size.load(Ordering::Relaxed)
+    }
 }
 
 fn remove_timeline_and_uninit_mark(timeline_dir: &Path, uninit_mark: &Path) -> anyhow::Result<()> {
@@ -2572,8 +2630,10 @@ where
 pub mod harness {
     use bytes::{Bytes, BytesMut};
     use once_cell::sync::Lazy;
+    use once_cell::sync::OnceCell;
     use std::sync::{Arc, RwLock, RwLockReadGuard, RwLockWriteGuard};
     use std::{fs, path::PathBuf};
+    use utils::logging;
     use utils::lsn::Lsn;
 
     use crate::{
@@ -2637,6 +2697,8 @@ pub mod harness {
         ),
     }
 
+    static LOG_HANDLE: OnceCell<()> = OnceCell::new();
+
     impl<'a> TenantHarness<'a> {
         pub fn create(test_name: &'static str) -> anyhow::Result<Self> {
             Self::create_internal(test_name, false)
@@ -2651,6 +2713,10 @@ pub mod harness {
                 (Some(LOCK.read().unwrap()), None)
             };
 
+            LOG_HANDLE.get_or_init(|| {
+                logging::init(logging::LogFormat::Test).expect("Failed to init test logging")
+            });
+
             let repo_dir = PageServerConf::test_repo_dir(test_name);
             let _ = fs::remove_dir_all(&repo_dir);
             fs::create_dir_all(&repo_dir)?;
@@ -2786,15 +2852,15 @@ mod tests {
         drop(writer);
 
         assert_eq!(
-            tline.get(*TEST_KEY, Lsn(0x10)).no_ondemand_download()?,
+            tline.get(*TEST_KEY, Lsn(0x10)).await?,
             TEST_IMG("foo at 0x10")
         );
         assert_eq!(
-            tline.get(*TEST_KEY, Lsn(0x1f)).no_ondemand_download()?,
+            tline.get(*TEST_KEY, Lsn(0x1f)).await?,
             TEST_IMG("foo at 0x10")
         );
         assert_eq!(
-            tline.get(*TEST_KEY, Lsn(0x20)).no_ondemand_download()?,
+            tline.get(*TEST_KEY, Lsn(0x20)).await?,
             TEST_IMG("foo at 0x20")
         );
 
@@ -2873,15 +2939,15 @@ mod tests {
 
         // Check page contents on both branches
         assert_eq!(
-            from_utf8(&tline.get(TEST_KEY_A, Lsn(0x40)).no_ondemand_download()?)?,
+            from_utf8(&tline.get(TEST_KEY_A, Lsn(0x40)).await?)?,
             "foo at 0x40"
         );
         assert_eq!(
-            from_utf8(&newtline.get(TEST_KEY_A, Lsn(0x40)).no_ondemand_download()?)?,
+            from_utf8(&newtline.get(TEST_KEY_A, Lsn(0x40)).await?)?,
             "bar at 0x40"
         );
         assert_eq!(
-            from_utf8(&newtline.get(TEST_KEY_B, Lsn(0x40)).no_ondemand_download()?)?,
+            from_utf8(&newtline.get(TEST_KEY_B, Lsn(0x40)).await?)?,
             "foobar at 0x20"
         );
 
@@ -3040,10 +3106,7 @@ mod tests {
         tenant
             .gc_iteration(Some(TIMELINE_ID), 0x10, Duration::ZERO)
             .await?;
-        assert!(newtline
-            .get(*TEST_KEY, Lsn(0x25))
-            .no_ondemand_download()
-            .is_ok());
+        assert!(newtline.get(*TEST_KEY, Lsn(0x25)).await.is_ok());
 
         Ok(())
     }
@@ -3073,7 +3136,7 @@ mod tests {
 
         // Check that the data is still accessible on the branch.
         assert_eq!(
-            newtline.get(*TEST_KEY, Lsn(0x50)).no_ondemand_download()?,
+            newtline.get(*TEST_KEY, Lsn(0x50)).await?,
             TEST_IMG(&format!("foo at {}", Lsn(0x40)))
         );
 
@@ -3221,23 +3284,23 @@ mod tests {
         tline.compact().await?;
 
         assert_eq!(
-            tline.get(*TEST_KEY, Lsn(0x10)).no_ondemand_download()?,
+            tline.get(*TEST_KEY, Lsn(0x10)).await?,
             TEST_IMG("foo at 0x10")
         );
         assert_eq!(
-            tline.get(*TEST_KEY, Lsn(0x1f)).no_ondemand_download()?,
+            tline.get(*TEST_KEY, Lsn(0x1f)).await?,
             TEST_IMG("foo at 0x10")
         );
         assert_eq!(
-            tline.get(*TEST_KEY, Lsn(0x20)).no_ondemand_download()?,
+            tline.get(*TEST_KEY, Lsn(0x20)).await?,
             TEST_IMG("foo at 0x20")
         );
         assert_eq!(
-            tline.get(*TEST_KEY, Lsn(0x30)).no_ondemand_download()?,
+            tline.get(*TEST_KEY, Lsn(0x30)).await?,
             TEST_IMG("foo at 0x30")
         );
         assert_eq!(
-            tline.get(*TEST_KEY, Lsn(0x40)).no_ondemand_download()?,
+            tline.get(*TEST_KEY, Lsn(0x40)).await?,
             TEST_IMG("foo at 0x40")
         );
 
@@ -3347,7 +3410,7 @@ mod tests {
             for (blknum, last_lsn) in updated.iter().enumerate() {
                 test_key.field6 = blknum as u32;
                 assert_eq!(
-                    tline.get(test_key, lsn).no_ondemand_download()?,
+                    tline.get(test_key, lsn).await?,
                     TEST_IMG(&format!("{} at {}", blknum, last_lsn))
                 );
             }
@@ -3433,7 +3496,7 @@ mod tests {
             for (blknum, last_lsn) in updated.iter().enumerate() {
                 test_key.field6 = blknum as u32;
                 assert_eq!(
-                    tline.get(test_key, lsn).no_ondemand_download()?,
+                    tline.get(test_key, lsn).await?,
                     TEST_IMG(&format!("{} at {}", blknum, last_lsn))
                 );
             }
@@ -3508,7 +3571,7 @@ mod tests {
                 println!("checking [{idx}][{blknum}] at {lsn}");
                 test_key.field6 = blknum as u32;
                 assert_eq!(
-                    tline.get(test_key, *lsn).no_ondemand_download()?,
+                    tline.get(test_key, *lsn).await?,
                     TEST_IMG(&format!("{idx} {blknum} at {lsn}"))
                 );
             }

pageserver/src/tenant/config.rs

@@ -30,7 +30,7 @@ pub mod defaults {
     pub const DEFAULT_GC_HORIZON: u64 = 64 * 1024 * 1024;
     pub const DEFAULT_GC_PERIOD: &str = "100 s";
     pub const DEFAULT_IMAGE_CREATION_THRESHOLD: usize = 3;
-    pub const DEFAULT_PITR_INTERVAL: &str = "30 days";
+    pub const DEFAULT_PITR_INTERVAL: &str = "7 days";
     pub const DEFAULT_WALRECEIVER_CONNECT_TIMEOUT: &str = "2 seconds";
     pub const DEFAULT_WALRECEIVER_LAGGING_WAL_TIMEOUT: &str = "3 seconds";
     pub const DEFAULT_MAX_WALRECEIVER_LSN_WAL_LAG: u64 = 10 * 1024 * 1024;

pageserver/src/tenant/layer_map.rs

@@ -202,6 +202,13 @@ impl<T: ?Sized> PartialEq for LayerRTreeObject<T> {
         // references.  Clippy complains about this. In practice it
         // seems to work, the assertion below would be triggered
         // otherwise but this ought to be fixed.
+
+        {
+            let left = Arc::as_ptr(&self.layer);
+            let right = Arc::as_ptr(&other.layer);
+            tracing::info!(?left, ?right, "comparing ptr_eq");
+        }
+
         #[allow(clippy::vtable_address_comparisons)]
         Arc::ptr_eq(&self.layer, &other.layer)
     }
@@ -250,18 +257,37 @@ where
     L: ?Sized + Layer,
 {
     ///
-    /// Find the latest layer that covers the given 'key', with lsn <
-    /// 'end_lsn'.
+    /// Find the latest layer (by lsn.end) that covers the given
+    /// 'key', with lsn.start < 'end_lsn'.
     ///
-    /// Returns the layer, if any, and an 'lsn_floor' value that
-    /// indicates which portion of the layer the caller should
-    /// check. 'lsn_floor' is normally the start-LSN of the layer, but
-    /// can be greater if there is an overlapping layer that might
-    /// contain the version, even if it's missing from the returned
-    /// layer.
+    /// The caller of this function is the page reconstruction
+    /// algorithm looking for the next relevant delta layer, or
+    /// the terminal image layer. The caller will pass the lsn_floor
+    /// value as end_lsn in the next call to search.
+    ///
+    /// If there's an image layer exactly below the given end_lsn,
+    /// search should return that layer regardless if there are
+    /// overlapping deltas.
+    ///
+    /// If the latest layer is a delta and there is an overlapping
+    /// image with it below, the lsn_floor returned should be right
+    /// above that image so we don't skip it in the search. Otherwise
+    /// the lsn_floor returned should be the bottom of the delta layer
+    /// because we should make as much progress down the lsn axis
+    /// as possible. It's fine if this way we skip some overlapping
+    /// deltas, because the delta we returned would contain the same
+    /// wal content.
+    ///
+    /// TODO: This API is convoluted and inefficient. If the caller
+    /// makes N search calls, we'll end up finding the same latest
+    /// image layer N times. We should either cache the latest image
+    /// layer result, or simplify the api to `get_latest_image` and
+    /// `get_latest_delta`, and only call `get_latest_image` once.
+    ///
+    /// NOTE: This only searches the 'historic' layers, *not* the
+    /// 'open' and 'frozen' layers!
     ///
     pub fn search(&self, key: Key, end_lsn: Lsn) -> Option<SearchResult<L>> {
-        // linear search
         // Find the latest image layer that covers the given key
         let mut latest_img: Option<Arc<L>> = None;
         let mut latest_img_lsn: Option<Lsn> = None;
@@ -399,7 +425,9 @@ where
         NUM_ONDISK_LAYERS.dec();
     }
 
-    /// Is there a newer image layer for given key- and LSN-range?
+    /// Is there a newer image layer for given key- and LSN-range? Or a set
+    /// of image layers within the specified lsn range that cover the entire
+    /// specified key range?
     ///
     /// This is used for garbage collection, to determine if an old layer can
     /// be deleted.
@@ -486,8 +514,8 @@ where
 
     ///
     /// Divide the whole given range of keys into sub-ranges based on the latest
-    /// image layer that covers each range. (This is used when creating  new
-    /// image layers)
+    /// image layer that covers each range at the specified lsn (inclusive).
+    /// This is used when creating  new image layers.
     ///
     // FIXME: clippy complains that the result type is very complex. She's probably
     // right...
@@ -539,8 +567,15 @@ where
         Ok(ranges)
     }
 
-    /// Count how many L1 delta layers there are that overlap with the
-    /// given key and LSN range.
+    /// Count the height of the tallest stack of deltas in this 2d region.
+    ///
+    /// This number is used to compute the largest number of deltas that
+    /// we'll need to visit for any page reconstruction in this region.
+    /// We use this heuristic to decide whether to create an image layer.
+    ///
+    /// TODO currently we just return the total number of deltas in the
+    ///      region, no matter if they're stacked on top of each other
+    ///      or next to each other.
     pub fn count_deltas(&self, key_range: &Range<Key>, lsn_range: &Range<Lsn>) -> Result<usize> {
         let mut result = 0;
         if lsn_range.start >= lsn_range.end {

pageserver/src/tenant/mgr.rs

@@ -430,7 +430,7 @@ where
         Err(e) => {
             let tenants_accessor = TENANTS.read().await;
             match tenants_accessor.get(&tenant_id) {
-                Some(tenant) => tenant.set_broken(),
+                Some(tenant) => tenant.set_broken(&e.to_string()),
                 None => warn!("Tenant {tenant_id} got removed from memory"),
             }
             Err(e)
@@ -492,3 +492,53 @@ pub async fn immediate_gc(
 
     Ok(wait_task_done)
 }
+
+#[cfg(feature = "testing")]
+pub async fn immediate_compact(
+    tenant_id: TenantId,
+    timeline_id: TimelineId,
+) -> Result<tokio::sync::oneshot::Receiver<anyhow::Result<()>>, ApiError> {
+    let guard = TENANTS.read().await;
+
+    let tenant = guard
+        .get(&tenant_id)
+        .map(Arc::clone)
+        .with_context(|| format!("Tenant {tenant_id} not found"))
+        .map_err(ApiError::NotFound)?;
+
+    let timeline = tenant
+        .get_timeline(timeline_id, true)
+        .map_err(ApiError::NotFound)?;
+
+    // Run in task_mgr to avoid race with detach operation
+    let (task_done, wait_task_done) = tokio::sync::oneshot::channel();
+    task_mgr::spawn(
+        &tokio::runtime::Handle::current(),
+        TaskKind::Compaction,
+        Some(tenant_id),
+        Some(timeline_id),
+        &format!(
+            "timeline_compact_handler compaction run for tenant {tenant_id} timeline {timeline_id}"
+        ),
+        false,
+        async move {
+            let result = timeline
+                .compact()
+                .instrument(
+                    info_span!("manual_compact", tenant = %tenant_id, timeline = %timeline_id),
+                )
+                .await;
+
+            match task_done.send(result) {
+                Ok(_) => (),
+                Err(result) => error!("failed to send compaction result: {result:?}"),
+            }
+            Ok(())
+        },
+    );
+
+    // drop the guard until after we've spawned the task so that timeline shutdown will wait for the task
+    drop(guard);
+
+    Ok(wait_task_done)
+}

pageserver/src/tenant/remote_timeline_client.rs

@@ -16,7 +16,7 @@
 //! unless the pageserver is configured without remote storage.
 //!
 //! We allocate the client instance in [Timeline][`crate::tenant::Timeline`], i.e.,
-//! either in [`crate::tenant_mgr`] during startup or when creating a new
+//! either in [`crate::tenant::mgr`] during startup or when creating a new
 //! timeline.
 //! However, the client does not become ready for use until we've initialized its upload queue:
 //!
@@ -135,7 +135,7 @@
 //! - Initiate upload queue with that [`IndexPart`].
 //! - Reschedule all lost operations by comparing the local filesystem state
 //!   and remote state as per [`IndexPart`]. This is done in
-//!   [`Timeline::setup_timeline`] and [`Timeline::reconcile_with_remote`].
+//!   [`Timeline::timeline_init_and_sync`] and [`Timeline::reconcile_with_remote`].
 //!
 //! Note that if we crash during file deletion between the index update
 //! that removes the file from the list of files, and deleting the remote file,
@@ -214,7 +214,7 @@ use anyhow::ensure;
 use remote_storage::{DownloadError, GenericRemoteStorage};
 use std::ops::DerefMut;
 use tokio::runtime::Runtime;
-use tracing::{info, warn};
+use tracing::{debug, info, warn};
 use tracing::{info_span, Instrument};
 use utils::lsn::Lsn;
 
@@ -298,8 +298,8 @@ impl RemoteTimelineClient {
         conf: &'static PageServerConf,
         tenant_id: TenantId,
         timeline_id: TimelineId,
-    ) -> anyhow::Result<RemoteTimelineClient> {
-        Ok(RemoteTimelineClient {
+    ) -> RemoteTimelineClient {
+        RemoteTimelineClient {
             conf,
             runtime: &BACKGROUND_RUNTIME,
             tenant_id,
@@ -307,7 +307,7 @@ impl RemoteTimelineClient {
             storage_impl: remote_storage,
             upload_queue: Mutex::new(UploadQueue::Uninitialized),
             metrics: Arc::new(RemoteTimelineClientMetrics::new(&tenant_id, &timeline_id)),
-        })
+        }
     }
 
     /// Initialize the upload queue for a remote storage that already received
@@ -367,6 +367,10 @@ impl RemoteTimelineClient {
 
     /// Download index file
     pub async fn download_index_file(&self) -> Result<IndexPart, DownloadError> {
+        let _unfinished_gauge_guard = self
+            .metrics
+            .call_begin(&RemoteOpFileKind::Index, &RemoteOpKind::Download);
+
         download::download_index_part(
             self.conf,
             &self.storage_impl,
@@ -393,22 +397,27 @@ impl RemoteTimelineClient {
         layer_file_name: &LayerFileName,
         layer_metadata: &LayerFileMetadata,
     ) -> anyhow::Result<u64> {
-        let downloaded_size = download::download_layer_file(
-            self.conf,
-            &self.storage_impl,
-            self.tenant_id,
-            self.timeline_id,
-            layer_file_name,
-            layer_metadata,
-        )
-        .measure_remote_op(
-            self.tenant_id,
-            self.timeline_id,
-            RemoteOpFileKind::Layer,
-            RemoteOpKind::Download,
-            Arc::clone(&self.metrics),
-        )
-        .await?;
+        let downloaded_size = {
+            let _unfinished_gauge_guard = self
+                .metrics
+                .call_begin(&RemoteOpFileKind::Layer, &RemoteOpKind::Download);
+            download::download_layer_file(
+                self.conf,
+                &self.storage_impl,
+                self.tenant_id,
+                self.timeline_id,
+                layer_file_name,
+                layer_metadata,
+            )
+            .measure_remote_op(
+                self.tenant_id,
+                self.timeline_id,
+                RemoteOpFileKind::Layer,
+                RemoteOpKind::Download,
+                Arc::clone(&self.metrics),
+            )
+            .await?
+        };
 
         // Update the metadata for given layer file. The remote index file
         // might be missing some information for the file; this allows us
@@ -517,7 +526,7 @@ impl RemoteTimelineClient {
             metadata_bytes,
         );
         let op = UploadOp::UploadMetadata(index_part, disk_consistent_lsn);
-        self.update_upload_queue_unfinished_metric(1, &op);
+        self.calls_unfinished_metric_begin(&op);
         upload_queue.queued_operations.push_back(op);
         upload_queue.latest_files_changes_since_metadata_upload_scheduled = 0;
 
@@ -549,7 +558,7 @@ impl RemoteTimelineClient {
         upload_queue.latest_files_changes_since_metadata_upload_scheduled += 1;
 
         let op = UploadOp::UploadLayer(layer_file_name.clone(), layer_metadata.clone());
-        self.update_upload_queue_unfinished_metric(1, &op);
+        self.calls_unfinished_metric_begin(&op);
         upload_queue.queued_operations.push_back(op);
 
         info!(
@@ -601,7 +610,7 @@ impl RemoteTimelineClient {
             // schedule the actual deletions
             for name in names {
                 let op = UploadOp::Delete(RemoteOpFileKind::Layer, name.clone());
-                self.update_upload_queue_unfinished_metric(1, &op);
+                self.calls_unfinished_metric_begin(&op);
                 upload_queue.queued_operations.push_back(op);
                 info!("scheduled layer file deletion {}", name.file_name());
             }
@@ -675,7 +684,7 @@ impl RemoteTimelineClient {
             // We can launch this task. Remove it from the queue first.
             let next_op = upload_queue.queued_operations.pop_front().unwrap();
 
-            info!("starting op: {}", next_op);
+            debug!("starting op: {}", next_op);
 
             // Update the counters
             match next_op {
@@ -747,13 +756,13 @@ impl RemoteTimelineClient {
             // Note: We only check for the shutdown requests between retries, so
             // if a shutdown request arrives while we're busy uploading, in the
             // upload::upload:*() call below, we will wait not exit until it has
-            // finisheed. We probably could cancel the upload by simply dropping
+            // finished. We probably could cancel the upload by simply dropping
             // the Future, but we're not 100% sure if the remote storage library
             // is cancellation safe, so we don't dare to do that. Hopefully, the
             // upload finishes or times out soon enough.
             if task_mgr::is_shutdown_requested() {
                 info!("upload task cancelled by shutdown request");
-                self.update_upload_queue_unfinished_metric(-1, &task.op);
+                self.calls_unfinished_metric_end(&task.op);
                 self.stop();
                 return;
             }
@@ -867,7 +876,7 @@ impl RemoteTimelineClient {
                 task.op, retries
             );
         } else {
-            info!("remote task {} completed successfully", task.op);
+            debug!("remote task {} completed successfully", task.op);
         }
 
         // The task has completed succesfully. Remove it from the in-progress list.
@@ -901,22 +910,40 @@ impl RemoteTimelineClient {
             // Launch any queued tasks that were unblocked by this one.
             self.launch_queued_tasks(upload_queue);
         }
-        self.update_upload_queue_unfinished_metric(-1, &task.op);
+        self.calls_unfinished_metric_end(&task.op);
     }
 
-    fn update_upload_queue_unfinished_metric(&self, delta: i64, op: &UploadOp) {
-        let (file_kind, op_kind) = match op {
+    fn calls_unfinished_metric_impl(
+        &self,
+        op: &UploadOp,
+    ) -> Option<(RemoteOpFileKind, RemoteOpKind)> {
+        let res = match op {
             UploadOp::UploadLayer(_, _) => (RemoteOpFileKind::Layer, RemoteOpKind::Upload),
             UploadOp::UploadMetadata(_, _) => (RemoteOpFileKind::Index, RemoteOpKind::Upload),
             UploadOp::Delete(file_kind, _) => (*file_kind, RemoteOpKind::Delete),
             UploadOp::Barrier(_) => {
                 // we do not account these
-                return;
+                return None;
             }
         };
-        self.metrics
-            .unfinished_tasks(&file_kind, &op_kind)
-            .add(delta)
+        Some(res)
+    }
+
+    fn calls_unfinished_metric_begin(&self, op: &UploadOp) {
+        let (file_kind, op_kind) = match self.calls_unfinished_metric_impl(op) {
+            Some(x) => x,
+            None => return,
+        };
+        let guard = self.metrics.call_begin(&file_kind, &op_kind);
+        guard.will_decrement_manually(); // in unfinished_ops_metric_end()
+    }
+
+    fn calls_unfinished_metric_end(&self, op: &UploadOp) {
+        let (file_kind, op_kind) = match self.calls_unfinished_metric_impl(op) {
+            Some(x) => x,
+            None => return,
+        };
+        self.metrics.call_end(&file_kind, &op_kind);
     }
 
     fn stop(&self) {
@@ -967,7 +994,7 @@ impl RemoteTimelineClient {
 
                 // Tear down queued ops
                 for op in qi.queued_operations.into_iter() {
-                    self.update_upload_queue_unfinished_metric(-1, &op);
+                    self.calls_unfinished_metric_end(&op);
                     // Dropping UploadOp::Barrier() here will make wait_completion() return with an Err()
                     // which is exactly what we want to happen.
                     drop(op);

pageserver/src/tenant/remote_timeline_client/download.rs

@@ -8,10 +8,9 @@ use std::future::Future;
 use std::path::Path;
 
 use anyhow::{anyhow, Context};
-use futures::stream::{FuturesUnordered, StreamExt};
 use tokio::fs;
 use tokio::io::AsyncWriteExt;
-use tracing::{debug, error, info, info_span, warn, Instrument};
+use tracing::{error, info, warn};
 
 use crate::config::PageServerConf;
 use crate::tenant::storage_layer::LayerFileName;
@@ -175,7 +174,7 @@ pub async fn list_remote_timelines<'a>(
     storage: &'a GenericRemoteStorage,
     conf: &'static PageServerConf,
     tenant_id: TenantId,
-) -> anyhow::Result<Vec<(TimelineId, IndexPart)>> {
+) -> anyhow::Result<HashSet<TimelineId>> {
     let tenant_path = conf.timelines_path(&tenant_id);
     let tenant_storage_path = conf.remote_path(&tenant_path)?;
 
@@ -194,7 +193,6 @@ pub async fn list_remote_timelines<'a>(
     }
 
     let mut timeline_ids = HashSet::new();
-    let mut part_downloads = FuturesUnordered::new();
 
     for timeline_remote_storage_key in timelines {
         let object_name = timeline_remote_storage_key.object_name().ok_or_else(|| {
@@ -205,35 +203,22 @@ pub async fn list_remote_timelines<'a>(
             format!("failed to parse object name into timeline id '{object_name}'")
         })?;
 
-        // list_prefixes returns all files with the prefix. If we haven't seen this timeline ID
-        // yet, launch a download task for it.
-        if !timeline_ids.contains(&timeline_id) {
-            timeline_ids.insert(timeline_id);
-            let storage_clone = storage.clone();
-            part_downloads.push(async move {
-                (
-                    timeline_id,
-                    download_index_part(conf, &storage_clone, tenant_id, timeline_id)
-                        .instrument(info_span!("download_index_part", timeline=%timeline_id))
-                        .await,
-                )
-            });
-        }
+        // list_prefixes is assumed to return unique names. Ensure this here.
+        // NB: it's safer to bail out than warn-log this because the pageserver
+        //     needs to absolutely know about _all_ timelines that exist, so that
+        //     GC knows all the branchpoints. If we skipped over a timeline instead,
+        //     GC could delete a layer that's still needed by that timeline.
+        anyhow::ensure!(
+            !timeline_ids.contains(&timeline_id),
+            "list_prefixes contains duplicate timeline id {timeline_id}"
+        );
+        timeline_ids.insert(timeline_id);
     }
 
-    // Wait for all the download tasks to complete.
-    let mut timeline_parts = Vec::new();
-    while let Some((timeline_id, part_upload_result)) = part_downloads.next().await {
-        let index_part = part_upload_result
-            .with_context(|| format!("Failed to fetch index part for timeline {timeline_id}"))?;
-
-        debug!("Successfully fetched index part for timeline {timeline_id}");
-        timeline_parts.push((timeline_id, index_part));
-    }
-    Ok(timeline_parts)
+    Ok(timeline_ids)
 }
 
-pub async fn download_index_part(
+pub(super) async fn download_index_part(
     conf: &'static PageServerConf,
     storage: &GenericRemoteStorage,
     tenant_id: TenantId,

pageserver/src/tenant/remote_timeline_client/index.rs

@@ -83,11 +83,6 @@ where
     /// Additional metadata can might exist in `layer_metadata`.
     pub timeline_layers: HashSet<L>,
 
-    /// FIXME: unused field. This should be removed, but that changes the on-disk format,
-    /// so we need to make sure we're backwards-` (and maybe forwards-) compatible
-    /// First pass is to move it to Optional and the next would be its removal
-    missing_layers: Option<HashSet<L>>,
-
     /// Per layer file name metadata, which can be present for a present or missing layer file.
     ///
     /// Older versions of `IndexPart` will not have this property or have only a part of metadata
@@ -167,8 +162,6 @@ impl IndexPartUnclean {
         let IndexPartUnclean {
             version,
             timeline_layers,
-            // this is an unused field, ignore it on cleaning
-            missing_layers: _,
             layer_metadata,
             disk_consistent_lsn,
             metadata_bytes,
@@ -189,7 +182,6 @@ impl IndexPartUnclean {
                     }
                 })
                 .collect(),
-            missing_layers: None,
             layer_metadata: layer_metadata
                 .into_iter()
                 .filter_map(|(l, m)| l.into_clean().map(|l| (l, m)))
@@ -225,7 +217,6 @@ impl IndexPart {
         Self {
             version: Self::LATEST_VERSION,
             timeline_layers,
-            missing_layers: Some(HashSet::new()),
             layer_metadata,
             disk_consistent_lsn,
             metadata_bytes,
@@ -259,7 +250,6 @@ mod tests {
     fn v0_indexpart_is_parsed() {
         let example = r#"{
             "timeline_layers":["000000000000000000000000000000000000-FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF__0000000001696070-00000000016960E9"],
-            "missing_layers":["LAYER_FILE_NAME::test/not_a_real_layer_but_adding_coverage"],
             "disk_consistent_lsn":"0/16960E8",
             "metadata_bytes":[113,11,159,210,0,54,0,4,0,0,0,0,1,105,96,232,1,0,0,0,0,1,105,96,112,0,0,0,0,0,0,0,0,0,0,0,0,0,1,105,96,112,0,0,0,0,1,105,96,112,0,0,0,14,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0]
         }"#;
@@ -267,7 +257,6 @@ mod tests {
         let expected = IndexPart {
             version: 0,
             timeline_layers: HashSet::from(["000000000000000000000000000000000000-FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF__0000000001696070-00000000016960E9".parse().unwrap()]),
-            missing_layers: None, // disabled fields should not carry unused values further
             layer_metadata: HashMap::default(),
             disk_consistent_lsn: "0/16960E8".parse::<Lsn>().unwrap(),
             metadata_bytes: [113,11,159,210,0,54,0,4,0,0,0,0,1,105,96,232,1,0,0,0,0,1,105,96,112,0,0,0,0,0,0,0,0,0,0,0,0,0,1,105,96,112,0,0,0,0,1,105,96,112,0,0,0,14,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0].to_vec(),
@@ -283,7 +272,6 @@ mod tests {
         let example = r#"{
             "version":1,
             "timeline_layers":["000000000000000000000000000000000000-FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF__0000000001696070-00000000016960E9"],
-            "missing_layers":["LAYER_FILE_NAME::test/not_a_real_layer_but_adding_coverage"],
             "layer_metadata":{
                 "000000000000000000000000000000000000-FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF__0000000001696070-00000000016960E9": { "file_size": 25600000 },
                 "LAYER_FILE_NAME::test/not_a_real_layer_but_adding_coverage": { "file_size": 9007199254741001 }
@@ -296,7 +284,6 @@ mod tests {
             // note this is not verified, could be anything, but exists for humans debugging.. could be the git version instead?
             version: 1,
             timeline_layers: HashSet::from(["000000000000000000000000000000000000-FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF__0000000001696070-00000000016960E9".parse().unwrap()]),
-            missing_layers: None,
             layer_metadata: HashMap::from([
                 ("000000000000000000000000000000000000-FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF__0000000001696070-00000000016960E9".parse().unwrap(), IndexLayerMetadata {
                     file_size: Some(25600000),
@@ -322,6 +309,7 @@ mod tests {
         let example = r#"{
             "version":1,
             "timeline_layers":["000000000000000000000000000000000000-FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF__0000000001696070-00000000016960E9"],
+            "missing_layers":["This shouldn't fail deserialization"],
             "layer_metadata":{
                 "000000000000000000000000000000000000-FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF__0000000001696070-00000000016960E9": { "file_size": 25600000 },
                 "LAYER_FILE_NAME::test/not_a_real_layer_but_adding_coverage": { "file_size": 9007199254741001 }
@@ -346,7 +334,6 @@ mod tests {
             ]),
             disk_consistent_lsn: "0/16960E8".parse::<Lsn>().unwrap(),
             metadata_bytes: [112,11,159,210,0,54,0,4,0,0,0,0,1,105,96,232,1,0,0,0,0,1,105,96,112,0,0,0,0,0,0,0,0,0,0,0,0,0,1,105,96,112,0,0,0,0,1,105,96,112,0,0,0,14,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0].to_vec(),
-            missing_layers: None,
         };
 
         let part = serde_json::from_str::<IndexPartUnclean>(example).unwrap();

pageserver/src/tenant/size.rs

@@ -23,7 +23,13 @@ use tracing::*;
 pub struct ModelInputs {
     updates: Vec<Update>,
     retention_period: u64,
+
+    /// Relevant lsns per timeline.
+    ///
+    /// This field is not required for deserialization purposes, which is mostly used in tests. The
+    /// LSNs explain the outcome (updates) but are not needed in size calculation.
     #[serde_as(as = "HashMap<serde_with::DisplayFromStr, _>")]
+    #[serde(default)]
     timeline_inputs: HashMap<TimelineId, TimelineInputs>,
 }
 
@@ -32,6 +38,8 @@ pub struct ModelInputs {
 #[serde_with::serde_as]
 #[derive(Debug, serde::Serialize, serde::Deserialize)]
 struct TimelineInputs {
+    #[serde_as(as = "serde_with::DisplayFromStr")]
+    ancestor_lsn: Lsn,
     #[serde_as(as = "serde_with::DisplayFromStr")]
     last_record: Lsn,
     #[serde_as(as = "serde_with::DisplayFromStr")]
@@ -44,6 +52,116 @@ struct TimelineInputs {
     next_gc_cutoff: Lsn,
 }
 
+// Adjust BranchFrom sorting so that we always process ancestor
+// before descendants. This is needed to correctly calculate size of
+// descendant timelines.
+//
+// Note that we may have multiple BranchFroms at the same LSN, so we
+// need to sort them in the tree order.
+//
+// see updates_sort_with_branches_at_same_lsn test below
+fn sort_updates_in_tree_order(updates: Vec<Update>) -> anyhow::Result<Vec<Update>> {
+    let mut sorted_updates = Vec::with_capacity(updates.len());
+    let mut known_timelineids = HashSet::new();
+    let mut i = 0;
+    while i < updates.len() {
+        let curr_upd = &updates[i];
+
+        if let Command::BranchFrom(parent_id) = curr_upd.command {
+            let parent_id = match parent_id {
+                Some(parent_id) if known_timelineids.contains(&parent_id) => {
+                    // we have already processed ancestor
+                    // process this BranchFrom Update normally
+                    known_timelineids.insert(curr_upd.timeline_id);
+                    sorted_updates.push(*curr_upd);
+                    i += 1;
+                    continue;
+                }
+                None => {
+                    known_timelineids.insert(curr_upd.timeline_id);
+                    sorted_updates.push(*curr_upd);
+                    i += 1;
+                    continue;
+                }
+                Some(parent_id) => parent_id,
+            };
+
+            let mut j = i;
+
+            // we have not processed ancestor yet.
+            // there is a chance that it is at the same Lsn
+            if !known_timelineids.contains(&parent_id) {
+                let mut curr_lsn_branchfroms: HashMap<TimelineId, Vec<(TimelineId, usize)>> =
+                    HashMap::new();
+
+                // inspect all branchpoints at the same lsn
+                while j < updates.len() && updates[j].lsn == curr_upd.lsn {
+                    let lookahead_upd = &updates[j];
+                    j += 1;
+
+                    if let Command::BranchFrom(lookahead_parent_id) = lookahead_upd.command {
+                        match lookahead_parent_id {
+                            Some(lookahead_parent_id)
+                                if !known_timelineids.contains(&lookahead_parent_id) =>
+                            {
+                                // we have not processed ancestor yet
+                                // store it for later
+                                let es =
+                                    curr_lsn_branchfroms.entry(lookahead_parent_id).or_default();
+                                es.push((lookahead_upd.timeline_id, j));
+                            }
+                            _ => {
+                                // we have already processed ancestor
+                                // process this BranchFrom Update normally
+                                known_timelineids.insert(lookahead_upd.timeline_id);
+                                sorted_updates.push(*lookahead_upd);
+                            }
+                        }
+                    }
+                }
+
+                // process BranchFroms in the tree order
+                // check that we don't have a cycle if somet entry is orphan
+                // (this should not happen, but better to be safe)
+                let mut processed_some_entry = true;
+                while processed_some_entry {
+                    processed_some_entry = false;
+
+                    curr_lsn_branchfroms.retain(|parent_id, branchfroms| {
+                        if known_timelineids.contains(parent_id) {
+                            for (timeline_id, j) in branchfroms {
+                                known_timelineids.insert(*timeline_id);
+                                sorted_updates.push(updates[*j - 1]);
+                            }
+                            processed_some_entry = true;
+                            false
+                        } else {
+                            true
+                        }
+                    });
+                }
+
+                if !curr_lsn_branchfroms.is_empty() {
+                    // orphans are expected to be rare and transient between tenant reloads
+                    // for example, an broken ancestor without the child branch being broken.
+                    anyhow::bail!(
+                        "orphan branch(es) detected in BranchFroms: {curr_lsn_branchfroms:?}"
+                    );
+                }
+            }
+
+            assert!(j > i);
+            i = j;
+        } else {
+            // not a BranchFrom, keep the same order
+            sorted_updates.push(*curr_upd);
+            i += 1;
+        }
+    }
+
+    Ok(sorted_updates)
+}
+
 /// Gathers the inputs for the tenant sizing model.
 ///
 /// Tenant size does not consider the latest state, but only the state until next_gc_cutoff, which
@@ -68,19 +186,20 @@ pub(super) async fn gather_inputs(
     // our advantage with `?` error handling.
     let mut joinset = tokio::task::JoinSet::new();
 
-    let timelines = tenant
+    // refresh is needed to update gc related pitr_cutoff and horizon_cutoff
+    tenant
         .refresh_gc_info()
         .await
         .context("Failed to refresh gc_info before gathering inputs")?;
 
+    let timelines = tenant.list_timelines();
+
     if timelines.is_empty() {
-        // All timelines are below tenant's gc_horizon; alternative would be to use
-        // Tenant::list_timelines but then those gc_info's would not be updated yet, possibly
-        // missing GcInfo::retain_lsns or having obsolete values for cutoff's.
+        // perhaps the tenant has just been created, and as such doesn't have any data yet
         return Ok(ModelInputs {
             updates: vec![],
             retention_period: 0,
-            timeline_inputs: HashMap::new(),
+            timeline_inputs: HashMap::default(),
         });
     }
 
@@ -91,13 +210,25 @@ pub(super) async fn gather_inputs(
 
     let mut updates = Vec::new();
 
-    // record the per timline values used to determine `retention_period`
+    // record the per timeline values useful to debug the model inputs, also used to track
+    // ancestor_lsn without keeping a hold of Timeline
     let mut timeline_inputs = HashMap::with_capacity(timelines.len());
 
     // used to determine the `retention_period` for the size model
     let mut max_cutoff_distance = None;
 
+    // mapping from (TimelineId, Lsn) => if this branch point has been handled already via
+    // GcInfo::retain_lsns or if it needs to have its logical_size calculated.
+    let mut referenced_branch_froms = HashMap::<(TimelineId, Lsn), bool>::new();
+
     for timeline in timelines {
+        if !timeline.is_active() {
+            anyhow::bail!(
+                "timeline {} is not active, cannot calculate tenant_size now",
+                timeline.timeline_id
+            );
+        }
+
         let last_record_lsn = timeline.get_last_record_lsn();
 
         let (interesting_lsns, horizon_cutoff, pitr_cutoff, next_gc_cutoff) = {
@@ -163,13 +294,30 @@ pub(super) async fn gather_inputs(
 
         // all timelines branch from something, because it might be impossible to pinpoint
         // which is the tenant_size_model's "default" branch.
+
+        let ancestor_lsn = timeline.get_ancestor_lsn();
+
         updates.push(Update {
-            lsn: timeline.get_ancestor_lsn(),
+            lsn: ancestor_lsn,
             command: Command::BranchFrom(timeline.get_ancestor_timeline_id()),
             timeline_id: timeline.timeline_id,
         });
 
+        if let Some(parent_timeline_id) = timeline.get_ancestor_timeline_id() {
+            // refresh_gc_info will update branchpoints and pitr_cutoff but only do it for branches
+            // which are over gc_horizon. for example, a "main" branch which never received any
+            // updates apart from initdb not have branch points recorded.
+            referenced_branch_froms
+                .entry((parent_timeline_id, timeline.get_ancestor_lsn()))
+                .or_default();
+        }
+
         for (lsn, _kind) in &interesting_lsns {
+            // mark this visited so don't need to re-process this parent
+            *referenced_branch_froms
+                .entry((timeline.timeline_id, *lsn))
+                .or_default() = true;
+
             if let Some(size) = logical_size_cache.get(&(timeline.timeline_id, *lsn)) {
                 updates.push(Update {
                     lsn: *lsn,
@@ -185,22 +333,10 @@ pub(super) async fn gather_inputs(
             }
         }
 
-        // all timelines also have an end point if they have made any progress
-        if last_record_lsn > timeline.get_ancestor_lsn()
-            && !interesting_lsns
-                .iter()
-                .any(|(lsn, _)| lsn == &last_record_lsn)
-        {
-            updates.push(Update {
-                lsn: last_record_lsn,
-                command: Command::EndOfBranch,
-                timeline_id: timeline.timeline_id,
-            });
-        }
-
         timeline_inputs.insert(
             timeline.timeline_id,
             TimelineInputs {
+                ancestor_lsn,
                 last_record: last_record_lsn,
                 // this is not used above, because it might not have updated recently enough
                 latest_gc_cutoff: *timeline.get_latest_gc_cutoff_lsn(),
@@ -211,6 +347,80 @@ pub(super) async fn gather_inputs(
         );
     }
 
+    // iterate over discovered branch points and make sure we are getting logical sizes at those
+    // points.
+    for ((timeline_id, lsn), handled) in referenced_branch_froms.iter() {
+        if *handled {
+            continue;
+        }
+
+        let timeline_id = *timeline_id;
+        let lsn = *lsn;
+
+        match timeline_inputs.get(&timeline_id) {
+            Some(inputs) if inputs.ancestor_lsn == lsn => {
+                // we don't need an update at this branch point which is also point where
+                // timeline_id branch was branched from.
+                continue;
+            }
+            Some(_) => {}
+            None => {
+                // we should have this because we have iterated through all of the timelines
+                anyhow::bail!("missing timeline_input for {timeline_id}")
+            }
+        }
+
+        if let Some(size) = logical_size_cache.get(&(timeline_id, lsn)) {
+            updates.push(Update {
+                lsn,
+                timeline_id,
+                command: Command::Update(*size),
+            });
+
+            needed_cache.insert((timeline_id, lsn));
+        } else {
+            let timeline = tenant
+                .get_timeline(timeline_id, false)
+                .context("find referenced ancestor timeline")?;
+            let parallel_size_calcs = Arc::clone(limit);
+            joinset.spawn(calculate_logical_size(
+                parallel_size_calcs,
+                timeline.clone(),
+                lsn,
+            ));
+
+            if let Some(parent_id) = timeline.get_ancestor_timeline_id() {
+                // we should not find new ones because we iterated tenants all timelines
+                anyhow::ensure!(
+                    timeline_inputs.contains_key(&parent_id),
+                    "discovered new timeline {parent_id} (parent of {timeline_id})"
+                );
+            }
+        };
+    }
+
+    // finally add in EndOfBranch for all timelines where their last_record_lsn is not a branch
+    // point. this is needed by the model.
+    for (timeline_id, inputs) in timeline_inputs.iter() {
+        let lsn = inputs.last_record;
+
+        if referenced_branch_froms.contains_key(&(*timeline_id, lsn)) {
+            // this means that the (timeline_id, last_record_lsn) represents a branch point
+            // we do not want to add EndOfBranch updates for these points because it doesn't fit
+            // into the current tenant_size_model.
+            continue;
+        }
+
+        if lsn > inputs.ancestor_lsn {
+            // all timelines also have an end point if they have made any progress
+            updates.push(Update {
+                lsn,
+                command: Command::EndOfBranch,
+                timeline_id: *timeline_id,
+            });
+        }
+    }
+
     let mut have_any_error = false;
 
     while let Some(res) = joinset.join_next().await {
@@ -267,8 +477,13 @@ pub(super) async fn gather_inputs(
     // for branch points, which come as multiple updates at the same LSN, the Command::Update
     // is needed before a branch is made out of that branch Command::BranchFrom. this is
     // handled by the variant order in `Command`.
+    //
     updates.sort_unstable();
 
+    // And another sort to handle Command::BranchFrom ordering
+    // in case when there are multiple branches at the same LSN.
+    let sorted_updates = sort_updates_in_tree_order(updates)?;
+
     let retention_period = match max_cutoff_distance {
         Some(max) => max.0,
         None => {
@@ -277,7 +492,7 @@ pub(super) async fn gather_inputs(
     };
 
     Ok(ModelInputs {
-        updates,
+        updates: sorted_updates,
         retention_period,
         timeline_inputs,
     })
@@ -295,21 +510,23 @@ impl ModelInputs {
                 command: op,
                 timeline_id,
             } = update;
+
             let Lsn(now) = *lsn;
             match op {
                 Command::Update(sz) => {
-                    storage.insert_point(&Some(*timeline_id), "".into(), now, Some(*sz));
+                    storage.insert_point(&Some(*timeline_id), "".into(), now, Some(*sz))?;
                 }
                 Command::EndOfBranch => {
-                    storage.insert_point(&Some(*timeline_id), "".into(), now, None);
+                    storage.insert_point(&Some(*timeline_id), "".into(), now, None)?;
                 }
                 Command::BranchFrom(parent) => {
-                    storage.branch(parent, Some(*timeline_id));
+                    // This branch command may fail if it cannot find a parent to branch from.
+                    storage.branch(parent, Some(*timeline_id))?;
                 }
             }
         }
 
-        Ok(storage.calculate(self.retention_period).total_children())
+        Ok(storage.calculate(self.retention_period)?.total_children())
     }
 }
 
@@ -372,6 +589,7 @@ async fn calculate_logical_size(
 
     let size_res = timeline
         .spawn_ondemand_logical_size_calculation(lsn)
+        .instrument(info_span!("spawn_ondemand_logical_size_calculation"))
         .await?;
     Ok(TimelineAtLsnSizeResult(timeline, lsn, size_res))
 }
@@ -457,9 +675,146 @@ fn updates_sort() {
 fn verify_size_for_multiple_branches() {
     // this is generated from integration test test_tenant_size_with_multiple_branches, but this way
     // it has the stable lsn's
-    let doc = r#"{"updates":[{"lsn":"0/0","command":{"branch_from":null},"timeline_id":"cd9d9409c216e64bf580904facedb01b"},{"lsn":"0/176FA40","command":{"update":25763840},"timeline_id":"cd9d9409c216e64bf580904facedb01b"},{"lsn":"0/176FA40","command":{"branch_from":"cd9d9409c216e64bf580904facedb01b"},"timeline_id":"10b532a550540bc15385eac4edde416a"},{"lsn":"0/1819818","command":{"update":26075136},"timeline_id":"10b532a550540bc15385eac4edde416a"},{"lsn":"0/18B5E40","command":{"update":26427392},"timeline_id":"cd9d9409c216e64bf580904facedb01b"},{"lsn":"0/18D3DF0","command":{"update":26492928},"timeline_id":"cd9d9409c216e64bf580904facedb01b"},{"lsn":"0/18D3DF0","command":{"branch_from":"cd9d9409c216e64bf580904facedb01b"},"timeline_id":"230fc9d756f7363574c0d66533564dcc"},{"lsn":"0/220F438","command":{"update":25239552},"timeline_id":"230fc9d756f7363574c0d66533564dcc"}],"retention_period":131072,"timeline_inputs":{"cd9d9409c216e64bf580904facedb01b":{"last_record":"0/18D5E40","latest_gc_cutoff":"0/169ACF0","horizon_cutoff":"0/18B5E40","pitr_cutoff":"0/18B5E40","next_gc_cutoff":"0/18B5E40"},"10b532a550540bc15385eac4edde416a":{"last_record":"0/1839818","latest_gc_cutoff":"0/169ACF0","horizon_cutoff":"0/1819818","pitr_cutoff":"0/1819818","next_gc_cutoff":"0/1819818"},"230fc9d756f7363574c0d66533564dcc":{"last_record":"0/222F438","latest_gc_cutoff":"0/169ACF0","horizon_cutoff":"0/220F438","pitr_cutoff":"0/220F438","next_gc_cutoff":"0/220F438"}}}"#;
+    //
+    // timelineinputs have been left out, because those explain the inputs, but don't participate
+    // in further size calculations.
+    let doc = r#"{"updates":[{"lsn":"0/0","command":{"branch_from":null},"timeline_id":"cd9d9409c216e64bf580904facedb01b"},{"lsn":"0/176FA40","command":{"update":25763840},"timeline_id":"cd9d9409c216e64bf580904facedb01b"},{"lsn":"0/176FA40","command":{"branch_from":"cd9d9409c216e64bf580904facedb01b"},"timeline_id":"10b532a550540bc15385eac4edde416a"},{"lsn":"0/1819818","command":{"update":26075136},"timeline_id":"10b532a550540bc15385eac4edde416a"},{"lsn":"0/18B5E40","command":{"update":26427392},"timeline_id":"cd9d9409c216e64bf580904facedb01b"},{"lsn":"0/18D3DF0","command":{"update":26492928},"timeline_id":"cd9d9409c216e64bf580904facedb01b"},{"lsn":"0/18D3DF0","command":{"branch_from":"cd9d9409c216e64bf580904facedb01b"},"timeline_id":"230fc9d756f7363574c0d66533564dcc"},{"lsn":"0/220F438","command":{"update":25239552},"timeline_id":"230fc9d756f7363574c0d66533564dcc"}],"retention_period":131072}"#;
 
     let inputs: ModelInputs = serde_json::from_str(doc).unwrap();
 
     assert_eq!(inputs.calculate().unwrap(), 36_409_872);
 }
+
+#[test]
+fn updates_sort_with_branches_at_same_lsn() {
+    use std::str::FromStr;
+    use Command::{BranchFrom, EndOfBranch};
+
+    macro_rules! lsn {
+        ($e:expr) => {
+            Lsn::from_str($e).unwrap()
+        };
+    }
+
+    let ids = [
+        TimelineId::from_str("00000000000000000000000000000000").unwrap(),
+        TimelineId::from_str("11111111111111111111111111111111").unwrap(),
+        TimelineId::from_str("22222222222222222222222222222222").unwrap(),
+        TimelineId::from_str("33333333333333333333333333333333").unwrap(),
+        TimelineId::from_str("44444444444444444444444444444444").unwrap(),
+    ];
+
+    // issue https://github.com/neondatabase/neon/issues/3179
+    let commands = vec![
+        Update {
+            lsn: lsn!("0/0"),
+            command: BranchFrom(None),
+            timeline_id: ids[0],
+        },
+        Update {
+            lsn: lsn!("0/169AD58"),
+            command: Command::Update(25387008),
+            timeline_id: ids[0],
+        },
+        // next three are wrongly sorted, because
+        // ids[1] is branched from before ids[1] exists
+        // and ids[2] is branched from before ids[2] exists
+        Update {
+            lsn: lsn!("0/169AD58"),
+            command: BranchFrom(Some(ids[1])),
+            timeline_id: ids[3],
+        },
+        Update {
+            lsn: lsn!("0/169AD58"),
+            command: BranchFrom(Some(ids[0])),
+            timeline_id: ids[2],
+        },
+        Update {
+            lsn: lsn!("0/169AD58"),
+            command: BranchFrom(Some(ids[2])),
+            timeline_id: ids[1],
+        },
+        Update {
+            lsn: lsn!("0/1CA85B8"),
+            command: Command::Update(28925952),
+            timeline_id: ids[1],
+        },
+        Update {
+            lsn: lsn!("0/1CD85B8"),
+            command: Command::Update(29024256),
+            timeline_id: ids[1],
+        },
+        Update {
+            lsn: lsn!("0/1CD85B8"),
+            command: BranchFrom(Some(ids[1])),
+            timeline_id: ids[4],
+        },
+        Update {
+            lsn: lsn!("0/22DCE70"),
+            command: Command::Update(32546816),
+            timeline_id: ids[3],
+        },
+        Update {
+            lsn: lsn!("0/230CE70"),
+            command: EndOfBranch,
+            timeline_id: ids[3],
+        },
+    ];
+
+    let expected = vec![
+        Update {
+            lsn: lsn!("0/0"),
+            command: BranchFrom(None),
+            timeline_id: ids[0],
+        },
+        Update {
+            lsn: lsn!("0/169AD58"),
+            command: Command::Update(25387008),
+            timeline_id: ids[0],
+        },
+        Update {
+            lsn: lsn!("0/169AD58"),
+            command: BranchFrom(Some(ids[0])),
+            timeline_id: ids[2],
+        },
+        Update {
+            lsn: lsn!("0/169AD58"),
+            command: BranchFrom(Some(ids[2])),
+            timeline_id: ids[1],
+        },
+        Update {
+            lsn: lsn!("0/169AD58"),
+            command: BranchFrom(Some(ids[1])),
+            timeline_id: ids[3],
+        },
+        Update {
+            lsn: lsn!("0/1CA85B8"),
+            command: Command::Update(28925952),
+            timeline_id: ids[1],
+        },
+        Update {
+            lsn: lsn!("0/1CD85B8"),
+            command: Command::Update(29024256),
+            timeline_id: ids[1],
+        },
+        Update {
+            lsn: lsn!("0/1CD85B8"),
+            command: BranchFrom(Some(ids[1])),
+            timeline_id: ids[4],
+        },
+        Update {
+            lsn: lsn!("0/22DCE70"),
+            command: Command::Update(32546816),
+            timeline_id: ids[3],
+        },
+        Update {
+            lsn: lsn!("0/230CE70"),
+            command: EndOfBranch,
+            timeline_id: ids[3],
+        },
+    ];
+
+    let sorted_commands = sort_updates_in_tree_order(commands).unwrap();
+
+    assert_eq!(sorted_commands, expected);
+}

pageserver/src/tenant/storage_layer.rs

@@ -109,7 +109,7 @@ pub trait Layer: Send + Sync {
     /// See PageReconstructResult for possible return values. The collected data
     /// is appended to reconstruct_data; the caller should pass an empty struct
     /// on first call, or a struct with a cached older image of the page if one
-    /// is available. If this returns PageReconstructResult::Continue, look up
+    /// is available. If this returns ValueReconstructResult::Continue, look up
     /// the predecessor layer and call again with the same 'reconstruct_data' to
     /// collect more data.
     fn get_value_reconstruct_data(

pageserver/src/tenant/tasks.rs

@@ -52,6 +52,8 @@ async fn compaction_loop(tenant_id: TenantId) {
     info!("starting");
     TENANT_TASK_EVENTS.with_label_values(&["start"]).inc();
     async {
+        let mut first = true;
+
         loop {
             trace!("waking up");
 
@@ -68,10 +70,14 @@ async fn compaction_loop(tenant_id: TenantId) {
 
             let mut sleep_duration = tenant.get_compaction_period();
             if sleep_duration == Duration::ZERO {
-                info!("automatic compaction is disabled");
+                if first {
+                    info!("automatic compaction is disabled");
+                }
+                first = false;
                 // check again in 10 seconds, in case it's been enabled again.
                 sleep_duration = Duration::from_secs(10);
             } else {
+                first = true;
                 // Run compaction
                 if let Err(e) = tenant.compaction_iteration().await {
                     sleep_duration = wait_duration;
@@ -83,7 +89,7 @@ async fn compaction_loop(tenant_id: TenantId) {
             tokio::select! {
                 _ = task_mgr::shutdown_watcher() => {
                     info!("received cancellation request during idling");
-                    break ;
+                    break;
                 },
                 _ = tokio::time::sleep(sleep_duration) => {},
             }
@@ -103,6 +109,7 @@ async fn gc_loop(tenant_id: TenantId) {
     info!("starting");
     TENANT_TASK_EVENTS.with_label_values(&["start"]).inc();
     async {
+        let mut first = true;
         loop {
             trace!("waking up");
 
@@ -121,10 +128,14 @@ async fn gc_loop(tenant_id: TenantId) {
             let gc_horizon = tenant.get_gc_horizon();
             let mut sleep_duration = gc_period;
             if sleep_duration == Duration::ZERO {
-                info!("automatic GC is disabled");
+                if first {
+                    info!("automatic GC is disabled");
+                }
+                first = false;
                 // check again in 10 seconds, in case it's been enabled again.
                 sleep_duration = Duration::from_secs(10);
             } else {
+                first = true;
                 // Run gc
                 if gc_horizon > 0 {
                     if let Err(e) = tenant.gc_iteration(None, gc_horizon, tenant.get_pitr_interval()).await