Silence test failure caused by expected error in log.

Fix test_ondemand_download_large_rel if uploads are slow.
If the uploads after compaction happen slowly, they might have finished before the pageserver is shut down. The L0 files have been uploaded, so no data is lost, but then the query later in the test will need to download all the L0 files, and causes the test to fail because it specifically checks that download happens on-demand, not all at once.
2026-08-14 02:01:31 +00:00 · 2023-01-13 10:26:15 +02:00 · 2023-01-12 23:40:46 +02:00 · 2023-01-12 22:49:00 +02:00 · 2023-01-12 20:39:04 +02:00 · 2023-01-12 19:24:30 +02:00
85 changed files with 5087 additions and 3056 deletions
@@ -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.
+
@@ -123,8 +123,8 @@ runs:
          exit 1
        fi
        if [[ "${{ inputs.run_in_parallel }}" == "true" ]]; then
-          # -n8 uses eight processes to run tests via pytest-xdist
-          EXTRA_PARAMS="-n8 $EXTRA_PARAMS"
+          # -n4 uses four processes to run tests via pytest-xdist
+          EXTRA_PARAMS="-n4 $EXTRA_PARAMS"

          # --dist=loadgroup points tests marked with @pytest.mark.xdist_group
          # to the same worker to make @pytest.mark.order work with xdist
@@ -9,6 +9,7 @@ settings:
  authEndpoint: "http://console-staging.local/management/api/v2"
  domain: "*.eu-west-1.aws.neon.build"
  sentryEnvironment: "development"
+  wssPort: 8443

 # -- Additional labels for neon-proxy pods
 podLabels:
@@ -23,6 +24,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
@@ -9,6 +9,7 @@ settings:
  authEndpoint: "http://console-staging.local/management/api/v2"
  domain: "*.cloud.stage.neon.tech"
  sentryEnvironment: "development"
+  wssPort: 8443

 # -- Additional labels for neon-proxy pods
 podLabels:
@@ -23,6 +24,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
@@ -9,6 +9,7 @@ settings:
  authEndpoint: "http://console-staging.local/management/api/v2"
  domain: "*.us-east-2.aws.neon.build"
  sentryEnvironment: "development"
+  wssPort: 8443

 # -- Additional labels for neon-proxy pods
 podLabels:
@@ -23,6 +24,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
@@ -9,6 +9,7 @@ settings:
  authEndpoint: "http://console-release.local/management/api/v2"
  domain: "*.ap-southeast-1.aws.neon.tech"
  sentryEnvironment: "production"
+  wssPort: 8443

 # -- Additional labels for neon-proxy pods
 podLabels:
@@ -23,6 +24,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
@@ -9,6 +9,7 @@ settings:
  authEndpoint: "http://console-release.local/management/api/v2"
  domain: "*.eu-central-1.aws.neon.tech"
  sentryEnvironment: "production"
+  wssPort: 8443

 # -- Additional labels for neon-proxy pods
 podLabels:
@@ -23,6 +24,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
@@ -9,6 +9,7 @@ settings:
  authEndpoint: "http://console-release.local/management/api/v2"
  domain: "*.us-east-2.aws.neon.tech"
  sentryEnvironment: "production"
+  wssPort: 8443

 # -- Additional labels for neon-proxy pods
 podLabels:
@@ -23,6 +24,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
@@ -9,6 +9,7 @@ settings:
  authEndpoint: "http://console-release.local/management/api/v2"
  domain: "*.us-west-2.aws.neon.tech"
  sentryEnvironment: "production"
+  wssPort: 8443

 # -- Additional labels for neon-proxy pods
 podLabels:
@@ -23,6 +24,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
@@ -3,6 +3,7 @@ settings:
  authEndpoint: "http://console-release.local/management/api/v2"
  domain: "*.cloud.neon.tech"
  sentryEnvironment: "production"
+  wssPort: 8443

 podLabels:
  zenith_service: proxy-scram
@@ -16,6 +17,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
@@ -794,6 +794,8 @@ jobs:
    strategy:
      matrix:
        include: ${{fromJSON(needs.calculate-deploy-targets.outputs.matrix-include)}}
+    environment:
+      name: prod-old
    steps:
      - name: Checkout
        uses: actions/checkout@v3
@@ -839,7 +841,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
@@ -911,6 +915,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
@@ -950,6 +956,8 @@ jobs:
    strategy:
      matrix:
        include: ${{fromJSON(needs.calculate-deploy-targets.outputs.matrix-include)}}
+    environment:
+      name: prod-old
    env:
      KUBECONFIG: .kubeconfig
    steps:
@@ -975,8 +983,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
@@ -993,6 +1001,8 @@ jobs:
    strategy:
      matrix:
        include: ${{fromJSON(needs.calculate-deploy-targets.outputs.matrix-include)}}
+    environment:
+      name: prod-old
    env:
      KUBECONFIG: .kubeconfig
    steps:
@@ -1041,6 +1051,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
@@ -1056,19 +1068,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 ]
@@ -1088,6 +1100,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
@@ -1126,6 +1140,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
@@ -1141,7 +1157,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
@@ -1165,6 +1181,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
@@ -1700,6 +1700,19 @@ dependencies = [
 "tokio-io-timeout",
 ]

+[[package]]
+name = "hyper-tungstenite"
+version = "0.8.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d62004bcd4f6f85d9e2aa4206f1466ee67031f5ededcb6c6e62d48f9306ad879"
+dependencies = [
+ "hyper",
+ "pin-project",
+ "tokio",
+ "tokio-tungstenite",
+ "tungstenite",
+]
+
 [[package]]
 name = "iana-time-zone"
 version = "0.1.53"
@@ -2502,6 +2515,7 @@ dependencies = [
 "postgres-protocol",
 "rand",
 "serde",
+ "thiserror",
 "tokio",
 "tracing",
 "workspace_hack",
@@ -2657,6 +2671,7 @@ dependencies = [
 "hex",
 "hmac",
 "hyper",
+ "hyper-tungstenite",
 "itertools",
 "md5",
 "metrics",
@@ -2666,6 +2681,7 @@ dependencies = [
 "pq_proto",
 "rand",
 "rcgen",
+ "regex",
 "reqwest",
 "routerify",
 "rstest",
@@ -2677,6 +2693,7 @@ dependencies = [
 "sha2",
 "socket2",
 "thiserror",
+ "tls-listener",
 "tokio",
 "tokio-postgres",
 "tokio-postgres-rustls",
@@ -2686,6 +2703,7 @@ dependencies = [
 "url",
 "utils",
 "uuid",
+ "webpki-roots",
 "workspace_hack",
 "x509-parser",
 ]
@@ -3323,6 +3341,17 @@ dependencies = [
 "syn",
 ]

+[[package]]
+name = "sha-1"
+version = "0.10.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f5058ada175748e33390e40e872bd0fe59a19f265d0158daa551c5a88a76009c"
+dependencies = [
+ "cfg-if",
+ "cpufeatures",
+ "digest",
+]
+
 [[package]]
 name = "sha1"
 version = "0.10.5"
@@ -3687,10 +3716,24 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "cda74da7e1a664f795bb1f8a87ec406fb89a02522cf6e50620d016add6dbbf5c"

 [[package]]
-name = "tokio"
-version = "1.21.1"
+name = "tls-listener"
+version = "0.5.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "0020c875007ad96677dcc890298f4b942882c5d4eb7cc8f439fc3bf813dc9c95"
+checksum = "c9d4ff21187d434ac7709bfc7441ca88f63681247e5ad99f0f08c8c91ddc103d"
+dependencies = [
+ "futures-util",
+ "hyper",
+ "pin-project-lite",
+ "thiserror",
+ "tokio",
+ "tokio-rustls",
+]
+
+[[package]]
+name = "tokio"
+version = "1.24.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1d9f76183f91ecfb55e1d7d5602bd1d979e38a3a522fe900241cf195624d67ae"
 dependencies = [
 "autocfg",
 "bytes",
@@ -3698,12 +3741,11 @@ dependencies = [
 "memchr",
 "mio",
 "num_cpus",
- "once_cell",
 "pin-project-lite",
 "signal-hook-registry",
 "socket2",
 "tokio-macros",
- "winapi",
+ "windows-sys 0.42.0",
 ]

 [[package]]
@@ -3800,6 +3842,18 @@ dependencies = [
 "xattr",
 ]

+[[package]]
+name = "tokio-tungstenite"
+version = "0.17.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f714dd15bead90401d77e04243611caec13726c2408afd5b31901dfcdcb3b181"
+dependencies = [
+ "futures-util",
+ "log",
+ "tokio",
+ "tungstenite",
+]
+
 [[package]]
 name = "tokio-util"
 version = "0.7.4"
@@ -4026,6 +4080,25 @@ version = "0.2.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "59547bce71d9c38b83d9c0e92b6066c4253371f15005def0c30d9657f50c7642"

+[[package]]
+name = "tungstenite"
+version = "0.17.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e27992fd6a8c29ee7eef28fc78349aa244134e10ad447ce3b9f0ac0ed0fa4ce0"
+dependencies = [
+ "base64 0.13.1",
+ "byteorder",
+ "bytes",
+ "http",
+ "httparse",
+ "log",
+ "rand",
+ "sha-1",
+ "thiserror",
+ "url",
+ "utf-8",
+]
+
 [[package]]
 name = "typenum"
 version = "1.16.0"
@@ -4114,6 +4187,12 @@ version = "2.1.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "e8db7427f936968176eaa7cdf81b7f98b980b18495ec28f1b5791ac3bfe3eea9"

+[[package]]
+name = "utf-8"
+version = "0.7.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "09cc8ee72d2a9becf2f2febe0205bbed8fc6615b7cb429ad062dc7b7ddd036a9"
+
 [[package]]
 name = "utils"
 version = "0.1.0"
@@ -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",
@@ -29,7 +29,12 @@ 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

 #########################################################################################
 #
@@ -55,7 +60,10 @@ 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 && \
+    echo 'trusted = true' >> /usr/local/pgsql/share/extension/postgis_sfcgal.control

 #########################################################################################
 #
@@ -29,7 +29,12 @@ 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

 #########################################################################################
 #
@@ -55,7 +60,10 @@ 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 && \
+    echo 'trusted = true' >> /usr/local/pgsql/share/extension/postgis_sfcgal.control

 #########################################################################################
 #
@@ -118,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
@@ -131,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 ...
@@ -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)),
@@ -1,5 +1,6 @@
 use std::path::Path;
 use std::str::FromStr;
+use std::time::Instant;

 use anyhow::Result;
 use log::{info, log_enabled, warn, Level};
@@ -197,22 +198,18 @@ pub fn handle_roles(spec: &ComputeSpec, client: &mut Client) -> Result<()> {

 /// Reassign all dependent objects and delete requested roles.
 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");
        for op in ops {
            if op.action == "delete_role" {
                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 +220,7 @@ pub fn handle_role_deletions(node: &ComputeNode, client: &mut Client) -> Result<
                xact.execute(query.as_str(), &[])?;
            }
        }
+        xact.commit()?;
    }

    Ok(())
@@ -317,6 +315,7 @@ pub fn handle_databases(spec: &ComputeSpec, client: &mut Client) -> Result<()> {
        // 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);

+        let start_time = Instant::now();
        if let Some(r) = pg_db {
            // XXX: db owner name is returned as quoted string from Postgres,
            // when quoting is needed.
@@ -335,6 +334,8 @@ pub fn handle_databases(spec: &ComputeSpec, client: &mut Client) -> Result<()> {
                info_print!(" -> update");

                client.execute(query.as_str(), &[])?;
+                let elapsed = start_time.elapsed().as_millis();
+                info_print!(" ({} ms)", elapsed);
            }
        } else {
            let mut query: String = format!("CREATE DATABASE {} ", name.pg_quote());
@@ -342,6 +343,9 @@ pub fn handle_databases(spec: &ComputeSpec, client: &mut Client) -> Result<()> {

            query.push_str(&db.to_pg_options());
            client.execute(query.as_str(), &[])?;
+
+            let elapsed = start_time.elapsed().as_millis();
+            info_print!(" ({} ms)", elapsed);
        }

        info_print!("\n");
@@ -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)
@@ -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")
@@ -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)
@@ -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 {
@@ -13,5 +13,6 @@ rand = "0.8.3"
 serde = { version = "1.0", features = ["derive"] }
 tokio = { version = "1.17", features = ["macros"] }
 tracing = "0.1"
+thiserror = "1.0"

 workspace_hack = { version = "0.1", path = "../../workspace_hack" }
@@ -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,7 +685,7 @@ fn read_cstr(buf: &mut Bytes) -> anyhow::Result<Bytes> {
    Ok(result)
 }

-const SQLSTATE_INTERNAL_ERROR: &str = "XX000\0";
+pub const SQLSTATE_INTERNAL_ERROR: &[u8; 5] = b"XX000";

 impl<'a> BeMessage<'a> {
    /// Write message to the given buf.
@@ -767,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| {
@@ -778,7 +834,9 @@ impl<'a> BeMessage<'a> {
                    buf.put_slice(b"ERROR\0");

                    buf.put_u8(b'C'); // SQLSTATE error code
-                    buf.put_slice(SQLSTATE_INTERNAL_ERROR.as_bytes());
+                    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)?;
@@ -801,7 +859,7 @@ impl<'a> BeMessage<'a> {
                    buf.put_slice(b"NOTICE\0");

                    buf.put_u8(b'C'); // SQLSTATE error code
-                    buf.put_slice(SQLSTATE_INTERNAL_ERROR.as_bytes());
+                    buf.put_slice(&terminate_code(SQLSTATE_INTERNAL_ERROR));

                    buf.put_u8(b'M'); // the message
                    write_cstr(error_msg.as_bytes(), buf)?;
@@ -1089,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
+}
@@ -3,8 +3,9 @@
 //! 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 serde::{Deserialize, Serialize};
@@ -21,20 +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 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 {
@@ -66,7 +79,7 @@ impl FromStr for AuthType {
        match s {
            "Trust" => Ok(Self::Trust),
            "NeonJWT" => Ok(Self::NeonJWT),
-            _ => bail!("invalid value \"{s}\" for auth type"),
+            _ => anyhow::bail!("invalid value \"{s}\" for auth type"),
        }
    }
 }
@@ -154,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())
 }
@@ -188,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"),
        }
    }

@@ -215,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::*;
@@ -223,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.
@@ -246,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);
@@ -254,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();
@@ -263,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,
@@ -274,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);
                }
            }
        }
@@ -295,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");
            }
        }
    }
@@ -305,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();
@@ -339,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
@@ -379,8 +399,11 @@ impl PostgresBackend {
                        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);
                        }
                    }
                }
@@ -394,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)?;
            }
@@ -445,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
@@ -468,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:?}"
+                )));
            }
        }

@@ -4,39 +4,87 @@
 //! is rather narrow, but we can extend it once required.

 use crate::postgres_backend::AuthType;
-use anyhow::{bail, Context, Result};
+use anyhow::Context;
 use bytes::{Buf, Bytes, BytesMut};
-use pq_proto::{BeMessage, FeMessage, FeStartupPacket};
+use pq_proto::{BeMessage, ConnectionError, FeMessage, FeStartupPacket, SQLSTATE_INTERNAL_ERROR};
 use std::future::Future;
+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 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 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")))
    }
 }

@@ -70,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),
@@ -90,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),
@@ -103,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),
@@ -139,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())
 }
@@ -149,7 +194,7 @@ 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 {
@@ -167,17 +212,18 @@ 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<()> {
+    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);
@@ -187,7 +233,7 @@ impl PostgresBackend {
    }

    /// 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)
    }
@@ -223,7 +269,11 @@ impl PostgresBackend {
    }

    // 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,
@@ -237,7 +287,7 @@ impl PostgresBackend {
        &mut self,
        handler: &mut impl Handler,
        shutdown_watcher: F,
-    ) -> Result<()>
+    ) -> Result<(), QueryError>
    where
        F: Fn() -> S,
        S: Future,
@@ -273,7 +323,7 @@ impl PostgresBackend {
                        return Ok(());
                    }
                }
-                Ok::<(), anyhow::Error>(())
+                Ok::<(), QueryError>(())
            } => {
                // Handshake complete.
                result?;
@@ -318,14 +368,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 {
@@ -348,8 +398,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
@@ -389,8 +444,11 @@ impl PostgresBackend {
                        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);
                        }
                    }
                }
@@ -413,33 +471,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)?;
            }
@@ -464,11 +517,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
@@ -487,7 +542,10 @@ 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
+                )));
            }
        }

@@ -555,3 +613,28 @@ impl<'a> AsyncWrite for CopyDataWriter<'a> {
        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:?}");
+        }
+    }
+}
@@ -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!()
        }
    }
@@ -10,11 +10,10 @@
 //! 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, ensure, Context, Result};
 use bytes::{BufMut, BytesMut};
 use fail::fail_point;
 use std::fmt::Write as FmtWrite;
-use std::sync::Arc;
 use std::time::SystemTime;
 use tokio::io;
 use tokio::io::AsyncWrite;
@@ -28,7 +27,8 @@ use tracing::*;
 ///
 use tokio_tar::{Builder, EntryType, Header};

-use crate::tenant::{with_ondemand_download, Timeline};
+use crate::tenant::TimelineRequestContext;
+use crate::tenant::{PageReconstructError, Timeline};
 use pageserver_api::reltag::{RelTag, SlruKind};

 use postgres_ffi::pg_constants::{DEFAULTTABLESPACE_OID, GLOBALTABLESPACE_OID};
@@ -39,114 +39,135 @@ 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,
+    ctx: &'a TimelineRequestContext,
+) -> Result<(), PageReconstructError>
+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) && backup_prev != provided_prev_lsn {
+            return Err(PageReconstructError::Other(anyhow!(
+                "prev LSN doesn't match"
+            )));
+        }
+        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,
+        ctx,
+    };
+    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: AsyncWrite + Send + Sync + Unpin,
 {
    ar: Builder<&'a mut W>,
-    timeline: &'a Arc<Timeline>,
-    pub lsn: Lsn,
+    timeline: &'a Timeline,
+    lsn: Lsn,
    prev_record_lsn: Lsn,
    full_backup: bool,
+    ctx: &'a TimelineRequestContext,
 }

-// 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: AsyncWrite + Send + Sync + Unpin,
 {
-    pub fn new(
-        write: &'a mut 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_non_terminated(write),
-            timeline,
-            lsn: backup_lsn,
-            prev_record_lsn: prev_lsn,
-            full_backup,
-        })
-    }
-
-    pub async fn send_tarball(mut self) -> anyhow::Result<()> {
+    async fn send_tarball(mut self) -> Result<(), PageReconstructError> {
        // 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()).await?;
+            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).await?;
+                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()).await?;
+                self.ar
+                    .append(&header, &mut io::empty())
+                    .await
+                    .context("could not add config file to basebackup tarball")?;
            }
        }

@@ -156,8 +177,10 @@ where
            SlruKind::MultiXactOffsets,
            SlruKind::MultiXactMembers,
        ] {
-            for segno in
-                with_ondemand_download(|| self.timeline.list_slru_segments(kind, self.lsn)).await?
+            for segno in self
+                .timeline
+                .list_slru_segments(kind, self.lsn, self.ctx)
+                .await?
            {
                self.add_slru_segment(kind, segno).await?;
            }
@@ -165,44 +188,56 @@ where

        // Create tablespace directories
        for ((spcnode, dbnode), has_relmap_file) in
-            with_ondemand_download(|| self.timeline.list_dbdirs(self.lsn)).await?
+            self.timeline.list_dbdirs(self.lsn, self.ctx).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(|| self.timeline.list_rels(spcnode, dbnode, self.lsn))
-                        .await?
+                for rel in self
+                    .timeline
+                    .list_rels(spcnode, dbnode, self.lsn, self.ctx)
+                    .await?
                {
                    self.add_rel(rel).await?;
                }
            }
        }
-        for xid in with_ondemand_download(|| self.timeline.list_twophase_files(self.lsn)).await? {
+        for xid in self
+            .timeline
+            .list_twophase_files(self.lsn, self.ctx)
+            .await?
+        {
            self.add_twophase_file(xid).await?;
        }

        fail_point!("basebackup-before-control-file", |_| {
-            bail!("failpoint basebackup-before-control-file")
+            Err(PageReconstructError::from(anyhow!(
+                "failpoint basebackup-before-control-file"
+            )))
        });

        // Generate pg_control and bootstrap WAL segment.
        self.add_pgcontrol_file().await?;
-        self.ar.finish().await?;
+        self.ar.finish().await.context("could not finish tarball")?;
        debug!("all tarred up!");
        Ok(())
    }

-    async fn add_rel(&mut self, tag: RelTag) -> anyhow::Result<()> {
-        let nblocks =
-            with_ondemand_download(|| self.timeline.get_rel_size(tag, self.lsn, false)).await?;
+    async fn add_rel(&mut self, tag: RelTag) -> Result<(), PageReconstructError> {
+        let nblocks = self
+            .timeline
+            .get_rel_size(tag, self.lsn, false, self.ctx)
+            .await?;

        // If the relation is empty, create an empty file
        if nblocks == 0 {
            let file_name = tag.to_segfile_name(0);
            let header = new_tar_header(&file_name, 0)?;
-            self.ar.append(&header, &mut io::empty()).await?;
+            self.ar
+                .append(&header, &mut io::empty())
+                .await
+                .context("could not write empty relfile to tar stream")?;
            return Ok(());
        }

@@ -211,20 +246,21 @@ where
        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 startblk..endblk {
-                let img = with_ondemand_download(|| {
-                    self.timeline
-                        .get_rel_page_at_lsn(tag, blknum, self.lsn, false)
-                })
-                .await?;
+                let img = self
+                    .timeline
+                    .get_rel_page_at_lsn(tag, blknum, self.lsn, false, self.ctx)
+                    .await?;
                segment_data.extend_from_slice(&img[..]);
            }

            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?;
+            self.ar
+                .append(&header, segment_data.as_slice())
+                .await
+                .context("could not write relfile segment to tar stream")?;

            seg += 1;
            startblk = endblk;
@@ -237,17 +273,17 @@ where
    // Generate SLRU segment files from repository.
    //
    async fn add_slru_segment(&mut self, slru: SlruKind, segno: u32) -> anyhow::Result<()> {
-        let nblocks =
-            with_ondemand_download(|| self.timeline.get_slru_segment_size(slru, segno, self.lsn))
-                .await?;
+        let nblocks = self
+            .timeline
+            .get_slru_segment_size(slru, segno, self.lsn, self.ctx)
+            .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(|| {
-                self.timeline
-                    .get_slru_page_at_lsn(slru, segno, blknum, self.lsn)
-            })
-            .await?;
+            let img = self
+                .timeline
+                .get_slru_page_at_lsn(slru, segno, blknum, self.lsn, self.ctx)
+                .await?;

            if slru == SlruKind::Clog {
                ensure!(img.len() == BLCKSZ as usize || img.len() == BLCKSZ as usize + 8);
@@ -279,9 +315,10 @@ where
        has_relmap_file: bool,
    ) -> anyhow::Result<()> {
        let relmap_img = if has_relmap_file {
-            let img =
-                with_ondemand_download(|| self.timeline.get_relmap_file(spcnode, dbnode, self.lsn))
-                    .await?;
+            let img = self
+                .timeline
+                .get_relmap_file(spcnode, dbnode, self.lsn, self.ctx)
+                .await?;
            ensure!(img.len() == 512);
            Some(img)
        } else {
@@ -314,7 +351,9 @@ where
            // XLOG_TBLSPC_DROP records. But we probably should just
            // throw an error on CREATE TABLESPACE in the first place.
            if !has_relmap_file
-                && with_ondemand_download(|| self.timeline.list_rels(spcnode, dbnode, self.lsn))
+                && self
+                    .timeline
+                    .list_rels(spcnode, dbnode, self.lsn, self.ctx)
                    .await?
                    .is_empty()
            {
@@ -347,7 +386,10 @@ where
    // Extract twophase state files
    //
    async fn add_twophase_file(&mut self, xid: TransactionId) -> anyhow::Result<()> {
-        let img = with_ondemand_download(|| self.timeline.get_twophase_file(xid, self.lsn)).await?;
+        let img = self
+            .timeline
+            .get_twophase_file(xid, self.lsn, self.ctx)
+            .await?;

        let mut buf = BytesMut::new();
        buf.extend_from_slice(&img[..]);
@@ -383,10 +425,14 @@ where
            )
            .await?;

-        let checkpoint_bytes = with_ondemand_download(|| self.timeline.get_checkpoint(self.lsn))
+        let checkpoint_bytes = self
+            .timeline
+            .get_checkpoint(self.lsn, self.ctx)
            .await
            .context("failed to get checkpoint bytes")?;
-        let pg_control_bytes = with_ondemand_download(|| self.timeline.get_control_file(self.lsn))
+        let pg_control_bytes = self
+            .timeline
+            .get_control_file(self.lsn, self.ctx)
            .await
            .context("failed get control bytes")?;

@@ -13,8 +13,8 @@ use tracing::*;
 use metrics::set_build_info_metric;
 use pageserver::{
    config::{defaults::*, PageServerConf},
+    context::{DownloadBehavior, RequestContext, TaskKind},
    http, page_cache, page_service, task_mgr,
-    task_mgr::TaskKind,
    task_mgr::{
        BACKGROUND_RUNTIME, COMPUTE_REQUEST_RUNTIME, MGMT_REQUEST_RUNTIME, WALRECEIVER_RUNTIME,
    },
@@ -303,62 +303,79 @@ fn start_pageserver(conf: &'static PageServerConf) -> anyhow::Result<()> {
    {
        let _rt_guard = MGMT_REQUEST_RUNTIME.enter();

+        let mgmt_ctx = RequestContext::new(TaskKind::HttpEndpointListener, DownloadBehavior::Error);
+        let cancellation_token = Box::leak(Box::new(mgmt_ctx.cancellation_token().clone()));
        let router = http::make_router(conf, auth.clone(), remote_storage)?
            .build()
            .map_err(|err| anyhow!(err))?;
        let service = utils::http::RouterService::new(router).unwrap();
        let server = hyper::Server::from_tcp(http_listener)?
            .serve(service)
-            .with_graceful_shutdown(task_mgr::shutdown_watcher());
+            .with_graceful_shutdown(cancellation_token.cancelled());

        task_mgr::spawn(
            MGMT_REQUEST_RUNTIME.handle(),
-            TaskKind::HttpEndpointListener,
-            None,
-            None,
            "http endpoint listener",
            true,
            async {
-                server.await?;
-                Ok(())
+                match server.await {
+                    Ok(()) => info!("HTTP endpoint listener shut down"),
+                    Err(err) => error!("HTTP endpoint listener shut down with error: {err:?}"),
+                }
            },
        );
+    }

-        if let Some(metric_collection_endpoint) = &conf.metric_collection_endpoint {
-            task_mgr::spawn(
-                MGMT_REQUEST_RUNTIME.handle(),
-                TaskKind::MetricsCollection,
-                None,
-                None,
-                "consumption metrics collection",
-                true,
-                async move {
-                    pageserver::consumption_metrics::collect_metrics(
-                        metric_collection_endpoint,
-                        conf.metric_collection_interval,
-                        conf.id,
-                    )
-                    .instrument(info_span!("metrics_collection"))
-                    .await?;
-                    Ok(())
-                },
-            );
-        }
+    if let Some(metric_collection_endpoint) = &conf.metric_collection_endpoint {
+        let metrics_ctx = RequestContext::new(
+            TaskKind::MetricsCollection,
+            DownloadBehavior::Error, // metrics collector shouldn't be downloading anything
+        );
+        task_mgr::spawn(
+            MGMT_REQUEST_RUNTIME.handle(),
+            "consumption metrics collection",
+            true,
+            pageserver::consumption_metrics::collect_metrics(
+                metric_collection_endpoint,
+                conf.metric_collection_interval,
+                conf.id,
+                metrics_ctx,
+            )
+            .instrument(info_span!("metrics_collection")),
+        );
    }

    // Spawn a task to listen for libpq connections. It will spawn further tasks
    // for each connection. We created the listener earlier already.
-    task_mgr::spawn(
-        COMPUTE_REQUEST_RUNTIME.handle(),
-        TaskKind::LibpqEndpointListener,
-        None,
-        None,
-        "libpq endpoint listener",
-        true,
-        async move {
-            page_service::libpq_listener_main(conf, auth, pageserver_listener, conf.auth_type).await
-        },
-    );
+    {
+        let libpq_ctx = RequestContext::new(
+            TaskKind::LibpqEndpointListener,
+            // listener task shouldn't need to download anything. (We will
+            // create a separate sub-contexts for each connection, with their
+            // own download behavior. This context is used only to listen and
+            // accept connections.)
+            DownloadBehavior::Error,
+        );
+        task_mgr::spawn(
+            COMPUTE_REQUEST_RUNTIME.handle(),
+            "libpq endpoint listener",
+            true,
+            async move {
+                match page_service::libpq_listener_main(
+                    conf,
+                    auth,
+                    pageserver_listener,
+                    conf.auth_type,
+                    libpq_ctx,
+                )
+                .await
+                {
+                    Ok(()) => info!("libpq endpoint listener shut down"),
+                    Err(err) => error!("libpq endpoint listener shut down with error: {err:?}"),
+                }
+            },
+        );
+    }

    // All started up! Now just sit and wait for shutdown signal.
    signals.handle(|signal| match signal {
@@ -375,7 +392,7 @@ fn start_pageserver(conf: &'static PageServerConf) -> anyhow::Result<()> {
                "Got {}. Terminating gracefully in fast shutdown mode",
                signal.name()
            );
-            BACKGROUND_RUNTIME.block_on(pageserver::shutdown_pageserver(0));
+            BACKGROUND_RUNTIME.block_on(task_mgr::shutdown_pageserver(0));
            unreachable!()
        }
    })
@@ -9,9 +9,8 @@ use tracing::*;
 use utils::id::NodeId;
 use utils::id::TimelineId;

-use crate::task_mgr;
+use crate::context::RequestContext;
 use crate::tenant::mgr;
-use pageserver_api::models::TenantState;
 use utils::id::TenantId;

 use serde::{Deserialize, Serialize};
@@ -138,12 +137,13 @@ struct EventChunk<'a> {
    events: &'a [ConsumptionMetric],
 }

-/// Main thread that serves metrics collection
+/// Main task that serves metrics collection
 pub async fn collect_metrics(
    metric_collection_endpoint: &Url,
    metric_collection_interval: Duration,
    node_id: NodeId,
-) -> anyhow::Result<()> {
+    metrics_ctx: RequestContext,
+) {
    let mut ticker = tokio::time::interval(metric_collection_interval);

    info!("starting collect_metrics");
@@ -154,12 +154,15 @@ pub async fn collect_metrics(

    loop {
        tokio::select! {
-            _ = task_mgr::shutdown_watcher() => {
+            _ = metrics_ctx.cancelled() => {
                info!("collect_metrics received cancellation request");
-                return Ok(());
+                return;
            },
            _ = ticker.tick() => {
-                collect_metrics_task(&client, &mut cached_metrics, metric_collection_endpoint, node_id).await?;
+                if let Err(err) = collect_metrics_task(&client, &mut cached_metrics, metric_collection_endpoint, node_id, &metrics_ctx).await {
+                    // Log the error and continue
+                    error!("metrics collection failed: {err:?}");
+                }
            }
        }
    }
@@ -174,6 +177,7 @@ pub async fn collect_metrics_task(
    cached_metrics: &mut HashMap<ConsumptionMetricsKey, u64>,
    metric_collection_endpoint: &reqwest::Url,
    node_id: NodeId,
+    ctx: &RequestContext,
 ) -> anyhow::Result<()> {
    let mut current_metrics: Vec<(ConsumptionMetricsKey, u64)> = Vec::new();
    trace!(
@@ -186,18 +190,27 @@ pub async fn collect_metrics_task(

    // iterate through list of Active tenants and collect metrics
    for (tenant_id, tenant_state) in tenants {
-        if tenant_state != TenantState::Active {
+        if ctx.is_cancelled() {
            continue;
        }
-
-        let tenant = mgr::get_tenant(tenant_id, true).await?;
+        let tenant = mgr::get_tenant(tenant_id).await?;
+        // If the tenant was shut down while while we were looking elsewhere, skip it.
+        let tenant_ctx = match tenant.get_context(ctx) {
+            Ok(ctx) => ctx,
+            Err(_state) => {
+                debug!(
+                    "skipping metrics collection for tenant {tenant_id} because it is not active"
+                );
+                continue;
+            }
+        };

        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() {
+            if let Ok(timeline_ctx) = timeline.get_context(&tenant_ctx) {
                let timeline_written_size = u64::from(timeline.get_last_record_lsn());

                current_metrics.push((
@@ -209,7 +222,8 @@ pub async fn collect_metrics_task(
                    timeline_written_size,
                ));

-                let (timeline_logical_size, is_exact) = timeline.get_current_logical_size()?;
+                let (timeline_logical_size, is_exact) =
+                    timeline.get_current_logical_size(&timeline_ctx)?;
                // Only send timeline logical size when it is fully calculated.
                if is_exact {
                    current_metrics.push((
@@ -0,0 +1,348 @@
+//!
+//! Most async functions throughout the pageserver take a `ctx: &RequestContext`
+//! argument. It is used to control desired behaviour of the operation, and to
+//! allow cancelling the operation gracefully.
+//!
+//! # Context hierarchy
+//!
+//! RequestContext's form a hierarchy. For example:
+//!
+//!  listener context (LibpqEndpointListener)
+//!    connection context (PageRequestHandler)
+//!      per-request context (PageRequestHandler)
+//!
+//! The top "listener context" is created at pageserver startup. The tokio
+//! task that listens on the libpq protocol TCP port holds that context. When
+//! it accepts a connection, it spawns a new task to handle that connection
+//! and creates a new per-connection context for it. The mgmt API listener,
+//! background jobs, and other things form separate but similar hierarchies.
+//!
+//! Usually, each tokio task has its own context, but it's not a strict
+//! requirement and some tasks can hold multiple contexts, and converesely,
+//! some contexts are shared by multiple tasks that work together to perform
+//! some operation.
+//!
+//! The hierarchy is not explictly tracked in the RequestContext struct
+//! itself, but only by their cancellation tokens. It's entirely possible for
+//! the parent context to be dropped before its children.
+//!
+//! # Tenant and Timeline registration
+//!
+//! Most operations are performed on a particular Tenant or Timeline. When
+//! operating on a Tenant or Timeline, it's important that the Tenant/Timeline
+//! isn't detached or deleted while there are tasks working on it. To ensure
+//! that, a RequestContext can be registered with a Tenant or Timeline. See
+//! `Tenant::register_context` and `Timeline::register_context` When
+//! shutting down a Tenant or Timeline, the shutdown routine cancels all the
+//! registered contexts, and waits for them to be dropped before completing
+//! the shutdown.
+//!
+//! To enforce that you hold a registered context when operating on a Tenant
+//! or Timeline, most functions take a TimelineRequestContext or
+//! TenantRequestContext reference as argument.
+//!
+//! NOTE: The Tenant / Timeline registration is separate from the context
+//! hierarchy. You can create a new RequestContext with TimelineRequestContext
+//! as the parent, and register it with a different timeline, for example.
+//!
+//! # Notes
+//!
+//! All RequestContexts in the system have a unique ID, and are also tracked
+//! in a global hash table, CONTEXTS.
+//!
+//! - Futures are normally not assumed to be async cancellation-safe. Pass a
+//!   RequestContext as argument and use cancel() on it instead.
+//!
+//! - If you perform an operation that depends on some external actor or the
+//!   network, use the cancellation token to check for cancellation
+//!
+//! - By convention, the appropriate context for current operation is carried in
+//!   a variable called 'ctx'. If a function handles multiple contexts, it's
+//!   best to *not* have a variable called 'ctx', to force you to think which
+//!   one to use in each call.
+//!
+//! # TODO
+//! - include a unique request ID for tracing
+//!
+
+use once_cell::sync::Lazy;
+use tokio_util::sync::CancellationToken;
+use tracing::{info, warn};
+
+use std::collections::HashMap;
+use std::sync::atomic::{AtomicU64, Ordering};
+use std::sync::Mutex;
+
+/// Each RequestContext has a unique context ID. It's just an increasing
+/// number that we assign.
+static NEXT_CONTEXT_ID: AtomicU64 = AtomicU64::new(1);
+
+/// Global registry of contexts
+static CONTEXTS: Lazy<Mutex<HashMap<RequestContextId, (TaskKind, CancellationToken)>>> =
+    Lazy::new(|| Mutex::new(HashMap::new()));
+
+#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
+pub struct RequestContextId(u64);
+
+///
+pub struct RequestContext {
+    context_id: RequestContextId,
+    task_kind: TaskKind,
+
+    download_behavior: DownloadBehavior,
+    cancellation_token: CancellationToken,
+}
+
+/// DownloadBehavior option specifies the behavior if completing the operation
+/// would require downloading a layer file from remote storage.
+#[derive(Clone, Copy, PartialEq, Eq)]
+pub enum DownloadBehavior {
+    /// Download the layer file. It can take a while.
+    Download,
+
+    /// Download the layer file, but print a warning to the log. This should be used
+    /// in code where the layer file is expected to already exist locally.
+    Warn,
+
+    /// Return a PageReconstructError::NeedsDownload error
+    Error,
+}
+
+///
+/// There are many kinds of tasks in the system. Some are associated with a particular
+/// tenant or timeline, while others are global.
+///
+/// The task kind affects the shutdown sequence on pageserver shutdown and on detach
+/// of an individual tenant. For example, when shutting down the pageserver, we shut
+/// down the LibpqEndpointListeners first, so that we don't accept any more client
+/// connections while we perform the rest of the shutdown duties. See
+/// [`Timeline::graceful_shutdown and`] and [`tenant_mgr::shutdown_pageserver`]
+/// for details.
+///
+/// Note that we don't try to limit how many task of a certain kind can be running
+/// at the same time.
+///
+#[derive(Debug, PartialEq, Eq, Clone, Copy)]
+pub enum TaskKind {
+    // libpq listener task. It just accepts connection and spawns a
+    // PageRequestHandler task for each connection.
+    LibpqEndpointListener,
+
+    // HTTP endpoint listener.
+    HttpEndpointListener,
+
+    // Task that handles a single connection. A PageRequestHandler task
+    // starts detached from any particular tenant or timeline, but it can be
+    // associated with one later, after receiving a command from the client.
+    PageRequestHandler,
+
+    // Context for one management API request
+    MgmtRequest,
+
+    // Manages the WAL receiver connection for one timeline. It subscribes to
+    // events from storage_broker, decides which safekeeper to connect to. It spawns a
+    // separate WalReceiverConnection task to handle each connection.
+    WalReceiverManager,
+
+    // Handles a connection to a safekeeper, to stream WAL to a timeline.
+    WalReceiverConnection,
+
+    // Garbage collection worker. One per tenant
+    GarbageCollector,
+
+    // Compaction. One per tenant.
+    Compaction,
+
+    // Initial logical size calculation
+    InitialLogicalSizeCalculation,
+
+    // Task that flushes frozen in-memory layers to disk
+    LayerFlush,
+
+    // Task that uploads a file to remote storage
+    RemoteUploadTask,
+
+    // Task that downloads a file from remote storage
+    RemoteDownloadTask,
+
+    // task that handles the initial downloading of all tenants
+    InitialLoad,
+
+    // task that handles attaching a tenant
+    Attach,
+
+    // task that handles metrics collection
+    MetricsCollection,
+
+    // task that drives downloading layers
+    DownloadAllRemoteLayers,
+
+    // Only used in unit tests
+    UnitTest,
+}
+
+impl Drop for RequestContext {
+    fn drop(&mut self) {
+        CONTEXTS
+            .lock()
+            .unwrap()
+            .remove(&self.context_id)
+            .expect("context is not in global registry");
+    }
+}
+
+impl RequestContext {
+    /// Create a new RequestContext
+    pub fn new(task_kind: TaskKind, download_behavior: DownloadBehavior) -> Self {
+        let cancellation_token = CancellationToken::new();
+        let context_id = RequestContextId(NEXT_CONTEXT_ID.fetch_add(1, Ordering::Relaxed));
+        CONTEXTS
+            .lock()
+            .unwrap()
+            .insert(context_id, (task_kind, cancellation_token.clone()));
+
+        RequestContext {
+            task_kind,
+            context_id,
+            download_behavior,
+            cancellation_token,
+        }
+    }
+
+    /// Create a new RequestContext, as a child of 'parent'.
+    pub fn with_parent(
+        task_kind: TaskKind,
+        download_behavior: DownloadBehavior,
+        parent: &RequestContext,
+    ) -> Self {
+        let cancellation_token = parent.cancellation_token.child_token();
+        let context_id = RequestContextId(NEXT_CONTEXT_ID.fetch_add(1, Ordering::Relaxed));
+        CONTEXTS
+            .lock()
+            .unwrap()
+            .insert(context_id, (task_kind, cancellation_token.clone()));
+
+        RequestContext {
+            task_kind,
+            context_id,
+            download_behavior,
+            cancellation_token,
+        }
+    }
+
+    pub fn context_id(&self) -> RequestContextId {
+        self.context_id
+    }
+
+    pub fn task_kind(&self) -> TaskKind {
+        self.task_kind
+    }
+
+    pub fn download_behavior(&self) -> DownloadBehavior {
+        self.download_behavior
+    }
+
+    pub fn cancellation_token(&self) -> &CancellationToken {
+        &self.cancellation_token
+    }
+
+    pub fn is_cancelled(&self) -> bool {
+        self.cancellation_token.is_cancelled()
+    }
+
+    pub async fn cancelled(&self) {
+        self.cancellation_token.cancelled().await
+    }
+}
+
+///
+/// Cancel all the contexts in 'context_ids' and wait for them to finish.
+///
+/// Whenever we notice that one of the contexts has finished, it is removed
+/// from 'context_ids'. On return, it is empty.
+///
+pub async fn cancel_and_wait(context_ids: &mut Vec<RequestContextId>) {
+    {
+        let contexts = CONTEXTS.lock().unwrap();
+        context_ids.retain(|context_id| {
+            if let Some((task_kind, cancellation_token)) = contexts.get(context_id) {
+                info!("cancelling task {task_kind:?} with ID {context_id:?}");
+                cancellation_token.cancel();
+                true
+            } else {
+                // Already gone
+                false
+            }
+        });
+    }
+    wait_contexts_to_finish(context_ids).await
+}
+
+async fn wait_contexts_to_finish(context_ids: &mut Vec<RequestContextId>) {
+    let mut n = 0;
+    while !context_ids.is_empty() {
+        {
+            let contexts = CONTEXTS.lock().unwrap();
+            while let Some(context_id) = context_ids.last() {
+                if let Some((task_kind, _cancellation_token)) = contexts.get(context_id) {
+                    info!("waiting for task {task_kind:?} with ID {context_id:?} to finish");
+                    break;
+                } else {
+                    context_ids.pop();
+                }
+            }
+        }
+        if !context_ids.is_empty() {
+            crate::exponential_backoff(
+                n,
+                crate::DEFAULT_BASE_BACKOFF_SECONDS,
+                crate::DEFAULT_MAX_BACKOFF_SECONDS,
+            )
+            .await;
+            n += 1;
+        }
+    }
+}
+
+/// Cancel and wait for all tasks of given 'kind' to finish
+pub async fn shutdown_tasks(kind: TaskKind) {
+    let mut context_ids = Vec::new();
+    {
+        let contexts = CONTEXTS.lock().unwrap();
+        for (&context_id, (task_kind, cancellation_token)) in contexts.iter() {
+            if *task_kind == kind {
+                cancellation_token.cancel();
+                context_ids.push(context_id);
+            }
+        }
+    }
+    wait_contexts_to_finish(&mut context_ids).await
+}
+
+/// Cancel all remaining contexts.
+///
+/// This is used as part of pageserver shutdown. We have already shut down all
+/// tasks / contexts, this is just a backstop or sanity check to make sure we
+/// didn't miss anything. Hence, also print a warning for any remaining tasks.
+pub async fn shutdown_all_tasks() {
+    loop {
+        let mut context_ids = Vec::new();
+        {
+            let contexts = CONTEXTS.lock().unwrap();
+
+            if contexts.is_empty() {
+                return;
+            }
+
+            for (&context_id, (task_kind, cancellation_token)) in contexts.iter() {
+                cancellation_token.cancel();
+                context_ids.push(context_id);
+                warn!(
+                    "unexpected task of kind {:?} with ID {:?} still running",
+                    *task_kind, context_id
+                );
+            }
+        }
+        wait_contexts_to_finish(&mut context_ids).await
+    }
+}
@@ -4,16 +4,16 @@ use anyhow::{anyhow, Context, Result};
 use hyper::StatusCode;
 use hyper::{Body, Request, Response, Uri};
 use remote_storage::GenericRemoteStorage;
-use tokio_util::sync::CancellationToken;
 use tracing::*;

 use super::models::{
    StatusResponse, TenantConfigRequest, TenantCreateRequest, TenantCreateResponse, TenantInfo,
    TimelineCreateRequest, TimelineInfo,
 };
+use crate::context::{DownloadBehavior, RequestContext, TaskKind};
 use crate::pgdatadir_mapping::LsnForTimestamp;
 use crate::tenant::config::TenantConfOpt;
-use crate::tenant::{with_ondemand_download, Timeline};
+use crate::tenant::{PageReconstructError, Timeline, TimelineRequestContext};
 use crate::{config::PageServerConf, tenant::mgr};
 use utils::{
    auth::JwtAuth,
@@ -77,29 +77,50 @@ 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::NeedsDownload(_, _) => {
+            // This shouldn't happen, because we use a RequestContext that requests to
+            // download any missing layer files on-demand.
+            ApiError::InternalServerError(anyhow::anyhow!(
+                "would need to download remote layer file"
+            ))
+        }
+        PageReconstructError::Cancelled => {
+            ApiError::InternalServerError(anyhow::anyhow!("request was cancelled"))
+        }
+        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>,
    include_non_incremental_logical_size: bool,
+    ctx: Option<&TimelineRequestContext>,
 ) -> anyhow::Result<TimelineInfo> {
-    let mut info = build_timeline_info_common(timeline)?;
+    let mut info = build_timeline_info_common(timeline, ctx)?;
    if include_non_incremental_logical_size {
-        // XXX we should be using spawn_ondemand_logical_size_calculation here.
-        // Otherwise, if someone deletes the timeline / detaches the tenant while
-        // we're executing this function, we will outlive the timeline on-disk state.
-        info.current_logical_size_non_incremental = Some(
-            timeline
-                .get_current_logical_size_non_incremental(
-                    info.last_record_lsn,
-                    CancellationToken::new(),
-                )
-                .await?,
-        );
+        if let Some(ctx) = ctx {
+            info.current_logical_size_non_incremental = Some(
+                timeline
+                    .get_current_logical_size_non_incremental(info.last_record_lsn, ctx)
+                    .await?,
+            );
+        } else {
+            info!("could not calculate non-incremental size for timeline because it is not active");
+        }
    }
    Ok(info)
 }

-fn build_timeline_info_common(timeline: &Arc<Timeline>) -> anyhow::Result<TimelineInfo> {
+fn build_timeline_info_common(
+    timeline: &Arc<Timeline>,
+    ctx: Option<&TimelineRequestContext>,
+) -> anyhow::Result<TimelineInfo> {
    let last_record_lsn = timeline.get_last_record_lsn();
    let (wal_source_connstr, last_received_msg_lsn, last_received_msg_ts) = {
        let guard = timeline.last_received_wal.lock().unwrap();
@@ -119,12 +140,16 @@ fn build_timeline_info_common(timeline: &Arc<Timeline>) -> anyhow::Result<Timeli
        Lsn(0) => None,
        lsn @ Lsn(_) => Some(lsn),
    };
-    let current_logical_size = match timeline.get_current_logical_size() {
-        Ok((size, _)) => Some(size),
-        Err(err) => {
-            error!("Timeline info creation failed to get current logical size: {err:?}");
-            None
+    let current_logical_size = if let Some(ctx) = ctx {
+        match timeline.get_current_logical_size(ctx) {
+            Ok((size, _)) => Some(size),
+            Err(err) => {
+                error!("Timeline info creation failed to get current logical size: {err:?}");
+                None
+            }
        }
+    } else {
+        None
    };
    let current_physical_size = Some(timeline.layer_size_sum().approximate_is_ok());
    let state = timeline.current_state();
@@ -170,20 +195,23 @@ async fn timeline_create_handler(mut request: Request<Body>) -> Result<Response<
        .new_timeline_id
        .unwrap_or_else(TimelineId::generate);

-    let tenant = mgr::get_tenant(tenant_id, true)
+    let ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Download);
+
+    let (tenant, tenant_ctx) = mgr::get_active_tenant(tenant_id, &ctx)
        .await
        .map_err(ApiError::NotFound)?;
    match tenant.create_timeline(
        new_timeline_id,
        request_data.ancestor_timeline_id.map(TimelineId::from),
        request_data.ancestor_start_lsn,
-        request_data.pg_version.unwrap_or(crate::DEFAULT_PG_VERSION)
+        request_data.pg_version.unwrap_or(crate::DEFAULT_PG_VERSION),
+        &tenant_ctx,
    )
    .instrument(info_span!("timeline_create", tenant = %tenant_id, new_timeline = ?request_data.new_timeline_id, timeline_id = %new_timeline_id, lsn=?request_data.ancestor_start_lsn, pg_version=?request_data.pg_version))
    .await {
-        Ok(Some(new_timeline)) => {
+        Ok(Some((new_timeline, timeline_ctx))) => {
            // Created. Construct a TimelineInfo for it.
-            let timeline_info = build_timeline_info_common(&new_timeline)
+            let timeline_info = build_timeline_info_common(&new_timeline, Some(&timeline_ctx))
                .map_err(ApiError::InternalServerError)?;
            json_response(StatusCode::CREATED, timeline_info)
        }
@@ -198,21 +226,25 @@ async fn timeline_list_handler(request: Request<Body>) -> Result<Response<Body>,
        query_param_present(&request, "include-non-incremental-logical-size");
    check_permission(&request, Some(tenant_id))?;

+    let top_ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Download);
+
    let response_data = async {
-        let tenant = mgr::get_tenant(tenant_id, true)
+        let (tenant, tenant_ctx) = mgr::get_active_tenant(tenant_id, &top_ctx)
            .await
            .map_err(ApiError::NotFound)?;
        let timelines = tenant.list_timelines();

        let mut response_data = Vec::with_capacity(timelines.len());
        for timeline in timelines {
-            let timeline_info =
-                build_timeline_info(&timeline, include_non_incremental_logical_size)
-                    .await
-                    .context(
-                        "Failed to convert tenant timeline {timeline_id} into the local one: {e:?}",
-                    )
-                    .map_err(ApiError::InternalServerError)?;
+            let timeline_ctx = timeline.get_context(&tenant_ctx).ok();
+            let timeline_info = build_timeline_info(
+                &timeline,
+                include_non_incremental_logical_size,
+                timeline_ctx.as_ref(),
+            )
+            .await
+            .context("Failed to convert tenant timeline {timeline_id} into the local one: {e:?}")
+            .map_err(ApiError::InternalServerError)?;

            response_data.push(timeline_info);
        }
@@ -261,19 +293,26 @@ async fn timeline_detail_handler(request: Request<Body>) -> Result<Response<Body
        query_param_present(&request, "include-non-incremental-logical-size");
    check_permission(&request, Some(tenant_id))?;

+    let top_ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Download);
+
    let timeline_info = async {
-        let tenant = mgr::get_tenant(tenant_id, true)
+        let (tenant, tenant_ctx) = mgr::get_active_tenant(tenant_id, &top_ctx)
            .await
            .map_err(ApiError::NotFound)?;

        let timeline = tenant
-            .get_timeline(timeline_id, false)
+            .get_timeline(timeline_id)
            .map_err(ApiError::NotFound)?;
+        let timeline_ctx = timeline.get_context(&tenant_ctx).ok();

-        let timeline_info = build_timeline_info(&timeline, include_non_incremental_logical_size)
-            .await
-            .context("Failed to get local timeline info: {e:#}")
-            .map_err(ApiError::InternalServerError)?;
+        let timeline_info = build_timeline_info(
+            &timeline,
+            include_non_incremental_logical_size,
+            timeline_ctx.as_ref(),
+        )
+        .await
+        .context("Failed to get local timeline info: {e:#}")
+        .map_err(ApiError::InternalServerError)?;

        Ok::<_, ApiError>(timeline_info)
    }
@@ -294,13 +333,19 @@ async fn get_lsn_by_timestamp_handler(request: Request<Body>) -> Result<Response
        .map_err(ApiError::BadRequest)?;
    let timestamp_pg = postgres_ffi::to_pg_timestamp(timestamp);

-    let timeline = mgr::get_tenant(tenant_id, true)
+    let ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Download);
+
+    let (tenant, ctx) = mgr::get_active_tenant(tenant_id, &ctx)
        .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 (timeline, ctx) = tenant
+        .get_active_timeline(timeline_id, &ctx)
+        .map_err(ApiError::NotFound)?;
+    let result = timeline
+        .find_lsn_for_timestamp(timestamp_pg, &ctx)
        .await
-        .map_err(ApiError::InternalServerError)?;
+        .map_err(apierror_from_prerror)?;

    let result = match result {
        LsnForTimestamp::Present(lsn) => format!("{lsn}"),
@@ -340,7 +385,10 @@ async fn timeline_delete_handler(request: Request<Body>) -> Result<Response<Body
    let timeline_id: TimelineId = parse_request_param(&request, "timeline_id")?;
    check_permission(&request, Some(tenant_id))?;

-    mgr::delete_timeline(tenant_id, timeline_id)
+    // deleting shouldn't require downloading anything
+    let ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Warn);
+
+    mgr::delete_timeline(tenant_id, timeline_id, &ctx)
        .instrument(info_span!("timeline_delete", tenant = %tenant_id, timeline = %timeline_id))
        .await
        // FIXME: Errors from `delete_timeline` can occur for a number of reasons, incuding both
@@ -418,8 +466,10 @@ async fn tenant_status(request: Request<Body>) -> Result<Response<Body>, ApiErro
    let tenant_id: TenantId = parse_request_param(&request, "tenant_id")?;
    check_permission(&request, Some(tenant_id))?;

+    let mut _req_ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Download);
+
    let tenant_info = async {
-        let tenant = mgr::get_tenant(tenant_id, false).await?;
+        let tenant = mgr::get_tenant(tenant_id).await?;

        // Calculate total physical size of all timelines
        let mut current_physical_size = 0;
@@ -446,13 +496,15 @@ async fn tenant_size_handler(request: Request<Body>) -> Result<Response<Body>, A
    let tenant_id: TenantId = parse_request_param(&request, "tenant_id")?;
    check_permission(&request, Some(tenant_id))?;

-    let tenant = mgr::get_tenant(tenant_id, true)
+    let ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Download);
+
+    let (tenant, ctx) = mgr::get_active_tenant(tenant_id, &ctx)
        .await
        .map_err(ApiError::InternalServerError)?;

    // this can be long operation, it currently is not backed by any request coalescing or similar
    let inputs = tenant
-        .gather_size_inputs()
+        .gather_size_inputs(&ctx)
        .await
        .map_err(ApiError::InternalServerError)?;

@@ -495,6 +547,8 @@ fn bad_duration<'a>(field_name: &'static str, value: &'a str) -> impl 'a + Fn()
 async fn tenant_create_handler(mut request: Request<Body>) -> Result<Response<Body>, ApiError> {
    check_permission(&request, None)?;

+    let ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Download);
+
    let request_data: TenantCreateRequest = json_request(&mut request).await?;

    let mut tenant_conf = TenantConfOpt::default();
@@ -583,9 +637,9 @@ async fn tenant_create_handler(mut request: Request<Body>) -> Result<Response<Bo
        Some(tenant) => {
            // We created the tenant. Existing API semantics are that the tenant
            // is Active when this function returns.
-            if let res @ Err(_) = tenant.wait_to_become_active().await {
+            if let res @ Err(_) = tenant.wait_to_become_active(ctx).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
@@ -607,6 +661,8 @@ async fn tenant_config_handler(mut request: Request<Body>) -> Result<Response<Bo
    let tenant_id = request_data.tenant_id;
    check_permission(&request, Some(tenant_id))?;

+    let ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Download);
+
    let mut tenant_conf: TenantConfOpt = Default::default();
    if let Some(gc_period) = request_data.gc_period {
        tenant_conf.gc_period = Some(
@@ -669,7 +725,7 @@ async fn tenant_config_handler(mut request: Request<Body>) -> Result<Response<Bo
    }

    let state = get_state(&request);
-    mgr::update_tenant_config(state.conf, tenant_conf, tenant_id)
+    mgr::update_tenant_config(state.conf, tenant_conf, tenant_id, &ctx)
        .instrument(info_span!("tenant_config", tenant = ?tenant_id))
        .await
        // FIXME: `update_tenant_config` can fail because of both user and internal errors.
@@ -721,11 +777,21 @@ async fn timeline_gc_handler(mut request: Request<Body>) -> Result<Response<Body

    let gc_req: TimelineGcRequest = json_request(&mut request).await?;

-    let wait_task_done = mgr::immediate_gc(tenant_id, timeline_id, gc_req).await?;
-    let gc_result = wait_task_done
+    let ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Download);
+
+    let (tenant, ctx) = mgr::get_active_tenant(tenant_id, &ctx)
+        .await
+        .map_err(ApiError::NotFound)?;
+
+    let gc_horizon = gc_req.gc_horizon.unwrap_or_else(|| tenant.get_gc_horizon());
+    // Use tenant's pitr setting
+    let pitr = tenant.get_pitr_interval();
+
+    fail::fail_point!("immediate_gc_task_pre");
+    let gc_result = tenant
+        .gc_iteration(Some(timeline_id), gc_horizon, pitr, &ctx)
+        .instrument(info_span!("manual_gc", tenant = %tenant_id, timeline = %timeline_id))
        .await
-        .context("wait for gc task")
-        .map_err(ApiError::InternalServerError)?
        .map_err(ApiError::InternalServerError)?;

    json_response(StatusCode::OK, gc_result)
@@ -738,14 +804,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)
+    let ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Download);
+
+    let (tenant, ctx) = mgr::get_active_tenant(tenant_id, &ctx)
        .await
        .map_err(ApiError::NotFound)?;
-    let timeline = tenant
-        .get_timeline(timeline_id, true)
+    let (timeline, ctx) = tenant
+        .get_active_timeline(timeline_id, &ctx)
        .map_err(ApiError::NotFound)?;
    timeline
-        .compact()
+        .compact(&ctx)
+        .instrument(info_span!("manual_compact", tenant = %tenant_id, timeline = %timeline_id))
        .await
        .map_err(ApiError::InternalServerError)?;

@@ -759,18 +828,21 @@ async fn timeline_checkpoint_handler(request: Request<Body>) -> Result<Response<
    let timeline_id: TimelineId = parse_request_param(&request, "timeline_id")?;
    check_permission(&request, Some(tenant_id))?;

-    let tenant = mgr::get_tenant(tenant_id, true)
+    let ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Download);
+
+    let (tenant, ctx) = mgr::get_active_tenant(tenant_id, &ctx)
        .await
        .map_err(ApiError::NotFound)?;
-    let timeline = tenant
-        .get_timeline(timeline_id, true)
+    let (timeline, ctx) = tenant
+        .get_active_timeline(timeline_id, &ctx)
        .map_err(ApiError::NotFound)?;
    timeline
        .freeze_and_flush()
        .await
        .map_err(ApiError::InternalServerError)?;
    timeline
-        .compact()
+        .compact(&ctx)
+        .instrument(info_span!("manual_compact", tenant = %tenant_id, timeline = %timeline_id))
        .await
        .map_err(ApiError::InternalServerError)?;

@@ -784,13 +856,15 @@ async fn timeline_download_remote_layers_handler_post(
    let timeline_id: TimelineId = parse_request_param(&request, "timeline_id")?;
    check_permission(&request, Some(tenant_id))?;

-    let tenant = mgr::get_tenant(tenant_id, true)
+    let ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Download);
+
+    let (tenant, ctx) = mgr::get_active_tenant(tenant_id, &ctx)
        .await
        .map_err(ApiError::NotFound)?;
-    let timeline = tenant
-        .get_timeline(timeline_id, true)
+    let (timeline, ctx) = tenant
+        .get_active_timeline(timeline_id, &ctx)
        .map_err(ApiError::NotFound)?;
-    match timeline.spawn_download_all_remote_layers().await {
+    match timeline.spawn_download_all_remote_layers(&ctx).await {
        Ok(st) => json_response(StatusCode::ACCEPTED, st),
        Err(st) => json_response(StatusCode::CONFLICT, st),
    }
@@ -803,11 +877,13 @@ async fn timeline_download_remote_layers_handler_get(
    let timeline_id: TimelineId = parse_request_param(&request, "timeline_id")?;
    check_permission(&request, Some(tenant_id))?;

-    let tenant = mgr::get_tenant(tenant_id, true)
+    let ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Download);
+
+    let (tenant, ctx) = mgr::get_active_tenant(tenant_id, &ctx)
        .await
        .map_err(ApiError::NotFound)?;
-    let timeline = tenant
-        .get_timeline(timeline_id, true)
+    let (timeline, _ctx) = tenant
+        .get_active_timeline(timeline_id, &ctx)
        .map_err(ApiError::NotFound)?;
    let info = timeline
        .get_download_all_remote_layers_task_info()
@@ -13,7 +13,7 @@ use tracing::*;
 use walkdir::WalkDir;

 use crate::pgdatadir_mapping::*;
-use crate::tenant::Timeline;
+use crate::tenant::{Timeline, TimelineRequestContext};
 use crate::walingest::WalIngest;
 use crate::walrecord::DecodedWALRecord;
 use pageserver_api::reltag::{RelTag, SlruKind};
@@ -47,6 +47,7 @@ pub async fn import_timeline_from_postgres_datadir(
    tline: &Timeline,
    pgdata_path: &Path,
    pgdata_lsn: Lsn,
+    ctx: &TimelineRequestContext,
 ) -> Result<()> {
    let mut pg_control: Option<ControlFileData> = None;

@@ -69,7 +70,7 @@ pub async fn import_timeline_from_postgres_datadir(
            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, &mut file, len).await?
+                import_file(&mut modification, relative_path, &mut file, len, ctx).await?
            {
                pg_control = Some(control_file);
            }
@@ -99,6 +100,7 @@ pub async fn import_timeline_from_postgres_datadir(
        tline,
        Lsn(pg_control.checkPointCopy.redo),
        pgdata_lsn,
+        ctx,
    )
    .await?;

@@ -113,6 +115,7 @@ async fn import_rel(
    dboid: Oid,
    reader: &mut (impl AsyncRead + Send + Sync + Unpin),
    len: usize,
+    ctx: &TimelineRequestContext,
 ) -> anyhow::Result<()> {
    // Does it look like a relation file?
    trace!("importing rel file {}", path.display());
@@ -143,7 +146,14 @@ async fn import_rel(
    // 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, ctx)
+        .await
+    {
        if e.to_string().contains("already exists") {
            debug!("relation {} already exists. we must be extending it", rel);
        } else {
@@ -178,7 +188,7 @@ async fn import_rel(
    //
    // 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, ctx).await?;

    Ok(())
 }
@@ -191,6 +201,7 @@ async fn import_slru(
    path: &Path,
    reader: &mut (impl AsyncRead + Send + Sync + Unpin),
    len: usize,
+    ctx: &TimelineRequestContext,
 ) -> anyhow::Result<()> {
    info!("importing slru file {path:?}");

@@ -206,7 +217,9 @@ async fn import_slru(

    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, ctx)
+        .await?;

    let mut rpageno = 0;
    loop {
@@ -246,15 +259,15 @@ async fn import_wal(
    tline: &Timeline,
    startpoint: Lsn,
    endpoint: Lsn,
+    ctx: &TimelineRequestContext,
 ) -> 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).await?;
+    let mut walingest = WalIngest::new(tline, startpoint, ctx).await?;

    while last_lsn <= endpoint {
        // FIXME: assume postgresql tli 1 for now
@@ -277,6 +290,7 @@ async fn import_wal(
            file.seek(std::io::SeekFrom::Start(offset as u64))?;
        }

+        use std::io::Read;
        let nread = file.read_to_end(&mut buf)?;
        if nread != WAL_SEGMENT_SIZE - offset {
            // Maybe allow this for .partial files?
@@ -291,7 +305,7 @@ async fn import_wal(
        while last_lsn <= endpoint {
            if let Some((lsn, recdata)) = waldecoder.poll_decode()? {
                walingest
-                    .ingest_record(recdata, lsn, &mut modification, &mut decoded)
+                    .ingest_record(recdata, lsn, &mut modification, &mut decoded, ctx)
                    .await?;
                last_lsn = lsn;

@@ -320,6 +334,7 @@ pub async fn import_basebackup_from_tar(
    tline: &Timeline,
    reader: &mut (impl AsyncRead + Send + Sync + Unpin),
    base_lsn: Lsn,
+    ctx: &TimelineRequestContext,
 ) -> Result<()> {
    info!("importing base at {base_lsn}");
    let mut modification = tline.begin_modification(base_lsn);
@@ -338,7 +353,7 @@ pub async fn import_basebackup_from_tar(
        match header.entry_type() {
            tokio_tar::EntryType::Regular => {
                if let Some(res) =
-                    import_file(&mut modification, file_path.as_ref(), &mut entry, len).await?
+                    import_file(&mut modification, file_path.as_ref(), &mut entry, len, ctx).await?
                {
                    // We found the pg_control file.
                    pg_control = Some(res);
@@ -370,13 +385,14 @@ pub async fn import_wal_from_tar(
    reader: &mut (impl AsyncRead + Send + Sync + Unpin),
    start_lsn: Lsn,
    end_lsn: Lsn,
+    ctx: &TimelineRequestContext,
 ) -> Result<()> {
    // Set up walingest mutable state
    let mut waldecoder = WalStreamDecoder::new(start_lsn, tline.pg_version);
    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).await?;
+    let mut walingest = WalIngest::new(tline, start_lsn, ctx).await?;

    // Ingest wal until end_lsn
    info!("importing wal until {}", end_lsn);
@@ -425,7 +441,7 @@ pub async fn import_wal_from_tar(
        while last_lsn <= end_lsn {
            if let Some((lsn, recdata)) = waldecoder.poll_decode()? {
                walingest
-                    .ingest_record(recdata, lsn, &mut modification, &mut decoded)
+                    .ingest_record(recdata, lsn, &mut modification, &mut decoded, ctx)
                    .await?;
                last_lsn = lsn;

@@ -460,6 +476,7 @@ async fn import_file(
    file_path: &Path,
    reader: &mut (impl AsyncRead + Send + Sync + Unpin),
    len: usize,
+    ctx: &TimelineRequestContext,
 ) -> Result<Option<ControlFileData>> {
    let file_name = match file_path.file_name() {
        Some(name) => name.to_string_lossy(),
@@ -492,14 +509,16 @@ async fn import_file(
            }
            "pg_filenode.map" => {
                let bytes = read_all_bytes(reader).await?;
-                modification.put_relmap_file(spcnode, dbnode, bytes)?;
+                modification
+                    .put_relmap_file(spcnode, dbnode, bytes, ctx)
+                    .await?;
                debug!("imported relmap file")
            }
            "PG_VERSION" => {
                debug!("ignored PG_VERSION file");
            }
            _ => {
-                import_rel(modification, file_path, spcnode, dbnode, reader, len).await?;
+                import_rel(modification, file_path, spcnode, dbnode, reader, len, ctx).await?;
                debug!("imported rel creation");
            }
        }
@@ -515,37 +534,41 @@ async fn import_file(
        match file_name.as_ref() {
            "pg_filenode.map" => {
                let bytes = read_all_bytes(reader).await?;
-                modification.put_relmap_file(spcnode, dbnode, bytes)?;
+                modification
+                    .put_relmap_file(spcnode, dbnode, bytes, ctx)
+                    .await?;
                debug!("imported relmap file")
            }
            "PG_VERSION" => {
                debug!("ignored PG_VERSION file");
            }
            _ => {
-                import_rel(modification, file_path, spcnode, dbnode, reader, len).await?;
+                import_rel(modification, file_path, spcnode, dbnode, reader, len, ctx).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).await?;
+        import_slru(modification, slru, file_path, reader, len, ctx).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).await?;
+        import_slru(modification, slru, file_path, reader, len, ctx).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).await?;
+        import_slru(modification, slru, file_path, reader, len, ctx).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).await?;
-        modification.put_twophase_file(xid, Bytes::copy_from_slice(&bytes[..]))?;
+        modification
+            .put_twophase_file(xid, Bytes::copy_from_slice(&bytes[..]), ctx)
+            .await?;
        debug!("imported twophase file");
    } else if file_path.starts_with("pg_wal") {
        debug!("found wal file in base section. ignore it");
@@ -2,6 +2,7 @@ mod auth;
 pub mod basebackup;
 pub mod config;
 pub mod consumption_metrics;
+pub mod context;
 pub mod http;
 pub mod import_datadir;
 pub mod keyspace;
@@ -21,7 +22,6 @@ pub mod walredo;

 use std::path::Path;

-use crate::task_mgr::TaskKind;
 use tracing::info;

 /// Current storage format version
@@ -41,35 +41,6 @@ pub const DELTA_FILE_MAGIC: u16 = 0x5A61;

 static ZERO_PAGE: bytes::Bytes = bytes::Bytes::from_static(&[0u8; 8192]);

-pub async fn shutdown_pageserver(exit_code: i32) {
-    // Shut down the libpq endpoint task. This prevents new connections from
-    // being accepted.
-    task_mgr::shutdown_tasks(Some(TaskKind::LibpqEndpointListener), None, None).await;
-
-    // Shut down any page service tasks.
-    task_mgr::shutdown_tasks(Some(TaskKind::PageRequestHandler), None, None).await;
-
-    // Shut down all the tenants. This flushes everything to disk and kills
-    // the checkpoint and GC tasks.
-    tenant::mgr::shutdown_all_tenants().await;
-
-    // Stop syncing with remote storage.
-    //
-    // FIXME: Does this wait for the sync tasks to finish syncing what's queued up?
-    // Should it?
-    task_mgr::shutdown_tasks(Some(TaskKind::RemoteUploadTask), None, None).await;
-
-    // Shut down the HTTP endpoint last, so that you can still check the server's
-    // status while it's shutting down.
-    // FIXME: We should probably stop accepting commands like attach/detach earlier.
-    task_mgr::shutdown_tasks(Some(TaskKind::HttpEndpointListener), None, None).await;
-
-    // There should be nothing left, but let's be sure
-    task_mgr::shutdown_tasks(None, None, None).await;
-    info!("Shut down successfully completed");
-    std::process::exit(exit_code);
-}
-
 const DEFAULT_BASE_BACKOFF_SECONDS: f64 = 0.1;
 const DEFAULT_MAX_BACKOFF_SECONDS: f64 = 3.0;

@@ -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,
@@ -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;
@@ -28,6 +30,7 @@ use std::sync::Arc;
 use std::time::Duration;
 use tracing::*;
 use utils::id::ConnectionId;
+use utils::postgres_backend_async::QueryError;
 use utils::{
    auth::{Claims, JwtAuth, Scope},
    id::{TenantId, TimelineId},
@@ -40,28 +43,31 @@ use utils::{
 use crate::auth::check_permission;
 use crate::basebackup;
 use crate::config::PageServerConf;
+use crate::context::{DownloadBehavior, RequestContext, TaskKind};
 use crate::import_datadir::import_wal_from_tar;
 use crate::metrics::{LIVE_CONNECTIONS_COUNT, SMGR_QUERY_TIME};
 use crate::task_mgr;
-use crate::task_mgr::TaskKind;
 use crate::tenant::mgr;
-use crate::tenant::{Tenant, Timeline};
+use crate::tenant::{Tenant, TenantRequestContext, Timeline, TimelineRequestContext};
 use crate::trace::Tracer;

 use postgres_ffi::pg_constants::DEFAULTTABLESPACE_OID;
 use postgres_ffi::BLCKSZ;

-fn copyin_stream(pgb: &mut PostgresBackend) -> impl Stream<Item = io::Result<Bytes>> + '_ {
+fn copyin_stream<'a>(
+    pgb: &'a mut PostgresBackend,
+    ctx: &'a RequestContext,
+) -> impl Stream<Item = io::Result<Bytes>> + 'a {
    async_stream::try_stream! {
        loop {
            let msg = tokio::select! {
                biased;

-                _ = task_mgr::shutdown_watcher() => {
+                _ = ctx.cancelled() => {
                    // 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 }
@@ -74,14 +80,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;
                        }
@@ -91,12 +98,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))?;
                }
            };
        }
@@ -115,6 +126,7 @@ pub async fn libpq_listener_main(
    auth: Option<Arc<JwtAuth>>,
    listener: TcpListener,
    auth_type: AuthType,
+    listener_ctx: RequestContext,
 ) -> anyhow::Result<()> {
    listener.set_nonblocking(true)?;
    let tokio_listener = tokio::net::TcpListener::from_std(listener)?;
@@ -123,8 +135,9 @@ pub async fn libpq_listener_main(
    while let Some(res) = tokio::select! {
        biased;

-        _ = task_mgr::shutdown_watcher() => {
+        _ = listener_ctx.cancelled() => {
            // We were requested to shut down.
+            info!("libpq listener shutting down");
            None
        }

@@ -138,18 +151,33 @@ pub async fn libpq_listener_main(
                debug!("accepted connection from {}", peer_addr);
                let local_auth = auth.clone();

+                let connection_ctx = RequestContext::with_parent(
+                    TaskKind::PageRequestHandler,
+                    DownloadBehavior::Download,
+                    &listener_ctx,
+                );
+
                // PageRequestHandler tasks are not associated with any particular
                // timeline in the task manager. In practice most connections will
                // only deal with a particular timeline, but we don't know which one
                // yet.
                task_mgr::spawn(
                    &tokio::runtime::Handle::current(),
-                    TaskKind::PageRequestHandler,
-                    None,
-                    None,
                    "serving compute connection task",
                    false,
-                    page_service_conn_main(conf, local_auth, socket, auth_type),
+                    async move {
+                        if let Err(err) = page_service_conn_main(
+                            conf,
+                            local_auth,
+                            socket,
+                            auth_type,
+                            connection_ctx,
+                        )
+                        .await
+                        {
+                            error!("connection handler exited with error: {err:?}");
+                        }
+                    },
                );
            }
            Err(err) => {
@@ -169,6 +197,7 @@ async fn page_service_conn_main(
    auth: Option<Arc<JwtAuth>>,
    socket: tokio::net::TcpStream,
    auth_type: AuthType,
+    connection_ctx: RequestContext,
 ) -> anyhow::Result<()> {
    // Immediately increment the gauge, then create a job to decrement it on task exit.
    // One of the pros of `defer!` is that this will *most probably*
@@ -183,34 +212,32 @@ async fn page_service_conn_main(
        .set_nodelay(true)
        .context("could not set TCP_NODELAY")?;

-    let mut conn_handler = PageServerHandler::new(conf, auth);
+    let cancellation_token = connection_ctx.cancellation_token().clone();
+
+    let mut conn_handler = PageServerHandler::new(conf, auth, connection_ctx);
    let pgbackend = PostgresBackend::new(socket, auth_type, None)?;

    let result = pgbackend
-        .run(&mut conn_handler, task_mgr::shutdown_watcher)
+        .run(&mut conn_handler, || cancellation_token.cancelled())
        .await;
    match result {
        Ok(()) => {
            // 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"),
    }
 }

@@ -247,35 +274,38 @@ impl PageRequestMetrics {
    }
 }

-#[derive(Debug)]
 struct PageServerHandler {
    _conf: &'static PageServerConf,
    auth: Option<Arc<JwtAuth>>,
    claims: Option<Claims>,
+
+    connection_ctx: RequestContext,
 }

 impl PageServerHandler {
-    pub fn new(conf: &'static PageServerConf, auth: Option<Arc<JwtAuth>>) -> Self {
+    pub fn new(
+        conf: &'static PageServerConf,
+        auth: Option<Arc<JwtAuth>>,
+        connection_ctx: RequestContext,
+    ) -> Self {
        PageServerHandler {
            _conf: conf,
            auth,
            claims: None,
+            connection_ctx,
        }
    }

    #[instrument(skip(self, pgb))]
    async fn handle_pagerequests(
-        &self,
+        &mut self,
        pgb: &mut PostgresBackend,
        tenant_id: TenantId,
        timeline_id: TimelineId,
    ) -> anyhow::Result<()> {
-        // NOTE: pagerequests handler exits when connection is closed,
-        //       so there is no need to reset the association
-        task_mgr::associate_with(Some(tenant_id), Some(timeline_id));
+        let (tenant, ctx) = get_active_tenant_with_timeout(tenant_id, &self.connection_ctx).await?;

        // Make request tracer if needed
-        let tenant = get_active_tenant_with_timeout(tenant_id).await?;
        let mut tracer = if tenant.get_trace_read_requests() {
            let connection_id = ConnectionId::generate();
            let path = tenant
@@ -287,7 +317,7 @@ impl PageServerHandler {
        };

        // Check that the timeline exists
-        let timeline = tenant.get_timeline(timeline_id, true)?;
+        let (timeline, ctx) = tenant.get_active_timeline(timeline_id, &ctx)?;

        // switch client to COPYBOTH
        pgb.write_message(&BeMessage::CopyBothResponse)?;
@@ -299,7 +329,7 @@ impl PageServerHandler {
            let msg = tokio::select! {
                biased;

-                _ = task_mgr::shutdown_watcher() => {
+                _ = ctx.cancelled() => {
                    // We were requested to shut down.
                    info!("shutdown request received in page handler");
                    break;
@@ -312,7 +342,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
            };
@@ -326,22 +356,27 @@ impl PageServerHandler {

            let neon_fe_msg = PagestreamFeMessage::parse(&mut copy_data_bytes.reader())?;

+            // TODO: We could create a new per-request context here, with unique ID.
+            // Currently we use the same per-timeline context for all requests
+
            let response = match neon_fe_msg {
                PagestreamFeMessage::Exists(req) => {
                    let _timer = metrics.get_rel_exists.start_timer();
-                    self.handle_get_rel_exists_request(&timeline, &req).await
+                    self.handle_get_rel_exists_request(&timeline, &req, &ctx)
+                        .await
                }
                PagestreamFeMessage::Nblocks(req) => {
                    let _timer = metrics.get_rel_size.start_timer();
-                    self.handle_get_nblocks_request(&timeline, &req).await
+                    self.handle_get_nblocks_request(&timeline, &req, &ctx).await
                }
                PagestreamFeMessage::GetPage(req) => {
                    let _timer = metrics.get_page_at_lsn.start_timer();
-                    self.handle_get_page_at_lsn_request(&timeline, &req).await
+                    self.handle_get_page_at_lsn_request(&timeline, &req, &ctx)
+                        .await
                }
                PagestreamFeMessage::DbSize(req) => {
                    let _timer = metrics.get_db_size.start_timer();
-                    self.handle_db_size_request(&timeline, &req).await
+                    self.handle_db_size_request(&timeline, &req, &ctx).await
                }
            };

@@ -362,19 +397,20 @@ impl PageServerHandler {

    #[instrument(skip(self, pgb))]
    async fn handle_import_basebackup(
-        &self,
+        &mut self,
        pgb: &mut PostgresBackend,
        tenant_id: TenantId,
        timeline_id: TimelineId,
        base_lsn: Lsn,
        _end_lsn: Lsn,
        pg_version: u32,
-    ) -> anyhow::Result<()> {
-        task_mgr::associate_with(Some(tenant_id), Some(timeline_id));
+    ) -> Result<(), QueryError> {
        // Create empty timeline
        info!("creating new timeline");
-        let tenant = get_active_tenant_with_timeout(tenant_id).await?;
-        let timeline = tenant.create_empty_timeline(timeline_id, base_lsn, pg_version)?;
+        let (tenant, tenant_ctx) =
+            get_active_tenant_with_timeout(tenant_id, &self.connection_ctx).await?;
+        let (timeline, ctx) =
+            tenant.create_empty_timeline(timeline_id, base_lsn, pg_version, &tenant_ctx)?;

        // TODO mark timeline as not ready until it reaches end_lsn.
        // We might have some wal to import as well, and we should prevent compute
@@ -391,9 +427,9 @@ impl PageServerHandler {
        pgb.write_message(&BeMessage::CopyInResponse)?;
        pgb.flush().await?;

-        let mut copyin_stream = Box::pin(copyin_stream(pgb));
+        let mut copyin_stream = Box::pin(copyin_stream(pgb, &ctx));
        timeline
-            .import_basebackup_from_tar(&mut copyin_stream, base_lsn)
+            .import_basebackup_from_tar(&mut copyin_stream, base_lsn, &ctx)
            .await?;

        // Drain the rest of the Copy data
@@ -417,17 +453,21 @@ impl PageServerHandler {

    #[instrument(skip(self, pgb))]
    async fn handle_import_wal(
-        &self,
+        &mut self,
        pgb: &mut PostgresBackend,
        tenant_id: TenantId,
        timeline_id: TimelineId,
        start_lsn: Lsn,
        end_lsn: Lsn,
-    ) -> anyhow::Result<()> {
-        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);
+    ) -> Result<(), QueryError> {
+        let (timeline, ctx) =
+            get_active_timeline_with_timeout(tenant_id, timeline_id, &self.connection_ctx).await?;
+        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.
@@ -436,9 +476,9 @@ impl PageServerHandler {
        info!("importing wal");
        pgb.write_message(&BeMessage::CopyInResponse)?;
        pgb.flush().await?;
-        let mut copyin_stream = Box::pin(copyin_stream(pgb));
+        let mut copyin_stream = Box::pin(copyin_stream(pgb, &ctx));
        let mut reader = tokio_util::io::StreamReader::new(&mut copyin_stream);
-        import_wal_from_tar(&timeline, &mut reader, start_lsn, end_lsn).await?;
+        import_wal_from_tar(&timeline, &mut reader, start_lsn, end_lsn, &ctx).await?;
        info!("wal import complete");

        // Drain the rest of the Copy data
@@ -451,7 +491,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
@@ -480,7 +524,8 @@ impl PageServerHandler {
        mut lsn: Lsn,
        latest: bool,
        latest_gc_cutoff_lsn: &RcuReadGuard<Lsn>,
-    ) -> Result<Lsn> {
+        ctx: &TimelineRequestContext,
+    ) -> 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
@@ -503,7 +548,7 @@ impl PageServerHandler {
            if lsn <= last_record_lsn {
                lsn = last_record_lsn;
            } else {
-                timeline.wait_lsn(lsn).await?;
+                timeline.wait_lsn(lsn, ctx).await?;
                // Since we waited for 'lsn' to arrive, that is now the last
                // record LSN. (Or close enough for our purposes; the
                // last-record LSN can advance immediately after we return
@@ -511,11 +556,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?;
+            timeline.wait_lsn(lsn, ctx).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
@@ -523,60 +568,61 @@ impl PageServerHandler {
        Ok(lsn)
    }

-    #[instrument(skip(self, timeline, req), fields(rel = %req.rel, req_lsn = %req.lsn))]
+    #[instrument(skip(self, timeline, req, ctx), fields(rel = %req.rel, req_lsn = %req.lsn))]
    async fn handle_get_rel_exists_request(
        &self,
        timeline: &Timeline,
        req: &PagestreamExistsRequest,
-    ) -> Result<PagestreamBeMessage> {
+        ctx: &TimelineRequestContext,
+    ) -> 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 lsn =
+            Self::wait_or_get_last_lsn(timeline, req.lsn, req.latest, &latest_gc_cutoff_lsn, ctx)
+                .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, ctx)
+            .await?;

        Ok(PagestreamBeMessage::Exists(PagestreamExistsResponse {
            exists,
        }))
    }

-    #[instrument(skip(self, timeline, req), fields(rel = %req.rel, req_lsn = %req.lsn))]
+    #[instrument(skip(self, timeline, req, ctx), fields(rel = %req.rel, req_lsn = %req.lsn))]
    async fn handle_get_nblocks_request(
        &self,
        timeline: &Timeline,
        req: &PagestreamNblocksRequest,
-    ) -> Result<PagestreamBeMessage> {
+        ctx: &TimelineRequestContext,
+    ) -> 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 lsn =
+            Self::wait_or_get_last_lsn(timeline, req.lsn, req.latest, &latest_gc_cutoff_lsn, ctx)
+                .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, ctx).await?;

        Ok(PagestreamBeMessage::Nblocks(PagestreamNblocksResponse {
            n_blocks,
        }))
    }

-    #[instrument(skip(self, timeline, req), fields(dbnode = %req.dbnode, req_lsn = %req.lsn))]
+    #[instrument(skip(self, timeline, req, ctx), fields(dbnode = %req.dbnode, req_lsn = %req.lsn))]
    async fn handle_db_size_request(
        &self,
        timeline: &Timeline,
        req: &PagestreamDbSizeRequest,
-    ) -> Result<PagestreamBeMessage> {
+        ctx: &TimelineRequestContext,
+    ) -> 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 lsn =
+            Self::wait_or_get_last_lsn(timeline, req.lsn, req.latest, &latest_gc_cutoff_lsn, ctx)
+                .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, ctx)
+            .await?;
        let db_size = total_blocks as i64 * BLCKSZ as i64;

        Ok(PagestreamBeMessage::DbSize(PagestreamDbSizeResponse {
@@ -584,15 +630,17 @@ impl PageServerHandler {
        }))
    }

-    #[instrument(skip(self, timeline, req), fields(rel = %req.rel, blkno = %req.blkno, req_lsn = %req.lsn))]
+    #[instrument(skip(self, timeline, req, ctx), fields(rel = %req.rel, blkno = %req.blkno, req_lsn = %req.lsn))]
    async fn handle_get_page_at_lsn_request(
        &self,
        timeline: &Timeline,
        req: &PagestreamGetPageRequest,
-    ) -> Result<PagestreamBeMessage> {
+        ctx: &TimelineRequestContext,
+    ) -> 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 lsn =
+            Self::wait_or_get_last_lsn(timeline, req.lsn, req.latest, &latest_gc_cutoff_lsn, ctx)
+                .await?;
        /*
        // Add a 1s delay to some requests. The delay helps the requests to
        // hit the race condition from github issue #1047 more easily.
@@ -602,10 +650,9 @@ impl PageServerHandler {
        }
        */

-        let page = crate::tenant::with_ondemand_download(|| {
-            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, ctx)
+            .await?;

        Ok(PagestreamBeMessage::GetPage(PagestreamGetPageResponse {
            page,
@@ -614,7 +661,7 @@ impl PageServerHandler {

    #[instrument(skip(self, pgb))]
    async fn handle_basebackup_request(
-        &self,
+        &mut self,
        pgb: &mut PostgresBackend,
        tenant_id: TenantId,
        timeline_id: TimelineId,
@@ -623,12 +670,14 @@ impl PageServerHandler {
        full_backup: bool,
    ) -> anyhow::Result<()> {
        // check that the timeline exists
-        let timeline = get_active_timeline_with_timeout(tenant_id, timeline_id).await?;
+        let (timeline, ctx) =
+            get_active_timeline_with_timeout(tenant_id, timeline_id, &self.connection_ctx).await?;
+
        let latest_gc_cutoff_lsn = timeline.get_latest_gc_cutoff_lsn();
        if let Some(lsn) = lsn {
            // Backup was requested at a particular LSN. Wait for it to arrive.
            info!("waiting for {}", lsn);
-            timeline.wait_lsn(lsn).await?;
+            timeline.wait_lsn(lsn, &ctx).await?;
            timeline
                .check_lsn_is_in_scope(lsn, &latest_gc_cutoff_lsn)
                .context("invalid basebackup lsn")?;
@@ -638,13 +687,18 @@ impl PageServerHandler {
        pgb.write_message(&BeMessage::CopyOutResponse)?;
        pgb.flush().await?;

-        /* Send a tarball of the latest layer on the timeline */
+        // Send a tarball of the latest layer on the timeline
        {
            let mut writer = pgb.copyout_writer();
-            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().await?;
+            basebackup::send_basebackup_tarball(
+                &mut writer,
+                &timeline,
+                lsn,
+                prev_lsn,
+                full_backup,
+                &ctx,
+            )
+            .await?;
        }

        pgb.write_message(&BeMessage::CopyDone)?;
@@ -656,7 +710,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(());
@@ -678,20 +732,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!(
@@ -703,22 +756,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))?;

@@ -728,18 +792,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
            };
@@ -754,16 +824,21 @@ 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?;
+            let (timeline, _ctx) =
+                get_active_timeline_with_timeout(tenant_id, timeline_id, &self.connection_ctx)
+                    .await?;

            let end_of_timeline = timeline.get_last_record_rlsn();

@@ -782,22 +857,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
            };
@@ -822,12 +906,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))?;

@@ -845,7 +938,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 ") {
@@ -855,11 +951,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))?;

@@ -870,7 +974,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 ") {
@@ -881,12 +988,18 @@ 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))?;

-            let tenant = get_active_tenant_with_timeout(tenant_id).await?;
+            let (tenant, _ctx) =
+                get_active_tenant_with_timeout(tenant_id, &self.connection_ctx).await?;
            pgb.write_message(&BeMessage::RowDescription(&[
                RowDescriptor::int8_col(b"checkpoint_distance"),
                RowDescriptor::int8_col(b"checkpoint_timeout"),
@@ -923,7 +1036,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(())
@@ -935,12 +1050,25 @@ 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>> {
-    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
-            // no .context(), the error message is good enough and some tests depend on it
-            .map(move |()| tenant),
+async fn get_active_tenant_with_timeout(
+    tenant_id: TenantId,
+    parent_ctx: &RequestContext,
+) -> anyhow::Result<(Arc<Tenant>, TenantRequestContext)> {
+    let child_ctx = RequestContext::with_parent(
+        parent_ctx.task_kind(),
+        parent_ctx.download_behavior(),
+        parent_ctx,
+    );
+
+    let tenant = mgr::get_tenant(tenant_id).await?;
+    match tokio::time::timeout(
+        Duration::from_secs(30),
+        tenant.wait_to_become_active(child_ctx),
+    )
+    .await
+    {
+        Ok(Ok(ctx)) => Ok((tenant, ctx)),
+        Ok(Err(err)) => Err(err),
        Err(_) => anyhow::bail!("Timeout waiting for tenant {tenant_id} to become Active"),
    }
 }
@@ -949,8 +1077,9 @@ 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>> {
-    get_active_tenant_with_timeout(tenant_id)
+    ctx: &RequestContext,
+) -> anyhow::Result<(Arc<Timeline>, TimelineRequestContext)> {
+    get_active_tenant_with_timeout(tenant_id, ctx)
        .await
-        .and_then(|tenant| tenant.get_timeline(timeline_id, true))
+        .and_then(|(tenant, ctx)| tenant.get_active_timeline(timeline_id, &ctx))
 }
@@ -6,11 +6,10 @@
 //! walingest.rs handles a few things like implicit relation creation and extension.
 //! Clarify that)
 //!
-use super::tenant::PageReconstructResult;
 use crate::keyspace::{KeySpace, KeySpaceAccum};
-use crate::tenant::{with_ondemand_download, Timeline};
+use crate::repository::*;
+use crate::tenant::{PageReconstructError, Timeline, TimelineRequestContext};
 use crate::walrecord::NeonWalRecord;
-use crate::{repository::*, try_no_ondemand_download};
 use anyhow::Context;
 use bytes::{Buf, Bytes};
 use pageserver_api::reltag::{RelTag, SlruKind};
@@ -20,7 +19,6 @@ use postgres_ffi::{Oid, TimestampTz, TransactionId};
 use serde::{Deserialize, Serialize};
 use std::collections::{hash_map, HashMap, HashSet};
 use std::ops::Range;
-use tokio_util::sync::CancellationToken;
 use tracing::{debug, trace, warn};
 use utils::{bin_ser::BeSer, lsn::Lsn};

@@ -35,14 +33,6 @@ pub enum LsnForTimestamp {
    NoData(Lsn),
 }

-#[derive(Debug, thiserror::Error)]
-pub enum CalculateLogicalSizeError {
-    #[error("cancelled")]
-    Cancelled,
-    #[error(transparent)]
-    Other(#[from] anyhow::Error),
-}
-
 ///
 /// This impl provides all the functionality to store PostgreSQL relations, SLRUs,
 /// and other special kinds of files, in a versioned key-value store. The
@@ -92,76 +82,83 @@ 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> {
+        ctx: &TimelineRequestContext,
+    ) -> 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, ctx).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, ctx).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> {
+        ctx: &TimelineRequestContext,
+    ) -> 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, ctx).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, ctx).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> {
+        ctx: &TimelineRequestContext,
+    ) -> 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, ctx).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, ctx).await?;
        let nblocks = buf.get_u32_le();

        if latest {
@@ -174,47 +171,51 @@ 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> {
+        ctx: &TimelineRequestContext,
+    ) -> 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, ctx).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>> {
+        ctx: &TimelineRequestContext,
+    ) -> 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, ctx).await?;

        match RelDirectory::des(&buf).context("deserialization failure") {
            Ok(dir) => {
@@ -226,53 +227,56 @@ 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> {
+        ctx: &TimelineRequestContext,
+    ) -> Result<Bytes, PageReconstructError> {
        let key = slru_block_to_key(kind, segno, blknum);
-        self.get(key, lsn)
+        self.get(key, lsn, ctx).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> {
+        ctx: &TimelineRequestContext,
+    ) -> 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, ctx).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> {
+        ctx: &TimelineRequestContext,
+    ) -> 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, ctx).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 +287,11 @@ 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> {
+        ctx: &TimelineRequestContext,
+    ) -> 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,15 @@ 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,
+                    ctx,
+                )
+                .await?;

            if cmp {
                high = mid;
@@ -319,15 +327,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 +345,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 +357,25 @@ 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
-            ));
+        ctx: &TimelineRequestContext,
+    ) -> Result<bool, PageReconstructError> {
+        for segno in self
+            .list_slru_segments(SlruKind::Clog, probe_lsn, ctx)
+            .await?
+        {
+            let nblocks = self
+                .get_slru_segment_size(SlruKind::Clog, segno, probe_lsn, ctx)
+                .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, ctx)
+                    .await?;

                if clog_page.len() == BLCKSZ as usize + 8 {
                    let mut timestamp_bytes = [0u8; 8];
@@ -377,76 +384,99 @@ 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>> {
+        ctx: &TimelineRequestContext,
+    ) -> 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, ctx).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> {
+        ctx: &TimelineRequestContext,
+    ) -> Result<Bytes, PageReconstructError> {
        let key = relmap_file_key(spcnode, dbnode);

-        let buf = try_no_ondemand_download!(self.get(key, lsn));
-        PageReconstructResult::Success(buf)
+        let buf = self.get(key, lsn, ctx).await?;
+        Ok(buf)
    }

-    pub fn list_dbdirs(&self, lsn: Lsn) -> PageReconstructResult<HashMap<(Oid, Oid), bool>> {
+    pub async fn list_dbdirs(
+        &self,
+        lsn: Lsn,
+        ctx: &TimelineRequestContext,
+    ) -> 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, ctx).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,
+        ctx: &TimelineRequestContext,
+    ) -> 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, ctx).await?;
+        Ok(buf)
    }

-    pub fn list_twophase_files(&self, lsn: Lsn) -> PageReconstructResult<HashSet<TransactionId>> {
+    pub async fn list_twophase_files(
+        &self,
+        lsn: Lsn,
+        ctx: &TimelineRequestContext,
+    ) -> 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, ctx).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,
+        ctx: &TimelineRequestContext,
+    ) -> Result<Bytes, PageReconstructError> {
+        self.get(CONTROLFILE_KEY, lsn, ctx).await
    }

-    pub fn get_checkpoint(&self, lsn: Lsn) -> PageReconstructResult<Bytes> {
-        self.get(CHECKPOINT_KEY, lsn)
+    pub async fn get_checkpoint(
+        &self,
+        lsn: Lsn,
+        ctx: &TimelineRequestContext,
+    ) -> Result<Bytes, PageReconstructError> {
+        self.get(CHECKPOINT_KEY, lsn, ctx).await
    }

    /// Does the same as get_current_logical_size but counted on demand.
@@ -457,23 +487,20 @@ impl Timeline {
    pub async fn get_current_logical_size_non_incremental(
        &self,
        lsn: Lsn,
-        cancel: CancellationToken,
-    ) -> Result<u64, CalculateLogicalSizeError> {
+        ctx: &TimelineRequestContext,
+    ) -> Result<u64, PageReconstructError> {
        // Fetch list of database dirs and iterate them
-        let buf = self.get_download(DBDIR_KEY, lsn).await?;
+        let buf = self.get(DBDIR_KEY, lsn, ctx).await?;
        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?
-            {
-                if cancel.is_cancelled() {
-                    return Err(CalculateLogicalSizeError::Cancelled);
+            for rel in self.list_rels(*spcnode, *dbnode, lsn, ctx).await? {
+                if ctx.is_cancelled() {
+                    return Err(PageReconstructError::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, ctx).await?;
                let relsize = buf.get_u32_le();

                total_size += relsize as u64;
@@ -486,7 +513,11 @@ impl Timeline {
    /// Get a KeySpace that covers all the Keys that are in use at the given LSN.
    /// Anything that's not listed maybe removed from the underlying storage (from
    /// that LSN forwards).
-    pub async fn collect_keyspace(&self, lsn: Lsn) -> anyhow::Result<KeySpace> {
+    pub async fn collect_keyspace(
+        &self,
+        lsn: Lsn,
+        ctx: &TimelineRequestContext,
+    ) -> anyhow::Result<KeySpace> {
        // Iterate through key ranges, greedily packing them into partitions
        let mut result = KeySpaceAccum::new();

@@ -494,7 +525,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, ctx).await?;
        let dbdir = DbDirectory::des(&buf).context("deserialization failure")?;

        let mut dbs: Vec<(Oid, Oid)> = dbdir.dbdirs.keys().cloned().collect();
@@ -503,15 +534,16 @@ 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, ctx)
+                .await?
+                .iter()
+                .cloned()
+                .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, ctx).await?;
                let relsize = buf.get_u32_le();

                result.add_range(rel_block_to_key(rel, 0)..rel_block_to_key(rel, relsize));
@@ -527,13 +559,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, ctx).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, ctx).await?;
                let segsize = buf.get_u32_le();

                result.add_range(
@@ -545,7 +577,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, ctx).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 +735,15 @@ 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,
+        ctx: &TimelineRequestContext,
+    ) -> 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, ctx).await?;
        let mut dbdir = DbDirectory::des(&buf)?;

        let r = dbdir.dbdirs.insert((spcnode, dbnode), true);
@@ -731,9 +769,14 @@ 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,
+        ctx: &TimelineRequestContext,
+    ) -> anyhow::Result<()> {
        // Add it to the directory entry
-        let buf = self.get(TWOPHASEDIR_KEY).no_ondemand_download()?;
+        let buf = self.get(TWOPHASEDIR_KEY, ctx).await?;
        let mut dir = TwoPhaseDirectory::des(&buf)?;
        if !dir.xids.insert(xid) {
            anyhow::bail!("twophase file for xid {} already exists", xid);
@@ -757,16 +800,21 @@ 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,
+        ctx: &TimelineRequestContext,
+    ) -> 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()?;
+            .get_db_size(spcnode, dbnode, req_lsn, true, ctx)
+            .await?;

        // Remove entry from dbdir
-        let buf = self.get(DBDIR_KEY).no_ondemand_download()?;
+        let buf = self.get(DBDIR_KEY, ctx).await?;
        let mut dir = DbDirectory::des(&buf)?;
        if dir.dbdirs.remove(&(spcnode, dbnode)).is_some() {
            let buf = DbDirectory::ser(&dir)?;
@@ -789,11 +837,16 @@ 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,
+        ctx: &TimelineRequestContext,
+    ) -> 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, ctx).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 +858,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, ctx).await?)?
        };

        // Add the new relation to the rel directory entry, and write it back
@@ -833,17 +886,18 @@ 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,
+        ctx: &TimelineRequestContext,
+    ) -> 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, ctx).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, ctx).await?.get_u32_le();

            // Update the entry with the new size.
            let buf = nblocks.to_le_bytes();
@@ -863,12 +917,17 @@ 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,
+        ctx: &TimelineRequestContext,
+    ) -> 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, ctx).await?.get_u32_le();

        // only extend relation here. never decrease the size
        if nblocks > old_size {
@@ -884,12 +943,16 @@ 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,
+        ctx: &TimelineRequestContext,
+    ) -> 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, ctx).await?;
        let mut dir = RelDirectory::des(&buf)?;

        if dir.rels.remove(&(rel.relnode, rel.forknum)) {
@@ -900,7 +963,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, ctx).await?.get_u32_le();
        self.pending_nblocks -= old_size as i64;

        // Remove enty from relation size cache
@@ -912,15 +975,16 @@ impl<'a> DatadirModification<'a> {
        Ok(())
    }

-    pub fn put_slru_segment_creation(
+    pub async fn put_slru_segment_creation(
        &mut self,
        kind: SlruKind,
        segno: u32,
        nblocks: BlockNumber,
+        ctx: &TimelineRequestContext,
    ) -> 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, ctx).await?;
        let mut dir = SlruSegmentDirectory::des(&buf)?;

        if !dir.segments.insert(segno) {
@@ -956,10 +1020,15 @@ 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,
+        ctx: &TimelineRequestContext,
+    ) -> 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, ctx).await?;
        let mut dir = SlruSegmentDirectory::des(&buf)?;

        if !dir.segments.remove(&segno) {
@@ -983,9 +1052,13 @@ 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,
+        ctx: &TimelineRequestContext,
+    ) -> anyhow::Result<()> {
        // Remove it from the directory entry
-        let buf = self.get(TWOPHASEDIR_KEY).no_ondemand_download()?;
+        let buf = self.get(TWOPHASEDIR_KEY, ctx).await?;
        let mut dir = TwoPhaseDirectory::des(&buf)?;

        if !dir.xids.remove(&xid) {
@@ -1079,7 +1152,11 @@ impl<'a> DatadirModification<'a> {

    // Internal helper functions to batch the modifications

-    fn get(&self, key: Key) -> PageReconstructResult<Bytes> {
+    async fn get(
+        &self,
+        key: Key,
+        ctx: &TimelineRequestContext,
+    ) -> Result<Bytes, PageReconstructError> {
        // Have we already updated the same key? Read the pending updated
        // version in that case.
        //
@@ -1087,18 +1164,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, ctx).await
        }
    }

@@ -1505,17 +1584,18 @@ fn is_slru_block_key(key: Key) -> bool {

 #[cfg(test)]
 pub fn create_test_timeline(
-    tenant: &crate::tenant::Tenant,
+    tenant: &std::sync::Arc<crate::tenant::Tenant>,
    timeline_id: utils::id::TimelineId,
    pg_version: u32,
-) -> anyhow::Result<std::sync::Arc<Timeline>> {
-    let tline = tenant
-        .create_empty_timeline(timeline_id, Lsn(8), pg_version)?
-        .initialize()?;
+    tenant_ctx: &crate::tenant::TenantRequestContext,
+) -> anyhow::Result<(std::sync::Arc<Timeline>, TimelineRequestContext)> {
+    let (tline, timeline_ctx) =
+        tenant.create_empty_timeline(timeline_id, Lsn(8), pg_version, tenant_ctx)?;
+    let tline = tline.initialize(&timeline_ctx)?;
    let mut m = tline.begin_modification(Lsn(8));
    m.init_empty()?;
    m.commit()?;
-    Ok(tline)
+    Ok((tline, timeline_ctx))
 }

 #[allow(clippy::bool_assert_comparison)]
@@ -1,59 +1,21 @@
 //!
-//! This module provides centralized handling of tokio tasks in the Page Server.
+//! This module provides some helpers for spawning tokio tasks in the pageserver.
 //!
-//! We provide a few basic facilities:
-//! - A global registry of tasks that lists what kind of tasks they are, and
-//!   which tenant or timeline they are working on
-//!
-//! - The ability to request a task to shut down.
-//!
-//!
-//! # How it works?
-//!
-//! There is a global hashmap of all the tasks (`TASKS`). Whenever a new
-//! task is spawned, a PageServerTask entry is added there, and when a
-//! task dies, it removes itself from the hashmap. If you want to kill a
-//! task, you can scan the hashmap to find it.
-//!
-//! # Task shutdown
-//!
-//! To kill a task, we rely on co-operation from the victim. Each task is
-//! expected to periodically call the `is_shutdown_requested()` function, and
-//! if it returns true, exit gracefully. In addition to that, when waiting for
-//! the network or other long-running operation, you can use
-//! `shutdown_watcher()` function to get a Future that will become ready if
-//! the current task has been requested to shut down. You can use that with
-//! Tokio select!().
-//!
-//! TODO: This would be a good place to also handle panics in a somewhat sane way.
-//! Depending on what task panics, we might want to kill the whole server, or
-//! only a single tenant or timeline.
+//! Mostly just a wrapper around tokio::spawn, with some code to handle panics.
 //!

-// Clippy 1.60 incorrectly complains about the tokio::task_local!() macro.
-// Silence it. See https://github.com/rust-lang/rust-clippy/issues/9224.
-#![allow(clippy::declare_interior_mutable_const)]
-
-use std::collections::HashMap;
-use std::fmt;
 use std::future::Future;
-use std::panic::AssertUnwindSafe;
-use std::sync::atomic::{AtomicU64, Ordering};
-use std::sync::{Arc, Mutex};
+use std::panic::{resume_unwind, AssertUnwindSafe};

 use futures::FutureExt;
 use tokio::runtime::Runtime;
 use tokio::task::JoinHandle;
-use tokio::task_local;
-use tokio_util::sync::CancellationToken;

-use tracing::{debug, error, info, warn};
+use tracing::{debug, error, info};

 use once_cell::sync::Lazy;

-use utils::id::{TenantId, TimelineId};
-
-use crate::shutdown_pageserver;
+use crate::context::{self, TaskKind};

 //
 // There are four runtimes:
@@ -92,10 +54,6 @@ use crate::shutdown_pageserver;
 // runtime. If a GetPage request comes in before the load of a tenant has finished, the
 // GetPage request will wait for the tenant load to finish.
 //
-// The core Timeline code is synchronous, and uses a bunch of std Mutexes and RWLocks to
-// protect data structures. Let's keep it that way. Synchronous code is easier to debug
-// and analyze, and there's a lot of hairy, low-level, performance critical code there.
-//
 // It's nice to have different runtimes, so that you can quickly eyeball how much CPU
 // time each class of operations is taking, with 'top -H' or similar.
 //
@@ -135,355 +93,81 @@ pub static BACKGROUND_RUNTIME: Lazy<Runtime> = Lazy::new(|| {
        .expect("Failed to create background op runtime")
 });

-#[derive(Debug, Clone, Copy)]
-pub struct PageserverTaskId(u64);
-
-impl fmt::Display for PageserverTaskId {
-    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        self.0.fmt(f)
-    }
-}
-
-/// Each task that we track is associated with a "task ID". It's just an
-/// increasing number that we assign. Note that it is different from tokio::task::Id.
-static NEXT_TASK_ID: AtomicU64 = AtomicU64::new(1);
-
-/// Global registry of tasks
-static TASKS: Lazy<Mutex<HashMap<u64, Arc<PageServerTask>>>> =
-    Lazy::new(|| Mutex::new(HashMap::new()));
-
-task_local! {
-    // This is a cancellation token which will be cancelled when a task needs to shut down. The
-    // root token is kept in the global registry, so that anyone can send the signal to request
-    // task shutdown.
-    static SHUTDOWN_TOKEN: CancellationToken;
-
-    // Each task holds reference to its own PageServerTask here.
-    static CURRENT_TASK: Arc<PageServerTask>;
-}
-
-///
-/// There are many kinds of tasks in the system. Some are associated with a particular
-/// tenant or timeline, while others are global.
-///
-/// Note that we don't try to limit how many task of a certain kind can be running
-/// at the same time.
-///
-#[derive(Debug, PartialEq, Eq, Clone, Copy)]
-pub enum TaskKind {
-    // libpq listener task. It just accepts connection and spawns a
-    // PageRequestHandler task for each connection.
-    LibpqEndpointListener,
-
-    // HTTP endpoint listener.
-    HttpEndpointListener,
-
-    // Task that handles a single connection. A PageRequestHandler task
-    // starts detached from any particular tenant or timeline, but it can be
-    // associated with one later, after receiving a command from the client.
-    PageRequestHandler,
-
-    // Manages the WAL receiver connection for one timeline. It subscribes to
-    // events from storage_broker, decides which safekeeper to connect to. It spawns a
-    // separate WalReceiverConnection task to handle each connection.
-    WalReceiverManager,
-
-    // Handles a connection to a safekeeper, to stream WAL to a timeline.
-    WalReceiverConnection,
-
-    // Garbage collection worker. One per tenant
-    GarbageCollector,
-
-    // Compaction. One per tenant.
-    Compaction,
-
-    // Initial logical size calculation
-    InitialLogicalSizeCalculation,
-
-    // Task that flushes frozen in-memory layers to disk
-    LayerFlushTask,
-
-    // Task that uploads a file to remote storage
-    RemoteUploadTask,
-
-    // Task that downloads a file from remote storage
-    RemoteDownloadTask,
-
-    // task that handles the initial downloading of all tenants
-    InitialLoad,
-
-    // task that handles attaching a tenant
-    Attach,
-
-    // task that handhes metrics collection
-    MetricsCollection,
-
-    // task that drives downloading layers
-    DownloadAllRemoteLayers,
-}
-
-#[derive(Default)]
-struct MutableTaskState {
-    /// Tenant and timeline that this task is associated with.
-    tenant_id: Option<TenantId>,
-    timeline_id: Option<TimelineId>,
-
-    /// Handle for waiting for the task to exit. It can be None, if the
-    /// the task has already exited.
-    join_handle: Option<JoinHandle<()>>,
-}
-
-struct PageServerTask {
-    #[allow(dead_code)] // unused currently
-    task_id: PageserverTaskId,
-
-    kind: TaskKind,
-
-    name: String,
-
-    // To request task shutdown, just cancel this token.
-    cancel: CancellationToken,
-
-    mutable: Mutex<MutableTaskState>,
-}
-
 /// Launch a new task
-/// Note: if shutdown_process_on_error is set to true failure
-///   of the task will lead to shutdown of entire process
+///
+/// This is a wrapper around tokio::spawn. One difference is that the Future
+/// is marked to return nothing to avoid silently swallowing errors. This
+/// forces the future to handle errors by itself. If you need the return
+/// value, you could create another function that passes it through, but we
+/// don't have a need for that currently.
+///
+/// If shutdown_process_on_panic is set to true, panic of the task will lead
+/// to shutdown of entire process. Otherwise we log the panic and continue.
 pub fn spawn<F>(
    runtime: &tokio::runtime::Handle,
-    kind: TaskKind,
-    tenant_id: Option<TenantId>,
-    timeline_id: Option<TimelineId>,
    name: &str,
-    shutdown_process_on_error: bool,
+    shutdown_process_on_panic: bool,
    future: F,
-) -> PageserverTaskId
+) -> JoinHandle<F::Output>
 where
-    F: Future<Output = anyhow::Result<()>> + Send + 'static,
+    F: Future<Output = ()> + Send + 'static,
 {
-    let cancel = CancellationToken::new();
-    let task_id = NEXT_TASK_ID.fetch_add(1, Ordering::Relaxed);
-    let task = Arc::new(PageServerTask {
-        task_id: PageserverTaskId(task_id),
-        kind,
-        name: name.to_string(),
-        cancel: cancel.clone(),
-        mutable: Mutex::new(MutableTaskState {
-            tenant_id,
-            timeline_id,
-            join_handle: None,
-        }),
-    });
-
-    TASKS.lock().unwrap().insert(task_id, Arc::clone(&task));
-
-    let mut task_mut = task.mutable.lock().unwrap();
-
    let task_name = name.to_string();
-    let task_cloned = Arc::clone(&task);
-    let join_handle = runtime.spawn(task_wrapper(
-        task_name,
-        task_id,
-        task_cloned,
-        cancel,
-        shutdown_process_on_error,
-        future,
-    ));
-    task_mut.join_handle = Some(join_handle);
-    drop(task_mut);
-
-    // The task is now running. Nothing more to do here
-    PageserverTaskId(task_id)
+    runtime.spawn(task_wrapper(task_name, shutdown_process_on_panic, future))
 }

-/// This wrapper function runs in a newly-spawned task. It initializes the
-/// task-local variables and calls the payload function.
-async fn task_wrapper<F>(
-    task_name: String,
-    task_id: u64,
-    task: Arc<PageServerTask>,
-    shutdown_token: CancellationToken,
-    shutdown_process_on_error: bool,
-    future: F,
-) where
-    F: Future<Output = anyhow::Result<()>> + Send + 'static,
+/// This wrapper function runs in a newly-spawned task. To handle panics.
+async fn task_wrapper<F, R>(task_name: String, shutdown_process_on_panic: bool, future: F) -> R
+where
+    F: Future<Output = R> + Send + 'static,
 {
    debug!("Starting task '{}'", task_name);

-    let result = SHUTDOWN_TOKEN
-        .scope(
-            shutdown_token,
-            CURRENT_TASK.scope(task, {
-                // We use AssertUnwindSafe here so that the payload function
-                // doesn't need to be UnwindSafe. We don't do anything after the
-                // unwinding that would expose us to unwind-unsafe behavior.
-                AssertUnwindSafe(future).catch_unwind()
-            }),
-        )
-        .await;
-    task_finish(result, task_name, task_id, shutdown_process_on_error).await;
-}
+    // We use AssertUnwindSafe here so that the payload function
+    // doesn't need to be UnwindSafe. We don't do anything after the
+    // unwinding that would expose us to unwind-unsafe behavior.
+    let result = AssertUnwindSafe(future).catch_unwind().await;

-async fn task_finish(
-    result: std::result::Result<
-        anyhow::Result<()>,
-        std::boxed::Box<dyn std::any::Any + std::marker::Send>,
-    >,
-    task_name: String,
-    task_id: u64,
-    shutdown_process_on_error: bool,
-) {
-    // Remove our entry from the global hashmap.
-    let task = TASKS
-        .lock()
-        .unwrap()
-        .remove(&task_id)
-        .expect("no task in registry");
-
-    let mut shutdown_process = false;
-    {
-        let task_mut = task.mutable.lock().unwrap();
-
-        match result {
-            Ok(Ok(())) => {
-                debug!("Task '{}' exited normally", task_name);
-            }
-            Ok(Err(err)) => {
-                if shutdown_process_on_error {
-                    error!(
-                        "Shutting down: task '{}' tenant_id: {:?}, timeline_id: {:?} exited with error: {:?}",
-                        task_name, task_mut.tenant_id, task_mut.timeline_id, err
-                    );
-                    shutdown_process = true;
-                } else {
-                    error!(
-                        "Task '{}' tenant_id: {:?}, timeline_id: {:?} exited with error: {:?}",
-                        task_name, task_mut.tenant_id, task_mut.timeline_id, err
-                    );
-                }
-            }
-            Err(err) => {
-                if shutdown_process_on_error {
-                    error!(
-                        "Shutting down: task '{}' tenant_id: {:?}, timeline_id: {:?} panicked: {:?}",
-                        task_name, task_mut.tenant_id, task_mut.timeline_id, err
-                    );
-                    shutdown_process = true;
-                } else {
-                    error!(
-                        "Task '{}' tenant_id: {:?}, timeline_id: {:?} panicked: {:?}",
-                        task_name, task_mut.tenant_id, task_mut.timeline_id, err
-                    );
-                }
+    // Handle panics
+    match result {
+        Ok(result) => {
+            debug!("Task '{}' exited normally", task_name);
+            result
+        }
+        Err(err) => {
+            if shutdown_process_on_panic {
+                error!("Shutting down: task '{}' panicked: {:?}", task_name, err);
+                shutdown_pageserver(1).await;
+                unreachable!();
+            } else {
+                error!("Task '{}'  panicked: {:?}", task_name, err);
+                resume_unwind(err);
            }
        }
    }
-
-    if shutdown_process {
-        shutdown_pageserver(1).await;
-    }
 }

-// expected to be called from the task of the given id.
-pub fn associate_with(tenant_id: Option<TenantId>, timeline_id: Option<TimelineId>) {
-    CURRENT_TASK.with(|ct| {
-        let mut task_mut = ct.mutable.lock().unwrap();
-        task_mut.tenant_id = tenant_id;
-        task_mut.timeline_id = timeline_id;
-    });
-}
-
-/// Is there a task running that matches the criteria
-
-/// Signal and wait for tasks to shut down.
 ///
+/// Perform pageserver shutdown. This is called on receiving a signal,
+/// or if one of the tasks marked as 'shutdown_process_on_error' dies.
 ///
-/// The arguments are used to select the tasks to kill. Any None arguments are
-/// ignored. For example, to shut down all WalReceiver tasks:
-///
-///   shutdown_tasks(Some(TaskKind::WalReceiver), None, None)
-///
-/// Or to shut down all tasks for given timeline:
-///
-///   shutdown_tasks(None, Some(tenant_id), Some(timeline_id))
-///
-pub async fn shutdown_tasks(
-    kind: Option<TaskKind>,
-    tenant_id: Option<TenantId>,
-    timeline_id: Option<TimelineId>,
-) {
-    let mut victim_tasks = Vec::new();
+/// This never returns.
+pub async fn shutdown_pageserver(exit_code: i32) {
+    // Shut down the libpq endpoint task. This prevents new connections from
+    // being accepted.
+    context::shutdown_tasks(TaskKind::LibpqEndpointListener).await;

-    {
-        let tasks = TASKS.lock().unwrap();
-        for task in tasks.values() {
-            let task_mut = task.mutable.lock().unwrap();
-            if (kind.is_none() || Some(task.kind) == kind)
-                && (tenant_id.is_none() || task_mut.tenant_id == tenant_id)
-                && (timeline_id.is_none() || task_mut.timeline_id == timeline_id)
-            {
-                task.cancel.cancel();
-                victim_tasks.push(Arc::clone(task));
-            }
-        }
-    }
+    // Shut down all tenants gracefully
+    crate::tenant::mgr::shutdown_all_tenants().await;

-    for task in victim_tasks {
-        let join_handle = {
-            let mut task_mut = task.mutable.lock().unwrap();
-            info!("waiting for {} to shut down", task.name);
-            let join_handle = task_mut.join_handle.take();
-            drop(task_mut);
-            join_handle
-        };
-        if let Some(join_handle) = join_handle {
-            let _ = join_handle.await;
-        } else {
-            // Possibly one of:
-            //  * The task had not even fully started yet.
-            //  * It was shut down concurrently and already exited
-        }
-    }
-}
-
-pub fn current_task_kind() -> Option<TaskKind> {
-    CURRENT_TASK.try_with(|ct| ct.kind).ok()
-}
-
-pub fn current_task_id() -> Option<PageserverTaskId> {
-    CURRENT_TASK.try_with(|ct| ct.task_id).ok()
-}
-
-/// A Future that can be used to check if the current task has been requested to
-/// shut down.
-pub async fn shutdown_watcher() {
-    let token = SHUTDOWN_TOKEN
-        .try_with(|t| t.clone())
-        .expect("shutdown_requested() called in an unexpected task or thread");
-
-    token.cancelled().await;
-}
-
-/// Clone the current task's cancellation token, which can be moved across tasks.
-///
-/// When the task which is currently executing is shutdown, the cancellation token will be
-/// cancelled. It can however be moved to other tasks, such as `tokio::task::spawn_blocking` or
-/// `tokio::task::JoinSet::spawn`.
-pub fn shutdown_token() -> CancellationToken {
-    SHUTDOWN_TOKEN
-        .try_with(|t| t.clone())
-        .expect("shutdown_token() called in an unexpected task or thread")
-}
-
-/// Has the current task been requested to shut down?
-pub fn is_shutdown_requested() -> bool {
-    if let Ok(cancel) = SHUTDOWN_TOKEN.try_with(|t| t.clone()) {
-        cancel.is_cancelled()
-    } else {
-        if !cfg!(test) {
-            warn!("is_shutdown_requested() called in an unexpected task or thread");
-        }
-        false
-    }
+    // Shut down the HTTP endpoint last, so that you can still check the server's
+    // status while it's shutting down.
+    // FIXME: We should probably stop accepting commands like attach/detach earlier.
+    context::shutdown_tasks(TaskKind::HttpEndpointListener).await;
+
+    // There should be nothing left, but let's be sure
+    context::shutdown_all_tasks().await;
+
+    info!("Shut down successfully completed");
+    std::process::exit(exit_code);
 }
@@ -260,8 +260,10 @@ where
    /// contain the version, even if it's missing from the returned
    /// layer.
    ///
+    /// 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;
@@ -8,6 +8,8 @@ use std::sync::Arc;
 use tokio::fs;

 use anyhow::Context;
+use futures::stream::FuturesUnordered;
+use futures::StreamExt;
 use once_cell::sync::Lazy;
 use tokio::sync::RwLock;
 use tracing::*;
@@ -16,9 +18,9 @@ use remote_storage::GenericRemoteStorage;
 use utils::crashsafe;

 use crate::config::PageServerConf;
-use crate::task_mgr::{self, TaskKind};
+use crate::context::RequestContext;
 use crate::tenant::config::TenantConfOpt;
-use crate::tenant::{Tenant, TenantState};
+use crate::tenant::{Tenant, TenantRequestContext, TenantState};
 use crate::IGNORED_TENANT_FILE_NAME;

 use utils::fs_ext::PathExt;
@@ -181,25 +183,11 @@ pub async fn shutdown_all_tenants() {
        tenants_to_shut_down
    };

-    // Shut down all existing walreceiver connections and stop accepting the new ones.
-    task_mgr::shutdown_tasks(Some(TaskKind::WalReceiverManager), None, None).await;
-
-    // Ok, no background tasks running anymore. Flush any remaining data in
-    // memory to disk.
-    //
-    // We assume that any incoming connections that might request pages from
-    // the tenant have already been terminated by the caller, so there
-    // should be no more activity in any of the repositories.
-    //
-    // On error, log it but continue with the shutdown for other tenants.
-    for tenant in tenants_to_shut_down {
-        let tenant_id = tenant.tenant_id();
-        debug!("shutdown tenant {tenant_id}");
-
-        if let Err(err) = tenant.freeze_and_flush().await {
-            error!("Could not checkpoint tenant {tenant_id} during shutdown: {err:?}");
-        }
+    let mut shutdown_futures: FuturesUnordered<_> = FuturesUnordered::new();
+    for tenant in tenants_to_shut_down.iter() {
+        shutdown_futures.push(tenant.graceful_shutdown(true));
    }
+    while let Some(_result) = shutdown_futures.next().await {}
 }

 pub async fn create_tenant(
@@ -234,36 +222,47 @@ pub async fn update_tenant_config(
    conf: &'static PageServerConf,
    tenant_conf: TenantConfOpt,
    tenant_id: TenantId,
+    ctx: &RequestContext,
 ) -> anyhow::Result<()> {
    info!("configuring tenant {tenant_id}");
-    get_tenant(tenant_id, true)
-        .await?
-        .update_tenant_config(tenant_conf);
+    let (tenant, _ctx) = get_active_tenant(tenant_id, ctx).await?;
+
+    tenant.update_tenant_config(tenant_conf);
    Tenant::persist_tenant_config(&conf.tenant_config_path(tenant_id), tenant_conf, false)?;
    Ok(())
 }

 /// Gets the tenant from the in-memory data, erroring if it's absent or is not fitting to the query.
 /// `active_only = true` allows to query only tenants that are ready for operations, erroring on other kinds of tenants.
-pub async fn get_tenant(tenant_id: TenantId, active_only: bool) -> anyhow::Result<Arc<Tenant>> {
+pub async fn get_active_tenant(
+    tenant_id: TenantId,
+    parent_ctx: &RequestContext,
+) -> anyhow::Result<(Arc<Tenant>, TenantRequestContext)> {
+    let tenant = get_tenant(tenant_id).await?;
+    let tenant_ctx = match tenant.get_context(parent_ctx) {
+        Ok(ctx) => ctx,
+        Err(state) => anyhow::bail!("Tenant {} is not active, state: {:?}", tenant_id, state,),
+    };
+    Ok((tenant, tenant_ctx))
+}
+
+pub async fn get_tenant(tenant_id: TenantId) -> anyhow::Result<Arc<Tenant>> {
    let m = TENANTS.read().await;
    let tenant = m
        .get(&tenant_id)
        .with_context(|| format!("Tenant {tenant_id} not found in the local state"))?;
-    if active_only && !tenant.is_active() {
-        anyhow::bail!(
-            "Tenant {tenant_id} is not active. Current state: {:?}",
-            tenant.current_state()
-        )
-    } else {
-        Ok(Arc::clone(tenant))
-    }
+
+    Ok(Arc::clone(tenant))
 }

-pub async fn delete_timeline(tenant_id: TenantId, timeline_id: TimelineId) -> anyhow::Result<()> {
-    match get_tenant(tenant_id, true).await {
-        Ok(tenant) => {
-            tenant.delete_timeline(timeline_id).await?;
+pub async fn delete_timeline(
+    tenant_id: TenantId,
+    timeline_id: TimelineId,
+    ctx: &RequestContext,
+) -> anyhow::Result<()> {
+    match get_active_tenant(tenant_id, ctx).await {
+        Ok((tenant, ctx)) => {
+            tenant.delete_timeline(timeline_id, &ctx).await?;
        }
        Err(e) => anyhow::bail!("Cannot access tenant {tenant_id} in local tenant state: {e:?}"),
    }
@@ -395,27 +394,31 @@ where
    // The exclusive lock here ensures we don't miss the tenant state updates before trying another removal.
    // tenant-wde cleanup operations may take some time (removing the entire tenant directory), we want to
    // avoid holding the lock for the entire process.
-    {
+    let tenant = {
        let tenants_accessor = TENANTS.write().await;
        match tenants_accessor.get(&tenant_id) {
            Some(tenant) => match tenant.current_state() {
                TenantState::Attaching
                | TenantState::Loading
                | TenantState::Broken
-                | TenantState::Active => tenant.set_stopping(),
+                | TenantState::Active => {
+                    tenant.set_stopping();
+                    Arc::clone(tenant)
+                }
                TenantState::Stopping => {
                    anyhow::bail!("Tenant {tenant_id} is stopping already")
                }
            },
            None => anyhow::bail!("Tenant not found for id {tenant_id}"),
        }
-    }
+    };

-    // shutdown all tenant and timeline tasks: gc, compaction, page service)
-    // No new tasks will be started for this tenant because it's in `Stopping` state.
-    // Hence, once we're done here, the `tenant_cleanup` callback can mutate tenant on-disk state freely.
-    task_mgr::shutdown_tasks(None, Some(tenant_id), None).await;
+    // Shut down all tenant and timeline tasks.
+    tenant.graceful_shutdown(true).await;

+    // All tasks that operated on the tenant or any of its timelines have no finished,
+    // and they are in Stopped state so that new ones cannot appear anymore. Proceed
+    // with the cleanup.
    match tenant_cleanup
        .await
        .with_context(|| format!("Failed to run cleanup for tenant {tenant_id}"))
@@ -430,65 +433,10 @@ 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)
        }
    }
 }
-
-#[cfg(feature = "testing")]
-use {
-    crate::repository::GcResult, pageserver_api::models::TimelineGcRequest,
-    utils::http::error::ApiError,
-};
-
-#[cfg(feature = "testing")]
-pub async fn immediate_gc(
-    tenant_id: TenantId,
-    timeline_id: TimelineId,
-    gc_req: TimelineGcRequest,
-) -> Result<tokio::sync::oneshot::Receiver<Result<GcResult, anyhow::Error>>, 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 gc_horizon = gc_req.gc_horizon.unwrap_or_else(|| tenant.get_gc_horizon());
-    // Use tenant's pitr setting
-    let pitr = tenant.get_pitr_interval();
-
-    // 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::GarbageCollector,
-        Some(tenant_id),
-        Some(timeline_id),
-        &format!("timeline_gc_handler garbage collection run for tenant {tenant_id} timeline {timeline_id}"),
-        false,
-        async move {
-            fail::fail_point!("immediate_gc_task_pre");
-            let result = tenant
-                .gc_iteration(Some(timeline_id), gc_horizon, pitr)
-                .instrument(info_span!("manual_gc", tenant = %tenant_id, timeline = %timeline_id))
-                .await;
-                // FIXME: `gc_iteration` can return an error for multiple reasons; we should handle it
-                // better once the types support it.
-            match task_done.send(result) {
-                Ok(_) => (),
-                Err(result) => error!("failed to send gc 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)
-}
@@ -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:
 //!
@@ -214,6 +214,7 @@ use anyhow::ensure;
 use remote_storage::{DownloadError, GenericRemoteStorage};
 use std::ops::DerefMut;
 use tokio::runtime::Runtime;
+use tokio_util::sync::CancellationToken;
 use tracing::{debug, info, warn};
 use tracing::{info_span, Instrument};
 use utils::lsn::Lsn;
@@ -225,12 +226,12 @@ use crate::tenant::remote_timeline_client::index::LayerFileMetadata;
 use crate::{
    config::PageServerConf,
    task_mgr,
-    task_mgr::TaskKind,
    task_mgr::BACKGROUND_RUNTIME,
    tenant::metadata::TimelineMetadata,
    tenant::upload_queue::{
        UploadOp, UploadQueue, UploadQueueInitialized, UploadQueueStopped, UploadTask,
    },
+    tenant::TimelineRequestContext,
    {exponential_backoff, DEFAULT_BASE_BACKOFF_SECONDS, DEFAULT_MAX_BACKOFF_SECONDS},
 };

@@ -298,8 +299,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,31 +308,56 @@ 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
    /// an index file upload, i.e., it's not empty.
    /// The given `index_part` must be the one on the remote.
-    pub fn init_upload_queue(&self, index_part: &IndexPart) -> anyhow::Result<()> {
+    pub fn init_upload_queue(
+        self: &Arc<Self>,
+        index_part: &IndexPart,
+        upload_ctx: TimelineRequestContext,
+    ) -> anyhow::Result<()> {
+        let cancellation_token = upload_ctx.cancellation_token().clone();
        let mut upload_queue = self.upload_queue.lock().unwrap();
-        upload_queue.initialize_with_current_remote_index_part(index_part)?;
+        upload_queue.initialize_with_current_remote_index_part(index_part, upload_ctx)?;
        self.update_remote_physical_size_gauge(Some(index_part));
+        self.spawn_cancellation_watch(cancellation_token);
        Ok(())
    }

    /// Initialize the upload queue for the case where the remote storage is empty,
    /// i.e., it doesn't have an `IndexPart`.
    pub fn init_upload_queue_for_empty_remote(
-        &self,
+        self: &Arc<Self>,
        local_metadata: &TimelineMetadata,
+        upload_ctx: TimelineRequestContext,
    ) -> anyhow::Result<()> {
+        let cancellation_token = upload_ctx.cancellation_token().clone();
        let mut upload_queue = self.upload_queue.lock().unwrap();
-        upload_queue.initialize_empty_remote(local_metadata)?;
+        upload_queue.initialize_empty_remote(local_metadata, upload_ctx)?;
        self.update_remote_physical_size_gauge(None);
+        self.spawn_cancellation_watch(cancellation_token);
        Ok(())
    }

+    /// Spawn a task that calls `stop` on cancellation. It's important that we
+    /// stop the upload queue promptly, because it holds onto the RequestContext,
+    /// which in turn prevents the Timeline from shutting down.
+    fn spawn_cancellation_watch(self: &Arc<Self>, cancellation_token: CancellationToken) {
+        let self_rc = Arc::clone(self);
+        task_mgr::spawn(
+            self.runtime.handle(),
+            "remote upload queue cancellation watch",
+            false,
+            async move {
+                cancellation_token.cancelled().await;
+                self_rc.stop();
+            },
+        );
+    }
+
    pub fn last_uploaded_consistent_lsn(&self) -> Option<Lsn> {
        match &*self.upload_queue.lock().unwrap() {
            UploadQueue::Uninitialized => None,
@@ -367,6 +393,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 +423,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 +552,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 +584,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 +636,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());
            }
@@ -616,7 +651,10 @@ impl RemoteTimelineClient {
    ///
    /// Wait for all previously scheduled uploads/deletions to complete
    ///
-    pub async fn wait_completion(self: &Arc<Self>) -> anyhow::Result<()> {
+    pub async fn wait_completion(
+        self: &Arc<Self>,
+        ctx: &TimelineRequestContext,
+    ) -> anyhow::Result<()> {
        let (sender, mut receiver) = tokio::sync::watch::channel(());
        let barrier_op = UploadOp::Barrier(sender);

@@ -630,9 +668,16 @@ impl RemoteTimelineClient {
            self.launch_queued_tasks(upload_queue);
        }

-        if receiver.changed().await.is_err() {
-            anyhow::bail!("wait_completion aborted because upload queue was stopped");
-        }
+        tokio::select! {
+            result = receiver.changed() => {
+                if result.is_err() {
+                    anyhow::bail!("wait_completion aborted because upload queue was stopped");
+                }
+            },
+            _ = ctx.cancelled() => {
+                anyhow::bail!("request cancelled while waiting on uploads to finish");
+            },
+        };
        Ok(())
    }

@@ -710,16 +755,15 @@ impl RemoteTimelineClient {

            // Spawn task to perform the task
            let self_rc = Arc::clone(self);
+
+            let cancellation_token = upload_queue.upload_ctx.cancellation_token().clone();
+
            task_mgr::spawn(
                self.runtime.handle(),
-                TaskKind::RemoteUploadTask,
-                Some(self.tenant_id),
-                Some(self.timeline_id),
                "remote upload",
                false,
                async move {
-                    self_rc.perform_upload_task(task).await;
-                    Ok(())
+                    self_rc.perform_upload_task(task, cancellation_token).await;
                }
                .instrument(info_span!(parent: None, "remote_upload", tenant = %self.tenant_id, timeline = %self.timeline_id, upload_task_id = %task_id)),
            );
@@ -739,7 +783,11 @@ impl RemoteTimelineClient {
    /// The task can be shut down, however. That leads to stopping the whole
    /// queue.
    ///
-    async fn perform_upload_task(self: &Arc<Self>, task: Arc<UploadTask>) {
+    async fn perform_upload_task(
+        self: &Arc<Self>,
+        task: Arc<UploadTask>,
+        cancellation_token: CancellationToken,
+    ) {
        // Loop to retry until it completes.
        loop {
            // If we're requested to shut down, close up shop and exit.
@@ -747,13 +795,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() {
+            if cancellation_token.is_cancelled() {
                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;
            }
@@ -849,7 +897,7 @@ impl RemoteTimelineClient {

                    // sleep until it's time to retry, or we're cancelled
                    tokio::select! {
-                        _ = task_mgr::shutdown_watcher() => { },
+                        _ = cancellation_token.cancelled() => { },
                        _ = exponential_backoff(
                            retries,
                            DEFAULT_BASE_BACKOFF_SECONDS,
@@ -901,22 +949,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 +1033,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);
@@ -983,7 +1049,9 @@ impl RemoteTimelineClient {
 #[cfg(test)]
 mod tests {
    use super::*;
+    use crate::context::{DownloadBehavior, RequestContext, TaskKind};
    use crate::tenant::harness::{TenantHarness, TIMELINE_ID};
+    use crate::DEFAULT_PG_VERSION;
    use remote_storage::{RemoteStorageConfig, RemoteStorageKind};
    use std::{collections::HashSet, path::Path};
    use utils::lsn::Lsn;
@@ -1002,7 +1070,7 @@ mod tests {
            Lsn(0),
            // Any version will do
            // but it should be consistent with the one in the tests
-            crate::DEFAULT_PG_VERSION,
+            DEFAULT_PG_VERSION,
        );

        // go through serialize + deserialize to fix the header, including checksum
@@ -1037,9 +1105,19 @@ mod tests {
    // Test scheduling
    #[test]
    fn upload_scheduling() -> anyhow::Result<()> {
+        // Use a current-thread runtime in the test
+        let runtime = Box::leak(Box::new(
+            tokio::runtime::Builder::new_current_thread()
+                .enable_all()
+                .build()?,
+        ));
+        let _entered = runtime.enter();
+
        let harness = TenantHarness::create("upload_scheduling")?;
+        let (tenant, tenant_ctx) = runtime.block_on(harness.load());
+        let (_timeline, timeline_ctx) =
+            tenant.create_empty_timeline(TIMELINE_ID, Lsn(0), DEFAULT_PG_VERSION, &tenant_ctx)?;
        let timeline_path = harness.timeline_path(&TIMELINE_ID);
-        std::fs::create_dir_all(&timeline_path)?;

        let remote_fs_dir = harness.conf.workdir.join("remote_fs");
        std::fs::create_dir_all(remote_fs_dir)?;
@@ -1057,14 +1135,6 @@ mod tests {
            storage: RemoteStorageKind::LocalFs(remote_fs_dir.clone()),
        };

-        // Use a current-thread runtime in the test
-        let runtime = Box::leak(Box::new(
-            tokio::runtime::Builder::new_current_thread()
-                .enable_all()
-                .build()?,
-        ));
-        let _entered = runtime.enter();
-
        // Test outline:
        //
        // Schedule upload of a bunch of layers. Check that they are started immediately, not queued
@@ -1100,7 +1170,11 @@ mod tests {
        println!("remote_timeline_dir: {}", remote_timeline_dir.display());

        let metadata = dummy_metadata(Lsn(0x10));
-        client.init_upload_queue_for_empty_remote(&metadata)?;
+        let upload_ctx = timeline_ctx.register_another(RequestContext::new(
+            TaskKind::RemoteUploadTask,
+            DownloadBehavior::Error,
+        ));
+        client.init_upload_queue_for_empty_remote(&metadata, upload_ctx)?;

        // Create a couple of dummy files,  schedule upload for them
        let content_foo = dummy_contents("foo");
@@ -1140,7 +1214,7 @@ mod tests {
        }

        // Wait for the uploads to finish
-        runtime.block_on(client.wait_completion())?;
+        runtime.block_on(client.wait_completion(&timeline_ctx))?;
        {
            let mut guard = client.upload_queue.lock().unwrap();
            let upload_queue = guard.initialized_mut().unwrap();
@@ -1177,7 +1251,7 @@ mod tests {
        assert_remote_files(&["foo", "bar", "index_part.json"], &remote_timeline_dir);

        // Finish them
-        runtime.block_on(client.wait_completion())?;
+        runtime.block_on(client.wait_completion(&timeline_ctx))?;

        assert_remote_files(&["bar", "baz", "index_part.json"], &remote_timeline_dir);

@@ -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,
@@ -3,10 +3,9 @@ use std::collections::{HashMap, HashSet};
 use std::sync::Arc;

 use anyhow::Context;
-use tokio::sync::oneshot::error::RecvError;
 use tokio::sync::Semaphore;

-use crate::pgdatadir_mapping::CalculateLogicalSizeError;
+use crate::tenant::{PageReconstructError, TenantRequestContext, TimelineRequestContext};

 use super::Tenant;
 use utils::id::TimelineId;
@@ -63,13 +62,14 @@ pub(super) async fn gather_inputs(
    tenant: &Tenant,
    limit: &Arc<Semaphore>,
    logical_size_cache: &mut HashMap<(TimelineId, Lsn), u64>,
+    tenant_ctx: &TenantRequestContext,
 ) -> anyhow::Result<ModelInputs> {
    // with joinset, on drop, all of the tasks will just be de-scheduled, which we can use to
    // our advantage with `?` error handling.
    let mut joinset = tokio::task::JoinSet::new();

    let timelines = tenant
-        .refresh_gc_info()
+        .refresh_gc_info(tenant_ctx)
        .await
        .context("Failed to refresh gc_info before gathering inputs")?;

@@ -97,9 +97,21 @@ pub(super) async fn gather_inputs(
    // used to determine the `retention_period` for the size model
    let mut max_cutoff_distance = None;

+    let mut ctx_dropguards: Vec<tokio_util::sync::DropGuard> = Vec::new();
+
    for timeline in timelines {
        let last_record_lsn = timeline.get_last_record_lsn();

+        let ctx = match timeline.get_context(tenant_ctx) {
+            Ok(ctx) => ctx,
+            Err(state) => {
+                info!("skipping tenant size calculation for timeline because it is in {state:?} state");
+                continue;
+            }
+        };
+        ctx_dropguards.push(ctx.cancellation_token().clone().drop_guard());
+        let ctx = Arc::new(ctx);
+
        let (interesting_lsns, horizon_cutoff, pitr_cutoff, next_gc_cutoff) = {
            // there's a race between the update (holding tenant.gc_lock) and this read but it
            // might not be an issue, because it's not for Timeline::gc
@@ -169,19 +181,23 @@ pub(super) async fn gather_inputs(
            timeline_id: timeline.timeline_id,
        });

-        for (lsn, _kind) in &interesting_lsns {
-            if let Some(size) = logical_size_cache.get(&(timeline.timeline_id, *lsn)) {
+        for (lsn, _kind) in interesting_lsns.iter() {
+            let lsn = *lsn;
+            if let Some(size) = logical_size_cache.get(&(timeline.timeline_id, lsn)) {
                updates.push(Update {
-                    lsn: *lsn,
+                    lsn,
                    timeline_id: timeline.timeline_id,
                    command: Command::Update(*size),
                });

-                needed_cache.insert((timeline.timeline_id, *lsn));
+                needed_cache.insert((timeline.timeline_id, lsn));
            } else {
                let timeline = Arc::clone(&timeline);
                let parallel_size_calcs = Arc::clone(limit);
-                joinset.spawn(calculate_logical_size(parallel_size_calcs, timeline, *lsn));
+                let ctx_clone = Arc::clone(&ctx);
+                joinset.spawn(async move {
+                    calculate_logical_size(parallel_size_calcs, timeline, lsn, &ctx_clone).await
+                });
            }
        }

@@ -357,7 +373,7 @@ enum LsnKind {
 struct TimelineAtLsnSizeResult(
    Arc<crate::tenant::Timeline>,
    utils::lsn::Lsn,
-    Result<u64, CalculateLogicalSizeError>,
+    Result<u64, PageReconstructError>,
 );

 #[instrument(skip_all, fields(timeline_id=%timeline.timeline_id, lsn=%lsn))]
@@ -365,14 +381,13 @@ async fn calculate_logical_size(
    limit: Arc<tokio::sync::Semaphore>,
    timeline: Arc<crate::tenant::Timeline>,
    lsn: utils::lsn::Lsn,
-) -> Result<TimelineAtLsnSizeResult, RecvError> {
+    ctx: &TimelineRequestContext,
+) -> Result<TimelineAtLsnSizeResult, PageReconstructError> {
    let _permit = tokio::sync::Semaphore::acquire_owned(limit)
        .await
-        .expect("global semaphore should not had been closed");
+        .expect("global semaphore should not have been closed");

-    let size_res = timeline
-        .spawn_ondemand_logical_size_calculation(lsn)
-        .await?;
+    let size_res = timeline.calculate_logical_size(lsn, ctx).await;
    Ok(TimelineAtLsnSizeResult(timeline, lsn, size_res))
 }

@@ -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(
@@ -1,45 +1,39 @@
 //! This module contains functions to serve per-tenant background processes,
 //! such as compaction and GC

-use std::ops::ControlFlow;
 use std::sync::Arc;
 use std::time::Duration;

+use crate::context::{DownloadBehavior, RequestContext, TaskKind};
 use crate::metrics::TENANT_TASK_EVENTS;
 use crate::task_mgr;
-use crate::task_mgr::{TaskKind, BACKGROUND_RUNTIME};
-use crate::tenant::mgr;
-use crate::tenant::{Tenant, TenantState};
+use crate::task_mgr::BACKGROUND_RUNTIME;
+use crate::tenant::Tenant;
 use tracing::*;
-use utils::id::TenantId;

-pub fn start_background_loops(tenant_id: TenantId) {
+pub fn start_background_loops(tenant: &Arc<Tenant>) {
+    let tenant_id = tenant.tenant_id;
+
+    let tenant_clone = Arc::clone(tenant);
    task_mgr::spawn(
        BACKGROUND_RUNTIME.handle(),
-        TaskKind::Compaction,
-        Some(tenant_id),
-        None,
        &format!("compactor for tenant {tenant_id}"),
        false,
        async move {
-            compaction_loop(tenant_id)
+            compaction_loop(&tenant_clone)
                .instrument(info_span!("compaction_loop", tenant_id = %tenant_id))
                .await;
-            Ok(())
        },
    );
+    let tenant_clone = Arc::clone(tenant);
    task_mgr::spawn(
        BACKGROUND_RUNTIME.handle(),
-        TaskKind::GarbageCollector,
-        Some(tenant_id),
-        None,
        &format!("garbage collector for tenant {tenant_id}"),
        false,
        async move {
-            gc_loop(tenant_id)
+            gc_loop(&tenant_clone)
                .instrument(info_span!("gc_loop", tenant_id = %tenant_id))
                .await;
-            Ok(())
        },
    );
 }
@@ -47,25 +41,27 @@ pub fn start_background_loops(tenant_id: TenantId) {
 ///
 /// Compaction task's main loop
 ///
-async fn compaction_loop(tenant_id: TenantId) {
+async fn compaction_loop(tenant: &Arc<Tenant>) {
    let wait_duration = Duration::from_secs(2);
    info!("starting");
    TENANT_TASK_EVENTS.with_label_values(&["start"]).inc();
    async {
+        let top_ctx = RequestContext::new(TaskKind::Compaction, DownloadBehavior::Download);
+
+        let tenant_ctx = match tenant.get_context(&top_ctx) {
+            Ok(ctx) => ctx,
+            Err(state) => {
+                // This could happen if the tenant is detached or the pageserver is shut
+                // down immediately after loading or attaching completed and the tenant
+                // was activated. It seems unlikely enough in practice that we better print
+                // a warning, as it could also be a bug.
+                error!("Not running compaction loop, tenant is not active: {state:?}");
+                return;
+            }
+        };
        loop {
            trace!("waking up");

-            let tenant = tokio::select! {
-                _ = task_mgr::shutdown_watcher() => {
-                    info!("received cancellation request");
-                    return;
-                },
-                tenant_wait_result = wait_for_active_tenant(tenant_id, wait_duration) => match tenant_wait_result {
-                    ControlFlow::Break(()) => return,
-                    ControlFlow::Continue(tenant) => tenant,
-                },
-            };
-
            let mut sleep_duration = tenant.get_compaction_period();
            if sleep_duration == Duration::ZERO {
                info!("automatic compaction is disabled");
@@ -73,7 +69,7 @@ async fn compaction_loop(tenant_id: TenantId) {
                sleep_duration = Duration::from_secs(10);
            } else {
                // Run compaction
-                if let Err(e) = tenant.compaction_iteration().await {
+                if let Err(e) = tenant.compaction_iteration(&tenant_ctx).await {
                    sleep_duration = wait_duration;
                    error!("Compaction failed, retrying in {:?}: {e:?}", sleep_duration);
                }
@@ -81,9 +77,9 @@ async fn compaction_loop(tenant_id: TenantId) {

            // Sleep
            tokio::select! {
-                _ = task_mgr::shutdown_watcher() => {
+                _ = tenant_ctx.cancelled() => {
                    info!("received cancellation request during idling");
-                    break ;
+                    break;
                },
                _ = tokio::time::sleep(sleep_duration) => {},
            }
@@ -98,25 +94,28 @@ async fn compaction_loop(tenant_id: TenantId) {
 ///
 /// GC task's main loop
 ///
-async fn gc_loop(tenant_id: TenantId) {
+async fn gc_loop(tenant: &Arc<Tenant>) {
    let wait_duration = Duration::from_secs(2);
    info!("starting");
    TENANT_TASK_EVENTS.with_label_values(&["start"]).inc();
    async {
+        // GC might require downloading, to find the cutoff LSN that corresponds to the
+        // cutoff specified as time.
+        let top_ctx = RequestContext::new(TaskKind::GarbageCollector, DownloadBehavior::Download);
+        let tenant_ctx = match tenant.get_context(&top_ctx) {
+            Ok(ctx) => ctx,
+            Err(state) => {
+                // This could happen if the tenant is detached or the pageserver is shut
+                // down immediately after loading or attaching completed and the tenant
+                // was activated. It seems unlikely enough in practice that we better print
+                // a warning, as it could also be a bug.
+                error!("Not running GC loop, tenant is not active: {state:?}");
+                return;
+            }
+        };
        loop {
            trace!("waking up");

-            let tenant = tokio::select! {
-                _ = task_mgr::shutdown_watcher() => {
-                    info!("received cancellation request");
-                    return;
-                },
-                tenant_wait_result = wait_for_active_tenant(tenant_id, wait_duration) => match tenant_wait_result {
-                    ControlFlow::Break(()) => return,
-                    ControlFlow::Continue(tenant) => tenant,
-                },
-            };
-
            let gc_period = tenant.get_gc_period();
            let gc_horizon = tenant.get_gc_horizon();
            let mut sleep_duration = gc_period;
@@ -127,7 +126,10 @@ async fn gc_loop(tenant_id: TenantId) {
            } else {
                // Run gc
                if gc_horizon > 0 {
-                    if let Err(e) = tenant.gc_iteration(None, gc_horizon, tenant.get_pitr_interval()).await
+                    // Run compaction
+                    if let Err(e) = tenant
+                        .gc_iteration(None, gc_horizon, tenant.get_pitr_interval(), &tenant_ctx)
+                        .await
                    {
                        sleep_duration = wait_duration;
                        error!("Gc failed, retrying in {:?}: {e:?}", sleep_duration);
@@ -137,7 +139,7 @@ async fn gc_loop(tenant_id: TenantId) {

            // Sleep
            tokio::select! {
-                _ = task_mgr::shutdown_watcher() => {
+                _ = tenant_ctx.cancelled() => {
                    info!("received cancellation request during idling");
                    break;
                },
@@ -149,46 +151,3 @@ async fn gc_loop(tenant_id: TenantId) {
    TENANT_TASK_EVENTS.with_label_values(&["stop"]).inc();
    trace!("GC loop stopped.");
 }
-
-async fn wait_for_active_tenant(
-    tenant_id: TenantId,
-    wait: Duration,
-) -> ControlFlow<(), Arc<Tenant>> {
-    let tenant = loop {
-        match mgr::get_tenant(tenant_id, false).await {
-            Ok(tenant) => break tenant,
-            Err(e) => {
-                error!("Failed to get a tenant {tenant_id}: {e:#}");
-                tokio::time::sleep(wait).await;
-            }
-        }
-    };
-
-    // if the tenant has a proper status already, no need to wait for anything
-    if tenant.current_state() == TenantState::Active {
-        ControlFlow::Continue(tenant)
-    } else {
-        let mut tenant_state_updates = tenant.subscribe_for_state_updates();
-        loop {
-            match tenant_state_updates.changed().await {
-                Ok(()) => {
-                    let new_state = *tenant_state_updates.borrow();
-                    match new_state {
-                        TenantState::Active => {
-                            debug!("Tenant state changed to active, continuing the task loop");
-                            return ControlFlow::Continue(tenant);
-                        }
-                        state => {
-                            debug!("Not running the task loop, tenant is not active: {state:?}");
-                            continue;
-                        }
-                    }
-                }
-                Err(_sender_dropped_error) => {
-                    info!("Tenant dropped the state updates sender, quitting waiting for tenant and the task loop");
-                    return ControlFlow::Break(());
-                }
-            }
-        }
-    }
-}
@@ -4,6 +4,7 @@ use super::storage_layer::LayerFileName;
 use crate::tenant::metadata::TimelineMetadata;
 use crate::tenant::remote_timeline_client::index::IndexPart;
 use crate::tenant::remote_timeline_client::index::LayerFileMetadata;
+use crate::tenant::TimelineRequestContext;
 use std::collections::{HashMap, VecDeque};
 use std::fmt::Debug;

@@ -73,6 +74,13 @@ pub(crate) struct UploadQueueInitialized {
    /// tasks to finish. For example, metadata upload cannot be performed before all
    /// preceding layer file uploads have completed.
    pub(crate) queued_operations: VecDeque<UploadOp>,
+
+    /// Context used for the upload tasks. Note that this is associated with the
+    /// Timeline, so this prevents the Timeline from being shut down. To ensure quick
+    /// shutdown, RemoteTimelineClient spawns a task to wait for cancellation on the
+    /// context and stop the queue. Otherwise we woudn't notice the cancellation
+    /// until next upload attempt.
+    pub(crate) upload_ctx: TimelineRequestContext,
 }

 pub(crate) struct UploadQueueStopped {
@@ -83,6 +91,7 @@ impl UploadQueue {
    pub(crate) fn initialize_empty_remote(
        &mut self,
        metadata: &TimelineMetadata,
+        upload_ctx: TimelineRequestContext,
    ) -> anyhow::Result<&mut UploadQueueInitialized> {
        match self {
            UploadQueue::Uninitialized => (),
@@ -108,6 +117,7 @@ impl UploadQueue {
            num_inprogress_deletions: 0,
            inprogress_tasks: HashMap::new(),
            queued_operations: VecDeque::new(),
+            upload_ctx,
        };

        *self = UploadQueue::Initialized(state);
@@ -117,6 +127,7 @@ impl UploadQueue {
    pub(crate) fn initialize_with_current_remote_index_part(
        &mut self,
        index_part: &IndexPart,
+        upload_ctx: TimelineRequestContext,
    ) -> anyhow::Result<&mut UploadQueueInitialized> {
        match self {
            UploadQueue::Uninitialized => (),
@@ -153,6 +164,7 @@ impl UploadQueue {
            num_inprogress_deletions: 0,
            inprogress_tasks: HashMap::new(),
            queued_operations: VecDeque::new(),
+            upload_ctx,
        };

        *self = UploadQueue::Initialized(state);
@@ -31,6 +31,7 @@ use once_cell::sync::OnceCell;
 use std::future::Future;
 use storage_broker::BrokerClientChannel;
 use tokio::sync::watch;
+use tokio_util::sync::CancellationToken;
 use tracing::*;

 pub use connection_manager::spawn_connection_manager_task;
@@ -76,7 +77,7 @@ pub fn is_broker_client_initialized() -> bool {

 /// A handle of an asynchronous task.
 /// The task has a channel that it can use to communicate its lifecycle events in a certain form, see [`TaskEvent`]
-/// and a cancellation channel that it can listen to for earlier interrupts.
+/// and a cancellation token that it can listen to for earlier interrupts.
 ///
 /// Note that the communication happens via the `watch` channel, that does not accumulate the events, replacing the old one with the never one on submission.
 /// That may lead to certain events not being observed by the listener.
@@ -84,7 +85,7 @@ pub fn is_broker_client_initialized() -> bool {
 pub struct TaskHandle<E> {
    join_handle: Option<tokio::task::JoinHandle<anyhow::Result<()>>>,
    events_receiver: watch::Receiver<TaskStateUpdate<E>>,
-    cancellation: watch::Sender<()>,
+    cancellation: CancellationToken,
 }

 pub enum TaskEvent<E> {
@@ -102,20 +103,18 @@ pub enum TaskStateUpdate<E> {
 impl<E: Clone> TaskHandle<E> {
    /// Initializes the task, starting it immediately after the creation.
    pub fn spawn<Fut>(
-        task: impl FnOnce(watch::Sender<TaskStateUpdate<E>>, watch::Receiver<()>) -> Fut
-            + Send
-            + 'static,
+        task: impl FnOnce(watch::Sender<TaskStateUpdate<E>>) -> Fut + Send + 'static,
+        cancellation: CancellationToken,
    ) -> Self
    where
        Fut: Future<Output = anyhow::Result<()>> + Send,
        E: Send + Sync + 'static,
    {
-        let (cancellation, cancellation_receiver) = watch::channel(());
        let (events_sender, events_receiver) = watch::channel(TaskStateUpdate::Started);

        let join_handle = WALRECEIVER_RUNTIME.spawn(async move {
            events_sender.send(TaskStateUpdate::Started).ok();
-            task(events_sender, cancellation_receiver).await
+            task(events_sender).await
        });

        TaskHandle {
@@ -157,7 +156,7 @@ impl<E: Clone> TaskHandle<E> {
    /// Aborts current task, waiting for it to finish.
    pub async fn shutdown(self) {
        if let Some(jh) = self.join_handle {
-            self.cancellation.send(()).ok();
+            self.cancellation.cancel();
            match jh.await {
                Ok(Ok(())) => debug!("Shutdown success"),
                Ok(Err(e)) => error!("Shutdown task error: {e:?}"),
@@ -11,9 +11,9 @@

 use std::{collections::HashMap, num::NonZeroU64, ops::ControlFlow, sync::Arc, time::Duration};

-use crate::task_mgr::TaskKind;
+use crate::context::{DownloadBehavior, RequestContext, TaskKind};
 use crate::task_mgr::WALRECEIVER_RUNTIME;
-use crate::tenant::Timeline;
+use crate::tenant::{Timeline, TimelineRequestContext};
 use crate::{task_mgr, walreceiver::TaskStateUpdate};
 use anyhow::Context;
 use chrono::{NaiveDateTime, Utc};
@@ -46,6 +46,7 @@ pub fn spawn_connection_manager_task(
    lagging_wal_timeout: Duration,
    max_lsn_wal_lag: NonZeroU64,
    auth_token: Option<Arc<String>>,
+    ctx: TimelineRequestContext,
 ) {
    let mut broker_client = get_broker_client().clone();

@@ -54,9 +55,6 @@ pub fn spawn_connection_manager_task(

    task_mgr::spawn(
        WALRECEIVER_RUNTIME.handle(),
-        TaskKind::WalReceiverManager,
-        Some(tenant_id),
-        Some(timeline_id),
        &format!("walreceiver for timeline {tenant_id}/{timeline_id}"),
        false,
        async move {
@@ -70,20 +68,21 @@ pub fn spawn_connection_manager_task(
            );
            loop {
                select! {
-                    _ = task_mgr::shutdown_watcher() => {
+                    _ = ctx.cancelled() => {
                        info!("WAL receiver shutdown requested, shutting down");
                        walreceiver_state.shutdown().await;
-                        return Ok(());
+                        return;
                    },
                    loop_step_result = connection_manager_loop_step(
                        &mut broker_client,
                        &mut walreceiver_state,
+                        &ctx,
                    ) => match loop_step_result {
                        ControlFlow::Continue(()) => continue,
                        ControlFlow::Break(()) => {
                            info!("Connection manager loop ended, shutting down");
                            walreceiver_state.shutdown().await;
-                            return Ok(());
+                            return;
                        }
                    },
                }
@@ -101,6 +100,7 @@ pub fn spawn_connection_manager_task(
 async fn connection_manager_loop_step(
    broker_client: &mut BrokerClientChannel,
    walreceiver_state: &mut WalreceiverState,
+    ctx: &TimelineRequestContext,
 ) -> ControlFlow<(), ()> {
    let mut timeline_state_updates = walreceiver_state.timeline.subscribe_for_state_updates();

@@ -226,6 +226,7 @@ async fn connection_manager_loop_step(
                .change_connection(
                    new_candidate.safekeeper_id,
                    new_candidate.wal_source_connconf,
+                    ctx,
                )
                .await
        }
@@ -389,26 +390,38 @@ impl WalreceiverState {
        &mut self,
        new_sk_id: NodeId,
        new_wal_source_connconf: PgConnectionConfig,
+        ctx: &TimelineRequestContext,
    ) {
        self.drop_old_connection(true).await;

        let id = self.id;
        let connect_timeout = self.wal_connect_timeout;
        let timeline = Arc::clone(&self.timeline);
-        let connection_handle = TaskHandle::spawn(move |events_sender, cancellation| {
-            async move {
-                super::walreceiver_connection::handle_walreceiver_connection(
-                    timeline,
-                    new_wal_source_connconf,
-                    events_sender,
-                    cancellation,
-                    connect_timeout,
-                )
-                .await
-                .context("walreceiver connection handling failure")
-            }
-            .instrument(info_span!("walreceiver_connection", id = %id, node_id = %new_sk_id))
-        });
+
+        let child_ctx = ctx.register_another(RequestContext::with_parent(
+            TaskKind::WalReceiverConnection,
+            DownloadBehavior::Download,
+            ctx,
+        ));
+        let cancellation_token = child_ctx.cancellation_token().clone();
+
+        let connection_handle = TaskHandle::spawn(
+            move |events_sender| {
+                async move {
+                    super::walreceiver_connection::handle_walreceiver_connection(
+                        timeline,
+                        new_wal_source_connconf,
+                        events_sender,
+                        connect_timeout,
+                        child_ctx,
+                    )
+                    .await
+                    .context("walreceiver connection handling failure")
+                }
+                .instrument(info_span!("walreceiver_connection", id = %id, node_id = %new_sk_id))
+            },
+            cancellation_token,
+        );

        let now = Utc::now().naive_utc();
        self.wal_connection = Some(WalConnection {
@@ -820,6 +833,7 @@ fn wal_stream_connection_config(
 mod tests {
    use super::*;
    use crate::tenant::harness::{TenantHarness, TIMELINE_ID};
+    use tokio_util::sync::CancellationToken;
    use url::Host;

    fn dummy_broker_sk_timeline(
@@ -900,12 +914,15 @@ mod tests {
            started_at: now,
            sk_id: connected_sk_id,
            status: connection_status,
-            connection_task: TaskHandle::spawn(move |sender, _| async move {
-                sender
-                    .send(TaskStateUpdate::Progress(connection_status))
-                    .ok();
-                Ok(())
-            }),
+            connection_task: TaskHandle::spawn(
+                move |sender| async move {
+                    sender
+                        .send(TaskStateUpdate::Progress(connection_status))
+                        .ok();
+                    Ok(())
+                },
+                CancellationToken::new(),
+            ),
            discovered_new_wal: None,
        });
        state.wal_stream_candidates = HashMap::from([
@@ -1062,12 +1079,15 @@ mod tests {
            started_at: now,
            sk_id: connected_sk_id,
            status: connection_status,
-            connection_task: TaskHandle::spawn(move |sender, _| async move {
-                sender
-                    .send(TaskStateUpdate::Progress(connection_status))
-                    .ok();
-                Ok(())
-            }),
+            connection_task: TaskHandle::spawn(
+                move |sender| async move {
+                    sender
+                        .send(TaskStateUpdate::Progress(connection_status))
+                        .ok();
+                    Ok(())
+                },
+                CancellationToken::new(),
+            ),
            discovered_new_wal: None,
        });
        state.wal_stream_candidates = HashMap::from([
@@ -1127,12 +1147,15 @@ mod tests {
            started_at: now,
            sk_id: NodeId(1),
            status: connection_status,
-            connection_task: TaskHandle::spawn(move |sender, _| async move {
-                sender
-                    .send(TaskStateUpdate::Progress(connection_status))
-                    .ok();
-                Ok(())
-            }),
+            connection_task: TaskHandle::spawn(
+                move |sender| async move {
+                    sender
+                        .send(TaskStateUpdate::Progress(connection_status))
+                        .ok();
+                    Ok(())
+                },
+                CancellationToken::new(),
+            ),
            discovered_new_wal: None,
        });
        state.wal_stream_candidates = HashMap::from([(
@@ -1189,7 +1212,10 @@ mod tests {
            started_at: now,
            sk_id: NodeId(1),
            status: connection_status,
-            connection_task: TaskHandle::spawn(move |_, _| async move { Ok(()) }),
+            connection_task: TaskHandle::spawn(
+                move |_| async move { Ok(()) },
+                CancellationToken::new(),
+            ),
            discovered_new_wal: Some(NewCommittedWAL {
                discovered_at: time_over_threshold,
                lsn: new_lsn,
@@ -1233,18 +1259,18 @@ mod tests {
    const DUMMY_SAFEKEEPER_HOST: &str = "safekeeper_connstr";

    async fn dummy_state(harness: &TenantHarness<'_>) -> WalreceiverState {
+        let (tenant, tenant_ctx) = harness.load().await;
+        let (timeline, timeline_ctx) = tenant
+            .create_empty_timeline(TIMELINE_ID, Lsn(0), crate::DEFAULT_PG_VERSION, &tenant_ctx)
+            .expect("Failed to create an empty timeline for dummy wal connection manager");
+        let timeline = timeline.initialize(&timeline_ctx).unwrap();
+
        WalreceiverState {
            id: TenantTimelineId {
                tenant_id: harness.tenant_id,
                timeline_id: TIMELINE_ID,
            },
-            timeline: harness
-                .load()
-                .await
-                .create_empty_timeline(TIMELINE_ID, Lsn(0), crate::DEFAULT_PG_VERSION)
-                .expect("Failed to create an empty timeline for dummy wal connection manager")
-                .initialize()
-                .unwrap(),
+            timeline,
            wal_connect_timeout: Duration::from_secs(1),
            lagging_wal_timeout: Duration::from_secs(1),
            max_lsn_wal_lag: NonZeroU64::new(1024 * 1024).unwrap(),
@@ -1,6 +1,7 @@
 //! Actual Postgres connection handler to stream WAL to the server.

 use std::{
+    error::Error,
    str::FromStr,
    sync::Arc,
    time::{Duration, SystemTime},
@@ -11,7 +12,7 @@ use bytes::BytesMut;
 use chrono::{NaiveDateTime, Utc};
 use fail::fail_point;
 use futures::StreamExt;
-use postgres::{SimpleQueryMessage, SimpleQueryRow};
+use postgres::{error::SqlState, SimpleQueryMessage, SimpleQueryRow};
 use postgres_ffi::v14::xlog_utils::normalize_lsn;
 use postgres_ffi::WAL_SEGMENT_SIZE;
 use postgres_protocol::message::backend::ReplicationMessage;
@@ -23,16 +24,15 @@ use tracing::{debug, error, info, trace, warn};
 use crate::{metrics::LIVE_CONNECTIONS_COUNT, walreceiver::TaskStateUpdate};
 use crate::{
    task_mgr,
-    task_mgr::TaskKind,
    task_mgr::WALRECEIVER_RUNTIME,
-    tenant::{Timeline, WalReceiverInfo},
+    tenant::{Timeline, TimelineRequestContext, WalReceiverInfo},
    walingest::WalIngest,
    walrecord::DecodedWALRecord,
 };
 use postgres_connection::PgConnectionConfig;
 use postgres_ffi::waldecoder::WalStreamDecoder;
 use pq_proto::ReplicationFeedback;
-use utils::lsn::Lsn;
+use utils::{lsn::Lsn, postgres_backend_async::is_expected_io_error};

 /// Status of the connection.
 #[derive(Debug, Clone, Copy)]
@@ -58,8 +58,8 @@ pub async fn handle_walreceiver_connection(
    timeline: Arc<Timeline>,
    wal_source_connconf: PgConnectionConfig,
    events_sender: watch::Sender<TaskStateUpdate<WalConnectionStatus>>,
-    mut cancellation: watch::Receiver<()>,
    connect_timeout: Duration,
+    ctx: TimelineRequestContext,
 ) -> anyhow::Result<()> {
    // Connect to the database in replication mode.
    info!("connecting to {wal_source_connconf:?}");
@@ -68,10 +68,17 @@ pub async fn handle_walreceiver_connection(
        let mut config = wal_source_connconf.to_tokio_postgres_config();
        config.application_name("pageserver");
        config.replication_mode(tokio_postgres::config::ReplicationMode::Physical);
-        time::timeout(connect_timeout, config.connect(postgres::NoTls))
-            .await
-            .context("Timed out while waiting for walreceiver connection to open")?
-            .context("Failed to open walreceiver connection")?
+        match time::timeout(connect_timeout, config.connect(postgres::NoTls)).await {
+            Ok(Ok(client_and_conn)) => client_and_conn,
+            Ok(Err(conn_err)) => {
+                let expected_error = ignore_expected_errors(conn_err)?;
+                info!("DB connection stream finished: {expected_error}");
+                return Ok(());
+            }
+            Err(elapsed) => anyhow::bail!(
+                "Timed out while waiting {elapsed} for walreceiver connection to open"
+            ),
+        }
    };

    info!("connected!");
@@ -90,12 +97,9 @@ pub async fn handle_walreceiver_connection(

    // The connection object performs the actual communication with the database,
    // so spawn it off to run on its own.
-    let mut connection_cancellation = cancellation.clone();
+    let cancellation_token = ctx.cancellation_token().clone();
    task_mgr::spawn(
        WALRECEIVER_RUNTIME.handle(),
-        TaskKind::WalReceiverConnection,
-        Some(timeline.tenant_id),
-        Some(timeline.timeline_id),
        "walreceiver connection",
        false,
        async move {
@@ -103,17 +107,14 @@ pub async fn handle_walreceiver_connection(
                connection_result = connection => match connection_result{
                    Ok(()) => info!("Walreceiver db connection closed"),
                    Err(connection_error) => {
-                        if connection_error.is_closed() {
-                            info!("Connection closed regularly: {connection_error}")
-                        } else {
-                            warn!("Connection aborted: {connection_error}")
+                        if let Err(e) = ignore_expected_errors(connection_error) {
+                            warn!("Connection aborted: {e:#}")
                        }
                    }
                },

-                _ = connection_cancellation.changed() => info!("Connection cancelled"),
+                _ = cancellation_token.cancelled() => info!("Connection cancelled"),
            }
-            Ok(())
        },
    );

@@ -173,11 +174,13 @@ pub async fn handle_walreceiver_connection(

    let mut waldecoder = WalStreamDecoder::new(startpoint, timeline.pg_version);

-    let mut walingest = WalIngest::new(timeline.as_ref(), startpoint).await?;
+    let mut walingest = WalIngest::new(timeline.as_ref(), startpoint, &ctx).await?;
+
+    let cancellation = ctx.cancellation_token().clone();

    while let Some(replication_message) = {
        select! {
-            _ = cancellation.changed() => {
+            _ = cancellation.cancelled() => {
                info!("walreceiver interrupted");
                None
            }
@@ -187,14 +190,9 @@ pub async fn handle_walreceiver_connection(
        let replication_message = match replication_message {
            Ok(message) => message,
            Err(replication_error) => {
-                if replication_error.is_closed() {
-                    info!("Replication stream got closed");
-                    return Ok(());
-                } else {
-                    return Err(
-                        anyhow::Error::new(replication_error).context("replication stream error")
-                    );
-                }
+                let expected_error = ignore_expected_errors(replication_error)?;
+                info!("Replication stream finished: {expected_error}");
+                return Ok(());
            }
        };

@@ -249,7 +247,7 @@ pub async fn handle_walreceiver_connection(
                        ensure!(lsn.is_aligned());

                        walingest
-                            .ingest_record(recdata.clone(), lsn, &mut modification, &mut decoded)
+                            .ingest_record(recdata, lsn, &mut modification, &mut decoded, &ctx)
                            .await
                            .with_context(|| format!("could not ingest record at {lsn}"))?;

@@ -327,7 +325,7 @@ pub async fn handle_walreceiver_connection(
            // Send the replication feedback message.
            // Regular standby_status_update fields are put into this message.
            let (timeline_logical_size, _) = timeline
-                .get_current_logical_size()
+                .get_current_logical_size(&ctx)
                .context("Status update creation failed to get current logical size")?;
            let status_update = ReplicationFeedback {
                current_timeline_size: timeline_logical_size,
@@ -400,3 +398,32 @@ async fn identify_system(client: &mut Client) -> anyhow::Result<IdentifySystem>
        Err(IdentifyError.into())
    }
 }
+
+/// We don't want to report connectivity problems as real errors towards connection manager because
+/// 1. they happen frequently enough to make server logs hard to read and
+/// 2. the connection manager can retry other safekeeper.
+///
+/// If this function returns `Ok(pg_error)`, it's such an error.
+/// The caller should log it at info level and then report to connection manager that we're done handling this connection.
+/// Connection manager will then handle reconnections.
+///
+/// If this function returns an `Err()`, the caller can bubble it up using `?`.
+/// The connection manager will log the error at ERROR level.
+fn ignore_expected_errors(pg_error: postgres::Error) -> anyhow::Result<postgres::Error> {
+    if pg_error.is_closed()
+        || pg_error
+            .source()
+            .and_then(|source| source.downcast_ref::<std::io::Error>())
+            .map(is_expected_io_error)
+            .unwrap_or(false)
+    {
+        return Ok(pg_error);
+    } else if let Some(db_error) = pg_error.as_db_error() {
+        if db_error.code() == &SqlState::CONNECTION_FAILURE
+            && db_error.message().contains("end streaming")
+        {
+            return Ok(pg_error);
+        }
+    }
+    Err(pg_error).context("connection error")
+}
@@ -111,6 +111,7 @@ pageserver_connect()
 				PQfinish(pageserver_conn);
 				pageserver_conn = NULL;
 				FreeWaitEventSet(pageserver_conn_wes);
+				pageserver_conn_wes = NULL;

 				neon_log(ERROR, "could not complete handshake with pageserver: %s",
 						 msg);
@@ -179,7 +180,10 @@ pageserver_disconnect(void)
 		prefetch_on_ps_disconnect();
 	}
 	if (pageserver_conn_wes != NULL)
+	{
 		FreeWaitEventSet(pageserver_conn_wes);
+		pageserver_conn_wes = NULL;
+	}
 }

 static void
@@ -206,7 +210,7 @@ pageserver_send(NeonRequest * request)
 	 */
 	if (PQputCopyData(pageserver_conn, req_buff.data, req_buff.len) <= 0)
 	{
-		char	   *msg = PQerrorMessage(pageserver_conn);
+		char	   *msg = pchomp(PQerrorMessage(pageserver_conn));

 		pageserver_disconnect();
 		neon_log(ERROR, "failed to send page request: %s", msg);
@@ -239,29 +243,33 @@ pageserver_receive(void)
 	PG_TRY();
 	{
 		/* read response */
-		resp_buff.len = call_PQgetCopyData(&resp_buff.data);
-		resp_buff.cursor = 0;
+		int			rc;

-		if (resp_buff.len < 0)
+		rc = call_PQgetCopyData(&resp_buff.data);
+		if (rc >= 0)
 		{
-			if (resp_buff.len == -1)
+			resp_buff.len = rc;
+			resp_buff.cursor = 0;
+			resp = nm_unpack_response(&resp_buff);
+			PQfreemem(resp_buff.data);
+
+			if (message_level_is_interesting(PageStoreTrace))
 			{
-				pageserver_disconnect();
-				return NULL;
+				char	   *msg = nm_to_string((NeonMessage *) resp);
+
+				neon_log(PageStoreTrace, "got response: %s", msg);
+				pfree(msg);
 			}
-			else if (resp_buff.len == -2)
-				neon_log(ERROR, "could not read COPY data: %s", PQerrorMessage(pageserver_conn));
 		}
-		resp = nm_unpack_response(&resp_buff);
-		PQfreemem(resp_buff.data);
-
-		if (message_level_is_interesting(PageStoreTrace))
+		else if (rc == -1)
 		{
-			char	   *msg = nm_to_string((NeonMessage *) resp);
-
-			neon_log(PageStoreTrace, "got response: %s", msg);
-			pfree(msg);
+			pageserver_disconnect();
+			resp = NULL;
 		}
+		else if (rc == -2)
+			neon_log(ERROR, "could not read COPY data: %s", PQerrorMessage(pageserver_conn));
+		else
+			neon_log(ERROR, "unexpected PQgetCopyData return value: %d", rc);
 	}
 	PG_CATCH();
 	{
@@ -52,6 +52,7 @@
 #include "access/xlogdefs.h"
 #include "catalog/pg_class.h"
 #include "common/hashfn.h"
+#include "executor/instrument.h"
 #include "pagestore_client.h"
 #include "postmaster/interrupt.h"
 #include "postmaster/autovacuum.h"
@@ -250,11 +251,6 @@ PrefetchState *MyPState;
 	) \
 )

-int			n_prefetch_hits = 0;
-int			n_prefetch_misses = 0;
-int			n_prefetch_missed_caches = 0;
-int			n_prefetch_dupes = 0;
-
 XLogRecPtr	prefetch_lsn = 0;

 static bool compact_prefetch_buffers(void);
@@ -770,7 +766,7 @@ prefetch_register_buffer(BufferTag tag, bool *force_latest, XLogRecPtr *force_ls
 		else
 		{
 			/* The buffered request is good enough, return that index */
-			n_prefetch_dupes++;
+			pgBufferUsage.prefetch.duplicates++;
 			return ring_index;
 		}
 	}
@@ -1845,7 +1841,7 @@ neon_read_at_lsn(RelFileNode rnode, ForkNumber forkNum, BlockNumber blkno,
 		if (slot->effective_request_lsn >= request_lsn)
 		{
 			ring_index = slot->my_ring_index;
-			n_prefetch_hits += 1;
+			pgBufferUsage.prefetch.hits += 1;
 		}
 		else /* the current prefetch LSN is not large enough, so drop the prefetch */
 		{
@@ -1860,7 +1856,7 @@ neon_read_at_lsn(RelFileNode rnode, ForkNumber forkNum, BlockNumber blkno,
 			}
 			/* drop caches */
 			prefetch_set_unused(slot->my_ring_index);
-			n_prefetch_missed_caches += 1;
+			pgBufferUsage.prefetch.expired += 1;
 			/* make it look like a prefetch cache miss */
 			entry = NULL;
 		}
@@ -1870,7 +1866,7 @@ neon_read_at_lsn(RelFileNode rnode, ForkNumber forkNum, BlockNumber blkno,
 	{
 		if (entry == NULL)
 		{
-			n_prefetch_misses += 1;
+			pgBufferUsage.prefetch.misses += 1;

 			ring_index = prefetch_register_buffer(buftag, &request_latest,
 												  &request_lsn);
@@ -17,12 +17,14 @@ hashbrown = "0.12"
 hex = "0.4.3"
 hmac = "0.12.1"
 hyper = "0.14"
+hyper-tungstenite = "0.8.1"
 itertools = "0.10.3"
 md5 = "0.7.0"
 once_cell = "1.13.0"
 parking_lot = "0.12"
 pin-project-lite = "0.2.7"
 rand = "0.8.3"
+regex = "1.4.5"
 reqwest = { version = "0.11", default-features = false, features = [ "json", "rustls-tls" ] }
 routerify = "3"
 rustls = "0.20.0"
@@ -36,10 +38,12 @@ thiserror = "1.0.30"
 tokio = { version = "1.17", features = ["macros"] }
 tokio-postgres = { git = "https://github.com/neondatabase/rust-postgres.git", rev="43e6db254a97fdecbce33d8bc0890accfd74495e" }
 tokio-rustls = "0.23.0"
+tls-listener = { version = "0.5.1", features = ["rustls", "hyper-h1"] }
 tracing = "0.1.36"
 tracing-subscriber = { version = "0.3", features = ["env-filter"] }
 url = "2.2.2"
 uuid = { version = "1.2", features = ["v4", "serde"] }
+webpki-roots = "0.22.5"
 x509-parser = "0.14"

 metrics = { path = "../libs/metrics" }
@@ -149,7 +149,7 @@ impl BackendType<'_, ClientCredentials<'_>> {
        // If there's no project so far, that entails that client doesn't
        // support SNI or other means of passing the project name.
        // We now expect to see a very specific payload in the place of password.
-        let fetch_magic_payload = async {
+        let fetch_magic_payload = |client| async {
            warn!("project name not specified, resorting to the password hack auth flow");
            let payload = AuthFlow::new(client)
                .begin(auth::PasswordHack)
@@ -161,10 +161,26 @@ impl BackendType<'_, ClientCredentials<'_>> {
            auth::Result::Ok(payload)
        };

+        // If we want to use cleartext password flow, we can read the password
+        // from the client and pretend that it's a magic payload (PasswordHack hack).
+        let fetch_plaintext_password = |client| async {
+            info!("using cleartext password flow");
+            let payload = AuthFlow::new(client)
+                .begin(auth::CleartextPassword)
+                .await?
+                .authenticate()
+                .await?;
+
+            auth::Result::Ok(auth::password_hack::PasswordHackPayload {
+                project: String::new(),
+                password: payload,
+            })
+        };
+
        // TODO: find a proper way to merge those very similar blocks.
        let (mut node, payload) = match self {
            Console(endpoint, creds) if creds.project.is_none() => {
-                let payload = fetch_magic_payload.await?;
+                let payload = fetch_magic_payload(client).await?;

                let mut creds = creds.as_ref();
                creds.project = Some(payload.project.as_str().into());
@@ -174,8 +190,18 @@ impl BackendType<'_, ClientCredentials<'_>> {

                (node, payload)
            }
+            Console(endpoint, creds) if creds.use_cleartext_password_flow => {
+                // This is a hack to allow cleartext password in secure connections (wss).
+                let payload = fetch_plaintext_password(client).await?;
+                let creds = creds.as_ref();
+                let node = console::Api::new(endpoint, extra, &creds)
+                    .wake_compute()
+                    .await?;
+
+                (node, payload)
+            }
            Postgres(endpoint, creds) if creds.project.is_none() => {
-                let payload = fetch_magic_payload.await?;
+                let payload = fetch_magic_payload(client).await?;

                let mut creds = creds.as_ref();
                creds.project = Some(payload.project.as_str().into());
@@ -34,6 +34,9 @@ pub struct ClientCredentials<'a> {
    pub user: &'a str,
    pub dbname: &'a str,
    pub project: Option<Cow<'a, str>>,
+    /// If `True`, we'll use the old cleartext password flow. This is used for
+    /// websocket connections, which want to minimize the number of round trips.
+    pub use_cleartext_password_flow: bool,
 }

 impl ClientCredentials<'_> {
@@ -50,6 +53,7 @@ impl<'a> ClientCredentials<'a> {
            user: self.user,
            dbname: self.dbname,
            project: self.project().map(Cow::Borrowed),
+            use_cleartext_password_flow: self.use_cleartext_password_flow,
        }
    }
 }
@@ -59,6 +63,7 @@ impl<'a> ClientCredentials<'a> {
        params: &'a StartupMessageParams,
        sni: Option<&str>,
        common_name: Option<&str>,
+        use_cleartext_password_flow: bool,
    ) -> Result<Self, ClientCredsParseError> {
        use ClientCredsParseError::*;

@@ -108,6 +113,7 @@ impl<'a> ClientCredentials<'a> {
            user = user,
            dbname = dbname,
            project = project.as_deref(),
+            use_cleartext_password_flow = use_cleartext_password_flow,
            "credentials"
        );

@@ -115,6 +121,7 @@ impl<'a> ClientCredentials<'a> {
            user,
            dbname,
            project,
+            use_cleartext_password_flow,
        })
    }
 }
@@ -141,7 +148,7 @@ mod tests {
        let options = StartupMessageParams::new([("user", "john_doe")]);

        // TODO: check that `creds.dbname` is None.
-        let creds = ClientCredentials::parse(&options, None, None)?;
+        let creds = ClientCredentials::parse(&options, None, None, false)?;
        assert_eq!(creds.user, "john_doe");

        Ok(())
@@ -151,7 +158,7 @@ mod tests {
    fn parse_missing_project() -> anyhow::Result<()> {
        let options = StartupMessageParams::new([("user", "john_doe"), ("database", "world")]);

-        let creds = ClientCredentials::parse(&options, None, None)?;
+        let creds = ClientCredentials::parse(&options, None, None, false)?;
        assert_eq!(creds.user, "john_doe");
        assert_eq!(creds.dbname, "world");
        assert_eq!(creds.project, None);
@@ -166,7 +173,7 @@ mod tests {
        let sni = Some("foo.localhost");
        let common_name = Some("localhost");

-        let creds = ClientCredentials::parse(&options, sni, common_name)?;
+        let creds = ClientCredentials::parse(&options, sni, common_name, false)?;
        assert_eq!(creds.user, "john_doe");
        assert_eq!(creds.dbname, "world");
        assert_eq!(creds.project.as_deref(), Some("foo"));
@@ -182,7 +189,7 @@ mod tests {
            ("options", "-ckey=1 project=bar -c geqo=off"),
        ]);

-        let creds = ClientCredentials::parse(&options, None, None)?;
+        let creds = ClientCredentials::parse(&options, None, None, false)?;
        assert_eq!(creds.user, "john_doe");
        assert_eq!(creds.dbname, "world");
        assert_eq!(creds.project.as_deref(), Some("bar"));
@@ -201,7 +208,7 @@ mod tests {
        let sni = Some("baz.localhost");
        let common_name = Some("localhost");

-        let creds = ClientCredentials::parse(&options, sni, common_name)?;
+        let creds = ClientCredentials::parse(&options, sni, common_name, false)?;
        assert_eq!(creds.user, "john_doe");
        assert_eq!(creds.dbname, "world");
        assert_eq!(creds.project.as_deref(), Some("baz"));
@@ -220,7 +227,8 @@ mod tests {
        let sni = Some("second.localhost");
        let common_name = Some("localhost");

-        let err = ClientCredentials::parse(&options, sni, common_name).expect_err("should fail");
+        let err =
+            ClientCredentials::parse(&options, sni, common_name, false).expect_err("should fail");
        match err {
            InconsistentProjectNames { domain, option } => {
                assert_eq!(option, "first");
@@ -237,7 +245,8 @@ mod tests {
        let sni = Some("project.localhost");
        let common_name = Some("example.com");

-        let err = ClientCredentials::parse(&options, sni, common_name).expect_err("should fail");
+        let err =
+            ClientCredentials::parse(&options, sni, common_name, false).expect_err("should fail");
        match err {
            InconsistentSni { sni, cn } => {
                assert_eq!(sni, "project.localhost");
@@ -37,6 +37,17 @@ impl AuthMethod for PasswordHack {
    }
 }

+/// Use clear-text password auth called `password` in docs
+/// <https://www.postgresql.org/docs/current/auth-password.html>
+pub struct CleartextPassword;
+
+impl AuthMethod for CleartextPassword {
+    #[inline(always)]
+    fn first_message(&self) -> BeMessage<'_> {
+        Be::AuthenticationCleartextPassword
+    }
+}
+
 /// This wrapper for [`PqStream`] performs client authentication.
 #[must_use]
 pub struct AuthFlow<'a, Stream, State> {
@@ -86,6 +97,18 @@ impl<S: AsyncRead + AsyncWrite + Unpin> AuthFlow<'_, S, PasswordHack> {
    }
 }

+impl<S: AsyncRead + AsyncWrite + Unpin> AuthFlow<'_, S, CleartextPassword> {
+    /// Perform user authentication. Raise an error in case authentication failed.
+    pub async fn authenticate(self) -> super::Result<Vec<u8>> {
+        let msg = self.stream.read_password_message().await?;
+        let password = msg
+            .strip_suffix(&[0])
+            .ok_or(AuthErrorImpl::MalformedPassword("missing terminator"))?;
+
+        Ok(password.to_vec())
+    }
+}
+
 /// Stream wrapper for handling [SCRAM](crate::scram) auth.
 impl<S: AsyncRead + AsyncWrite + Unpin> AuthFlow<'_, S, Scram<'_>> {
    /// Perform user authentication. Raise an error in case authentication failed.
@@ -1,4 +1,5 @@
 pub mod server;
+pub mod websocket;

 use crate::url::ApiUrl;

@@ -0,0 +1,263 @@
+use bytes::{Buf, Bytes};
+use futures::{Sink, Stream, StreamExt};
+use hyper::server::accept::{self};
+use hyper::server::conn::AddrIncoming;
+use hyper::upgrade::Upgraded;
+use hyper::{Body, Request, Response, StatusCode};
+use hyper_tungstenite::{tungstenite, WebSocketStream};
+use hyper_tungstenite::{tungstenite::Message, HyperWebsocket};
+use pin_project_lite::pin_project;
+use tokio::net::TcpListener;
+
+use std::convert::Infallible;
+use std::future::ready;
+use std::pin::Pin;
+use std::sync::Arc;
+use std::task::{Context, Poll};
+use tls_listener::TlsListener;
+
+use tokio::io::{self, AsyncBufRead, AsyncRead, AsyncWrite, ReadBuf};
+
+use tracing::{error, info, info_span, warn, Instrument};
+use utils::http::{error::ApiError, json::json_response};
+
+use crate::cancellation::CancelMap;
+use crate::config::ProxyConfig;
+use crate::proxy::handle_ws_client;
+
+pin_project! {
+    /// This is a wrapper around a WebSocketStream that implements AsyncRead and AsyncWrite.
+    pub struct WebSocketRW {
+        #[pin]
+        stream: WebSocketStream<Upgraded>,
+        chunk: Option<bytes::Bytes>,
+    }
+}
+
+// FIXME: explain why this is safe or try to remove `unsafe impl`.
+unsafe impl Sync for WebSocketRW {}
+
+impl WebSocketRW {
+    pub fn new(stream: WebSocketStream<Upgraded>) -> Self {
+        Self {
+            stream,
+            chunk: None,
+        }
+    }
+
+    fn has_chunk(&self) -> bool {
+        if let Some(ref chunk) = self.chunk {
+            chunk.remaining() > 0
+        } else {
+            false
+        }
+    }
+}
+
+fn ws_err_into(e: tungstenite::Error) -> io::Error {
+    io::Error::new(io::ErrorKind::Other, e.to_string())
+}
+
+impl AsyncWrite for WebSocketRW {
+    fn poll_write(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &[u8],
+    ) -> Poll<Result<usize, io::Error>> {
+        let mut this = self.project();
+        match this.stream.as_mut().poll_ready(cx) {
+            Poll::Ready(Ok(())) => {
+                if let Err(e) = this
+                    .stream
+                    .as_mut()
+                    .start_send(Message::Binary(buf.to_vec()))
+                {
+                    Poll::Ready(Err(ws_err_into(e)))
+                } else {
+                    Poll::Ready(Ok(buf.len()))
+                }
+            }
+            Poll::Ready(Err(e)) => Poll::Ready(Err(ws_err_into(e))),
+            Poll::Pending => {
+                cx.waker().wake_by_ref();
+                Poll::Pending
+            }
+        }
+    }
+
+    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
+        self.project().stream.poll_flush(cx).map_err(ws_err_into)
+    }
+
+    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
+        self.project().stream.poll_close(cx).map_err(ws_err_into)
+    }
+}
+
+impl AsyncRead for WebSocketRW {
+    fn poll_read(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &mut ReadBuf<'_>,
+    ) -> Poll<io::Result<()>> {
+        if buf.remaining() == 0 {
+            return Poll::Ready(Ok(()));
+        }
+
+        let inner_buf = match self.as_mut().poll_fill_buf(cx) {
+            Poll::Ready(Ok(buf)) => buf,
+            Poll::Ready(Err(err)) => return Poll::Ready(Err(err)),
+            Poll::Pending => return Poll::Pending,
+        };
+        let len = std::cmp::min(inner_buf.len(), buf.remaining());
+        buf.put_slice(&inner_buf[..len]);
+
+        self.consume(len);
+        Poll::Ready(Ok(()))
+    }
+}
+
+impl AsyncBufRead for WebSocketRW {
+    fn poll_fill_buf(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> {
+        loop {
+            if self.as_mut().has_chunk() {
+                let buf = self.project().chunk.as_ref().unwrap().chunk();
+                return Poll::Ready(Ok(buf));
+            } else {
+                match self.as_mut().project().stream.poll_next(cx) {
+                    Poll::Ready(Some(Ok(message))) => match message {
+                        Message::Text(_) => {}
+                        Message::Binary(chunk) => {
+                            *self.as_mut().project().chunk = Some(Bytes::from(chunk));
+                        }
+                        Message::Ping(_) => {
+                            // No need to send a reply: tungstenite takes care of this for you.
+                        }
+                        Message::Pong(_) => {}
+                        Message::Close(_) => {
+                            // No need to send a reply: tungstenite takes care of this for you.
+                            return Poll::Ready(Ok(&[]));
+                        }
+                        Message::Frame(_) => {
+                            unreachable!();
+                        }
+                    },
+                    Poll::Ready(Some(Err(err))) => return Poll::Ready(Err(ws_err_into(err))),
+                    Poll::Ready(None) => return Poll::Ready(Ok(&[])),
+                    Poll::Pending => return Poll::Pending,
+                }
+            }
+        }
+    }
+
+    fn consume(self: Pin<&mut Self>, amt: usize) {
+        if amt > 0 {
+            self.project()
+                .chunk
+                .as_mut()
+                .expect("No chunk present")
+                .advance(amt);
+        }
+    }
+}
+
+async fn serve_websocket(
+    websocket: HyperWebsocket,
+    config: &ProxyConfig,
+    cancel_map: &CancelMap,
+    session_id: uuid::Uuid,
+    hostname: Option<String>,
+) -> anyhow::Result<()> {
+    let websocket = websocket.await?;
+    handle_ws_client(
+        config,
+        cancel_map,
+        session_id,
+        WebSocketRW::new(websocket),
+        hostname,
+    )
+    .await?;
+    Ok(())
+}
+
+async fn ws_handler(
+    mut request: Request<Body>,
+    config: &'static ProxyConfig,
+    cancel_map: Arc<CancelMap>,
+    session_id: uuid::Uuid,
+) -> Result<Response<Body>, ApiError> {
+    let host = request
+        .headers()
+        .get("host")
+        .and_then(|h| h.to_str().ok())
+        .and_then(|h| h.split(':').next())
+        .map(|s| s.to_string());
+
+    // Check if the request is a websocket upgrade request.
+    if hyper_tungstenite::is_upgrade_request(&request) {
+        let (response, websocket) = hyper_tungstenite::upgrade(&mut request, None)
+            .map_err(|e| ApiError::BadRequest(e.into()))?;
+
+        tokio::spawn(async move {
+            if let Err(e) = serve_websocket(websocket, config, &cancel_map, session_id, host).await
+            {
+                error!("error in websocket connection: {:?}", e);
+            }
+        });
+
+        // Return the response so the spawned future can continue.
+        Ok(response)
+    } else {
+        json_response(StatusCode::OK, "Connect with a websocket client")
+    }
+}
+
+pub async fn task_main(
+    ws_listener: TcpListener,
+    config: &'static ProxyConfig,
+) -> anyhow::Result<()> {
+    scopeguard::defer! {
+        info!("websocket server has shut down");
+    }
+
+    let tls_config = config.tls_config.as_ref().map(|cfg| cfg.to_server_config());
+    let tls_acceptor: tokio_rustls::TlsAcceptor = match tls_config {
+        Some(config) => config.into(),
+        None => {
+            warn!("TLS config is missing, WebSocket Secure server will not be started");
+            return Ok(());
+        }
+    };
+
+    let addr_incoming = AddrIncoming::from_listener(ws_listener)?;
+
+    let tls_listener = TlsListener::new(tls_acceptor, addr_incoming).filter(|conn| {
+        if let Err(err) = conn {
+            error!("failed to accept TLS connection for websockets: {:?}", err);
+            ready(false)
+        } else {
+            ready(true)
+        }
+    });
+
+    let make_svc = hyper::service::make_service_fn(|_stream| async move {
+        Ok::<_, Infallible>(hyper::service::service_fn(
+            move |req: Request<Body>| async move {
+                let cancel_map = Arc::new(CancelMap::default());
+                let session_id = uuid::Uuid::new_v4();
+                ws_handler(req, config, cancel_map, session_id)
+                    .instrument(info_span!(
+                        "ws-client",
+                        session = format_args!("{session_id}")
+                    ))
+                    .await
+            },
+        ))
+    });
+
+    hyper::Server::builder(accept::from_stream(tls_listener))
+        .serve(make_svc)
+        .await?;
+
+    Ok(())
+}
@@ -110,12 +110,23 @@ async fn main() -> anyhow::Result<()> {
    info!("Starting proxy on {proxy_address}");
    let proxy_listener = TcpListener::bind(proxy_address).await?;

-    let tasks = [
+    let mut tasks = vec![
        tokio::spawn(http::server::task_main(http_listener)),
        tokio::spawn(proxy::task_main(config, proxy_listener)),
        tokio::task::spawn_blocking(move || mgmt::thread_main(mgmt_listener)),
-    ]
-    .map(flatten_err);
+    ];
+
+    if let Some(wss_address) = arg_matches.get_one::<String>("wss") {
+        let wss_address: SocketAddr = wss_address.parse()?;
+        info!("Starting wss on {}", wss_address);
+        let wss_listener = TcpListener::bind(wss_address).await?;
+        tasks.push(tokio::spawn(http::websocket::task_main(
+            wss_listener,
+            config,
+        )));
+    }
+
+    let tasks = tasks.into_iter().map(flatten_err);

    set_build_info_metric(GIT_VERSION);
    // This will block until all tasks have completed.
@@ -155,6 +166,11 @@ fn cli() -> clap::Command {
                .help("listen for incoming http connections (metrics, etc) on ip:port")
                .default_value("127.0.0.1:7001"),
        )
+        .arg(
+            Arg::new("wss")
+                .long("wss")
+                .help("listen for incoming wss connections on ip:port"),
+        )
        .arg(
            Arg::new("uri")
                .short('u')
@@ -9,7 +9,10 @@ use std::{
    thread,
 };
 use tracing::{error, info, info_span};
-use utils::postgres_backend::{self, AuthType, PostgresBackend};
+use utils::{
+    postgres_backend::{self, AuthType, PostgresBackend},
+    postgres_backend_async::QueryError,
+};

 /// Console management API listener thread.
 /// It spawns console response handlers needed for the link auth.
@@ -47,7 +50,7 @@ pub fn thread_main(listener: TcpListener) -> anyhow::Result<()> {
    }
 }

-fn handle_connection(socket: TcpStream) -> anyhow::Result<()> {
+fn handle_connection(socket: TcpStream) -> Result<(), QueryError> {
    let pgbackend = PostgresBackend::new(socket, AuthType::Trust, None, true)?;
    pgbackend.run(&mut MgmtHandler)
 }
@@ -58,7 +61,7 @@ pub type ComputeReady = Result<DatabaseInfo, String>;
 // TODO: replace with an http-based protocol.
 struct MgmtHandler;
 impl postgres_backend::Handler for MgmtHandler {
-    fn process_query(&mut self, pgb: &mut PostgresBackend, query: &str) -> anyhow::Result<()> {
+    fn process_query(&mut self, pgb: &mut PostgresBackend, query: &str) -> Result<(), QueryError> {
        try_process_query(pgb, query).map_err(|e| {
            error!("failed to process response: {e:?}");
            e
@@ -66,8 +69,8 @@ impl postgres_backend::Handler for MgmtHandler {
    }
 }

-fn try_process_query(pgb: &mut PostgresBackend, query: &str) -> anyhow::Result<()> {
-    let resp: KickSession = serde_json::from_str(query)?;
+fn try_process_query(pgb: &mut PostgresBackend, query: &str) -> Result<(), QueryError> {
+    let resp: KickSession = serde_json::from_str(query).context("Failed to parse query as json")?;

    let span = info_span!("event", session_id = resp.session_id);
    let _enter = span.enter();
@@ -81,7 +84,7 @@ fn try_process_query(pgb: &mut PostgresBackend, query: &str) -> anyhow::Result<(
        }
        Err(e) => {
            error!("failed to deliver response to per-client task");
-            pgb.write_message(&BeMessage::ErrorResponse(&e.to_string()))?;
+            pgb.write_message(&BeMessage::ErrorResponse(&e.to_string(), None))?;
        }
    }

@@ -82,6 +82,47 @@ pub async fn task_main(
    }
 }

+pub async fn handle_ws_client(
+    config: &ProxyConfig,
+    cancel_map: &CancelMap,
+    session_id: uuid::Uuid,
+    stream: impl AsyncRead + AsyncWrite + Unpin + Send,
+    hostname: Option<String>,
+) -> anyhow::Result<()> {
+    // The `closed` counter will increase when this future is destroyed.
+    NUM_CONNECTIONS_ACCEPTED_COUNTER.inc();
+    scopeguard::defer! {
+        NUM_CONNECTIONS_CLOSED_COUNTER.inc();
+    }
+
+    let tls = config.tls_config.as_ref();
+    let hostname = hostname.as_deref();
+
+    // TLS is None here, because the connection is already encrypted.
+    let do_handshake = handshake(stream, None, cancel_map).instrument(info_span!("handshake"));
+    let (mut stream, params) = match do_handshake.await? {
+        Some(x) => x,
+        None => return Ok(()), // it's a cancellation request
+    };
+
+    // Extract credentials which we're going to use for auth.
+    let creds = {
+        let common_name = tls.and_then(|tls| tls.common_name.as_deref());
+        let result = config
+            .auth_backend
+            .as_ref()
+            .map(|_| auth::ClientCredentials::parse(&params, hostname, common_name, true))
+            .transpose();
+
+        async { result }.or_else(|e| stream.throw_error(e)).await?
+    };
+
+    let client = Client::new(stream, creds, &params, session_id);
+    cancel_map
+        .with_session(|session| client.connect_to_db(session))
+        .await
+}
+
 async fn handle_client(
    config: &ProxyConfig,
    cancel_map: &CancelMap,
@@ -108,7 +149,7 @@ async fn handle_client(
        let result = config
            .auth_backend
            .as_ref()
-            .map(|_| auth::ClientCredentials::parse(&params, sni, common_name))
+            .map(|_| auth::ClientCredentials::parse(&params, sni, common_name, false))
            .transpose();

        async { result }.or_else(|e| stream.throw_error(e)).await?
@@ -2,7 +2,7 @@ use crate::error::UserFacingError;
 use anyhow::bail;
 use bytes::BytesMut;
 use pin_project_lite::pin_project;
-use pq_proto::{BeMessage, FeMessage, FeStartupPacket};
+use pq_proto::{BeMessage, ConnectionError, FeMessage, FeStartupPacket};
 use rustls::ServerConfig;
 use std::pin::Pin;
 use std::sync::Arc;
@@ -47,18 +47,13 @@ fn err_connection() -> io::Error {
    io::Error::new(io::ErrorKind::ConnectionAborted, "connection is lost")
 }

-// TODO: change error type of `FeMessage::read_fut`
-fn from_anyhow(e: anyhow::Error) -> io::Error {
-    io::Error::new(io::ErrorKind::Other, e.to_string())
-}
-
 impl<S: AsyncRead + Unpin> PqStream<S> {
    /// Receive [`FeStartupPacket`], which is a first packet sent by a client.
    pub async fn read_startup_packet(&mut self) -> io::Result<FeStartupPacket> {
        // TODO: `FeStartupPacket::read_fut` should return `FeStartupPacket`
        let msg = FeStartupPacket::read_fut(&mut self.stream)
            .await
-            .map_err(from_anyhow)?
+            .map_err(ConnectionError::into_io_error)?
            .ok_or_else(err_connection)?;

        match msg {
@@ -80,7 +75,7 @@ impl<S: AsyncRead + Unpin> PqStream<S> {
    async fn read_message(&mut self) -> io::Result<FeMessage> {
        FeMessage::read_fut(&mut self.stream)
            .await
-            .map_err(from_anyhow)?
+            .map_err(ConnectionError::into_io_error)?
            .ok_or_else(err_connection)
    }
 }
@@ -112,7 +107,8 @@ impl<S: AsyncWrite + Unpin> PqStream<S> {
    /// This method exists due to `&str` not implementing `Into<anyhow::Error>`.
    pub async fn throw_error_str<T>(&mut self, error: &'static str) -> anyhow::Result<T> {
        tracing::info!("forwarding error to user: {error}");
-        self.write_message(&BeMessage::ErrorResponse(error)).await?;
+        self.write_message(&BeMessage::ErrorResponse(error, None))
+            .await?;
        bail!(error)
    }

@@ -124,7 +120,8 @@ impl<S: AsyncWrite + Unpin> PqStream<S> {
    {
        let msg = error.to_string_client();
        tracing::info!("forwarding error to user: {msg}");
-        self.write_message(&BeMessage::ErrorResponse(&msg)).await?;
+        self.write_message(&BeMessage::ErrorResponse(&msg, None))
+            .await?;
        bail!(error)
    }
 }
@@ -1,10 +1,5 @@
 [toolchain]
-# We try to stick to a toolchain version that is widely available on popular distributions, so that most people
-# can use the toolchain that comes with their operating system. But if there's a feature we miss badly from a later
-# version, we can consider updating.
-# See https://tracker.debian.org/pkg/rustc for more details on Debian rustc package,
-# we use "unstable" version number as the highest version used in the project by default.
-channel = "1.62.1"
+channel = "1.66.1"
 profile = "default"
 # The default profile includes rustc, rust-std, cargo, rust-docs, rustfmt and clippy.
 # https://rust-lang.github.io/rustup/concepts/profiles.html
@@ -229,11 +229,7 @@ fn start_safekeeper(conf: SafeKeeperConf) -> Result<()> {
    let conf_cloned = conf.clone();
    let safekeeper_thread = thread::Builder::new()
        .name("safekeeper thread".into())
-        .spawn(|| {
-            if let Err(e) = wal_service::thread_main(conf_cloned, pg_listener) {
-                info!("safekeeper thread terminated: {e}");
-            }
-        })
+        .spawn(|| wal_service::thread_main(conf_cloned, pg_listener))
        .unwrap();

    threads.push(safekeeper_thread);
@@ -8,7 +8,7 @@ use crate::receive_wal::ReceiveWalConn;
 use crate::send_wal::ReplicationConn;

 use crate::{GlobalTimelines, SafeKeeperConf};
-use anyhow::{bail, ensure, Context, Result};
+use anyhow::Context;

 use postgres_ffi::PG_TLI;
 use regex::Regex;
@@ -17,6 +17,7 @@ use pq_proto::{BeMessage, FeStartupPacket, RowDescriptor, INT4_OID, TEXT_OID};
 use std::str;
 use tracing::info;
 use utils::auth::{Claims, Scope};
+use utils::postgres_backend_async::QueryError;
 use utils::{
    id::{TenantId, TenantTimelineId, TimelineId},
    lsn::Lsn,
@@ -42,7 +43,7 @@ enum SafekeeperPostgresCommand {
    JSONCtrl { cmd: AppendLogicalMessage },
 }

-fn parse_cmd(cmd: &str) -> Result<SafekeeperPostgresCommand> {
+fn parse_cmd(cmd: &str) -> anyhow::Result<SafekeeperPostgresCommand> {
    if cmd.starts_with("START_WAL_PUSH") {
        Ok(SafekeeperPostgresCommand::StartWalPush)
    } else if cmd.starts_with("START_REPLICATION") {
@@ -62,13 +63,17 @@ fn parse_cmd(cmd: &str) -> Result<SafekeeperPostgresCommand> {
            cmd: serde_json::from_str(cmd)?,
        })
    } else {
-        bail!("unsupported command {}", cmd);
+        anyhow::bail!("unsupported command {cmd}");
    }
 }

 impl postgres_backend::Handler for SafekeeperPostgresHandler {
    // tenant_id and timeline_id are passed in connection string params
-    fn startup(&mut self, _pgb: &mut PostgresBackend, sm: &FeStartupPacket) -> Result<()> {
+    fn startup(
+        &mut self,
+        _pgb: &mut PostgresBackend,
+        sm: &FeStartupPacket,
+    ) -> Result<(), QueryError> {
        if let FeStartupPacket::StartupMessage { params, .. } = sm {
            if let Some(options) = params.options_raw() {
                for opt in options {
@@ -77,10 +82,14 @@ impl postgres_backend::Handler for SafekeeperPostgresHandler {
                    // https://github.com/neondatabase/neon/pull/2433#discussion_r970005064
                    match opt.split_once('=') {
                        Some(("ztenantid", value)) | Some(("tenant_id", value)) => {
-                            self.tenant_id = Some(value.parse()?);
+                            self.tenant_id = Some(value.parse().with_context(|| {
+                                format!("Failed to parse {value} as tenant id")
+                            })?);
                        }
                        Some(("ztimelineid", value)) | Some(("timeline_id", value)) => {
-                            self.timeline_id = Some(value.parse()?);
+                            self.timeline_id = Some(value.parse().with_context(|| {
+                                format!("Failed to parse {value} as timeline id")
+                            })?);
                        }
                        _ => continue,
                    }
@@ -93,7 +102,9 @@ impl postgres_backend::Handler for SafekeeperPostgresHandler {

            Ok(())
        } else {
-            bail!("Safekeeper received unexpected initial message: {:?}", sm);
+            Err(QueryError::Other(anyhow::anyhow!(
+                "Safekeeper received unexpected initial message: {sm:?}"
+            )))
        }
    }

@@ -101,7 +112,7 @@ impl postgres_backend::Handler for SafekeeperPostgresHandler {
        &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
@@ -109,13 +120,12 @@ impl postgres_backend::Handler for SafekeeperPostgresHandler {
            .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!(
@@ -127,7 +137,11 @@ impl postgres_backend::Handler for SafekeeperPostgresHandler {
        Ok(())
    }

-    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> {
        if query_string
            .to_ascii_lowercase()
            .starts_with("set datestyle to ")
@@ -148,19 +162,26 @@ impl postgres_backend::Handler for SafekeeperPostgresHandler {
        self.check_permission(Some(tenant_id))?;
        self.ttid = TenantTimelineId::new(tenant_id, timeline_id);

-        match cmd {
+        let res = match cmd {
            SafekeeperPostgresCommand::StartWalPush => ReceiveWalConn::new(pgb).run(self),
            SafekeeperPostgresCommand::StartReplication { start_lsn } => {
                ReplicationConn::new(pgb).run(self, pgb, start_lsn)
            }
            SafekeeperPostgresCommand::IdentifySystem => self.handle_identify_system(pgb),
            SafekeeperPostgresCommand::JSONCtrl { ref cmd } => handle_json_ctrl(self, pgb, cmd),
-        }
-        .context(format!(
-            "Failed to process query for timeline {timeline_id}"
-        ))?;
+        };

-        Ok(())
+        match res {
+            Ok(()) => Ok(()),
+            Err(QueryError::Disconnected(connection_error)) => {
+                info!("Timeline {tenant_id}/{timeline_id} query failed with connection error: {connection_error}");
+                Err(QueryError::Disconnected(connection_error))
+            }
+            Err(QueryError::Other(e)) => Err(QueryError::Other(e.context(format!(
+                "Failed to process query for timeline {}",
+                self.ttid
+            )))),
+        }
    }
 }

@@ -178,7 +199,7 @@ impl SafekeeperPostgresHandler {

    // 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.conf.auth.is_none() {
            // auth is set to Trust, nothing to check so just return ok
            return Ok(());
@@ -196,7 +217,7 @@ impl SafekeeperPostgresHandler {
    ///
    /// Handle IDENTIFY_SYSTEM replication command
    ///
-    fn handle_identify_system(&mut self, pgb: &mut PostgresBackend) -> Result<()> {
+    fn handle_identify_system(&mut self, pgb: &mut PostgresBackend) -> Result<(), QueryError> {
        let tli = GlobalTimelines::get(self.ttid)?;

        let lsn = if self.is_walproposer_recovery() {
@@ -8,11 +8,12 @@

 use std::sync::Arc;

-use anyhow::Result;
+use anyhow::Context;
 use bytes::Bytes;
 use serde::{Deserialize, Serialize};
 use tracing::*;
 use utils::id::TenantTimelineId;
+use utils::postgres_backend_async::QueryError;

 use crate::handler::SafekeeperPostgresHandler;
 use crate::safekeeper::{AcceptorProposerMessage, AppendResponse, ServerInfo};
@@ -47,7 +48,7 @@ pub struct AppendLogicalMessage {
    pg_version: u32,
 }

-#[derive(Serialize, Deserialize)]
+#[derive(Debug, Serialize, Deserialize)]
 struct AppendResult {
    // safekeeper state after append
    state: SafeKeeperState,
@@ -62,8 +63,8 @@ pub fn handle_json_ctrl(
    spg: &SafekeeperPostgresHandler,
    pgb: &mut PostgresBackend,
    append_request: &AppendLogicalMessage,
-) -> Result<()> {
-    info!("JSON_CTRL request: {:?}", append_request);
+) -> Result<(), QueryError> {
+    info!("JSON_CTRL request: {append_request:?}");

    // need to init safekeeper state before AppendRequest
    let tli = prepare_safekeeper(spg.ttid, append_request.pg_version)?;
@@ -78,7 +79,8 @@ pub fn handle_json_ctrl(
        state: tli.get_state().1,
        inserted_wal,
    };
-    let response_data = serde_json::to_vec(&response)?;
+    let response_data = serde_json::to_vec(&response)
+        .with_context(|| format!("Response {response:?} is not a json array"))?;

    pgb.write_message_noflush(&BeMessage::RowDescription(&[RowDescriptor {
        name: b"json",
@@ -93,7 +95,7 @@ pub fn handle_json_ctrl(

 /// Prepare safekeeper to process append requests without crashes,
 /// by sending ProposerGreeting with default server.wal_seg_size.
-fn prepare_safekeeper(ttid: TenantTimelineId, pg_version: u32) -> Result<Arc<Timeline>> {
+fn prepare_safekeeper(ttid: TenantTimelineId, pg_version: u32) -> anyhow::Result<Arc<Timeline>> {
    GlobalTimelines::create(
        ttid,
        ServerInfo {
@@ -106,7 +108,7 @@ fn prepare_safekeeper(ttid: TenantTimelineId, pg_version: u32) -> Result<Arc<Tim
    )
 }

-fn send_proposer_elected(tli: &Arc<Timeline>, term: Term, lsn: Lsn) -> Result<()> {
+fn send_proposer_elected(tli: &Arc<Timeline>, term: Term, lsn: Lsn) -> anyhow::Result<()> {
    // add new term to existing history
    let history = tli.get_state().1.acceptor_state.term_history;
    let history = history.up_to(lsn.checked_sub(1u64).unwrap());
@@ -125,7 +127,7 @@ fn send_proposer_elected(tli: &Arc<Timeline>, term: Term, lsn: Lsn) -> Result<()
    Ok(())
 }

-#[derive(Serialize, Deserialize)]
+#[derive(Debug, Serialize, Deserialize)]
 struct InsertedWAL {
    begin_lsn: Lsn,
    end_lsn: Lsn,
@@ -134,7 +136,10 @@ struct InsertedWAL {

 /// Extend local WAL with new LogicalMessage record. To do that,
 /// create AppendRequest with new WAL and pass it to safekeeper.
-fn append_logical_message(tli: &Arc<Timeline>, msg: &AppendLogicalMessage) -> Result<InsertedWAL> {
+fn append_logical_message(
+    tli: &Arc<Timeline>,
+    msg: &AppendLogicalMessage,
+) -> anyhow::Result<InsertedWAL> {
    let wal_data = encode_logical_message(&msg.lm_prefix, &msg.lm_message);
    let sk_state = tli.get_state().1;

@@ -2,11 +2,13 @@
 //! Gets messages from the network, passes them down to consensus module and
 //! sends replies back.

-use anyhow::{anyhow, bail, Result};
+use anyhow::anyhow;
+use anyhow::Context;

 use bytes::BytesMut;
 use tracing::*;
 use utils::lsn::Lsn;
+use utils::postgres_backend_async::QueryError;

 use crate::safekeeper::ServerInfo;
 use crate::timeline::Timeline;
@@ -43,7 +45,7 @@ impl<'pg> ReceiveWalConn<'pg> {
    }

    // Send message to the postgres
-    fn write_msg(&mut self, msg: &AcceptorProposerMessage) -> Result<()> {
+    fn write_msg(&mut self, msg: &AcceptorProposerMessage) -> anyhow::Result<()> {
        let mut buf = BytesMut::with_capacity(128);
        msg.serialize(&mut buf)?;
        self.pg_backend.write_message(&BeMessage::CopyData(&buf))?;
@@ -51,7 +53,7 @@ impl<'pg> ReceiveWalConn<'pg> {
    }

    /// Receive WAL from wal_proposer
-    pub fn run(&mut self, spg: &mut SafekeeperPostgresHandler) -> Result<()> {
+    pub fn run(&mut self, spg: &mut SafekeeperPostgresHandler) -> Result<(), QueryError> {
        let _enter = info_span!("WAL acceptor", ttid = %spg.ttid).entered();

        // Notify the libpq client that it's allowed to send `CopyData` messages
@@ -79,7 +81,11 @@ impl<'pg> ReceiveWalConn<'pg> {
                };
                GlobalTimelines::create(spg.ttid, server_info, Lsn::INVALID, Lsn::INVALID)?
            }
-            _ => bail!("unexpected message {:?} instead of greeting", next_msg),
+            _ => {
+                return Err(QueryError::Other(anyhow::anyhow!(
+                    "unexpected message {next_msg:?} instead of greeting"
+                )))
+            }
        };

        let mut next_msg = Some(next_msg);
@@ -134,25 +140,32 @@ impl<'pg> ReceiveWalConn<'pg> {

 struct ProposerPollStream {
    msg_rx: Receiver<ProposerAcceptorMessage>,
-    read_thread: Option<thread::JoinHandle<Result<()>>>,
+    read_thread: Option<thread::JoinHandle<Result<(), QueryError>>>,
 }

 impl ProposerPollStream {
-    fn new(mut r: ReadStream) -> Result<Self> {
+    fn new(mut r: ReadStream) -> anyhow::Result<Self> {
        let (msg_tx, msg_rx) = channel();

        let read_thread = thread::Builder::new()
            .name("Read WAL thread".into())
-            .spawn(move || -> Result<()> {
+            .spawn(move || -> Result<(), QueryError> {
                loop {
                    let copy_data = match FeMessage::read(&mut r)? {
-                        Some(FeMessage::CopyData(bytes)) => bytes,
-                        Some(msg) => bail!("expected `CopyData` message, found {:?}", msg),
-                        None => bail!("connection closed unexpectedly"),
-                    };
+                        Some(FeMessage::CopyData(bytes)) => Ok(bytes),
+                        Some(msg) => Err(QueryError::Other(anyhow::anyhow!(
+                            "expected `CopyData` message, found {msg:?}"
+                        ))),
+                        None => Err(QueryError::from(std::io::Error::new(
+                            std::io::ErrorKind::ConnectionAborted,
+                            "walproposer closed the connection",
+                        ))),
+                    }?;

                    let msg = ProposerAcceptorMessage::parse(copy_data)?;
-                    msg_tx.send(msg)?;
+                    msg_tx
+                        .send(msg)
+                        .context("Failed to send the proposer message")?;
                }
                // msg_tx will be dropped here, this will also close msg_rx
            })?;
@@ -163,17 +176,19 @@ impl ProposerPollStream {
        })
    }

-    fn recv_msg(&mut self) -> Result<ProposerAcceptorMessage> {
+    fn recv_msg(&mut self) -> Result<ProposerAcceptorMessage, QueryError> {
        self.msg_rx.recv().map_err(|_| {
            // return error from the read thread
            let res = match self.read_thread.take() {
                Some(thread) => thread.join(),
-                None => return anyhow!("read thread is gone"),
+                None => return QueryError::Other(anyhow::anyhow!("read thread is gone")),
            };

            match res {
-                Ok(Ok(())) => anyhow!("unexpected result from read thread"),
-                Err(err) => anyhow!("read thread panicked: {:?}", err),
+                Ok(Ok(())) => {
+                    QueryError::Other(anyhow::anyhow!("unexpected result from read thread"))
+                }
+                Err(err) => QueryError::Other(anyhow::anyhow!("read thread panicked: {err:?}")),
                Ok(Err(err)) => err,
            }
        })
@@ -5,7 +5,7 @@ use crate::handler::SafekeeperPostgresHandler;
 use crate::timeline::{ReplicaState, Timeline};
 use crate::wal_storage::WalReader;
 use crate::GlobalTimelines;
-use anyhow::{bail, Context, Result};
+use anyhow::Context;

 use bytes::Bytes;
 use postgres_ffi::get_current_timestamp;
@@ -15,7 +15,8 @@ use std::cmp::min;
 use std::net::Shutdown;
 use std::sync::Arc;
 use std::time::Duration;
-use std::{str, thread};
+use std::{io, str, thread};
+use utils::postgres_backend_async::QueryError;

 use pq_proto::{BeMessage, FeMessage, ReplicationFeedback, WalSndKeepAlive, XLogDataBody};
 use tokio::sync::watch::Receiver;
@@ -91,7 +92,7 @@ impl ReplicationConn {
    fn background_thread(
        mut stream_in: ReadStream,
        replica_guard: Arc<ReplicationConnGuard>,
-    ) -> Result<()> {
+    ) -> anyhow::Result<()> {
        let replica_id = replica_guard.replica;
        let timeline = &replica_guard.timeline;

@@ -140,7 +141,7 @@ impl ReplicationConn {
                    // Shutdown the connection, because rust-postgres client cannot be dropped
                    // when connection is alive.
                    let _ = stream_in.shutdown(Shutdown::Both);
-                    bail!("Copy failed");
+                    anyhow::bail!("Copy failed");
                }
                _ => {
                    // We only handle `CopyData`, 'Sync', 'CopyFail' messages. Anything else is ignored.
@@ -160,7 +161,7 @@ impl ReplicationConn {
        spg: &mut SafekeeperPostgresHandler,
        pgb: &mut PostgresBackend,
        mut start_pos: Lsn,
-    ) -> Result<()> {
+    ) -> Result<(), QueryError> {
        let _enter = info_span!("WAL sender", ttid = %spg.ttid).entered();

        let tli = GlobalTimelines::get(spg.ttid)?;
@@ -256,8 +257,10 @@ impl ReplicationConn {
                        // to right pageserver.
                        if tli.should_walsender_stop(replica_id) {
                            // Shut down, timeline is suspended.
-                            // TODO create proper error type for this
-                            bail!("end streaming to {:?}", spg.appname);
+                            return Err(QueryError::from(io::Error::new(
+                                io::ErrorKind::ConnectionAborted,
+                                format!("end streaming to {:?}", spg.appname),
+                            )));
                        }

                        // timeout expired: request pageserver status
@@ -265,8 +268,7 @@ impl ReplicationConn {
                            sent_ptr: end_pos.0,
                            timestamp: get_current_timestamp(),
                            request_reply: true,
-                        }))
-                        .context("Failed to send KeepAlive message")?;
+                        }))?;
                        continue;
                    }
                }
@@ -301,7 +303,7 @@ impl ReplicationConn {
 const POLL_STATE_TIMEOUT: Duration = Duration::from_secs(1);

 // Wait until we have commit_lsn > lsn or timeout expires. Returns latest commit_lsn.
-async fn wait_for_lsn(rx: &mut Receiver<Lsn>, lsn: Lsn) -> Result<Option<Lsn>> {
+async fn wait_for_lsn(rx: &mut Receiver<Lsn>, lsn: Lsn) -> anyhow::Result<Option<Lsn>> {
    let commit_lsn: Lsn = *rx.borrow();
    if commit_lsn > lsn {
        return Ok(Some(commit_lsn));
@@ -279,7 +279,9 @@ impl GlobalTimelines {
                let (dir_existed, was_active) = timeline.delete_from_disk(&mut shared_state)?;

                // Remove timeline from the map.
-                TIMELINES_STATE.lock().unwrap().timelines.remove(ttid);
+                // FIXME: re-enable it once we fix the issue with recreation of deleted timelines
+                // https://github.com/neondatabase/neon/issues/3146
+                // TIMELINES_STATE.lock().unwrap().timelines.remove(ttid);

                Ok(TimelineDeleteForceResult {
                    dir_existed,
@@ -346,15 +348,16 @@ impl GlobalTimelines {
                .tenant_dir(tenant_id),
        )?;

-        let tlis_after_delete = Self::get_all_for_tenant(*tenant_id);
-        if !tlis_after_delete.is_empty() {
-            // Some timelines were created while we were deleting them, returning error
-            // to the caller, so it can retry later.
-            bail!(
-                "failed to delete all timelines for tenant {}: some timelines were created while we were deleting them",
-                tenant_id
-            );
-        }
+        // FIXME: we temporarily disabled removing timelines from the map, see `delete_force`
+        // let tlis_after_delete = Self::get_all_for_tenant(*tenant_id);
+        // if !tlis_after_delete.is_empty() {
+        //     // Some timelines were created while we were deleting them, returning error
+        //     // to the caller, so it can retry later.
+        //     bail!(
+        //         "failed to delete all timelines for tenant {}: some timelines were created while we were deleting them",
+        //         tenant_id
+        //     );
+        // }

        Ok(deleted)
    }
@@ -2,18 +2,18 @@
 //!   WAL service listens for client connections and
 //!   receive WAL from wal_proposer and send it to WAL receivers
 //!
-use anyhow::Result;
 use regex::Regex;
 use std::net::{TcpListener, TcpStream};
 use std::thread;
 use tracing::*;
+use utils::postgres_backend_async::QueryError;

 use crate::handler::SafekeeperPostgresHandler;
 use crate::SafeKeeperConf;
 use utils::postgres_backend::{AuthType, PostgresBackend};

 /// Accept incoming TCP connections and spawn them into a background thread.
-pub fn thread_main(conf: SafeKeeperConf, listener: TcpListener) -> Result<()> {
+pub fn thread_main(conf: SafeKeeperConf, listener: TcpListener) -> ! {
    loop {
        match listener.accept() {
            Ok((socket, peer_addr)) => {
@@ -44,7 +44,7 @@ fn get_tid() -> u64 {

 /// This is run by `thread_main` above, inside a background thread.
 ///
-fn handle_socket(socket: TcpStream, conf: SafeKeeperConf) -> Result<()> {
+fn handle_socket(socket: TcpStream, conf: SafeKeeperConf) -> Result<(), QueryError> {
    let _enter = info_span!("", tid = ?get_tid()).entered();

    socket.set_nodelay(true)?;
@@ -40,10 +40,9 @@ def parse_metrics(text: str, name: str = "") -> Metrics:


 PAGESERVER_PER_TENANT_REMOTE_TIMELINE_CLIENT_METRICS: Tuple[str, ...] = (
-    "pageserver_remote_upload_queue_unfinished_tasks",
-    "pageserver_remote_operation_seconds_bucket",
-    "pageserver_remote_operation_seconds_count",
-    "pageserver_remote_operation_seconds_sum",
+    "pageserver_remote_timeline_client_calls_unfinished",
+    *[f"pageserver_remote_timeline_client_calls_started_{x}" for x in ["bucket", "count", "sum"]],
+    *[f"pageserver_remote_operation_seconds_{x}" for x in ["bucket", "count", "sum"]],
    "pageserver_remote_physical_size",
 )

@@ -18,6 +18,7 @@ from contextlib import closing, contextmanager
 from dataclasses import dataclass, field
 from enum import Flag, auto
 from functools import cached_property
+from itertools import chain, product
 from pathlib import Path
 from types import TracebackType
 from typing import Any, Dict, Iterator, List, Optional, Tuple, Type, Union, cast
@@ -34,6 +35,7 @@ from _pytest.config import Config
 from _pytest.config.argparsing import Parser
 from _pytest.fixtures import FixtureRequest
 from fixtures.log_helper import log
+from fixtures.metrics import parse_metrics
 from fixtures.types import Lsn, TenantId, TimelineId
 from fixtures.utils import (
    ATTACHMENT_NAME_REGEX,
@@ -595,6 +597,7 @@ class NeonEnvBuilder:
        rust_log_override: Optional[str] = None,
        default_branch_name: str = DEFAULT_BRANCH_NAME,
        preserve_database_files: bool = False,
+        initial_tenant: Optional[TenantId] = None,
    ):
        self.repo_dir = repo_dir
        self.rust_log_override = rust_log_override
@@ -617,8 +620,9 @@ class NeonEnvBuilder:
        self.pg_distrib_dir = pg_distrib_dir
        self.pg_version = pg_version
        self.preserve_database_files = preserve_database_files
+        self.initial_tenant = initial_tenant or TenantId.generate()

-    def init(self) -> NeonEnv:
+    def init_configs(self) -> NeonEnv:
        # Cannot create more than one environment from one builder
        assert self.env is None, "environment already initialized"
        self.env = NeonEnv(self)
@@ -629,8 +633,17 @@ class NeonEnvBuilder:
        self.env.start()

    def init_start(self) -> NeonEnv:
-        env = self.init()
+        env = self.init_configs()
        self.start()
+
+        # Prepare the default branch to start the postgres on later.
+        # Pageserver itself does not create tenants and timelines, until started first and asked via HTTP API.
+        log.info(
+            f"Services started, creating initial tenant {env.initial_tenant} and its initial timeline"
+        )
+        initial_tenant, initial_timeline = env.neon_cli.create_tenant(tenant_id=env.initial_tenant)
+        log.info(f"Initial timeline {initial_tenant}/{initial_timeline} created successfully")
+
        return env

    def enable_remote_storage(
@@ -889,12 +902,12 @@ class NeonEnv:

        # generate initial tenant ID here instead of letting 'neon init' generate it,
        # so that we don't need to dig it out of the config file afterwards.
-        self.initial_tenant = TenantId.generate()
+        self.initial_tenant = config.initial_tenant

        # Create a config file corresponding to the options
        toml = textwrap.dedent(
            f"""
-            default_tenant_id = '{self.initial_tenant}'
+            default_tenant_id = '{config.initial_tenant}'
        """
        )

@@ -1409,6 +1422,33 @@ class PageserverHttpClient(requests.Session):
        ]
        return sample.value

+    def get_remote_timeline_client_metric(
+        self,
+        metric_name: str,
+        tenant_id: TenantId,
+        timeline_id: TimelineId,
+        file_kind: str,
+        op_kind: str,
+    ) -> Optional[float]:
+        metrics = parse_metrics(self.get_metrics(), "pageserver")
+        matches = metrics.query_all(
+            name=metric_name,
+            filter={
+                "tenant_id": str(tenant_id),
+                "timeline_id": str(timeline_id),
+                "file_kind": str(file_kind),
+                "op_kind": str(op_kind),
+            },
+        )
+        if len(matches) == 0:
+            value = None
+        elif len(matches) == 1:
+            value = matches[0].value
+            assert value is not None
+        else:
+            assert len(matches) < 2, "above filter should uniquely identify metric"
+        return value
+
    def get_metric_value(self, name: str) -> Optional[str]:
        metrics = self.get_metrics()
        relevant = [line for line in metrics.splitlines() if line.startswith(name)]
@@ -1528,6 +1568,7 @@ class NeonCli(AbstractNeonCli):
        tenant_id: Optional[TenantId] = None,
        timeline_id: Optional[TimelineId] = None,
        conf: Optional[Dict[str, str]] = None,
+        set_default: bool = False,
    ) -> Tuple[TenantId, TimelineId]:
        """
        Creates a new tenant, returns its id and its initial timeline's id.
@@ -1536,47 +1577,51 @@ class NeonCli(AbstractNeonCli):
            tenant_id = TenantId.generate()
        if timeline_id is None:
            timeline_id = TimelineId.generate()
-        if conf is None:
-            res = self.raw_cli(
-                [
-                    "tenant",
-                    "create",
-                    "--tenant-id",
-                    str(tenant_id),
-                    "--timeline-id",
-                    str(timeline_id),
-                    "--pg-version",
-                    self.env.pg_version,
-                ]
-            )
-        else:
-            res = self.raw_cli(
-                [
-                    "tenant",
-                    "create",
-                    "--tenant-id",
-                    str(tenant_id),
-                    "--timeline-id",
-                    str(timeline_id),
-                    "--pg-version",
-                    self.env.pg_version,
-                ]
-                + sum(list(map(lambda kv: (["-c", kv[0] + ":" + kv[1]]), conf.items())), [])
+
+        args = [
+            "tenant",
+            "create",
+            "--tenant-id",
+            str(tenant_id),
+            "--timeline-id",
+            str(timeline_id),
+            "--pg-version",
+            self.env.pg_version,
+        ]
+        if conf is not None:
+            args.extend(
+                chain.from_iterable(
+                    product(["-c"], (f"{key}:{value}" for key, value in conf.items()))
+                )
            )
+        if set_default:
+            args.append("--set-default")
+
+        res = self.raw_cli(args)
        res.check_returncode()
        return tenant_id, timeline_id

+    def set_default(self, tenant_id: TenantId):
+        """
+        Update default tenant for future operations that require tenant_id.
+        """
+        res = self.raw_cli(["tenant", "set-default", "--tenant-id", str(tenant_id)])
+        res.check_returncode()
+
    def config_tenant(self, tenant_id: TenantId, conf: Dict[str, str]):
        """
        Update tenant config.
        """
-        if conf is None:
-            res = self.raw_cli(["tenant", "config", "--tenant-id", str(tenant_id)])
-        else:
-            res = self.raw_cli(
-                ["tenant", "config", "--tenant-id", str(tenant_id)]
-                + sum(list(map(lambda kv: (["-c", kv[0] + ":" + kv[1]]), conf.items())), [])
+
+        args = ["tenant", "config", "--tenant-id", str(tenant_id)]
+        if conf is not None:
+            args.extend(
+                chain.from_iterable(
+                    product(["-c"], (f"{key}:{value}" for key, value in conf.items()))
+                )
            )
+
+        res = self.raw_cli(args)
        res.check_returncode()

    def list_tenants(self) -> "subprocess.CompletedProcess[str]":
@@ -1611,36 +1656,6 @@ class NeonCli(AbstractNeonCli):

        return TimelineId(str(created_timeline_id))

-    def create_root_branch(
-        self,
-        branch_name: str,
-        tenant_id: Optional[TenantId] = None,
-    ):
-        cmd = [
-            "timeline",
-            "create",
-            "--branch-name",
-            branch_name,
-            "--tenant-id",
-            str(tenant_id or self.env.initial_tenant),
-            "--pg-version",
-            self.env.pg_version,
-        ]
-
-        res = self.raw_cli(cmd)
-        res.check_returncode()
-
-        matches = CREATE_TIMELINE_ID_EXTRACTOR.search(res.stdout)
-
-        created_timeline_id = None
-        if matches is not None:
-            created_timeline_id = matches.group("timeline_id")
-
-        if created_timeline_id is None:
-            raise Exception("could not find timeline id after `neon timeline create` invocation")
-        else:
-            return TimelineId(created_timeline_id)
-
    def create_branch(
        self,
        new_branch_name: str = DEFAULT_BRANCH_NAME,
@@ -1696,17 +1711,12 @@ class NeonCli(AbstractNeonCli):
    def init(
        self,
        config_toml: str,
-        initial_timeline_id: Optional[TimelineId] = None,
    ) -> "subprocess.CompletedProcess[str]":
        with tempfile.NamedTemporaryFile(mode="w+") as tmp:
            tmp.write(config_toml)
            tmp.flush()

-            cmd = ["init", f"--config={tmp.name}"]
-            if initial_timeline_id:
-                cmd.extend(["--timeline-id", str(initial_timeline_id)])
-
-            cmd.extend(["--pg-version", self.env.pg_version])
+            cmd = ["init", f"--config={tmp.name}", "--pg-version", self.env.pg_version]

            append_pageserver_param_overrides(
                params_to_update=cmd,
@@ -1903,14 +1913,17 @@ class NeonPageserver(PgProtocol):
            ".*wal receiver task finished with an error: walreceiver connection handling failure.*",
            ".*Shutdown task error: walreceiver connection handling failure.*",
            ".*wal_connection_manager.*tcp connect error: Connection refused.*",
-            ".*query handler for .* failed: Connection reset by peer.*",
-            ".*serving compute connection task.*exited with error: Broken pipe.*",
-            ".*Connection aborted: error communicating with the server: Broken pipe.*",
-            ".*Connection aborted: error communicating with the server: Transport endpoint is not connected.*",
-            ".*Connection aborted: error communicating with the server: Connection reset by peer.*",
+            ".*query handler for .* failed: Socket IO error: Connection reset by peer.*",
+            ".*connection handler exited with error: Postgres connection error.*",
+            ".*connection handler exited with error: Connection reset by peer.*",
+            ".*connection handler exited with error: Postgres query error.*",
+            ".*Connection aborted: connection error: error communicating with the server: Broken pipe.*",
+            ".*Connection aborted: connection error: error communicating with the server: Transport endpoint is not connected.*",
+            ".*Connection aborted: connection error: error communicating with the server: Connection reset by peer.*",
            ".*kill_and_wait_impl.*: wait successful.*",
-            ".*end streaming to Some.*",
+            ".*Replication stream finished: db error: ERROR: Socket IO error: end streaming to Some.*",
            ".*query handler for 'pagestream.*failed: Broken pipe.*",  # pageserver notices compute shut down
+            ".*query handler for 'pagestream.*failed: Connection reset by peer.*",  # pageserver notices compute shut down
            # safekeeper connection can fail with this, in the window between timeline creation
            # and streaming start
            ".*Failed to process query for timeline .*: state uninitialized, no data to read.*",
@@ -1932,6 +1945,7 @@ class NeonPageserver(PgProtocol):
            # Tenant::delete_timeline() can cause any of the four following errors.
            # FIXME: we shouldn't be considering it an error: https://github.com/neondatabase/neon/issues/2946
            ".*could not flush frozen layer.*queue is in state Stopped",  # when schedule layer upload fails because queued got closed before compaction got killed
+            ".*Compaction failed, .* queue is in state Stopped.*",  # similarly for compaction
            ".*wait for layer upload ops to complete.*",  # .*Caused by:.*wait_completion aborted because upload queue was stopped
            ".*gc_loop.*Gc failed, retrying in.*timeline is Stopping",  # When gc checks timeline state after acquiring layer_removal_cs
            ".*compaction_loop.*Compaction failed, retrying in.*timeline is Stopping",  # When compaction checks timeline state after acquiring layer_removal_cs
@@ -15,7 +15,7 @@ def test_broken_timeline(neon_env_builder: NeonEnvBuilder):

    env.pageserver.allowed_errors.extend(
        [
-            ".*Failed to reconstruct the page.*",
+            ".*Failed to load delta layer.*",
            ".*could not find data for key.*",
            ".*is not active. Current state: Broken.*",
            ".*will not become active. Current state: Broken.*",
@@ -99,7 +99,7 @@ def test_broken_timeline(neon_env_builder: NeonEnvBuilder):
    # Third timeline will also fail during basebackup, because the layer file is corrupt.
    # It will fail when we try to read (and reconstruct) a page from it, ergo the error message.
    # (We don't check layer file contents on startup, when loading the timeline)
-    with pytest.raises(Exception, match="Failed to reconstruct the page") as err:
+    with pytest.raises(Exception, match="Failed to load delta layer") as err:
        pg3.start()
    log.info(
        f"As expected, compute startup failed for timeline {tenant3}/{timeline3} with corrupt layers: {err}"
@@ -1,3 +1,5 @@
+import time
+
 import pytest
 from fixtures.log_helper import log
 from fixtures.metrics import parse_metrics
@@ -20,9 +22,19 @@ def httpserver_listen_address(port_distributor: PortDistributor):
    return ("localhost", port)


-num_metrics_received = 0
+initial_tenant = TenantId.generate()
 remote_uploaded = 0
-first_request = True
+checks = {
+    "written_size": lambda value: value > 0,
+    "resident_size": lambda value: value >= 0,
+    # >= 0 check here is to avoid race condition when we receive metrics before
+    # remote_uploaded is updated
+    "remote_storage_size": lambda value: value > 0 if remote_uploaded > 0 else value >= 0,
+    # logical size may lag behind the actual size, so allow 0 here
+    "timeline_logical_size": lambda value: value >= 0,
+}
+
+metric_kinds_checked = set([])


 #
@@ -36,38 +48,19 @@ def metrics_handler(request: Request) -> Response:
    log.info("received events:")
    log.info(events)

-    checks = {
-        "written_size": lambda value: value > 0,
-        "resident_size": lambda value: value >= 0,
-        # >= 0 check here is to avoid race condition when we receive metrics before
-        # remote_uploaded is updated
-        "remote_storage_size": lambda value: value > 0 if remote_uploaded > 0 else value >= 0,
-        # logical size may lag behind the actual size, so allow 0 here
-        "timeline_logical_size": lambda value: value >= 0,
-    }
-
-    events_received = 0
    for event in events:
-        check = checks.get(event["metric"])
+        assert event["tenant_id"] == str(
+            initial_tenant
+        ), "Expecting metrics only from the initial tenant"
+        metric_name = event["metric"]
+
+        check = checks.get(metric_name)
        # calm down mypy
        if check is not None:
-            assert check(event["value"]), f"{event['metric']} isn't valid"
-            events_received += 1
+            assert check(event["value"]), f"{metric_name} isn't valid"
+            global metric_kinds_checked
+            metric_kinds_checked.add(metric_name)

-    global first_request
-    # check that all checks were sent
-    # but only on the first request, because we don't send non-changed metrics
-    if first_request:
-        # we may receive more metrics than we check,
-        # because there are two timelines
-        # and we may receive per-timeline metrics from both
-        # if the test was slow enough for these metrics to be collected
-        # -1 because that is ok to not receive timeline_logical_size
-        assert events_received >= len(checks) - 1
-        first_request = False
-
-    global num_metrics_received
-    num_metrics_received += 1
    return Response(status=200)


@@ -83,11 +76,14 @@ def test_metric_collection(
    (host, port) = httpserver_listen_address
    metric_collection_endpoint = f"http://{host}:{port}/billing/api/v1/usage_events"

+    # Require collecting metrics frequently, since we change
+    # the timeline and want something to be logged about it.
+    #
    # Disable time-based pitr, we will use the manual GC calls
    # to trigger remote storage operations in a controlled way
    neon_env_builder.pageserver_config_override = (
        f"""
-        metric_collection_interval="60s"
+        metric_collection_interval="1s"
        metric_collection_endpoint="{metric_collection_endpoint}"
    """
        + "tenant_config={pitr_interval = '0 sec'}"
@@ -100,6 +96,9 @@ def test_metric_collection(

    log.info(f"test_metric_collection endpoint is {metric_collection_endpoint}")

+    # Set initial tenant of the test, that we expect the logs from
+    global initial_tenant
+    initial_tenant = neon_env_builder.initial_tenant
    # mock http server that returns OK for the metrics
    httpserver.expect_request("/billing/api/v1/usage_events", method="POST").respond_with_handler(
        metrics_handler
@@ -154,7 +153,11 @@ def test_metric_collection(
        remote_uploaded = get_num_remote_ops("index", "upload")
        assert remote_uploaded > 0

-    # check that all requests are served
+    # wait longer than collecting interval and check that all requests are served
+    time.sleep(3)
    httpserver.check()
-    global num_metrics_received
-    assert num_metrics_received > 0, "no metrics were received"
+    global metric_kinds_checked, checks
+    expected_checks = set(checks.keys())
+    assert len(metric_kinds_checked) == len(
+        checks
+    ), f"Expected to receive and check all kind of metrics, but {expected_checks - metric_kinds_checked} got uncovered"
@@ -1,10 +1,17 @@
-from fixtures.neon_fixtures import NeonEnvBuilder
+from fixtures.neon_fixtures import NeonEnvBuilder, PortDistributor


 # Test that neon cli is able to start and stop all processes with the user defaults.
-# def test_neon_cli_basics(neon_simple_env: NeonEnv):
-def test_neon_cli_basics(neon_env_builder: NeonEnvBuilder):
-    env = neon_env_builder.init()
+# Repeats the example from README.md as close as it can
+def test_neon_cli_basics(neon_env_builder: NeonEnvBuilder, port_distributor: PortDistributor):
+    env = neon_env_builder.init_configs()
+    # Skipping the init step that creates a local tenant in Pytest tests
+    try:
+        env.neon_cli.start()
+        env.neon_cli.create_tenant(tenant_id=env.initial_tenant, set_default=True)
+        env.neon_cli.pg_start(node_name="main", port=port_distributor.get_port())

-    env.neon_cli.start()
-    env.neon_cli.stop()
+        env.neon_cli.create_branch(new_branch_name="migration_check")
+        env.neon_cli.pg_start(node_name="migration_check", port=port_distributor.get_port())
+    finally:
+        env.neon_cli.stop()
@@ -87,8 +87,22 @@ def test_ondemand_download_large_rel(
    # run checkpoint manually to be sure that data landed in remote storage
    client.timeline_checkpoint(tenant_id, timeline_id)

-    # wait until pageserver successfully uploaded a checkpoint to remote storage
+    # wait until pageserver successfully uploaded all the data, and the result
+    # of the compaction, to remote storage
    wait_for_upload(client, tenant_id, timeline_id, current_lsn)
+    def assert_queue_empty(file_kind, op_kind):
+        val = client.get_remote_timeline_client_metric(
+            "pageserver_remote_timeline_client_calls_unfinished",
+            tenant_id,
+            timeline_id,
+            file_kind,
+            op_kind,
+        )
+        assert val is not None, "expecting metric to be present"
+        assert int(val) == 0
+    wait_until(60, 0.5, lambda: assert_queue_empty(file_kind="layer", op_kind="upload"))
+    wait_until(60, 0.5, lambda: assert_queue_empty(file_kind="index", op_kind="upload"))
+    wait_until(60, 0.5, lambda: assert_queue_empty(file_kind="layer", op_kind="delete"))
    log.info("uploads have finished")

    ##### Stop the first pageserver instance, erase all its data
@@ -112,7 +126,7 @@ def test_ondemand_download_large_rel(
    # relation, so they are likely already downloaded. But the middle of the
    # table should not have been needed by anything yet.
    with pg.cursor() as cur:
-        assert query_scalar(cur, "select count(*) from tbl where id = 500000") == 1
+        assert query_scalar(cur, f"select count(*) from tbl where id = {num_rows}/2") == 1

    after_downloads = get_num_downloaded_layers(client, tenant_id, timeline_id)
    log.info(f"layers downloaded before {before_downloads} and after {after_downloads}")
@@ -120,7 +134,7 @@ def test_ondemand_download_large_rel(


 #
-# If you have a relation with a long history of updates,the pageserver downloads the layer
+# If you have a relation with a long history of updates, the pageserver downloads the layer
 # files containing the history as needed by timetravel queries.
 #
@pytest.mark.parametrize("remote_storage_kind", available_remote_storages())
@@ -189,13 +203,10 @@ def test_ondemand_download_timetravel(
        # run checkpoint manually to be sure that data landed in remote storage
        client.timeline_checkpoint(tenant_id, timeline_id)

-    # wait until pageserver successfully uploaded a checkpoint to remote storage
-    wait_for_upload(client, tenant_id, timeline_id, current_lsn)
-    log.info("uploads have finished")
-
    ##### Stop the first pageserver instance, erase all its data
    env.postgres.stop_all()

+    # wait until pageserver has successfully uploaded all the data to remote storage
    wait_for_sk_commit_lsn_to_reach_remote_storage(
        tenant_id, timeline_id, env.safekeepers, env.pageserver
    )
@@ -227,11 +238,15 @@ def test_ondemand_download_timetravel(

    wait_until(10, 0.2, lambda: assert_tenant_status(client, tenant_id, "Active"))

-    # current_physical_size reports sum of layer file sizes, regardless of local or remote
+    # The current_physical_size reports the sum of layers loaded in the layer
+    # map, regardless of where the layer files are located. So even though we
+    # just removed the local files, they still count towards
+    # current_physical_size because they are loaded as `RemoteLayer`s.
    assert filled_current_physical == get_api_current_physical_size()

+    # Run queries at different points in time
    num_layers_downloaded = [0]
-    physical_size = [get_resident_physical_size()]
+    resident_size = [get_resident_physical_size()]
    for (checkpoint_number, lsn) in lsns:
        pg_old = env.postgres.create_start(
            branch_name="main", node_name=f"test_old_lsn_{checkpoint_number}", lsn=lsn
@@ -268,13 +283,15 @@ def test_ondemand_download_timetravel(
        if len(num_layers_downloaded) > 4:
            assert after_downloads > num_layers_downloaded[len(num_layers_downloaded) - 4]

-        # Likewise, assert that the physical_size metric grows as layers are downloaded
-        physical_size.append(get_resident_physical_size())
-        log.info(f"physical_size[-1]={physical_size[-1]}")
-        if len(physical_size) > 4:
-            assert physical_size[-1] > physical_size[len(physical_size) - 4]
+        # Likewise, assert that the resident_physical_size metric grows as layers are downloaded
+        resident_size.append(get_resident_physical_size())
+        log.info(f"resident_size[-1]={resident_size[-1]}")
+        if len(resident_size) > 4:
+            assert resident_size[-1] > resident_size[len(resident_size) - 4]

-        # current_physical_size reports sum of layer file sizes, regardless of local or remote
+        # current_physical_size reports the total size of all layer files, whether
+        # they are present only in the remote storage, only locally, or both.
+        # It should not change.
        assert filled_current_physical == get_api_current_physical_size()


@@ -368,7 +385,7 @@ def test_download_remote_layers_api(
    env.pageserver.allowed_errors.extend(
        [
            f".*download_all_remote_layers.*{tenant_id}.*{timeline_id}.*layer download failed.*remote-storage-download-pre-rename failpoint",
-            f".*initial size calculation.*{tenant_id}.*{timeline_id}.*Failed to calculate logical size",
+            f".*initial size calculation for {tenant_id}/{timeline_id} failed.*remote-storage-download-pre-rename failpoint",
        ]
    )

@@ -12,11 +12,9 @@ def test_pageserver_recovery(neon_env_builder: NeonEnvBuilder):
    # Override default checkpointer settings to run it more often
    neon_env_builder.pageserver_config_override = "tenant_config={checkpoint_distance = 1048576}"

-    env = neon_env_builder.init()
+    env = neon_env_builder.init_start()
    env.pageserver.is_testing_enabled_or_skip()

-    neon_env_builder.start()
-
    # These warnings are expected, when the pageserver is restarted abruptly
    env.pageserver.allowed_errors.append(".*found future delta layer.*")
    env.pageserver.allowed_errors.append(".*found future image layer.*")
@@ -2,11 +2,11 @@
 # env NEON_PAGESERVER_OVERRIDES="remote_storage={local_path='/tmp/neon_zzz/'}" poetry ......

 import os
-import re
 import shutil
 import threading
 import time
 from pathlib import Path
+from typing import Dict, List, Tuple

 import pytest
 from fixtures.log_helper import log
@@ -271,14 +271,15 @@ def test_remote_storage_upload_queue_retries(
        wait_for_last_flush_lsn(env, pg, tenant_id, timeline_id)

    def get_queued_count(file_kind, op_kind):
-        metrics = client.get_metrics()
-        matches = re.search(
-            f'^pageserver_remote_upload_queue_unfinished_tasks{{file_kind="{file_kind}",op_kind="{op_kind}",tenant_id="{tenant_id}",timeline_id="{timeline_id}"}} (\\S+)$',
-            metrics,
-            re.MULTILINE,
+        val = client.get_remote_timeline_client_metric(
+            "pageserver_remote_timeline_client_calls_unfinished",
+            tenant_id,
+            timeline_id,
+            file_kind,
+            op_kind,
        )
-        assert matches
-        return int(matches[1])
+        assert val is not None, "expecting metric to be present"
+        return int(val)

    # create some layers & wait for uploads to finish
    overwrite_data_and_wait_for_it_to_arrive_at_pageserver("a")
@@ -368,6 +369,168 @@ def test_remote_storage_upload_queue_retries(
        assert query_scalar(cur, "SELECT COUNT(*) FROM foo WHERE val = 'd'") == 10000


+@pytest.mark.parametrize("remote_storage_kind", [RemoteStorageKind.LOCAL_FS])
+def test_remote_timeline_client_calls_started_metric(
+    neon_env_builder: NeonEnvBuilder,
+    remote_storage_kind: RemoteStorageKind,
+):
+
+    neon_env_builder.enable_remote_storage(
+        remote_storage_kind=remote_storage_kind,
+        test_name="test_remote_timeline_client_metrics",
+    )
+
+    env = neon_env_builder.init_start()
+
+    # create tenant with config that will determinstically allow
+    # compaction and gc
+    tenant_id, timeline_id = env.neon_cli.create_tenant(
+        conf={
+            # small checkpointing and compaction targets to ensure we generate many upload operations
+            "checkpoint_distance": f"{128 * 1024}",
+            "compaction_threshold": "1",
+            "compaction_target_size": f"{128 * 1024}",
+            # no PITR horizon, we specify the horizon when we request on-demand GC
+            "pitr_interval": "0s",
+            # disable background compaction and GC. We invoke it manually when we want it to happen.
+            "gc_period": "0s",
+            "compaction_period": "0s",
+            # don't create image layers, that causes just noise
+            "image_creation_threshold": "10000",
+        }
+    )
+
+    client = env.pageserver.http_client()
+
+    pg = env.postgres.create_start("main", tenant_id=tenant_id)
+
+    pg.safe_psql("CREATE TABLE foo (id INTEGER PRIMARY KEY, val text)")
+
+    def overwrite_data_and_wait_for_it_to_arrive_at_pageserver(data):
+        # create initial set of layers & upload them with failpoints configured
+        pg.safe_psql_many(
+            [
+                f"""
+               INSERT INTO foo (id, val)
+               SELECT g, '{data}'
+               FROM generate_series(1, 10000) g
+               ON CONFLICT (id) DO UPDATE
+               SET val = EXCLUDED.val
+               """,
+                # to ensure that GC can actually remove some layers
+                "VACUUM foo",
+            ]
+        )
+        wait_for_last_flush_lsn(env, pg, tenant_id, timeline_id)
+
+    def get_queued_count(file_kind, op_kind):
+        val = client.get_remote_timeline_client_metric(
+            "pageserver_remote_timeline_client_calls_unfinished",
+            tenant_id,
+            timeline_id,
+            file_kind,
+            op_kind,
+        )
+        if val is None:
+            return val
+        return int(val)
+
+    def wait_upload_queue_empty():
+        wait_until(2, 1, lambda: get_queued_count(file_kind="layer", op_kind="upload") == 0)
+        wait_until(2, 1, lambda: get_queued_count(file_kind="index", op_kind="upload") == 0)
+        wait_until(2, 1, lambda: get_queued_count(file_kind="layer", op_kind="delete") == 0)
+
+    calls_started: Dict[Tuple[str, str], List[int]] = {
+        ("layer", "upload"): [0],
+        ("index", "upload"): [0],
+        ("layer", "delete"): [0],
+    }
+
+    def fetch_calls_started():
+        for (file_kind, op_kind), observations in calls_started.items():
+            val = client.get_remote_timeline_client_metric(
+                "pageserver_remote_timeline_client_calls_started_count",
+                tenant_id,
+                timeline_id,
+                file_kind,
+                op_kind,
+            )
+            assert val is not None, f"expecting metric to be present: {file_kind} {op_kind}"
+            val = int(val)
+            observations.append(val)
+
+    def ensure_calls_started_grew():
+        for (file_kind, op_kind), observations in calls_started.items():
+            log.info(f"ensure_calls_started_grew: {file_kind} {op_kind}: {observations}")
+            assert all(
+                x < y for x, y in zip(observations, observations[1:])
+            ), f"observations for {file_kind} {op_kind} did not grow monotonically: {observations}"
+
+    def churn(data_pass1, data_pass2):
+        overwrite_data_and_wait_for_it_to_arrive_at_pageserver(data_pass1)
+        client.timeline_checkpoint(tenant_id, timeline_id)
+        client.timeline_compact(tenant_id, timeline_id)
+        overwrite_data_and_wait_for_it_to_arrive_at_pageserver(data_pass2)
+        client.timeline_checkpoint(tenant_id, timeline_id)
+        client.timeline_compact(tenant_id, timeline_id)
+        gc_result = client.timeline_gc(tenant_id, timeline_id, 0)
+        print_gc_result(gc_result)
+        assert gc_result["layers_removed"] > 0
+
+    # create some layers & wait for uploads to finish
+    churn("a", "b")
+
+    wait_upload_queue_empty()
+
+    # ensure that we updated the calls_started metric
+    fetch_calls_started()
+    ensure_calls_started_grew()
+
+    # more churn to cause more operations
+    churn("c", "d")
+
+    # ensure that the calls_started metric continued to be updated
+    fetch_calls_started()
+    ensure_calls_started_grew()
+
+    ### now we exercise the download path
+    calls_started.clear()
+    calls_started.update(
+        {
+            ("index", "download"): [0],
+            ("layer", "download"): [0],
+        }
+    )
+
+    env.pageserver.stop(immediate=True)
+    env.postgres.stop_all()
+
+    dir_to_clear = Path(env.repo_dir) / "tenants"
+    shutil.rmtree(dir_to_clear)
+    os.mkdir(dir_to_clear)
+
+    env.pageserver.start()
+    client = env.pageserver.http_client()
+
+    client.tenant_attach(tenant_id)
+
+    def tenant_active():
+        all_states = client.tenant_list()
+        [tenant] = [t for t in all_states if TenantId(t["id"]) == tenant_id]
+        assert tenant["state"] == "Active"
+
+    wait_until(30, 1, tenant_active)
+
+    log.info("restarting postgres to validate")
+    pg = env.postgres.create_start("main", tenant_id=tenant_id)
+    with pg.cursor() as cur:
+        assert query_scalar(cur, "SELECT COUNT(*) FROM foo WHERE val = 'd'") == 10000
+
+    # ensure that we updated the calls_started download metric
+    fetch_calls_started()
+    ensure_calls_started_grew()
+
+
 # Test that we correctly handle timeline with layers stuck in upload queue
@pytest.mark.parametrize("remote_storage_kind", [RemoteStorageKind.LOCAL_FS])
 def test_timeline_deletion_with_files_stuck_in_upload_queue(
@@ -380,6 +543,7 @@ def test_timeline_deletion_with_files_stuck_in_upload_queue(
    )

    env = neon_env_builder.init_start()
+    env.pageserver.allowed_errors.append(".*request cancelled while waiting on uploads to finish.*")

    # create tenant with config that will determinstically allow
    # compaction and gc
@@ -401,15 +565,14 @@ def test_timeline_deletion_with_files_stuck_in_upload_queue(
    client = env.pageserver.http_client()

    def get_queued_count(file_kind, op_kind):
-        metrics = client.get_metrics()
-        matches = re.search(
-            f'^pageserver_remote_upload_queue_unfinished_tasks{{file_kind="{file_kind}",op_kind="{op_kind}",tenant_id="{tenant_id}",timeline_id="{timeline_id}"}} (\\S+)$',
-            metrics,
-            re.MULTILINE,
+        val = client.get_remote_timeline_client_metric(
+            "pageserver_remote_timeline_client_calls_unfinished",
+            tenant_id,
+            timeline_id,
+            file_kind,
+            op_kind,
        )
-        if matches is None:
-            return None
-        return int(matches[1])
+        return int(val) if val is not None else val

    pg = env.postgres.create_start("main", tenant_id=tenant_id)

@@ -1,9 +1,13 @@
+import asyncio
+import random
 import time
 from threading import Thread

+import asyncpg
 import pytest
 from fixtures.log_helper import log
 from fixtures.neon_fixtures import (
+    NeonEnv,
    NeonEnvBuilder,
    PageserverApiException,
    PageserverHttpClient,
@@ -12,6 +16,7 @@ from fixtures.neon_fixtures import (
    available_remote_storages,
    wait_for_last_record_lsn,
    wait_for_upload,
+    wait_until,
    wait_until_tenant_state,
 )
 from fixtures.types import Lsn, TenantId, TimelineId
@@ -84,6 +89,151 @@ def test_tenant_reattach(
    assert env.pageserver.log_contains(".*download.*failed, will retry.*")


+num_connections = 10
+num_rows = 100000
+updates_to_perform = 0
+
+updates_started = 0
+updates_finished = 0
+
+
+# Run random UPDATEs on test table. On failure, try again.
+async def update_table(pg_conn: asyncpg.Connection):
+    global updates_started, updates_finished, updates_to_perform
+
+    while updates_started < updates_to_perform or updates_to_perform == 0:
+        updates_started += 1
+        id = random.randrange(1, num_rows)
+
+        # Loop to retry until the UPDATE succeeds
+        while True:
+            try:
+                await pg_conn.fetchrow(f"UPDATE t SET counter = counter + 1 WHERE id = {id}")
+                updates_finished += 1
+                if updates_finished % 1000 == 0:
+                    log.info(f"update {updates_finished} / {updates_to_perform}")
+                break
+            except asyncpg.PostgresError as e:
+                # Received error from Postgres. Log it, sleep a little, and continue
+                log.info(f"UPDATE error: {e}")
+                await asyncio.sleep(0.1)
+
+
+async def sleep_and_reattach(pageserver_http: PageserverHttpClient, tenant_id: TenantId):
+    global updates_started, updates_finished, updates_to_perform
+
+    # Wait until we have performed some updates
+    wait_until(20, 0.5, lambda: updates_finished > 500)
+
+    log.info("Detaching tenant")
+    pageserver_http.tenant_detach(tenant_id)
+    await asyncio.sleep(1)
+    log.info("Re-attaching tenant")
+    pageserver_http.tenant_attach(tenant_id)
+    log.info("Re-attach finished")
+
+    # Continue with 5000 more updates
+    updates_to_perform = updates_started + 5000
+
+
+# async guts of test_tenant_reattach_while_bysy test
+async def reattach_while_busy(
+    env: NeonEnv, pg: Postgres, pageserver_http: PageserverHttpClient, tenant_id: TenantId
+):
+    workers = []
+    for worker_id in range(num_connections):
+        pg_conn = await pg.connect_async()
+        workers.append(asyncio.create_task(update_table(pg_conn)))
+
+    workers.append(asyncio.create_task(sleep_and_reattach(pageserver_http, tenant_id)))
+    await asyncio.gather(*workers)
+
+    assert updates_finished == updates_to_perform
+
+
+# Detach and re-attach tenant, while compute is busy running queries.
+#
+# Some of the queries may fail, in the window that the tenant has been
+# detached but not yet re-attached. But Postgres itself should keep
+# running, and when we retry the queries, they should start working
+# after the attach has finished.
+
+# FIXME:
+#
+# This is pretty unstable at the moment. I've seen it fail with a warning like this:
+#
+# AssertionError: assert not ['2023-01-05T13:09:40.708303Z  WARN remote_upload{tenant=c3fc41f6cf29a7626b90316e3518cd4b timeline=7978246f85faa71ab03...1282b/000000000000000000000000000000000000-FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF__0000000001716699-0000000001736681"\n']
+#
+# (https://neon-github-public-dev.s3.amazonaws.com/reports/pr-3232/debug/3846817847/index.html#suites/f9eba3cfdb71aa6e2b54f6466222829b/470fc62b5db7d7d7/)
+# I believe that failure happened because there is a race condition
+# between detach and starting remote upload tasks:
+#
+# 1. detach_timeline calls task_mgr::shutdown_tasks(), sending shutdown
+#    signal to all in-progress tasks associated with the tenant.
+# 2. Just after shutdown_tasks() has collected the list of tasks,
+#    a new remote-upload task is spawned.
+#
+# See https://github.com/neondatabase/neon/issues/3273
+#
+#
+# I also saw this failure:
+#
+# test_runner/regress/test_tenant_detach.py:194: in test_tenant_reattach_while_busy
+#     asyncio.run(reattach_while_busy(env, pg, pageserver_http, tenant_id))
+# /home/nonroot/.pyenv/versions/3.9.2/lib/python3.9/asyncio/runners.py:44: in run
+#     return loop.run_until_complete(main)
+# /home/nonroot/.pyenv/versions/3.9.2/lib/python3.9/asyncio/base_events.py:642: in run_until_complete
+#     return future.result()
+# test_runner/regress/test_tenant_detach.py:151: in reattach_while_busy
+#     assert updates_finished == updates_to_perform
+# E   assert 5010 == 10010
+# E     +5010
+# E     -10010
+#
+# I don't know what's causing that...
+@pytest.mark.skip(reason="fixme")
+@pytest.mark.parametrize("remote_storage_kind", available_remote_storages())
+def test_tenant_reattach_while_busy(
+    neon_env_builder: NeonEnvBuilder,
+    remote_storage_kind: RemoteStorageKind,
+):
+    neon_env_builder.enable_remote_storage(
+        remote_storage_kind=remote_storage_kind,
+        test_name="test_tenant_reattach_while_busy",
+    )
+    env = neon_env_builder.init_start()
+
+    # Attempts to connect from compute to pageserver while the tenant is
+    # temporarily detached produces these errors in the pageserver log.
+    env.pageserver.allowed_errors.append(".*Tenant .* not found in the local state.*")
+    env.pageserver.allowed_errors.append(
+        ".*Tenant .* will not become active\\. Current state: Stopping.*"
+    )
+    env.pageserver.allowed_errors.append(".*Tenant .* is not active, state: Stopping.*")
+
+    pageserver_http = env.pageserver.http_client()
+
+    # create new nenant
+    tenant_id, timeline_id = env.neon_cli.create_tenant(
+        # Create layers aggressively
+        conf={"checkpoint_distance": "100000"}
+    )
+
+    pg = env.postgres.create_start("main", tenant_id=tenant_id)
+
+    cur = pg.connect().cursor()
+
+    cur.execute("CREATE TABLE t(id int primary key, counter int)")
+    cur.execute(f"INSERT INTO t SELECT generate_series(1,{num_rows}), 0")
+
+    # Run the test
+    asyncio.run(reattach_while_busy(env, pg, pageserver_http, tenant_id))
+
+    # Verify table contents
+    assert query_scalar(cur, "SELECT count(*) FROM t") == num_rows
+    assert query_scalar(cur, "SELECT sum(counter) FROM t") == updates_to_perform
+
+
 def test_tenant_detach_smoke(neon_env_builder: NeonEnvBuilder):
    env = neon_env_builder.init_start()
    pageserver_http = env.pageserver.http_client()
@@ -567,6 +567,8 @@ def test_s3_wal_replay(neon_env_builder: NeonEnvBuilder, remote_storage_kind: Re
    for sk in env.safekeepers:
        cli = sk.http_client()
        cli.timeline_delete_force(tenant_id, timeline_id)
+        # restart safekeeper to clear its in-memory state
+        sk.stop().start()
        cli.timeline_create(tenant_id, timeline_id, pg_version, last_lsn)
        f_partial_path = (
            Path(sk.data_dir()) / str(tenant_id) / str(timeline_id) / f_partial_saved.name
@@ -1105,7 +1107,7 @@ def test_delete_force(neon_env_builder: NeonEnvBuilder, auth_enabled: bool):
    env.pageserver.allowed_errors.extend(
        [
            ".*Failed to process query for timeline .*: Timeline .* was not found in global map.*",
-            ".*end streaming to Some.*",
+            ".*Failed to process query for timeline .*: Timeline .* was cancelled and cannot be used anymore.*",
        ]
    )

@@ -1206,7 +1208,7 @@ def test_delete_force(neon_env_builder: NeonEnvBuilder, auth_enabled: bool):

    # Remove initial tenant again.
    response = sk_http.tenant_delete_force(tenant_id)
-    assert response == {}
+    # assert response == {}
    assert not (sk_data_dir / str(tenant_id)).exists()
    assert (sk_data_dir / str(tenant_id_other) / str(timeline_id_other)).is_dir()

@@ -40,7 +40,7 @@ scopeguard = { version = "1", features = ["use_std"] }
 serde = { version = "1", features = ["alloc", "derive", "serde_derive", "std"] }
 serde_json = { version = "1", features = ["raw_value", "std"] }
 socket2 = { version = "0.4", default-features = false, features = ["all"] }
-tokio = { version = "1", features = ["bytes", "fs", "io-std", "io-util", "libc", "macros", "memchr", "mio", "net", "num_cpus", "once_cell", "process", "rt", "rt-multi-thread", "signal-hook-registry", "socket2", "sync", "time", "tokio-macros"] }
+tokio = { version = "1", features = ["bytes", "fs", "io-std", "io-util", "libc", "macros", "memchr", "mio", "net", "num_cpus", "process", "rt", "rt-multi-thread", "signal-hook-registry", "socket2", "sync", "time", "tokio-macros"] }
 tokio-util = { version = "0.7", features = ["codec", "io", "io-util", "tracing"] }
 tower = { version = "0.4", features = ["__common", "balance", "buffer", "discover", "futures-core", "futures-util", "indexmap", "limit", "load", "log", "make", "pin-project", "pin-project-lite", "rand", "ready-cache", "retry", "slab", "timeout", "tokio", "tokio-util", "tracing", "util"] }
 tracing = { version = "0.1", features = ["attributes", "log", "std", "tracing-attributes"] }
Author	SHA1	Message	Date
Heikki Linnakangas	9b5792b9bf	Silence test failure caused by expected error in log.	2023-01-13 10:26:15 +02:00
Heikki Linnakangas	be0dfa9d3a	Fix test_ondemand_download_large_rel if uploads are slow. If the uploads after compaction happen slowly, they might have finished before the pageserver is shut down. The L0 files have been uploaded, so no data is lost, but then the query later in the test will need to download all the L0 files, and causes the test to fail because it specifically checks that download happens on-demand, not all at once.	2023-01-12 23:40:46 +02:00
Heikki Linnakangas	292c42731e	Put back spans that were lost along with 'immediate_compact'.	2023-01-12 22:49:00 +02:00
Heikki Linnakangas	867b35ce55	Try to fix regression failures.	2023-01-12 20:39:04 +02:00
Heikki Linnakangas	14ff793582	Add comment about the effect of TaskKind to shutdown sequence	2023-01-12 19:24:30 +02:00
Heikki Linnakangas	5aaa5302eb	Introduce RequestContexts. RequestContext is used to track each "operation" or "task" in a way that's not tied to tokio tasks. It provides support for fine-grained cancellation of individual operations, or all tasks working on an active tenant or timeline. Most async functions now take a RequestContext argument. RequestContexts form a hierarchy, so that you have a top-level context e.g. for a TCP listener task, a child context for each task handling a connection, and perhaps a grandchild context for each individual client request. In addition to the hierarchy, each RequestContext can be associated with a Tenant or Timeline object. This is used to prevent a Tenant or Timeline from being deleted or detached while there are still tasks accessing it. This fixes a long-standing race conditions between GC/compaction and deletion (see issues #2914 and compiler in any way, but the functions like `get_active_timeline` make it easy to do the right thing. This replaces most of the machinery in `task_mgr.rs`. We don't track running tasks as such anymore, only RequestContexts. In practice, every task holds onto a RequestContext. In addition to supporting cancellation, the RequestContext specifies the desired behavior if a remote layer is needed for the operation. This replaces the `with_ondemand_download_sync` and `no_ondemand_download` macros. The on-demand download now happens deep in the call stack, in get_reconstruct_data(), and the caller is no longer involved in the download, except by passing a RequestContext that specifies whether to do on-demand download or not. The PageReconstructResult type is gone but the PageReconstructError::NeedsDownload variant remains. It's now returned if the context specified "don't do on-demand download", and a layer is missing. TODO: - Enforce better that you hold a RequestContext associated with a Tenant or Timeline. - All the fields in RequestContext are currently 'pub', but things will break if you modify the tenant/timeline fields directly. Make that more safe. - When you create a subcontext, should it inherit the Tenant / Timeline of its parent? - Can the walreceiver::TaskHandle stuff be replaced with this? - Extract smaller patches: - What else could we extract?	2023-01-12 19:24:25 +02:00
Heikki Linnakangas	6a53b8fac6	Add placeholders for RequestContext and friends. This commit adds the 'ctx' parameter to all the functions that will need an active context. However, you can just create new contexts on the fly, there is no cross-checks that the tenant/timeline is still in active state. You can simply call Tenant::get_context or Timeline::get_context, and they always succee. In the next commit, we will change the functions for constructing contexts, so that you cannot create a new TenantRequestContext if the tenant is being stopped (and similarly for TimelineRequestContext). This commit isn't useful on its own, but splitting these fairly mechanical changes helps to make the next commit smaller, and thus easier to review. Because the contexts are merely passed through places, and not actually used for anything, this introduces a lot of "unused variable" warnings. They will go away with the next commit.	2023-01-12 19:17:44 +02:00
Heikki Linnakangas	57a6e931ea	Comment, formatting and other cosmetic cleanup.	2023-01-12 19:05:13 +02:00
Heikki Linnakangas	0cceb14e48	Add a FIXME on ugly error message parsing.	2023-01-12 19:05:13 +02:00
Konstantin Knizhnik	1983c4d4ad	Explain prefetch (#3002 ) Co-authored-by: Bojan Serafimov <bojan.serafimov7@gmail.com>	2023-01-12 18:12:40 +02:00
Heikki Linnakangas	d7c41cbbee	Replace tokio::watch with CancellationToken. PR #3228 starts to use CancellationTokens more widely, this is a small part extracted from that.	2023-01-12 17:37:15 +02:00
Vadim Kharitonov	29a2465276	Update rust version in toolchain	2023-01-12 15:16:52 +01:00
Arthur Petukhovsky	f49e923d87	Keep deleted timelines in memory of safekeeper (#3300 ) A temporal fix for https://github.com/neondatabase/neon/issues/3146, until we come up with a reliable way to create and delete timelines in all safekeepers.	2023-01-12 15:33:07 +03:00
Vadim Kharitonov	a0ee306c74	Enable safe contrib extensions	2023-01-12 12:41:53 +01:00
Heikki Linnakangas	c1731bc4f0	Push on-demand download into Timeline::get() function itself. This makes Timeline::get() async, and all functions that call it directly or indirectly with it. The with_ondemand_download() mechanism is gone, Timeline::get() now always downloads files, whether you want it or not. That is what all the current callers want, so even though this loses the capability to get a page only if it's already in the pageserver, without downloading, we were not using that capability. There were some places that used 'no_ondemand_download' in the WAL ingestion code that would error out if a layer file was not found locally, but those were dubious. We do actually want to on-demand download in all of those places. Per discussion at https://github.com/neondatabase/neon/pull/3233#issuecomment-1368032358	2023-01-12 11:53:10 +02:00
Sergey Melnikov	95bf19b85a	Add --atomic to all helm upgrade operations (#3299 ) When number of github actions workers is changed, some jobs get killed. When helm if killed during the upgrade, release stuck in pending-upgrade state. --atomic should initiate automatic rollback in this case.	2023-01-10 10:05:27 +00:00
Vadim Kharitonov	80d4afab0c	Update tokio version (RUSTSEC-2023-0001)	2023-01-10 09:02:00 +01:00
Arthur Petukhovsky	0807522a64	Enable wss proxy in all regions (#3292 ) Follow-up to https://github.com/neondatabase/helm-charts/pull/24 and #3247	2023-01-09 19:56:12 +00:00
Christian Schwarz	8eebd5f039	run on-demand compaction in a task_mgr task With this patch, tenant_detach and timeline_delete's task_mgr::shutdown_tasks() call will wait for on-demand compaction to finish. Before this patch, the on-demand compaction would grab the layer_removal_cs after tenant_detach / timeline_delete had removed the timeline directory. This resulted in error No such file or directory (os error 2) NB: I already implemented this pattern for ondemand GC a while back. fixes https://github.com/neondatabase/neon/issues/3136	2023-01-09 19:08:22 +01:00
Heikki Linnakangas	8c07ef413d	Minor cleanup of test_ondemand_download_timetravel test. - Fix and improve comments - Rename 'physical_size' local variable to 'resident_size' for clarity. - Remove one 'unnecessary wait_for_upload' call. The 'wait_for_sk_commit_lsn_to_reach_remote_storage' call after shutting down compute is sufficient.	2023-01-09 18:56:50 +02:00
Sergey Melnikov	14df37c108	Use GHA environments for gradual prod rollout (#3295 ) Each release will wait for manual approval for each region	2023-01-09 20:18:16 +04:00
Christian Schwarz	d4d0aa6ed6	gc_iteration_internal: better log message & debug log level if nothing to do fixes https://github.com/neondatabase/neon/issues/3107	2023-01-09 13:53:59 +01:00
Kirill Bulatov	a457256fef	Fix log message matching (#3291 ) Spotted https://neon-github-public-dev.s3.amazonaws.com/reports/main/debug/3871991071/index.html#suites/158be07438eb5188d40b466b6acfaeb3/22966d740e33b677/ failing on `main`, fixes that by using a proper regex match string. Also removes one clippy lint suppression.	2023-01-09 14:25:12 +02:00
Shany Pozin	3a22e1335d	Adding a PR template (#3288 ) ## Describe your changes Added a PR template ## Issue ticket number and link #3162 ## 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.	2023-01-09 12:15:53 +00:00
Sergey Melnikov	93c77b0383	Use GHA environment for per-region deploy approvals on staging (#3293 ) Each main deploy will wait for manual approval for each region	2023-01-09 15:40:14 +04:00
Shany Pozin	7920b39a27	Adding transition reason to the log when a tenant is moved to Broken state (#3289 ) #3160	2023-01-09 10:24:50 +02:00
Kirill Bulatov	23d5e2bdaa	Fix common pg port in the CLI basics test (#3283 ) Closes https://github.com/neondatabase/neon/issues/3282	2023-01-07 00:46:42 +02:00
Christian Schwarz	3526323bc4	prepare Timeline::get_reconstruct_data for becoming async (#3271 ) This patch restructures the code so that PR https://github.com/neondatabase/neon/pull/3228 can seamlessly replace the return PageReconstructResult::NeedsDownload with a download_remote_layer().await. Background: PR https://github.com/neondatabase/neon/pull/3228 will turn get_reconstruct_data() async and do the on-demand download right in place, instead of returning a PageReconstructResult::NeedsDownload. Current rustc requires that the layers lock guard be not in scope across an await point. For on-demand download inside get_reconstruct_data(), we need to do download_remote_layer().await. Supersedes https://github.com/neondatabase/neon/pull/3260 See my comment there: https://github.com/neondatabase/neon/pull/3260#issuecomment-1370752407 Co-authored-by: Heikki Linnakangas <heikki@neon.tech>	2023-01-06 19:42:25 +02:00
Heikki Linnakangas	af9425394f	Print time taken by CREATE/ALTER DATABASE at compute start. Trying to investigate why the "apply_config" stage is taking longer than expected. This proves or disproves that it's the CREATE DATABASE statement.	2023-01-06 17:50:44 +02:00
Arthur Petukhovsky	debd134b15	Implement wss support in proxy (#3247 ) This is a hacky implementation of WebSocket server, embedded into our postgres proxy. The server is used to allow https://github.com/neondatabase/serverless to connect to our postgres from browser and serverless javascript functions. How it will work (general schema): - browser opens a websocket connection to `wss://ep-abc-xyz-123.xx-central-1.aws.neon.tech/` - proxy accepts this connection and terminates TLS (https) - inside encrypted tunnel (HTTPS), browser initiates plain (non-encrypted) postgres connection - proxy performs auth as in usual plain pg connection and forwards connection to the compute Related issue: #3225	2023-01-06 18:34:18 +03:00
Heikki Linnakangas	df42213dbb	Fix missing COMMIT in handle_role_deletions. There was no COMMIT, so the DROP ROLE commands were always implicitly rolled back. Fixes issue #3279.	2023-01-06 17:07:46 +02:00
Kirill Bulatov	b6237474d2	Fix README and basic startup example (#3275 ) Follow-up of https://github.com/neondatabase/neon/pull/3270 which made an example from main README.md not working. Fixes that, by adding a way to specify a default tenant now and modifies the basic neon_local test to start postgres and check branching. Not all neon_local commands are implemented, so not all README.md contents is tested yet.	2023-01-06 12:26:14 +02:00
Heikki Linnakangas	8b710b9753	Fix segfault if pageserver connection is lost during backend startup. It's not OK to return early from within a PG_TRY-CATCH block. The PG_TRY macro sets the global PG_exception_stack variable, and PG_END_TRY restores it. If we jump out in between with "return NULL", the PG_exception_stack is left to point to garbage. (I'm surprised the comments in PG_TRY_CATCH don't warn about this.) Add test that re-attaches tenant in pageserver while Postgres is running. If the tenant is detached while compute is connected and busy running queries, those queries will fail if they try to fetch any pages. But when the tenant is re-attached, things should start working again, without disconnecting the client <-> postgres connections. Without this fix, this reproduced the segfault. Fixes issue #3231	2023-01-05 18:51:47 +02:00
Heikki Linnakangas	c187de1101	Copy error message before it's freed. pageserver_disconnect() call invalidates 'pageserver_conn', including the error message pointer we got from PQerrorMessage(pageserver_conn). Copy the message to a temporary variable before disconnecting, like we do in a few other places. In the passing, clear 'pageserver_conn_wes' variable in a few places where it was free'd. I didn't see any live bug from this, but since pageserver_disconnect() checks if it's NULL, let's not leave it dangling to already-free'd memory.	2023-01-05 18:51:47 +02:00
Kirill Bulatov	8712e1899e	Move initial timeline creation into pytest (#3270 ) For every Python test, we start the storage first, and expect that later, in the test, when we start a compute, it will work without specific timeline and tenant creation or their IDs specified. For that, we have a concept of "default" branch that was created on the control plane level first, but that's not needed at all, given that it's only Python tests that need it: let them create the initial timeline during set-up. Before, control plane started and stopped pageserver for timeline creation, now Python harness runs an extra tenant creation request on test env init. I had to adjust the metrics test, turns out it registered the metrics from the default tenant after an extra pageserver restart. New model does not sent the metrics before the collection time happens, and that was 30s before.	2023-01-05 17:48:27 +02:00
Christian Schwarz	d7f1e30112	remote_timeline_client: more metrics & metrics-related cleanups - Clean up redundant metric removal in TimelineMetrics::drop. RemoteTimelineClientMetrics is responsible for cleaning up REMOTE_OPERATION_TIME andREMOTE_UPLOAD_QUEUE_UNFINISHED_TASKS. - Rename `pageserver_remote_upload_queue_unfinished_tasks` to `pageserver_remote_timeline_client_calls_unfinished`. The new name reflects that the metric is with respect to the entire call to remote timeline client. This includes wait time in the upload queue and hence it's a longer span than what `pageserver_remote_OPERATION_seconds` measures. - Add the `pageserver_remote_timeline_client_calls_started` histogram. See the metric description for why we need it. - Add helper functions `call_begin` etc to `RemoteTimelineClientMetrics` to centralize the logic for updating the metrics above (they relate to each other, see comments in code). - Use these constructs to track ongoing downloads in `pageserver_remote_timeline_client_calls_unfinished` refs https://github.com/neondatabase/neon/issues/2029 fixes https://github.com/neondatabase/neon/issues/3249 closes https://github.com/neondatabase/neon/pull/3250	2023-01-05 11:50:17 +01:00
Christian Schwarz	6a9d1030a6	use RemoteTimelineClient for downloading index part during tenant_attach Before this change, we would not .measure_remote_op for index part downloads. And more generally, it's good to pass not just uploads but also downloads through RemoteTimelineClient, e.g., if we ever want to implement some timeline-scoped policies there. Found this while working on https://github.com/neondatabase/neon/pull/3250 where I add a metric to measure the degree of concurrent downloads. Layer download was missing in a test that I added there.	2023-01-05 11:08:50 +01:00
Heikki Linnakangas	8c6e607327	Refactor send_tarball() (#3259 ) The Basebackup struct is really just a convenient place to carry the various parameters around in send_tarball and its subroutines. Make it internal to the send_tarball function.	2023-01-04 23:03:16 +02:00
Vadim Kharitonov	f436fb2dfb	Fix panics at compute_ctl:monitor	2023-01-04 17:26:42 +01:00
Kirill Bulatov	8932d14d50	Revert "Run Python tests in 8 threads (#3206 )" (#3264 ) This reverts commit `56a4466d0a`. Seems that flackiness increased after this commit, while the time decrease was a couple of seconds. With every regular Python test spawing 1 etcd, 3 safekeepers, 1 pageserver, few CLI commands and post-run cleanup hooks, it might be hard to run many such tests in parallel. We could return to this later, after we consider alternative test structure and/or CI runner structure.	2023-01-04 17:31:51 +02:00
Kirill Bulatov	efad64bc7f	Expect compute shutdown test log error (#3262 ) https://neon-github-public-dev.s3.amazonaws.com/reports/pr-3261/debug/3833043374/index.html#suites/ffbb7f9930a77115316b58ff32b7c719/1f6ebaedc0a113a1/ Spotted a flacky test that appeared after https://github.com/neondatabase/neon/pull/3227 changes	2023-01-04 10:45:11 +00:00
Kirill Bulatov	10dae79c6d	Tone down safekeeper and pageserver walreceiver errors (#3227 ) Closes https://github.com/neondatabase/neon/issues/3114 Adds more typization into errors that appear during protocol messages (`FeMessage`), postgres and walreceiver connections. Socket IO errors are now better detected and logged with lesser (INFO, DEBUG) error level, without traces that they were logged before, when they were wrapped in anyhow context.	2023-01-03 20:42:04 +00:00