test

see
https://github.com/appsmithorg/appsmith/issus/826#issuecomment-703093426
for details

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

@@ -1,37 +1,39 @@
+# Helm chart values for neon-proxy-scram.
+# This is a YAML-formatted file.
+
+image:
+  repository: neondatabase/neon
+
 settings:
-  authBackend: "link"
-  authEndpoint: "https://console.neon.tech/authenticate_proxy_request/"
-  uri: "https://console.neon.tech/psql_session/"
+  authBackend: "console"
+  authEndpoint: "http://console-release.local/management/api/v2"
+  domain: "*.cloud.neon.tech"
   sentryEnvironment: "production"
+  wssPort: 8443
+  metricCollectionEndpoint: "http://console-release.local/billing/api/v1/usage_events"
+  metricCollectionInterval: "10min"
 
-# -- Additional labels for zenith-proxy pods
+# -- Additional labels for neon-proxy pods
 podLabels:
-  zenith_service: proxy
-  zenith_env: production
+  zenith_service: proxy-scram
+  zenith_env: prod
   zenith_region: us-west-2
-  zenith_region_slug: oregon
-
-service:
-  annotations:
-    service.beta.kubernetes.io/aws-load-balancer-type: external
-    service.beta.kubernetes.io/aws-load-balancer-nlb-target-type: ip
-    service.beta.kubernetes.io/aws-load-balancer-scheme: internal
-    external-dns.alpha.kubernetes.io/hostname: proxy-release.local
-  type: LoadBalancer
+  zenith_region_slug: us-west-2
 
 exposedService:
   annotations:
     service.beta.kubernetes.io/aws-load-balancer-type: external
     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: connect.neon.tech,pg.neon.tech
+    external-dns.alpha.kubernetes.io/hostname: neon-proxy-scram-legacy.eta.us-west-2.aws.neon.tech
+  httpsPort: 443
 
-metrics:
-  enabled: true
-  serviceMonitor:
-    enabled: true
-    selector:
-      release: kube-prometheus-stack
+#metrics:
+#  enabled: true
+#  serviceMonitor:
+#    enabled: true
+#    selector:
+#      release: kube-prometheus-stack
 
 extraManifests:
   - apiVersion: operator.victoriametrics.com/v1beta1

.github/workflows/build_and_test.yml

@@ -1,4 +1,4 @@
-name: Test and Deploy
+name: Build and Test
 
 on:
   push:
@@ -54,7 +54,7 @@ jobs:
   check-codestyle-python:
     runs-on: [ self-hosted, gen3, small ]
     container:
-      image: 369495373322.dkr.ecr.eu-central-1.amazonaws.com/cloud:pinned
+      image: 369495373322.dkr.ecr.eu-central-1.amazonaws.com/rust:pinned
       options: --init
 
     steps:
@@ -337,7 +337,7 @@ jobs:
         uses: ./.github/actions/save-coverage-data
 
   benchmarks:
-    runs-on: [ self-hosted, gen3, large ]
+    runs-on: [ self-hosted, gen3, small ]
     container:
       image: 369495373322.dkr.ecr.eu-central-1.amazonaws.com/rust:pinned
       options: --init
@@ -405,7 +405,7 @@ jobs:
           DATABASE_URL="$TEST_RESULT_CONNSTR" poetry run python3 scripts/ingest_regress_test_result.py --revision ${SHA} --reference ${GITHUB_REF} --build-type ${BUILD_TYPE} --ingest suites.json
 
   coverage-report:
-    runs-on: [ self-hosted, gen3, large ]
+    runs-on: [ self-hosted, gen3, small ]
     container:
       image: 369495373322.dkr.ecr.eu-central-1.amazonaws.com/rust:pinned
       options: --init
@@ -481,7 +481,7 @@ jobs:
             }"
 
   trigger-e2e-tests:
-    runs-on: [ self-hosted, gen3, large ]
+    runs-on: [ self-hosted, gen3, small ]
     container:
       image: 369495373322.dkr.ecr.eu-central-1.amazonaws.com/base:pinned
       options: --init
@@ -690,20 +690,37 @@ jobs:
   promote-images:
     runs-on: [ self-hosted, gen3, small ]
     needs: [ tag, test-images, vm-compute-node-image ]
+    container: golang:1.19-bullseye
     if: github.event_name != 'workflow_dispatch'
-    container: amazon/aws-cli
-    strategy:
-      fail-fast: false
-      matrix:
-        name: [ neon, compute-node-v14, vm-compute-node-v14, compute-node-v15, vm-compute-node-v15, compute-tools]
-    env:
-      AWS_DEFAULT_REGION: eu-central-1
 
     steps:
-      - name: Promote image to latest
+      - name: Install Crane & ECR helper
+        if: |
+          (github.ref_name == 'main' || github.ref_name == 'release') &&
+          github.event_name != 'workflow_dispatch'
         run: |
-          export MANIFEST=$(aws ecr batch-get-image --repository-name ${{ matrix.name }} --image-ids imageTag=${{needs.tag.outputs.build-tag}} --query 'images[].imageManifest' --output text)
-          aws ecr put-image --repository-name ${{ matrix.name }} --image-tag latest --image-manifest "$MANIFEST"
+          go install github.com/google/go-containerregistry/cmd/crane@31786c6cbb82d6ec4fb8eb79cd9387905130534e # v0.11.0
+          go install github.com/awslabs/amazon-ecr-credential-helper/ecr-login/cli/docker-credential-ecr-login@69c85dc22db6511932bbf119e1a0cc5c90c69a7f # v0.6.0
+
+      - name: Configure ECR login
+        run: |
+          mkdir /github/home/.docker/
+          echo "{\"credsStore\":\"ecr-login\"}" > /github/home/.docker/config.json
+
+      - name: Add latest tag to images
+        if: |
+          (github.ref_name == 'main' || github.ref_name == 'release') &&
+          github.event_name != 'workflow_dispatch'
+        run: |
+          crane tag 369495373322.dkr.ecr.eu-central-1.amazonaws.com/neon:${{needs.tag.outputs.build-tag}} latest
+          crane tag 369495373322.dkr.ecr.eu-central-1.amazonaws.com/compute-tools:${{needs.tag.outputs.build-tag}} latest
+          crane tag 369495373322.dkr.ecr.eu-central-1.amazonaws.com/compute-node-v14:${{needs.tag.outputs.build-tag}} latest
+          crane tag 369495373322.dkr.ecr.eu-central-1.amazonaws.com/vm-compute-node-v14:${{needs.tag.outputs.build-tag}} latest
+          crane tag 369495373322.dkr.ecr.eu-central-1.amazonaws.com/compute-node-v15:${{needs.tag.outputs.build-tag}} latest
+          crane tag 369495373322.dkr.ecr.eu-central-1.amazonaws.com/vm-compute-node-v15:${{needs.tag.outputs.build-tag}} latest
+
+      - name: Cleanup ECR folder
+        run: rm -rf ~/.ecr
 
   push-docker-hub:
     runs-on: [ self-hosted, dev, x64 ]
@@ -796,112 +813,6 @@ jobs:
       - name: Cleanup ECR folder
         run: rm -rf ~/.ecr
 
-  calculate-deploy-targets:
-    runs-on: [ self-hosted, gen3, small ]
-    if: |
-      github.ref_name == 'release' &&
-      github.event_name != 'workflow_dispatch'
-    outputs:
-      matrix-include: ${{ steps.set-matrix.outputs.include }}
-    steps:
-      - id: set-matrix
-        run: |
-          if [[ "$GITHUB_REF_NAME" == "release" ]]; then
-            PRODUCTION='{"env_name": "production", "proxy_job": "neon-proxy", "proxy_config": "production.proxy", "storage_broker_ns": "neon-storage-broker", "storage_broker_config": "production.neon-storage-broker", "kubeconfig_secret": "PRODUCTION_KUBECONFIG_DATA", "console_api_key_secret": "NEON_PRODUCTION_API_KEY"}'
-            echo "include=[$PRODUCTION]" >> $GITHUB_OUTPUT
-          else
-            echo "GITHUB_REF_NAME (value '$GITHUB_REF_NAME') is not set to 'release'"
-            exit 1
-          fi
-
-  deploy:
-    runs-on: [ self-hosted, gen3, small ]
-    container: 369495373322.dkr.ecr.eu-central-1.amazonaws.com/ansible:pinned
-    # We need both storage **and** compute images for deploy, because control plane picks the compute version based on the storage version.
-    # If it notices a fresh storage it may bump the compute version. And if compute image failed to build it may break things badly
-    needs: [ push-docker-hub, calculate-deploy-targets, tag, regress-tests ]
-    if: |
-      github.ref_name == 'release' &&
-      github.event_name != 'workflow_dispatch'
-    defaults:
-      run:
-        shell: bash
-    strategy:
-      matrix:
-        include: ${{fromJSON(needs.calculate-deploy-targets.outputs.matrix-include)}}
-    environment:
-      name: prod-old
-    steps:
-      - name: Checkout
-        uses: actions/checkout@v3
-        with:
-          submodules: true
-          fetch-depth: 0
-
-      - name: Redeploy
-        run: |
-          export DOCKER_TAG=${{needs.tag.outputs.build-tag}}
-          cd "$(pwd)/.github/ansible"
-
-          if [[ "$GITHUB_REF_NAME" == "main" ]]; then
-            ./get_binaries.sh
-          elif [[ "$GITHUB_REF_NAME" == "release" ]]; then
-            RELEASE=true ./get_binaries.sh
-          else
-            echo "GITHUB_REF_NAME (value '$GITHUB_REF_NAME') is not set to either 'main' or 'release'"
-            exit 1
-          fi
-
-          eval $(ssh-agent)
-          echo "${{ secrets.TELEPORT_SSH_KEY }}"  | tr -d '\n'| base64 --decode >ssh-key
-          echo "${{ secrets.TELEPORT_SSH_CERT }}" | tr -d '\n'| base64 --decode >ssh-key-cert.pub
-          chmod 0600 ssh-key
-          ssh-add ssh-key
-          rm -f ssh-key ssh-key-cert.pub
-          ANSIBLE_CONFIG=./ansible.cfg ansible-galaxy collection install sivel.toiletwater
-          ANSIBLE_CONFIG=./ansible.cfg ansible-playbook deploy.yaml -i ${{ matrix.env_name }}.hosts.yaml -e CONSOLE_API_TOKEN=${{ secrets[matrix.console_api_key_secret] }} -e SENTRY_URL_PAGESERVER=${{ secrets.SENTRY_URL_PAGESERVER }} -e SENTRY_URL_SAFEKEEPER=${{ secrets.SENTRY_URL_SAFEKEEPER }}
-          rm -f neon_install.tar.gz .neon_current_version
-
-  deploy-new:
-    runs-on: [ self-hosted, gen3, small ]
-    container: 369495373322.dkr.ecr.eu-central-1.amazonaws.com/ansible:pinned
-    # We need both storage **and** compute images for deploy, because control plane picks the compute version based on the storage version.
-    # If it notices a fresh storage it may bump the compute version. And if compute image failed to build it may break things badly
-    needs: [ push-docker-hub, tag, regress-tests ]
-    if: |
-      (github.ref_name == 'main') &&
-      github.event_name != 'workflow_dispatch'
-    defaults:
-      run:
-        shell: bash
-    strategy:
-      matrix:
-        target_region: [ eu-west-1, us-east-2 ]
-    environment:
-      name: dev-${{ matrix.target_region }}
-    steps:
-      - name: Checkout
-        uses: actions/checkout@v3
-        with:
-          submodules: true
-          fetch-depth: 0
-
-      - name: Redeploy
-        run: |
-          export DOCKER_TAG=${{needs.tag.outputs.build-tag}}
-          cd "$(pwd)/.github/ansible"
-          if [[ "$GITHUB_REF_NAME" == "main" ]]; then
-            ./get_binaries.sh
-          elif [[ "$GITHUB_REF_NAME" == "release" ]]; then
-            RELEASE=true ./get_binaries.sh
-          else
-            echo "GITHUB_REF_NAME (value '$GITHUB_REF_NAME') is not set to either 'main' or 'release'"
-            exit 1
-          fi
-          ansible-galaxy collection install sivel.toiletwater
-          ansible-playbook deploy.yaml -i staging.${{ matrix.target_region }}.hosts.yaml -e @ssm_config -e CONSOLE_API_TOKEN=${{ secrets.NEON_STAGING_API_KEY }} -e SENTRY_URL_PAGESERVER=${{ secrets.SENTRY_URL_PAGESERVER }} -e SENTRY_URL_SAFEKEEPER=${{ secrets.SENTRY_URL_SAFEKEEPER }}
-          rm -f neon_install.tar.gz .neon_current_version
-
   deploy-pr-test-new:
     runs-on: [ self-hosted, gen3, small ]
     container: 369495373322.dkr.ecr.eu-central-1.amazonaws.com/ansible:pinned
@@ -935,321 +846,40 @@ jobs:
           ansible-playbook deploy.yaml -i staging.${{ matrix.target_region }}.hosts.yaml -e @ssm_config -e CONSOLE_API_TOKEN=${{ secrets.NEON_STAGING_API_KEY }} -e SENTRY_URL_PAGESERVER=${{ secrets.SENTRY_URL_PAGESERVER }} -e SENTRY_URL_SAFEKEEPER=${{ secrets.SENTRY_URL_SAFEKEEPER }}
           rm -f neon_install.tar.gz .neon_current_version
 
-  deploy-prod-new:
-    runs-on: prod
-    container: 093970136003.dkr.ecr.eu-central-1.amazonaws.com/ansible:latest
-    # We need both storage **and** compute images for deploy, because control plane picks the compute version based on the storage version.
-    # If it notices a fresh storage it may bump the compute version. And if compute image failed to build it may break things badly
+      - name: Cleanup ansible folder
+        run: rm -rf ~/.ansible
+
+  deploy:
+    runs-on: [ self-hosted, gen3, small ]
+    container: 369495373322.dkr.ecr.eu-central-1.amazonaws.com/ansible:latest
     needs: [ push-docker-hub, tag, regress-tests ]
-    if: |
-      (github.ref_name == 'release') &&
-      github.event_name != 'workflow_dispatch'
-    defaults:
-      run:
-        shell: bash
-    strategy:
-      matrix:
-        target_region: [ us-east-2, us-west-2, eu-central-1, ap-southeast-1 ]
-    environment:
-      name: prod-${{ matrix.target_region }}
+    if: ( github.ref_name == 'main' || github.ref_name == 'release' ) && github.event_name != 'workflow_dispatch'
     steps:
       - name: Checkout
         uses: actions/checkout@v3
         with:
-          submodules: true
+          submodules: false
           fetch-depth: 0
 
-      - name: Redeploy
+      - name: Trigger deploy workflow
+        env:
+          GH_TOKEN: ${{ github.token }}
         run: |
-          export DOCKER_TAG=${{needs.tag.outputs.build-tag}}
-          cd "$(pwd)/.github/ansible"
-
           if [[ "$GITHUB_REF_NAME" == "main" ]]; then
-            ./get_binaries.sh
+            gh workflow run deploy-dev.yml --ref main -f branch=${{ github.sha }} -f dockerTag=${{needs.tag.outputs.build-tag}}
           elif [[ "$GITHUB_REF_NAME" == "release" ]]; then
-            RELEASE=true ./get_binaries.sh
+            gh workflow run deploy-prod.yml --ref release -f branch=${{ github.sha }} -f dockerTag=${{needs.tag.outputs.build-tag}} -f disclamerAcknowledged=true
           else
             echo "GITHUB_REF_NAME (value '$GITHUB_REF_NAME') is not set to either 'main' or 'release'"
             exit 1
           fi
 
-          ansible-galaxy collection install sivel.toiletwater
-          ansible-playbook deploy.yaml -i prod.${{ matrix.target_region }}.hosts.yaml -e @ssm_config -e CONSOLE_API_TOKEN=${{ secrets.NEON_PRODUCTION_API_KEY }} -e SENTRY_URL_PAGESERVER=${{ secrets.SENTRY_URL_PAGESERVER }} -e SENTRY_URL_SAFEKEEPER=${{ secrets.SENTRY_URL_SAFEKEEPER }}
-          rm -f neon_install.tar.gz .neon_current_version
-
-  deploy-proxy:
-    runs-on: [ self-hosted, gen3, small ]
-    container: 369495373322.dkr.ecr.eu-central-1.amazonaws.com/base:pinned
-    # Compute image isn't strictly required for proxy deploy, but let's still wait for it to run all deploy jobs consistently.
-    needs: [ push-docker-hub, calculate-deploy-targets, tag, regress-tests ]
-    if: |
-      github.ref_name == 'release' &&
-      github.event_name != 'workflow_dispatch'
-    defaults:
-      run:
-        shell: bash
-    strategy:
-      matrix:
-        include: ${{fromJSON(needs.calculate-deploy-targets.outputs.matrix-include)}}
-    environment:
-      name: prod-old
-    env:
-      KUBECONFIG: .kubeconfig
-    steps:
-      - name: Checkout
-        uses: actions/checkout@v3
-        with:
-          submodules: true
-          fetch-depth: 0
-
-      - name: Add curl
-        run: apt update && apt install curl -y
-
-      - name: Store kubeconfig file
-        run: |
-          echo "${{ secrets[matrix.kubeconfig_secret] }}" | base64 --decode > ${KUBECONFIG}
-          chmod 0600 ${KUBECONFIG}
-
-      - name: Setup helm v3
-        run: |
-          curl -s https://raw.githubusercontent.com/helm/helm/main/scripts/get-helm-3 | bash
-          helm repo add neondatabase https://neondatabase.github.io/helm-charts
-
-      - name: Re-deploy proxy
-        run: |
-          DOCKER_TAG=${{needs.tag.outputs.build-tag}}
-          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
-    runs-on: [ self-hosted, gen3, small ]
-    container: 369495373322.dkr.ecr.eu-central-1.amazonaws.com/base:pinned
-    # Compute image isn't strictly required for proxy deploy, but let's still wait for it to run all deploy jobs consistently.
-    needs: [ push-docker-hub, calculate-deploy-targets, tag, regress-tests ]
-    if: |
-      github.ref_name == 'release' &&
-      github.event_name != 'workflow_dispatch'
-    defaults:
-      run:
-        shell: bash
-    strategy:
-      matrix:
-        include: ${{fromJSON(needs.calculate-deploy-targets.outputs.matrix-include)}}
-    environment:
-      name: prod-old
-    env:
-      KUBECONFIG: .kubeconfig
-    steps:
-      - name: Checkout
-        uses: actions/checkout@v3
-        with:
-          submodules: true
-          fetch-depth: 0
-
-      - name: Add curl
-        run: apt update && apt install curl -y
-
-      - name: Store kubeconfig file
-        run: |
-          echo "${{ secrets[matrix.kubeconfig_secret] }}" | base64 --decode > ${KUBECONFIG}
-          chmod 0600 ${KUBECONFIG}
-
-      - name: Setup helm v3
-        run: |
-          curl -s https://raw.githubusercontent.com/helm/helm/main/scripts/get-helm-3 | bash
-          helm repo add neondatabase https://neondatabase.github.io/helm-charts
-
-      - name: Deploy storage-broker
-        run:
-          helm upgrade neon-storage-broker neondatabase/neon-storage-broker --namespace ${{ matrix.storage_broker_ns }} --create-namespace --install --atomic -f .github/helm-values/${{ matrix.storage_broker_config }}.yaml --set image.tag=${{ needs.tag.outputs.build-tag }} --set settings.sentryUrl=${{ secrets.SENTRY_URL_BROKER }} --wait --timeout 5m0s
-
-  deploy-proxy-new:
-    runs-on: [ self-hosted, gen3, small ]
-    container: 369495373322.dkr.ecr.eu-central-1.amazonaws.com/ansible:pinned
-    # Compute image isn't strictly required for proxy deploy, but let's still wait for it to run all deploy jobs consistently.
-    needs: [ push-docker-hub, tag, regress-tests ]
-    if: |
-      (github.ref_name == 'main') &&
-      github.event_name != 'workflow_dispatch'
-    defaults:
-      run:
-        shell: bash
-    strategy:
-      matrix:
-        include:
-          - target_region:  us-east-2
-            target_cluster: dev-us-east-2-beta
-            deploy_link_proxy: true
-            deploy_legacy_scram_proxy: true
-          - target_region:  eu-west-1
-            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
-        with:
-          submodules: true
-          fetch-depth: 0
-
-      - name: Configure environment
-        run: |
-          helm repo add neondatabase https://neondatabase.github.io/helm-charts
-          aws --region ${{ matrix.target_region }} eks update-kubeconfig --name  ${{ matrix.target_cluster }}
-
-      - 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 --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 --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 --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, gen3, small ]
-    container: 369495373322.dkr.ecr.eu-central-1.amazonaws.com/ansible:pinned
-    # Compute image isn't strictly required for proxy deploy, but let's still wait for it to run all deploy jobs consistently.
-    needs: [ push-docker-hub, tag, regress-tests ]
-    if: |
-      (github.ref_name == 'main') &&
-      github.event_name != 'workflow_dispatch'
-    defaults:
-      run:
-        shell: bash
-    strategy:
-      matrix:
-        include:
-          - target_region:  us-east-2
-            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
-        with:
-          submodules: true
-          fetch-depth: 0
-
-      - name: Configure environment
-        run: |
-          helm repo add neondatabase https://neondatabase.github.io/helm-charts
-          aws --region ${{ matrix.target_region }} eks update-kubeconfig --name  ${{ matrix.target_cluster }}
-
-      - name: Deploy storage-broker
-        run:
-          helm upgrade neon-storage-broker-lb neondatabase/neon-storage-broker --namespace neon-storage-broker-lb --create-namespace --install --atomic -f .github/helm-values/${{ matrix.target_cluster }}.neon-storage-broker.yaml --set image.tag=${{ needs.tag.outputs.build-tag }} --set settings.sentryUrl=${{ secrets.SENTRY_URL_BROKER }} --wait --timeout 5m0s
-
-  deploy-proxy-prod-new:
-    runs-on: prod
-    container: 093970136003.dkr.ecr.eu-central-1.amazonaws.com/ansible:latest
-    # Compute image isn't strictly required for proxy deploy, but let's still wait for it to run all deploy jobs consistently.
-    needs: [ push-docker-hub, tag, regress-tests ]
-    if: |
-      (github.ref_name == 'release') &&
-      github.event_name != 'workflow_dispatch'
-    defaults:
-      run:
-        shell: bash
-    strategy:
-      matrix:
-        include:
-          - target_region:  us-east-2
-            target_cluster: prod-us-east-2-delta
-            deploy_link_proxy: true
-          - target_region:  us-west-2
-            target_cluster: prod-us-west-2-eta
-            deploy_link_proxy: false
-          - target_region: eu-central-1
-            target_cluster: prod-eu-central-1-gamma
-            deploy_link_proxy: false
-          - target_region: ap-southeast-1
-            target_cluster: prod-ap-southeast-1-epsilon
-            deploy_link_proxy: false
-    environment:
-      name: prod-${{ matrix.target_region }}
-    steps:
-      - name: Checkout
-        uses: actions/checkout@v3
-        with:
-          submodules: true
-          fetch-depth: 0
-
-      - name: Configure environment
-        run: |
-          helm repo add neondatabase https://neondatabase.github.io/helm-charts
-          aws --region ${{ matrix.target_region }} eks update-kubeconfig --name  ${{ matrix.target_cluster }}
-
-      - 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 --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 --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
-
-  deploy-storage-broker-prod-new:
-    runs-on: prod
-    container: 093970136003.dkr.ecr.eu-central-1.amazonaws.com/ansible:latest
-    # Compute image isn't strictly required for proxy deploy, but let's still wait for it to run all deploy jobs consistently.
-    needs: [ push-docker-hub, tag, regress-tests ]
-    if: |
-      (github.ref_name == 'release') &&
-      github.event_name != 'workflow_dispatch'
-    defaults:
-      run:
-        shell: bash
-    strategy:
-      matrix:
-        include:
-          - target_region:  us-east-2
-            target_cluster: prod-us-east-2-delta
-          - target_region:  us-west-2
-            target_cluster: prod-us-west-2-eta
-          - target_region: eu-central-1
-            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
-        with:
-          submodules: true
-          fetch-depth: 0
-
-      - name: Configure environment
-        run: |
-          helm repo add neondatabase https://neondatabase.github.io/helm-charts
-          aws --region ${{ matrix.target_region }} eks update-kubeconfig --name  ${{ matrix.target_cluster }}
-
-      - name: Deploy storage-broker
-        run:
-          helm upgrade neon-storage-broker-lb neondatabase/neon-storage-broker --namespace neon-storage-broker-lb --create-namespace --install --atomic -f .github/helm-values/${{ matrix.target_cluster }}.neon-storage-broker.yaml --set image.tag=${{ needs.tag.outputs.build-tag }} --set settings.sentryUrl=${{ secrets.SENTRY_URL_BROKER }} --wait --timeout 5m0s
-
   promote-compatibility-data:
     runs-on: [ self-hosted, gen3, small ]
     container:
       image: 369495373322.dkr.ecr.eu-central-1.amazonaws.com/rust:pinned
       options: --init
-    needs: [ deploy, deploy-proxy ]
+    needs: [ push-docker-hub, tag, regress-tests ]
     if: github.ref_name == 'release' && github.event_name != 'workflow_dispatch'
     steps:
       - name: Promote compatibility snapshot for the release

.github/workflows/deploy-dev.yml

@@ -0,0 +1,179 @@
+name: Neon Deploy dev
+
+on:
+  workflow_dispatch:
+    inputs:
+      dockerTag:
+        description: 'Docker tag to deploy'
+        required: true
+        type: string
+      branch:
+        description: 'Branch or commit used for deploy scripts and configs'
+        required: true
+        type: string
+        default: 'main'
+      deployStorage:
+        description: 'Deploy storage'
+        required: true
+        type: boolean
+        default: true
+      deployProxy:
+        description: 'Deploy proxy'
+        required: true
+        type: boolean
+        default: true
+      deployStorageBroker:
+        description: 'Deploy storage-broker'
+        required: true
+        type: boolean
+        default: true
+
+env:
+  AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_DEV }}
+  AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_KEY_DEV }}
+
+concurrency:
+  group: deploy-dev
+  cancel-in-progress: false
+
+jobs:
+  deploy-storage-new:
+    runs-on: [ self-hosted, gen3, small ]
+    container:
+      image: 369495373322.dkr.ecr.eu-central-1.amazonaws.com/ansible:pinned
+      options: --user root --privileged
+    if: inputs.deployStorage
+    defaults:
+      run:
+        shell: bash
+    strategy:
+      matrix:
+        target_region: [ eu-west-1, us-east-2 ]
+    environment:
+      name: dev-${{ matrix.target_region }}
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v3
+        with:
+          submodules: true
+          fetch-depth: 0
+          ref: ${{ inputs.branch }}
+
+      - name: Redeploy
+        run: |
+          export DOCKER_TAG=${{ inputs.dockerTag }}
+          cd "$(pwd)/.github/ansible"
+
+          ./get_binaries.sh
+
+          ansible-galaxy collection install sivel.toiletwater
+          ansible-playbook deploy.yaml -i staging.${{ matrix.target_region }}.hosts.yaml -e @ssm_config -e CONSOLE_API_TOKEN=${{ secrets.NEON_STAGING_API_KEY }} -e SENTRY_URL_PAGESERVER=${{ secrets.SENTRY_URL_PAGESERVER }} -e SENTRY_URL_SAFEKEEPER=${{ secrets.SENTRY_URL_SAFEKEEPER }}
+          rm -f neon_install.tar.gz .neon_current_version
+
+      - name: Cleanup ansible folder
+        run: rm -rf ~/.ansible
+
+  deploy-proxy-new:
+    runs-on: [ self-hosted, gen3, small ]
+    container: 369495373322.dkr.ecr.eu-central-1.amazonaws.com/ansible:pinned
+    if: inputs.deployProxy
+    defaults:
+      run:
+        shell: bash
+    strategy:
+      matrix:
+        include:
+          - target_region:  us-east-2
+            target_cluster: dev-us-east-2-beta
+            deploy_link_proxy: true
+            deploy_legacy_scram_proxy: true
+          - target_region:  eu-west-1
+            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
+        with:
+          submodules: true
+          fetch-depth: 0
+          ref: ${{ inputs.branch }}
+  
+      - name: Configure AWS Credentials
+        uses: aws-actions/configure-aws-credentials@v1-node16
+        with:
+          role-to-assume: arn:aws:iam::369495373322:role/github-runner
+          aws-region: eu-central-1
+          role-skip-session-tagging: true
+          role-duration-seconds: 1800
+  
+      - name: Configure environment
+        run: |
+          helm repo add neondatabase https://neondatabase.github.io/helm-charts
+          aws --region ${{ matrix.target_region }} eks update-kubeconfig --name  ${{ matrix.target_cluster }}
+  
+      - name: Re-deploy scram proxy
+        run: |
+          DOCKER_TAG=${{ inputs.dockerTag }}
+          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=${{ inputs.dockerTag }}
+          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=${{ inputs.dockerTag }}
+          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
+  
+      - name: Cleanup helm folder
+        run: rm -rf ~/.cache
+  
+  deploy-storage-broker-new:
+    runs-on: [ self-hosted, gen3, small ]
+    container: 369495373322.dkr.ecr.eu-central-1.amazonaws.com/ansible:pinned
+    if: inputs.deployStorageBroker
+    defaults:
+      run:
+        shell: bash
+    strategy:
+      matrix:
+        include:
+          - target_region:  us-east-2
+            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
+        with:
+          submodules: true
+          fetch-depth: 0
+          ref: ${{ inputs.branch }}
+  
+      - name: Configure AWS Credentials
+        uses: aws-actions/configure-aws-credentials@v1-node16
+        with:
+          role-to-assume: arn:aws:iam::369495373322:role/github-runner
+          aws-region: eu-central-1
+          role-skip-session-tagging: true
+          role-duration-seconds: 1800
+  
+      - name: Configure environment
+        run: |
+          helm repo add neondatabase https://neondatabase.github.io/helm-charts
+          aws --region ${{ matrix.target_region }} eks update-kubeconfig --name  ${{ matrix.target_cluster }}
+  
+      - name: Deploy storage-broker
+        run:
+          helm upgrade neon-storage-broker-lb neondatabase/neon-storage-broker --namespace neon-storage-broker-lb --create-namespace --install --atomic -f .github/helm-values/${{ matrix.target_cluster }}.neon-storage-broker.yaml --set image.tag=${{ inputs.dockerTag }} --set settings.sentryUrl=${{ secrets.SENTRY_URL_BROKER }} --wait --timeout 5m0s
+  
+      - name: Cleanup helm folder
+        run: rm -rf ~/.cache

.github/workflows/deploy-prod.yml

@@ -0,0 +1,275 @@
+name: Neon Deploy prod
+
+on:
+  workflow_dispatch:
+    inputs:
+      dockerTag:
+        description: 'Docker tag to deploy'
+        required: true
+        type: string
+      branch:
+        description: 'Branch or commit used for deploy scripts and configs'
+        required: true
+        type: string
+        default: 'release'
+      deployStorage:
+        description: 'Deploy storage'
+        required: true
+        type: boolean
+        default: true
+      deployProxy:
+        description: 'Deploy proxy'
+        required: true
+        type: boolean
+        default: true
+      deployStorageBroker:
+        description: 'Deploy storage-broker'
+        required: true
+        type: boolean
+        default: true
+      disclamerAcknowledged:
+        description: 'I confirm that there is an emergency and I can not use regular release workflow'
+        required: true
+        type: boolean
+        default: false
+
+concurrency:
+  group: deploy-prod
+  cancel-in-progress: false
+
+jobs:
+  deploy-prod-new:
+    runs-on: prod
+    container: 093970136003.dkr.ecr.eu-central-1.amazonaws.com/ansible:latest
+    if: inputs.deployStorage && inputs.disclamerAcknowledged
+    defaults:
+      run:
+        shell: bash
+    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
+        with:
+          submodules: true
+          fetch-depth: 0
+          ref: ${{ inputs.branch }}
+
+      - name: Redeploy
+        run: |
+          export DOCKER_TAG=${{ inputs.dockerTag }}
+          cd "$(pwd)/.github/ansible"
+
+          ./get_binaries.sh
+
+          ansible-galaxy collection install sivel.toiletwater
+          ansible-playbook deploy.yaml -i prod.${{ matrix.target_region }}.hosts.yaml -e @ssm_config -e CONSOLE_API_TOKEN=${{ secrets.NEON_PRODUCTION_API_KEY }} -e SENTRY_URL_PAGESERVER=${{ secrets.SENTRY_URL_PAGESERVER }} -e SENTRY_URL_SAFEKEEPER=${{ secrets.SENTRY_URL_SAFEKEEPER }}
+          rm -f neon_install.tar.gz .neon_current_version
+
+  deploy-proxy-prod-new:
+    runs-on: prod
+    container: 093970136003.dkr.ecr.eu-central-1.amazonaws.com/ansible:latest
+    if: inputs.deployProxy && inputs.disclamerAcknowledged
+    defaults:
+      run:
+        shell: bash
+    strategy:
+      matrix:
+        include:
+          - target_region:  us-east-2
+            target_cluster: prod-us-east-2-delta
+            deploy_link_proxy: true
+            deploy_legacy_scram_proxy: false
+          - target_region:  us-west-2
+            target_cluster: prod-us-west-2-eta
+            deploy_link_proxy: false
+            deploy_legacy_scram_proxy: true
+          - target_region: eu-central-1
+            target_cluster: prod-eu-central-1-gamma
+            deploy_link_proxy: false
+            deploy_legacy_scram_proxy: false
+          - target_region: ap-southeast-1
+            target_cluster: prod-ap-southeast-1-epsilon
+            deploy_link_proxy: false
+            deploy_legacy_scram_proxy: false
+    environment:
+      name: prod-${{ matrix.target_region }}
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v3
+        with:
+          submodules: true
+          fetch-depth: 0
+          ref: ${{ inputs.branch }}
+
+      - name: Configure environment
+        run: |
+          helm repo add neondatabase https://neondatabase.github.io/helm-charts
+          aws --region ${{ matrix.target_region }} eks update-kubeconfig --name  ${{ matrix.target_cluster }}
+
+      - name: Re-deploy scram proxy
+        run: |
+          DOCKER_TAG=${{ inputs.dockerTag }}
+          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=${{ inputs.dockerTag }}
+          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=${{ inputs.dockerTag }}
+          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-prod-new:
+    runs-on: prod
+    container: 093970136003.dkr.ecr.eu-central-1.amazonaws.com/ansible:latest
+    if: inputs.deployStorageBroker && inputs.disclamerAcknowledged
+    defaults:
+      run:
+        shell: bash
+    strategy:
+      matrix:
+        include:
+          - target_region:  us-east-2
+            target_cluster: prod-us-east-2-delta
+          - target_region:  us-west-2
+            target_cluster: prod-us-west-2-eta
+          - target_region: eu-central-1
+            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
+        with:
+          submodules: true
+          fetch-depth: 0
+          ref: ${{ inputs.branch }}
+
+      - name: Configure environment
+        run: |
+          helm repo add neondatabase https://neondatabase.github.io/helm-charts
+          aws --region ${{ matrix.target_region }} eks update-kubeconfig --name  ${{ matrix.target_cluster }}
+
+      - name: Deploy storage-broker
+        run:
+          helm upgrade neon-storage-broker-lb neondatabase/neon-storage-broker --namespace neon-storage-broker-lb --create-namespace --install --atomic -f .github/helm-values/${{ matrix.target_cluster }}.neon-storage-broker.yaml --set image.tag=${{ inputs.dockerTag }} --set settings.sentryUrl=${{ secrets.SENTRY_URL_BROKER }} --wait --timeout 5m0s
+
+  # Deploy to old account below          
+
+  deploy:
+    runs-on: prod
+    container: 093970136003.dkr.ecr.eu-central-1.amazonaws.com/ansible:latest
+    if: inputs.deployStorage && inputs.disclamerAcknowledged
+    defaults:
+      run:
+        shell: bash
+    environment:
+      name: prod-old
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v3
+        with:
+          submodules: true
+          fetch-depth: 0
+          ref: ${{ inputs.branch }}
+
+      - name: Redeploy
+        run: |
+          export DOCKER_TAG=${{ inputs.dockerTag }}
+          cd "$(pwd)/.github/ansible"
+
+          ./get_binaries.sh
+
+          eval $(ssh-agent)
+          echo "${{ secrets.TELEPORT_SSH_KEY }}"  | tr -d '\n'| base64 --decode >ssh-key
+          echo "${{ secrets.TELEPORT_SSH_CERT }}" | tr -d '\n'| base64 --decode >ssh-key-cert.pub
+          chmod 0600 ssh-key
+          ssh-add ssh-key
+          rm -f ssh-key ssh-key-cert.pub
+          ANSIBLE_CONFIG=./ansible.cfg ansible-galaxy collection install sivel.toiletwater
+          ANSIBLE_CONFIG=./ansible.cfg ansible-playbook deploy.yaml -i production.hosts.yaml -e CONSOLE_API_TOKEN=${{ secrets.NEON_PRODUCTION_API_KEY }} -e SENTRY_URL_PAGESERVER=${{ secrets.SENTRY_URL_PAGESERVER }} -e SENTRY_URL_SAFEKEEPER=${{ secrets.SENTRY_URL_SAFEKEEPER }}
+          rm -f neon_install.tar.gz .neon_current_version
+
+      # Cleanup script fails otherwise - rm: cannot remove '/nvme/actions-runner/_work/_temp/_github_home/.ansible/collections': Permission denied
+      - name: Cleanup ansible folder
+        run: rm -rf ~/.ansible
+
+  deploy-proxy:
+    runs-on: [ self-hosted, gen3, small ]
+    container: 369495373322.dkr.ecr.eu-central-1.amazonaws.com/ansible:pinned
+    if: inputs.deployProxy && inputs.disclamerAcknowledged
+    defaults:
+      run:
+        shell: bash
+    environment:
+      name: prod-old
+    env:
+      KUBECONFIG: .kubeconfig
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v3
+        with:
+          submodules: true
+          fetch-depth: 0
+          ref: ${{ inputs.branch }}
+
+      - name: Store kubeconfig file
+        run: |
+          echo "${{ secrets.PRODUCTION_KUBECONFIG_DATA }}" | base64 --decode > ${KUBECONFIG}
+          chmod 0600 ${KUBECONFIG}
+
+      - name: Add neon helm chart
+        run: helm repo add neondatabase https://neondatabase.github.io/helm-charts
+
+      - name: Re-deploy proxy
+        run: |
+          DOCKER_TAG=${{ inputs.dockerTag }}
+          helm upgrade neon-proxy-scram neondatabase/neon-proxy --namespace neon-proxy --install --atomic -f .github/helm-values/production.proxy-scram.yaml --set image.tag=${DOCKER_TAG} --set settings.sentryUrl=${{ secrets.SENTRY_URL_PROXY }} --wait --timeout 15m0s
+
+      - name: Cleanup helm folder
+        run: rm -rf ~/.cache
+
+  deploy-storage-broker:
+    name: deploy storage broker on old staging and old prod
+    runs-on: [ self-hosted, gen3, small ]
+    container: 369495373322.dkr.ecr.eu-central-1.amazonaws.com/ansible:pinned
+    if: inputs.deployStorageBroker && inputs.disclamerAcknowledged
+    defaults:
+      run:
+        shell: bash
+    environment:
+      name: prod-old
+    env:
+      KUBECONFIG: .kubeconfig
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v3
+        with:
+          submodules: true
+          fetch-depth: 0
+          ref: ${{ inputs.branch }}
+
+      - name: Store kubeconfig file
+        run: |
+          echo "${{ secrets.PRODUCTION_KUBECONFIG_DATA }}" | base64 --decode > ${KUBECONFIG}
+          chmod 0600 ${KUBECONFIG}
+
+      - name: Add neon helm chart
+        run: helm repo add neondatabase https://neondatabase.github.io/helm-charts
+
+      - name: Deploy storage-broker
+        run:
+          helm upgrade neon-storage-broker neondatabase/neon-storage-broker --namespace neon-storage-broker --create-namespace --install --atomic -f .github/helm-values/production.neon-storage-broker.yaml --set image.tag=${{ inputs.dockerTag }} --set settings.sentryUrl=${{ secrets.SENTRY_URL_BROKER }} --wait --timeout 5m0s
+
+      - name: Cleanup helm folder
+        run: rm -rf ~/.cache

.github/workflows/release.yml

@@ -0,0 +1,33 @@
+name: Create Release Branch
+
+on:
+  schedule:
+    - cron: '0 10 * * 2'
+
+jobs:
+  create_release_branch:
+    runs-on: [ubuntu-latest]
+
+    steps:
+    - name: Check out code
+      uses: actions/checkout@v3
+      with:
+        ref: main
+
+    - name: Get current date
+      id: date
+      run: echo "date=$(date +'%Y-%m-%d')" >> $GITHUB_OUTPUT
+
+    - name: Create release branch
+      run: git checkout -b releases/${{ steps.date.outputs.date }}
+
+    - name: Push new branch
+      run: git push origin releases/${{ steps.date.outputs.date }}
+
+    - name: Create pull request into release
+      uses: thomaseizinger/create-pull-request@e3972219c86a56550fb70708d96800d8e24ba862 # 1.3.0
+      with:
+        GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+        head: releases/${{ steps.date.outputs.date }}
+        base: release
+        title: Release ${{ steps.date.outputs.date }}

Cargo.lock

@@ -17,6 +17,17 @@ version = "1.0.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "f26201604c87b1e01bd3d98f8d5d9a8fcbb815e8cedb41ffccbeb4bf593a35fe"
 
+[[package]]
+name = "ahash"
+version = "0.7.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fcb51a0695d8f838b1ee009b3fbf66bda078cd64590202a864a8f3e8c4315c47"
+dependencies = [
+ "getrandom",
+ "once_cell",
+ "version_check",
+]
+
 [[package]]
 name = "ahash"
 version = "0.8.2"
@@ -1515,6 +1526,9 @@ name = "hashbrown"
 version = "0.12.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "8a9ee70c43aaf417c914396645a0fa852624801b24ebb7ae78fe8272889ac888"
+dependencies = [
+ "ahash 0.7.6",
+]
 
 [[package]]
 name = "hashbrown"
@@ -1522,7 +1536,16 @@ version = "0.13.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "43a3c133739dddd0d2990f9a4bdf8eb4b21ef50e4851ca85ab661199821d510e"
 dependencies = [
- "ahash",
+ "ahash 0.8.2",
+]
+
+[[package]]
+name = "hashlink"
+version = "0.8.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "69fe1fcf8b4278d860ad0548329f892a3631fb63f82574df68275f34cdbe0ffa"
+dependencies = [
+ "hashbrown 0.12.3",
 ]
 
 [[package]]
@@ -2764,6 +2787,7 @@ dependencies = [
  "futures",
  "git-version",
  "hashbrown 0.13.2",
+ "hashlink",
  "hex",
  "hmac",
  "hostname",
@@ -4337,6 +4361,7 @@ dependencies = [
  "tokio-rustls",
  "tracing",
  "tracing-subscriber",
+ "url",
  "workspace_hack",
 ]
 
@@ -4623,6 +4648,7 @@ dependencies = [
  "futures-executor",
  "futures-task",
  "futures-util",
+ "hashbrown 0.12.3",
  "indexmap",
  "itertools",
  "libc",

Cargo.toml

@@ -44,6 +44,7 @@ futures-core = "0.3"
 futures-util = "0.3"
 git-version = "0.3"
 hashbrown = "0.13"
+hashlink = "0.8.1"
 hex = "0.4"
 hex-literal = "0.3"
 hmac = "0.12.1"

Dockerfile.compute-node

@@ -44,20 +44,30 @@ RUN cd postgres && \
 FROM build-deps AS postgis-build
 COPY --from=pg-build /usr/local/pgsql/ /usr/local/pgsql/
 RUN apt update && \
-    apt install -y gdal-bin libgdal-dev libprotobuf-c-dev protobuf-c-compiler xsltproc
+    apt install -y cmake gdal-bin libboost-dev libboost-thread-dev libboost-filesystem-dev \
+    libboost-system-dev libboost-iostreams-dev libboost-program-options-dev libboost-timer-dev \
+    libcgal-dev libgdal-dev libgmp-dev libmpfr-dev libopenscenegraph-dev libprotobuf-c-dev \
+    protobuf-c-compiler xsltproc
+
+RUN wget https://gitlab.com/Oslandia/SFCGAL/-/archive/v1.3.10/SFCGAL-v1.3.10.tar.gz && \
+    tar zxvf SFCGAL-v1.3.10.tar.gz && \
+    cd SFCGAL-v1.3.10 && cmake . && make -j $(getconf _NPROCESSORS_ONLN) && \
+    DESTDIR=/sfcgal make install -j $(getconf _NPROCESSORS_ONLN) && \
+    make clean && cp -R /sfcgal/* /
 
 RUN wget https://download.osgeo.org/postgis/source/postgis-3.3.1.tar.gz && \
     tar xvzf postgis-3.3.1.tar.gz && \
     cd postgis-3.3.1 && \
     ./autogen.sh && \
     export PATH="/usr/local/pgsql/bin:$PATH" && \
-    ./configure && \
+    ./configure --with-sfcgal=/usr/local/bin/sfcgal-config && \
     make -j $(getconf _NPROCESSORS_ONLN) install && \
     cd extensions/postgis && \
     make clean && \
     make -j $(getconf _NPROCESSORS_ONLN) install && \
     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_sfcgal.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/address_standardizer.control && \
@@ -162,6 +172,7 @@ RUN wget https://github.com/df7cb/postgresql-unit/archive/refs/tags/7.7.tar.gz &
 #########################################################################################
 FROM build-deps AS neon-pg-ext-build
 COPY --from=postgis-build /usr/local/pgsql/ /usr/local/pgsql/
+COPY --from=postgis-build /sfcgal/* /
 COPY --from=plv8-build /usr/local/pgsql/ /usr/local/pgsql/
 COPY --from=h3-pg-build /usr/local/pgsql/ /usr/local/pgsql/
 COPY --from=h3-pg-build /h3/usr /
@@ -224,7 +235,7 @@ COPY --from=compute-tools --chown=postgres /home/nonroot/target/release-line-deb
 # Install:
 # libreadline8 for psql
 # libossp-uuid16 for extension ossp-uuid
-# libgeos, libgdal, libproj and libprotobuf-c1 for PostGIS
+# libgeos, libgdal, libsfcgal1, libproj and libprotobuf-c1 for PostGIS
 RUN apt update &&  \
     apt install --no-install-recommends -y \
         libreadline8 \
@@ -233,6 +244,7 @@ RUN apt update &&  \
         libgdal28 \
         libproj19 \
         libprotobuf-c1 \
+        libsfcgal1 \
         gdb && \
     rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/*
 

README.md

@@ -108,7 +108,7 @@ make -j`sysctl -n hw.logicalcpu`
 To run the `psql` client, install the `postgresql-client` package or modify `PATH` and `LD_LIBRARY_PATH` to include `pg_install/bin` and `pg_install/lib`, respectively.
 
 To run the integration tests or Python scripts (not required to use the code), install
-Python (3.9 or higher), and install python3 packages using `./scripts/pysync` (requires [poetry](https://python-poetry.org/)) in the project directory.
+Python (3.9 or higher), and install python3 packages using `./scripts/pysync` (requires [poetry>=1.3](https://python-poetry.org/)) in the project directory.
 
 
 #### Running neon database
@@ -254,3 +254,6 @@ To get more familiar with this aspect, refer to:
 - Read `CONTRIBUTING.md` to learn about project code style and practices.
 - To get familiar with a source tree layout, use [/docs/sourcetree.md](/docs/sourcetree.md).
 - To learn more about PostgreSQL internals, check http://www.interdb.jp/pg/index.html
+
+
+dd

compute_tools/src/spec.rs

@@ -387,13 +387,13 @@ pub fn handle_databases(spec: &ComputeSpec, client: &mut Client) -> Result<()> {
                     name.pg_quote(),
                     db.owner.pg_quote()
                 );
-                let _ = info_span!("executing", query).entered();
+                let _guard = info_span!("executing", query).entered();
                 client.execute(query.as_str(), &[])?;
             }
             DatabaseAction::Create => {
                 let mut query: String = format!("CREATE DATABASE {} ", name.pg_quote());
                 query.push_str(&db.to_pg_options());
-                let _ = info_span!("executing", query).entered();
+                let _guard = info_span!("executing", query).entered();
                 client.execute(query.as_str(), &[])?;
             }
         };

libs/metrics/src/lib.rs

@@ -8,6 +8,7 @@ pub use prometheus::opts;
 pub use prometheus::register;
 pub use prometheus::{core, default_registry, proto};
 pub use prometheus::{exponential_buckets, linear_buckets};
+pub use prometheus::{register_counter_vec, Counter, CounterVec};
 pub use prometheus::{register_gauge, Gauge};
 pub use prometheus::{register_gauge_vec, GaugeVec};
 pub use prometheus::{register_histogram, Histogram};

libs/pageserver_api/src/models.rs

@@ -29,6 +29,14 @@ pub enum TenantState {
     Broken,
 }
 
+pub mod state {
+    pub const LOADING: &str = "loading";
+    pub const ATTACHING: &str = "attaching";
+    pub const ACTIVE: &str = "active";
+    pub const STOPPING: &str = "stopping";
+    pub const BROKEN: &str = "broken";
+}
+
 impl TenantState {
     pub fn has_in_progress_downloads(&self) -> bool {
         match self {
@@ -39,6 +47,16 @@ impl TenantState {
             Self::Broken => false,
         }
     }
+
+    pub fn as_str(&self) -> &'static str {
+        match self {
+            TenantState::Loading => state::LOADING,
+            TenantState::Attaching => state::ATTACHING,
+            TenantState::Active => state::ACTIVE,
+            TenantState::Stopping => state::STOPPING,
+            TenantState::Broken => state::BROKEN,
+        }
+    }
 }
 
 /// A state of a timeline in pageserver's memory.

libs/utils/Cargo.toml

@@ -37,6 +37,7 @@ metrics.workspace = true
 pq_proto.workspace = true
 
 workspace_hack.workspace = true
+url.workspace = true
 
 [dev-dependencies]
 byteorder.workspace = true

libs/utils/src/http/request.rs

@@ -1,4 +1,5 @@
-use std::str::FromStr;
+use core::fmt;
+use std::{borrow::Cow, str::FromStr};
 
 use super::error::ApiError;
 use anyhow::anyhow;
@@ -29,6 +30,50 @@ pub fn parse_request_param<T: FromStr>(
     }
 }
 
+fn get_query_param<'a>(
+    request: &'a Request<Body>,
+    param_name: &str,
+) -> Result<Option<Cow<'a, str>>, ApiError> {
+    let query = match request.uri().query() {
+        Some(q) => q,
+        None => return Ok(None),
+    };
+    let mut values = url::form_urlencoded::parse(query.as_bytes())
+        .filter_map(|(k, v)| if k == param_name { Some(v) } else { None })
+        // we call .next() twice below. If it's None the first time, .fuse() ensures it's None afterwards
+        .fuse();
+
+    let value1 = values.next();
+    if values.next().is_some() {
+        return Err(ApiError::BadRequest(anyhow!(
+            "param {param_name} specified more than once"
+        )));
+    }
+    Ok(value1)
+}
+
+pub fn must_get_query_param<'a>(
+    request: &'a Request<Body>,
+    param_name: &str,
+) -> Result<Cow<'a, str>, ApiError> {
+    get_query_param(request, param_name)?.ok_or_else(|| {
+        ApiError::BadRequest(anyhow!("no {param_name} specified in query parameters"))
+    })
+}
+
+pub fn parse_query_param<E: fmt::Display, T: FromStr<Err = E>>(
+    request: &Request<Body>,
+    param_name: &str,
+) -> Result<Option<T>, ApiError> {
+    get_query_param(request, param_name)?
+        .map(|v| {
+            v.parse().map_err(|e| {
+                ApiError::BadRequest(anyhow!("cannot parse query param {param_name}: {e}"))
+            })
+        })
+        .transpose()
+}
+
 pub async fn ensure_no_body(request: &mut Request<Body>) -> Result<(), ApiError> {
     match request.body_mut().data().await {
         Some(_) => Err(ApiError::BadRequest(anyhow!("Unexpected request body"))),

pageserver/src/bin/layer_map_analyzer.rs

@@ -0,0 +1,230 @@
+//! Tool for extracting content-dependent metadata about layers. Useful for scanning real project layer files and evaluating the effectiveness of different heuristics on them.
+//!
+//! Currently it only analyzes holes, which are regions within the layer range that the layer contains no updates for. In the future it might do more analysis (maybe key quantiles?) but it should never return sensitive data.
+
+use anyhow::Result;
+use std::cmp::Ordering;
+use std::collections::BinaryHeap;
+use std::ops::Range;
+use std::{env, fs, path::Path, path::PathBuf, str, str::FromStr};
+
+use pageserver::page_cache::PAGE_SZ;
+use pageserver::repository::{Key, KEY_SIZE};
+use pageserver::tenant::block_io::{BlockReader, FileBlockReader};
+use pageserver::tenant::disk_btree::{DiskBtreeReader, VisitDirection};
+use pageserver::tenant::storage_layer::delta_layer::{Summary, DELTA_KEY_SIZE};
+use pageserver::tenant::storage_layer::range_overlaps;
+use pageserver::virtual_file::VirtualFile;
+
+use utils::{bin_ser::BeSer, lsn::Lsn};
+
+const MIN_HOLE_LENGTH: i128 = (128 * 1024 * 1024 / PAGE_SZ) as i128;
+const DEFAULT_MAX_HOLES: usize = 10;
+
+/// Wrapper for key range to provide reverse ordering by range length for BinaryHeap
+#[derive(PartialEq, Eq)]
+struct Hole(Range<Key>);
+
+impl Ord for Hole {
+    fn cmp(&self, other: &Self) -> Ordering {
+        let other_len = other.0.end.to_i128() - other.0.start.to_i128();
+        let self_len = self.0.end.to_i128() - self.0.start.to_i128();
+        other_len.cmp(&self_len)
+    }
+}
+
+impl PartialOrd for Hole {
+    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
+        Some(self.cmp(other))
+    }
+}
+
+struct LayerFile {
+    key_range: Range<Key>,
+    lsn_range: Range<Lsn>,
+    is_delta: bool,
+    holes: Vec<Hole>,
+}
+
+impl LayerFile {
+    fn skips(&self, key_range: &Range<Key>) -> bool {
+        if !range_overlaps(&self.key_range, key_range) {
+            return false;
+        }
+        let start = match self
+            .holes
+            .binary_search_by_key(&key_range.start, |hole| hole.0.start)
+        {
+            Ok(index) => index,
+            Err(index) => {
+                if index == 0 {
+                    return false;
+                }
+                index - 1
+            }
+        };
+        self.holes[start].0.end >= key_range.end
+    }
+}
+
+fn parse_filename(name: &str) -> Option<LayerFile> {
+    let split: Vec<&str> = name.split("__").collect();
+    if split.len() != 2 {
+        return None;
+    }
+    let keys: Vec<&str> = split[0].split('-').collect();
+    let mut lsns: Vec<&str> = split[1].split('-').collect();
+    let is_delta = if lsns.len() == 1 {
+        lsns.push(lsns[0]);
+        false
+    } else {
+        true
+    };
+
+    let key_range = Key::from_hex(keys[0]).unwrap()..Key::from_hex(keys[1]).unwrap();
+    let lsn_range = Lsn::from_hex(lsns[0]).unwrap()..Lsn::from_hex(lsns[1]).unwrap();
+    let holes = Vec::new();
+    Some(LayerFile {
+        key_range,
+        lsn_range,
+        is_delta,
+        holes,
+    })
+}
+
+// Finds the max_holes largest holes, ignoring any that are smaller than MIN_HOLE_LENGTH"
+fn get_holes(path: &Path, max_holes: usize) -> Result<Vec<Hole>> {
+    let file = FileBlockReader::new(VirtualFile::open(path)?);
+    let summary_blk = file.read_blk(0)?;
+    let actual_summary = Summary::des_prefix(summary_blk.as_ref())?;
+    let tree_reader = DiskBtreeReader::<_, DELTA_KEY_SIZE>::new(
+        actual_summary.index_start_blk,
+        actual_summary.index_root_blk,
+        file,
+    );
+    // min-heap (reserve space for one more element added before eviction)
+    let mut heap: BinaryHeap<Hole> = BinaryHeap::with_capacity(max_holes + 1);
+    let mut prev_key: Option<Key> = None;
+    tree_reader.visit(
+        &[0u8; DELTA_KEY_SIZE],
+        VisitDirection::Forwards,
+        |key, _value| {
+            let curr = Key::from_slice(&key[..KEY_SIZE]);
+            if let Some(prev) = prev_key {
+                if curr.to_i128() - prev.to_i128() >= MIN_HOLE_LENGTH {
+                    heap.push(Hole(prev..curr));
+                    if heap.len() > max_holes {
+                        heap.pop(); // remove smallest hole
+                    }
+                }
+            }
+            prev_key = Some(curr.next());
+            true
+        },
+    )?;
+    let mut holes = heap.into_vec();
+    holes.sort_by_key(|hole| hole.0.start);
+    Ok(holes)
+}
+
+fn main() -> Result<()> {
+    let args: Vec<String> = env::args().collect();
+    if args.len() < 2 {
+        println!("Usage: layer_map_analyzer PAGESERVER_DATA_DIR [MAX_HOLES]");
+        return Ok(());
+    }
+    let storage_path = PathBuf::from_str(&args[1])?;
+    let max_holes = if args.len() > 2 {
+        args[2].parse::<usize>().unwrap()
+    } else {
+        DEFAULT_MAX_HOLES
+    };
+
+    // Initialize virtual_file (file desriptor cache) and page cache which are needed to access layer persistent B-Tree.
+    pageserver::virtual_file::init(10);
+    pageserver::page_cache::init(100);
+
+    let mut total_delta_layers = 0usize;
+    let mut total_image_layers = 0usize;
+    let mut total_excess_layers = 0usize;
+    for tenant in fs::read_dir(storage_path.join("tenants"))? {
+        let tenant = tenant?;
+        if !tenant.file_type()?.is_dir() {
+            continue;
+        }
+        for timeline in fs::read_dir(tenant.path().join("timelines"))? {
+            let timeline = timeline?;
+            if !timeline.file_type()?.is_dir() {
+                continue;
+            }
+            // Collect sorted vec of layers and count deltas
+            let mut layers = Vec::new();
+            let mut n_deltas = 0usize;
+
+            for layer in fs::read_dir(timeline.path())? {
+                let layer = layer?;
+                if let Some(mut layer_file) =
+                    parse_filename(&layer.file_name().into_string().unwrap())
+                {
+                    if layer_file.is_delta {
+                        layer_file.holes = get_holes(&layer.path(), max_holes)?;
+                        n_deltas += 1;
+                    }
+                    layers.push(layer_file);
+                }
+            }
+            layers.sort_by_key(|layer| layer.lsn_range.end);
+
+            // Count the number of holes and number of excess layers.
+            // Excess layer is image layer generated when holes in delta layers are not considered.
+            let mut n_excess_layers = 0usize;
+            let mut n_holes = 0usize;
+
+            for i in 0..layers.len() {
+                if !layers[i].is_delta {
+                    let mut n_deltas_since_last_image = 0usize;
+                    let mut n_skipped = 0usize;
+                    let img_key_range = &layers[i].key_range;
+                    for j in (0..i).rev() {
+                        if range_overlaps(img_key_range, &layers[j].key_range) {
+                            if layers[j].is_delta {
+                                n_deltas_since_last_image += 1;
+                                if layers[j].skips(img_key_range) {
+                                    n_skipped += 1;
+                                }
+                            } else {
+                                // Image layer is always dense, despite to the fact that it doesn't contain all possible
+                                // key values in the specified range: there are may be no keys in the storage belonging
+                                // to the image layer range but not present in the image layer.
+                                break;
+                            }
+                        }
+                    }
+                    if n_deltas_since_last_image >= 3 && n_deltas_since_last_image - n_skipped < 3 {
+                        // It is just approximation: it doesn't take in account all image coverage.
+                        // Moreover the new layer map doesn't count total deltas, but the max stack of overlapping deltas.
+                        n_excess_layers += 1;
+                    }
+                    n_holes += n_skipped;
+                }
+            }
+            println!(
+                "Tenant {} timeline {} delta layers {} image layers {} excess layers {} holes {}",
+                tenant.file_name().into_string().unwrap(),
+                timeline.file_name().into_string().unwrap(),
+                n_deltas,
+                layers.len() - n_deltas,
+                n_excess_layers,
+                n_holes
+            );
+            total_delta_layers += n_deltas;
+            total_image_layers += layers.len() - n_deltas;
+            total_excess_layers += n_excess_layers;
+        }
+    }
+    println!(
+        "Total delta layers {} image layers {} excess layers {}",
+        total_delta_layers, total_image_layers, total_excess_layers
+    );
+    Ok(())
+}

pageserver/src/bin/pageserver.rs

@@ -250,7 +250,7 @@ fn start_pageserver(conf: &'static PageServerConf) -> anyhow::Result<()> {
     let signals = signals::install_shutdown_handlers()?;
 
     // Launch broker client
-    WALRECEIVER_RUNTIME.block_on(pageserver::walreceiver::init_broker_client(conf))?;
+    WALRECEIVER_RUNTIME.block_on(pageserver::broker_client::init_broker_client(conf))?;
 
     // Initialize authentication for incoming connections
     let auth = match &conf.auth_type {

pageserver/src/bin/pageserver_binutils.rs

@@ -12,7 +12,9 @@ use anyhow::Context;
 use clap::{value_parser, Arg, Command};
 
 use pageserver::{
+    context::{DownloadBehavior, RequestContext},
     page_cache,
+    task_mgr::TaskKind,
     tenant::{dump_layerfile_from_path, metadata::TimelineMetadata},
     virtual_file,
 };
@@ -75,7 +77,8 @@ fn print_layerfile(path: &Path) -> anyhow::Result<()> {
     // Basic initialization of things that don't change after startup
     virtual_file::init(10);
     page_cache::init(100);
-    dump_layerfile_from_path(path, true)
+    let ctx = RequestContext::new(TaskKind::DebugTool, DownloadBehavior::Error);
+    dump_layerfile_from_path(path, true, &ctx)
 }
 
 fn handle_metadata(path: &Path, arg_matches: &clap::ArgMatches) -> Result<(), anyhow::Error> {

pageserver/src/broker_client.rs

@@ -0,0 +1,48 @@
+//! The broker client instance of the pageserver, created during pageserver startup.
+//! Used by each timelines' [`walreceiver`].
+
+use crate::config::PageServerConf;
+
+use anyhow::Context;
+use once_cell::sync::OnceCell;
+use storage_broker::BrokerClientChannel;
+use tracing::*;
+
+static BROKER_CLIENT: OnceCell<BrokerClientChannel> = OnceCell::new();
+
+///
+/// Initialize the broker client. This must be called once at page server startup.
+///
+pub async fn init_broker_client(conf: &'static PageServerConf) -> anyhow::Result<()> {
+    let broker_endpoint = conf.broker_endpoint.clone();
+
+    // Note: we do not attempt connecting here (but validate endpoints sanity).
+    let broker_client =
+        storage_broker::connect(broker_endpoint.clone(), conf.broker_keepalive_interval).context(
+            format!(
+                "Failed to create broker client to {}",
+                &conf.broker_endpoint
+            ),
+        )?;
+
+    if BROKER_CLIENT.set(broker_client).is_err() {
+        panic!("broker already initialized");
+    }
+
+    info!(
+        "Initialized broker client with endpoints: {}",
+        broker_endpoint
+    );
+    Ok(())
+}
+
+///
+/// Get a handle to the broker client
+///
+pub fn get_broker_client() -> &'static BrokerClientChannel {
+    BROKER_CLIENT.get().expect("broker client not initialized")
+}
+
+pub fn is_broker_client_initialized() -> bool {
+    BROKER_CLIENT.get().is_some()
+}

pageserver/src/http/routes.rs

@@ -7,6 +7,7 @@ use pageserver_api::models::DownloadRemoteLayersTaskSpawnRequest;
 use remote_storage::GenericRemoteStorage;
 use tokio_util::sync::CancellationToken;
 use tracing::*;
+use utils::http::request::{must_get_query_param, parse_query_param};
 
 use super::models::{
     StatusResponse, TenantConfigRequest, TenantCreateRequest, TenantCreateResponse, TenantInfo,
@@ -87,9 +88,7 @@ fn apierror_from_prerror(err: PageReconstructError) -> ApiError {
         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"
-            ))
+            ApiError::InternalServerError(anyhow::anyhow!("need to download remote layer file"))
         }
         PageReconstructError::Cancelled => {
             ApiError::InternalServerError(anyhow::anyhow!("request was cancelled"))
@@ -237,8 +236,8 @@ async fn timeline_create_handler(mut request: Request<Body>) -> Result<Response<
 
 async fn timeline_list_handler(request: Request<Body>) -> Result<Response<Body>, ApiError> {
     let tenant_id: TenantId = parse_request_param(&request, "tenant_id")?;
-    let include_non_incremental_logical_size =
-        query_param_present(&request, "include-non-incremental-logical-size");
+    let include_non_incremental_logical_size: Option<bool> =
+        parse_query_param(&request, "include-non-incremental-logical-size")?;
     check_permission(&request, Some(tenant_id))?;
 
     let ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Download);
@@ -251,13 +250,14 @@ async fn timeline_list_handler(request: Request<Body>) -> Result<Response<Body>,
 
         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, &ctx)
-                    .await
-                    .context(
-                        "Failed to convert tenant timeline {timeline_id} into the local one: {e:?}",
-                    )
-                    .map_err(ApiError::InternalServerError)?;
+            let timeline_info = build_timeline_info(
+                &timeline,
+                include_non_incremental_logical_size.unwrap_or(false),
+                &ctx,
+            )
+            .await
+            .context("Failed to convert tenant timeline {timeline_id} into the local one: {e:?}")
+            .map_err(ApiError::InternalServerError)?;
 
             response_data.push(timeline_info);
         }
@@ -269,36 +269,11 @@ async fn timeline_list_handler(request: Request<Body>) -> Result<Response<Body>,
     json_response(StatusCode::OK, response_data)
 }
 
-/// Checks if a query param is present in the request's URL
-fn query_param_present(request: &Request<Body>, param: &str) -> bool {
-    request
-        .uri()
-        .query()
-        .map(|v| url::form_urlencoded::parse(v.as_bytes()).any(|(p, _)| p == param))
-        .unwrap_or(false)
-}
-
-fn get_query_param(request: &Request<Body>, param_name: &str) -> Result<String, ApiError> {
-    request.uri().query().map_or(
-        Err(ApiError::BadRequest(anyhow!("empty query in request"))),
-        |v| {
-            url::form_urlencoded::parse(v.as_bytes())
-                .find(|(k, _)| k == param_name)
-                .map_or(
-                    Err(ApiError::BadRequest(anyhow!(
-                        "no {param_name} specified in query parameters"
-                    ))),
-                    |(_, v)| Ok(v.into_owned()),
-                )
-        },
-    )
-}
-
 async fn timeline_detail_handler(request: Request<Body>) -> Result<Response<Body>, ApiError> {
     let tenant_id: TenantId = parse_request_param(&request, "tenant_id")?;
     let timeline_id: TimelineId = parse_request_param(&request, "timeline_id")?;
-    let include_non_incremental_logical_size =
-        query_param_present(&request, "include-non-incremental-logical-size");
+    let include_non_incremental_logical_size: Option<bool> =
+        parse_query_param(&request, "include-non-incremental-logical-size")?;
     check_permission(&request, Some(tenant_id))?;
 
     // Logical size calculation needs downloading.
@@ -313,11 +288,14 @@ async fn timeline_detail_handler(request: Request<Body>) -> Result<Response<Body
             .get_timeline(timeline_id, false)
             .map_err(ApiError::NotFound)?;
 
-        let timeline_info =
-            build_timeline_info(&timeline, include_non_incremental_logical_size, &ctx)
-                .await
-                .context("get local timeline info")
-                .map_err(ApiError::InternalServerError)?;
+        let timeline_info = build_timeline_info(
+            &timeline,
+            include_non_incremental_logical_size.unwrap_or(false),
+            &ctx,
+        )
+        .await
+        .context("get local timeline info")
+        .map_err(ApiError::InternalServerError)?;
 
         Ok::<_, ApiError>(timeline_info)
     }
@@ -332,8 +310,8 @@ async fn get_lsn_by_timestamp_handler(request: Request<Body>) -> Result<Response
     check_permission(&request, Some(tenant_id))?;
 
     let timeline_id: TimelineId = parse_request_param(&request, "timeline_id")?;
-    let timestamp_raw = get_query_param(&request, "timestamp")?;
-    let timestamp = humantime::parse_rfc3339(timestamp_raw.as_str())
+    let timestamp_raw = must_get_query_param(&request, "timestamp")?;
+    let timestamp = humantime::parse_rfc3339(&timestamp_raw)
         .with_context(|| format!("Invalid time: {:?}", timestamp_raw))
         .map_err(ApiError::BadRequest)?;
     let timestamp_pg = postgres_ffi::to_pg_timestamp(timestamp);
@@ -505,13 +483,7 @@ 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 inputs_only = if query_param_present(&request, "inputs_only") {
-        get_query_param(&request, "inputs_only")?
-            .parse()
-            .map_err(|_| ApiError::BadRequest(anyhow!("failed to parse inputs_only")))?
-    } else {
-        false
-    };
+    let inputs_only: Option<bool> = parse_query_param(&request, "inputs_only")?;
 
     let ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Download);
     let tenant = mgr::get_tenant(tenant_id, true)
@@ -524,7 +496,7 @@ async fn tenant_size_handler(request: Request<Body>) -> Result<Response<Body>, A
         .await
         .map_err(ApiError::InternalServerError)?;
 
-    let size = if !inputs_only {
+    let size = if !inputs_only.unwrap_or(false) {
         Some(inputs.calculate().map_err(ApiError::InternalServerError)?)
     } else {
         None

pageserver/src/lib.rs

@@ -1,5 +1,6 @@
 mod auth;
 pub mod basebackup;
+pub mod broker_client;
 pub mod config;
 pub mod consumption_metrics;
 pub mod context;
@@ -16,7 +17,6 @@ pub mod tenant;
 pub mod trace;
 pub mod virtual_file;
 pub mod walingest;
-pub mod walreceiver;
 pub mod walrecord;
 pub mod walredo;
 

pageserver/src/metrics.rs

@@ -1,10 +1,12 @@
 use metrics::core::{AtomicU64, GenericCounter};
 use metrics::{
-    register_histogram, register_histogram_vec, register_int_counter, register_int_counter_vec,
-    register_int_gauge, register_int_gauge_vec, register_uint_gauge_vec, Histogram, HistogramVec,
-    IntCounter, IntCounterVec, IntGauge, IntGaugeVec, UIntGauge, UIntGaugeVec,
+    register_counter_vec, register_histogram, register_histogram_vec, register_int_counter,
+    register_int_counter_vec, register_int_gauge, register_int_gauge_vec, register_uint_gauge_vec,
+    Counter, CounterVec, Histogram, HistogramVec, IntCounter, IntCounterVec, IntGauge, IntGaugeVec,
+    UIntGauge, UIntGaugeVec,
 };
 use once_cell::sync::Lazy;
+use pageserver_api::models::state;
 use utils::id::{TenantId, TimelineId};
 
 /// Prometheus histogram buckets (in seconds) that capture the majority of
@@ -35,11 +37,29 @@ const STORAGE_TIME_OPERATIONS: &[&str] = &[
     "gc",
 ];
 
-pub static STORAGE_TIME: Lazy<HistogramVec> = Lazy::new(|| {
-    register_histogram_vec!(
-        "pageserver_storage_operations_seconds",
-        "Time spent on storage operations",
+pub static STORAGE_TIME_SUM_PER_TIMELINE: Lazy<CounterVec> = Lazy::new(|| {
+    register_counter_vec!(
+        "pageserver_storage_operations_seconds_sum",
+        "Total time spent on storage operations with operation, tenant and timeline dimensions",
         &["operation", "tenant_id", "timeline_id"],
+    )
+    .expect("failed to define a metric")
+});
+
+pub static STORAGE_TIME_COUNT_PER_TIMELINE: Lazy<IntCounterVec> = Lazy::new(|| {
+    register_int_counter_vec!(
+        "pageserver_storage_operations_seconds_count",
+        "Count of storage operations with operation, tenant and timeline dimensions",
+        &["operation", "tenant_id", "timeline_id"],
+    )
+    .expect("failed to define a metric")
+});
+
+pub static STORAGE_TIME_GLOBAL: Lazy<HistogramVec> = Lazy::new(|| {
+    register_histogram_vec!(
+        "pageserver_storage_operations_seconds_global",
+        "Time spent on storage operations",
+        &["operation"],
         get_buckets_for_critical_operations(),
     )
     .expect("failed to define a metric")
@@ -112,6 +132,24 @@ static CURRENT_LOGICAL_SIZE: Lazy<UIntGaugeVec> = Lazy::new(|| {
     .expect("failed to define current logical size metric")
 });
 
+// Metrics collected on tenant states.
+const TENANT_STATE_OPTIONS: &[&str] = &[
+    state::LOADING,
+    state::ATTACHING,
+    state::ACTIVE,
+    state::STOPPING,
+    state::BROKEN,
+];
+
+pub static TENANT_STATE_METRIC: Lazy<UIntGaugeVec> = Lazy::new(|| {
+    register_uint_gauge_vec!(
+        "pageserver_tenant_states_count",
+        "Count of tenants per state",
+        &["tenant_id", "state"]
+    )
+    .expect("Failed to register pageserver_tenant_states_count metric")
+});
+
 // Metrics for cloud upload. These metrics reflect data uploaded to cloud storage,
 // or in testing they estimate how much we would upload if we did.
 static NUM_PERSISTENT_FILES_CREATED: Lazy<IntCounterVec> = Lazy::new(|| {
@@ -375,18 +413,81 @@ pub static WAL_REDO_RECORD_COUNTER: Lazy<IntCounter> = Lazy::new(|| {
     .unwrap()
 });
 
+/// Similar to [`prometheus::HistogramTimer`] but does not record on drop.
+pub struct StorageTimeMetricsTimer {
+    metrics: StorageTimeMetrics,
+    start: Instant,
+}
+
+impl StorageTimeMetricsTimer {
+    fn new(metrics: StorageTimeMetrics) -> Self {
+        Self {
+            metrics,
+            start: Instant::now(),
+        }
+    }
+
+    /// Record the time from creation to now.
+    pub fn stop_and_record(self) {
+        let duration = self.start.elapsed().as_secs_f64();
+        self.metrics.timeline_sum.inc_by(duration);
+        self.metrics.timeline_count.inc();
+        self.metrics.global_histogram.observe(duration);
+    }
+}
+
+/// Timing facilities for an globally histogrammed metric, which is supported by per tenant and
+/// timeline total sum and count.
+#[derive(Clone, Debug)]
+pub struct StorageTimeMetrics {
+    /// Sum of f64 seconds, per operation, tenant_id and timeline_id
+    timeline_sum: Counter,
+    /// Number of oeprations, per operation, tenant_id and timeline_id
+    timeline_count: IntCounter,
+    /// Global histogram having only the "operation" label.
+    global_histogram: Histogram,
+}
+
+impl StorageTimeMetrics {
+    pub fn new(operation: &str, tenant_id: &str, timeline_id: &str) -> Self {
+        let timeline_sum = STORAGE_TIME_SUM_PER_TIMELINE
+            .get_metric_with_label_values(&[operation, tenant_id, timeline_id])
+            .unwrap();
+        let timeline_count = STORAGE_TIME_COUNT_PER_TIMELINE
+            .get_metric_with_label_values(&[operation, tenant_id, timeline_id])
+            .unwrap();
+        let global_histogram = STORAGE_TIME_GLOBAL
+            .get_metric_with_label_values(&[operation])
+            .unwrap();
+
+        StorageTimeMetrics {
+            timeline_sum,
+            timeline_count,
+            global_histogram,
+        }
+    }
+
+    /// Starts timing a new operation.
+    ///
+    /// Note: unlike [`prometheus::HistogramTimer`] the returned timer does not record on drop.
+    pub fn start_timer(&self) -> StorageTimeMetricsTimer {
+        StorageTimeMetricsTimer::new(self.clone())
+    }
+}
+
 #[derive(Debug)]
 pub struct TimelineMetrics {
     tenant_id: String,
     timeline_id: String,
     pub reconstruct_time_histo: Histogram,
     pub materialized_page_cache_hit_counter: GenericCounter<AtomicU64>,
-    pub flush_time_histo: Histogram,
-    pub compact_time_histo: Histogram,
-    pub create_images_time_histo: Histogram,
-    pub init_logical_size_histo: Histogram,
-    pub logical_size_histo: Histogram,
-    pub load_layer_map_histo: Histogram,
+    pub flush_time_histo: StorageTimeMetrics,
+    pub compact_time_histo: StorageTimeMetrics,
+    pub create_images_time_histo: StorageTimeMetrics,
+    pub init_logical_size_histo: StorageTimeMetrics,
+    pub logical_size_histo: StorageTimeMetrics,
+    pub load_layer_map_histo: StorageTimeMetrics,
+    pub garbage_collect_histo: StorageTimeMetrics,
     pub last_record_gauge: IntGauge,
     pub wait_lsn_time_histo: Histogram,
     pub resident_physical_size_gauge: UIntGauge,
@@ -406,24 +507,16 @@ impl TimelineMetrics {
         let materialized_page_cache_hit_counter = MATERIALIZED_PAGE_CACHE_HIT
             .get_metric_with_label_values(&[&tenant_id, &timeline_id])
             .unwrap();
-        let flush_time_histo = STORAGE_TIME
-            .get_metric_with_label_values(&["layer flush", &tenant_id, &timeline_id])
-            .unwrap();
-        let compact_time_histo = STORAGE_TIME
-            .get_metric_with_label_values(&["compact", &tenant_id, &timeline_id])
-            .unwrap();
-        let create_images_time_histo = STORAGE_TIME
-            .get_metric_with_label_values(&["create images", &tenant_id, &timeline_id])
-            .unwrap();
-        let init_logical_size_histo = STORAGE_TIME
-            .get_metric_with_label_values(&["init logical size", &tenant_id, &timeline_id])
-            .unwrap();
-        let logical_size_histo = STORAGE_TIME
-            .get_metric_with_label_values(&["logical size", &tenant_id, &timeline_id])
-            .unwrap();
-        let load_layer_map_histo = STORAGE_TIME
-            .get_metric_with_label_values(&["load layer map", &tenant_id, &timeline_id])
-            .unwrap();
+        let flush_time_histo = StorageTimeMetrics::new("layer flush", &tenant_id, &timeline_id);
+        let compact_time_histo = StorageTimeMetrics::new("compact", &tenant_id, &timeline_id);
+        let create_images_time_histo =
+            StorageTimeMetrics::new("create images", &tenant_id, &timeline_id);
+        let init_logical_size_histo =
+            StorageTimeMetrics::new("init logical size", &tenant_id, &timeline_id);
+        let logical_size_histo = StorageTimeMetrics::new("logical size", &tenant_id, &timeline_id);
+        let load_layer_map_histo =
+            StorageTimeMetrics::new("load layer map", &tenant_id, &timeline_id);
+        let garbage_collect_histo = StorageTimeMetrics::new("gc", &tenant_id, &timeline_id);
         let last_record_gauge = LAST_RECORD_LSN
             .get_metric_with_label_values(&[&tenant_id, &timeline_id])
             .unwrap();
@@ -453,6 +546,7 @@ impl TimelineMetrics {
             create_images_time_histo,
             init_logical_size_histo,
             logical_size_histo,
+            garbage_collect_histo,
             load_layer_map_histo,
             last_record_gauge,
             wait_lsn_time_histo,
@@ -478,7 +572,10 @@ impl Drop for TimelineMetrics {
         let _ = PERSISTENT_BYTES_WRITTEN.remove_label_values(&[tenant_id, timeline_id]);
 
         for op in STORAGE_TIME_OPERATIONS {
-            let _ = STORAGE_TIME.remove_label_values(&[op, tenant_id, timeline_id]);
+            let _ =
+                STORAGE_TIME_SUM_PER_TIMELINE.remove_label_values(&[op, tenant_id, timeline_id]);
+            let _ =
+                STORAGE_TIME_COUNT_PER_TIMELINE.remove_label_values(&[op, tenant_id, timeline_id]);
         }
         for op in STORAGE_IO_TIME_OPERATIONS {
             let _ = STORAGE_IO_TIME.remove_label_values(&[op, tenant_id, timeline_id]);
@@ -495,7 +592,10 @@ impl Drop for TimelineMetrics {
 }
 
 pub fn remove_tenant_metrics(tenant_id: &TenantId) {
-    let _ = STORAGE_TIME.remove_label_values(&["gc", &tenant_id.to_string(), "-"]);
+    let tid = tenant_id.to_string();
+    for state in TENANT_STATE_OPTIONS {
+        let _ = TENANT_STATE_METRIC.remove_label_values(&[&tid, state]);
+    }
 }
 
 use futures::Future;

pageserver/src/pgdatadir_mapping.rs

@@ -1653,7 +1653,7 @@ mod tests {
         assert!(tline.list_rels(0, TESTDB, Lsn(0x30))?.contains(&TESTREL_A));
 
         // Create a branch, check that the relation is visible there
-        repo.branch_timeline(TIMELINE_ID, NEW_TIMELINE_ID, Lsn(0x30))?;
+        repo.branch_timeline(&tline, NEW_TIMELINE_ID, Lsn(0x30))?;
         let newtline = match repo.get_timeline(NEW_TIMELINE_ID)?.local_timeline() {
             Some(timeline) => timeline,
             None => panic!("Should have a local timeline"),

pageserver/src/task_mgr.rs

@@ -255,6 +255,8 @@ pub enum TaskKind {
     // A request that comes in via the pageserver HTTP API.
     MgmtRequest,
 
+    DebugTool,
+
     #[cfg(test)]
     UnitTest,
 }

pageserver/src/tenant.rs

@@ -51,7 +51,7 @@ use crate::config::PageServerConf;
 use crate::context::{DownloadBehavior, RequestContext};
 use crate::import_datadir;
 use crate::is_uninit_mark;
-use crate::metrics::{remove_tenant_metrics, STORAGE_TIME};
+use crate::metrics::{remove_tenant_metrics, TENANT_STATE_METRIC};
 use crate::repository::GcResult;
 use crate::task_mgr;
 use crate::task_mgr::TaskKind;
@@ -77,7 +77,7 @@ use utils::{
 
 mod blob_io;
 pub mod block_io;
-mod disk_btree;
+pub mod disk_btree;
 pub(crate) mod ephemeral_file;
 pub mod layer_map;
 
@@ -1205,13 +1205,8 @@ impl Tenant {
                     ancestor_timeline.wait_lsn(*lsn, ctx).await?;
                 }
 
-                self.branch_timeline(
-                    ancestor_timeline_id,
-                    new_timeline_id,
-                    ancestor_start_lsn,
-                    ctx,
-                )
-                .await?
+                self.branch_timeline(&ancestor_timeline, new_timeline_id, ancestor_start_lsn, ctx)
+                    .await?
             }
             None => {
                 self.bootstrap_timeline(new_timeline_id, pg_version, ctx)
@@ -1247,17 +1242,11 @@ impl Tenant {
             "Cannot run GC iteration on inactive tenant"
         );
 
-        let timeline_str = target_timeline_id
-            .map(|x| x.to_string())
-            .unwrap_or_else(|| "-".to_string());
+        let gc_result = self
+            .gc_iteration_internal(target_timeline_id, horizon, pitr, ctx)
+            .await;
 
-        {
-            let _timer = STORAGE_TIME
-                .with_label_values(&["gc", &self.tenant_id.to_string(), &timeline_str])
-                .start_timer();
-            self.gc_iteration_internal(target_timeline_id, horizon, pitr, ctx)
-                .await
-        }
+        gc_result
     }
 
     /// Perform one compaction iteration.
@@ -1741,7 +1730,33 @@ impl Tenant {
         tenant_id: TenantId,
         remote_storage: Option<GenericRemoteStorage>,
     ) -> Tenant {
-        let (state, _) = watch::channel(state);
+        let (state, mut rx) = watch::channel(state);
+
+        tokio::spawn(async move {
+            let current_state = *rx.borrow_and_update();
+            let tid = tenant_id.to_string();
+            TENANT_STATE_METRIC
+                .with_label_values(&[&tid, current_state.as_str()])
+                .inc();
+            loop {
+                match rx.changed().await {
+                    Ok(()) => {
+                        let new_state = *rx.borrow();
+                        TENANT_STATE_METRIC
+                            .with_label_values(&[&tid, current_state.as_str()])
+                            .dec();
+                        TENANT_STATE_METRIC
+                            .with_label_values(&[&tid, new_state.as_str()])
+                            .inc();
+                    }
+                    Err(_sender_dropped_error) => {
+                        info!("Tenant dropped the state updates sender, quitting waiting for tenant state change");
+                        return;
+                    }
+                }
+            }
+        });
+
         Tenant {
             tenant_id,
             conf,
@@ -1800,69 +1815,70 @@ impl Tenant {
     }
 
     pub(super) fn persist_tenant_config(
+        tenant_id: &TenantId,
         target_config_path: &Path,
         tenant_conf: TenantConfOpt,
-        first_save: bool,
+        creating_tenant: bool,
     ) -> anyhow::Result<()> {
         let _enter = info_span!("saving tenantconf").entered();
-        info!("persisting tenantconf to {}", target_config_path.display());
 
-        // TODO this will prepend comments endlessly ?
-        let mut conf_content = r#"# This file contains a specific per-tenant's config.
-#  It is read in case of pageserver restart.
-
-[tenant_config]
-"#
-        .to_string();
-
-        // Convert the config to a toml file.
-        conf_content += &toml_edit::easy::to_string(&tenant_conf)?;
-
-        let mut target_config_file = VirtualFile::open_with_options(
-            target_config_path,
-            OpenOptions::new()
-                .truncate(true) // This needed for overwriting with small config files
-                .write(true)
-                .create_new(first_save),
-        )?;
-
-        target_config_file
-            .write(conf_content.as_bytes())
-            .context("Failed to write toml bytes into file")
-            .and_then(|_| {
-                target_config_file
-                    .sync_all()
-                    .context("Faile to fsync config file")
-            })
-            .with_context(|| {
+        // imitate a try-block with a closure
+        let do_persist = |target_config_path: &Path| -> anyhow::Result<()> {
+            let target_config_parent = target_config_path.parent().with_context(|| {
                 format!(
-                    "Failed to write config file into path '{}'",
+                    "Config path does not have a parent: {}",
                     target_config_path.display()
                 )
             })?;
 
-        // fsync the parent directory to ensure the directory entry is durable
-        if first_save {
-            target_config_path
-                .parent()
-                .context("Config file does not have a parent")
-                .and_then(|target_config_parent| {
-                    File::open(target_config_parent).context("Failed to open config parent")
-                })
-                .and_then(|tenant_dir| {
-                    tenant_dir
-                        .sync_all()
-                        .context("Failed to fsync config parent")
-                })
-                .with_context(|| {
-                    format!(
-                        "Failed to fsync on first save for config {}",
-                        target_config_path.display()
-                    )
-                })?;
-        }
+            info!("persisting tenantconf to {}", target_config_path.display());
 
-        Ok(())
+            let mut conf_content = r#"# This file contains a specific per-tenant's config.
+#  It is read in case of pageserver restart.
+
+[tenant_config]
+"#
+            .to_string();
+
+            // Convert the config to a toml file.
+            conf_content += &toml_edit::easy::to_string(&tenant_conf)?;
+
+            let mut target_config_file = VirtualFile::open_with_options(
+                target_config_path,
+                OpenOptions::new()
+                    .truncate(true) // This needed for overwriting with small config files
+                    .write(true)
+                    .create_new(creating_tenant)
+                    // when creating a new tenant, first_save will be true and `.create(true)` will be
+                    // ignored (per rust std docs).
+                    //
+                    // later when updating the config of created tenant, or persisting config for the
+                    // first time for attached tenant, the `.create(true)` is used.
+                    .create(true),
+            )?;
+
+            target_config_file
+                .write(conf_content.as_bytes())
+                .context("write toml bytes into file")
+                .and_then(|_| target_config_file.sync_all().context("fsync config file"))
+                .context("write config file")?;
+
+            // fsync the parent directory to ensure the directory entry is durable.
+            // before this was done conditionally on creating_tenant, but these management actions are rare
+            // enough to just fsync it always.
+
+            crashsafe::fsync(target_config_parent)?;
+            Ok(())
+        };
+
+        // this function is called from creating the tenant and updating the tenant config, which
+        // would otherwise share this context, so keep it here in one place.
+        do_persist(target_config_path).with_context(|| {
+            format!(
+                "write tenant {tenant_id} config to {}",
+                target_config_path.display()
+            )
+        })
     }
 
     //
@@ -2050,54 +2066,53 @@ impl Tenant {
     /// Branch an existing timeline
     async fn branch_timeline(
         &self,
-        src: TimelineId,
-        dst: TimelineId,
+        src_timeline: &Arc<Timeline>,
+        dst_id: TimelineId,
         start_lsn: Option<Lsn>,
         _ctx: &RequestContext,
     ) -> anyhow::Result<Arc<Timeline>> {
-        // We need to hold this lock to prevent GC from starting at the same time. GC scans the directory to learn
-        // about timelines, so otherwise a race condition is possible, where we create new timeline and GC
-        // concurrently removes data that is needed by the new timeline.
-        let _gc_cs = self.gc_cs.lock().await;
-        let timeline_uninit_mark = {
-            let timelines = self.timelines.lock().unwrap();
-            self.create_timeline_uninit_mark(dst, &timelines)?
-        };
-
-        // In order for the branch creation task to not wait for GC/compaction,
-        // we need to make sure that the starting LSN of the child branch is not out of scope midway by
-        //
-        // 1. holding the GC lock to prevent overwritting timeline's GC data
-        // 2. checking both the latest GC cutoff LSN and latest GC info of the source timeline
-        //
-        // Step 2 is to avoid initializing the new branch using data removed by past GC iterations
-        // or in-queue GC iterations.
-
-        let src_timeline = self.get_timeline(src, false).with_context(|| {
-            format!(
-                "No ancestor {} found for timeline {}/{}",
-                src, self.tenant_id, dst
-            )
-        })?;
-
-        let latest_gc_cutoff_lsn = src_timeline.get_latest_gc_cutoff_lsn();
+        let src_id = src_timeline.timeline_id;
 
         // If no start LSN is specified, we branch the new timeline from the source timeline's last record LSN
         let start_lsn = start_lsn.unwrap_or_else(|| {
             let lsn = src_timeline.get_last_record_lsn();
-            info!("branching timeline {dst} from timeline {src} at last record LSN: {lsn}");
+            info!("branching timeline {dst_id} from timeline {src_id} at last record LSN: {lsn}");
             lsn
         });
 
-        // Check if the starting LSN is out of scope because it is less than
-        // 1. the latest GC cutoff LSN or
-        // 2. the planned GC cutoff LSN, which is from an in-queue GC iteration.
+        // First acquire the GC lock so that another task cannot advance the GC
+        // cutoff in 'gc_info', and make 'start_lsn' invalid, while we are
+        // creating the branch.
+        let _gc_cs = self.gc_cs.lock().await;
+
+        // Create a placeholder for the new branch. This will error
+        // out if the new timeline ID is already in use.
+        let timeline_uninit_mark = {
+            let timelines = self.timelines.lock().unwrap();
+            self.create_timeline_uninit_mark(dst_id, &timelines)?
+        };
+
+        // Ensure that `start_lsn` is valid, i.e. the LSN is within the PITR
+        // horizon on the source timeline
+        //
+        // We check it against both the planned GC cutoff stored in 'gc_info',
+        // and the 'latest_gc_cutoff' of the last GC that was performed.  The
+        // planned GC cutoff in 'gc_info' is normally larger than
+        // 'latest_gc_cutoff_lsn', but beware of corner cases like if you just
+        // changed the GC settings for the tenant to make the PITR window
+        // larger, but some of the data was already removed by an earlier GC
+        // iteration.
+
+        // check against last actual 'latest_gc_cutoff' first
+        let latest_gc_cutoff_lsn = src_timeline.get_latest_gc_cutoff_lsn();
         src_timeline
             .check_lsn_is_in_scope(start_lsn, &latest_gc_cutoff_lsn)
             .context(format!(
                 "invalid branch start lsn: less than latest GC cutoff {}",
                 *latest_gc_cutoff_lsn,
             ))?;
+
+        // and then the planned GC cutoff
         {
             let gc_info = src_timeline.gc_info.read().unwrap();
             let cutoff = min(gc_info.pitr_cutoff, gc_info.horizon_cutoff);
@@ -2108,6 +2123,12 @@ impl Tenant {
             }
         }
 
+        //
+        // The branch point is valid, and we are still holding the 'gc_cs' lock
+        // so that GC cannot advance the GC cutoff until we are finished.
+        // Proceed with the branch creation.
+        //
+
         // Determine prev-LSN for the new timeline. We can only determine it if
         // the timeline was branched at the current end of the source timeline.
         let RecordLsn {
@@ -2126,7 +2147,7 @@ impl Tenant {
         let metadata = TimelineMetadata::new(
             start_lsn,
             dst_prev,
-            Some(src),
+            Some(src_id),
             start_lsn,
             *src_timeline.latest_gc_cutoff_lsn.read(), // FIXME: should we hold onto this guard longer?
             src_timeline.initdb_lsn,
@@ -2135,15 +2156,15 @@ impl Tenant {
         let mut timelines = self.timelines.lock().unwrap();
         let new_timeline = self
             .prepare_timeline(
-                dst,
+                dst_id,
                 metadata,
                 timeline_uninit_mark,
                 false,
-                Some(src_timeline),
+                Some(Arc::clone(src_timeline)),
             )?
             .initialize_with_lock(&mut timelines, true, true)?;
         drop(timelines);
-        info!("branched timeline {dst} from {src} at {start_lsn}");
+        info!("branched timeline {dst_id} from {src_id} at {start_lsn}");
 
         Ok(new_timeline)
     }
@@ -2512,26 +2533,19 @@ fn try_create_target_tenant_dir(
         target_tenant_directory,
         temporary_tenant_dir,
     )
-    .with_context(|| format!("Failed to resolve tenant {tenant_id} temporary timelines dir"))?;
+    .with_context(|| format!("resolve tenant {tenant_id} temporary timelines dir"))?;
     let temporary_tenant_config_path = rebase_directory(
         &conf.tenant_config_path(tenant_id),
         target_tenant_directory,
         temporary_tenant_dir,
     )
-    .with_context(|| format!("Failed to resolve tenant {tenant_id} temporary config path"))?;
+    .with_context(|| format!("resolve tenant {tenant_id} temporary config path"))?;
+
+    Tenant::persist_tenant_config(&tenant_id, &temporary_tenant_config_path, tenant_conf, true)?;
 
-    Tenant::persist_tenant_config(&temporary_tenant_config_path, tenant_conf, true).with_context(
-        || {
-            format!(
-                "Failed to write tenant {} config to {}",
-                tenant_id,
-                temporary_tenant_config_path.display()
-            )
-        },
-    )?;
     crashsafe::create_dir(&temporary_tenant_timelines_dir).with_context(|| {
         format!(
-            "could not create tenant {} temporary timelines directory {}",
+            "create tenant {} temporary timelines directory {}",
             tenant_id,
             temporary_tenant_timelines_dir.display()
         )
@@ -2542,7 +2556,7 @@ fn try_create_target_tenant_dir(
 
     fs::rename(temporary_tenant_dir, target_tenant_directory).with_context(|| {
         format!(
-            "failed to move tenant {} temporary directory {} into the permanent one {}",
+            "move tenant {} temporary directory {} into the permanent one {}",
             tenant_id,
             temporary_tenant_dir.display(),
             target_tenant_directory.display()
@@ -2550,14 +2564,14 @@ fn try_create_target_tenant_dir(
     })?;
     let target_dir_parent = target_tenant_directory.parent().with_context(|| {
         format!(
-            "Failed to get tenant {} dir parent for {}",
+            "get tenant {} dir parent for {}",
             tenant_id,
             target_tenant_directory.display()
         )
     })?;
     crashsafe::fsync(target_dir_parent).with_context(|| {
         format!(
-            "Failed to fsync renamed directory's parent {} for tenant {}",
+            "fsync renamed directory's parent {} for tenant {}",
             target_dir_parent.display(),
             tenant_id,
         )
@@ -2629,7 +2643,11 @@ impl Drop for Tenant {
     }
 }
 /// Dump contents of a layer file to stdout.
-pub fn dump_layerfile_from_path(path: &Path, verbose: bool) -> anyhow::Result<()> {
+pub fn dump_layerfile_from_path(
+    path: &Path,
+    verbose: bool,
+    ctx: &RequestContext,
+) -> anyhow::Result<()> {
     use std::os::unix::fs::FileExt;
 
     // All layer files start with a two-byte "magic" value, to identify the kind of
@@ -2639,8 +2657,8 @@ pub fn dump_layerfile_from_path(path: &Path, verbose: bool) -> anyhow::Result<()
     file.read_exact_at(&mut header_buf, 0)?;
 
     match u16::from_be_bytes(header_buf) {
-        crate::IMAGE_FILE_MAGIC => ImageLayer::new_for_path(path, file)?.dump(verbose)?,
-        crate::DELTA_FILE_MAGIC => DeltaLayer::new_for_path(path, file)?.dump(verbose)?,
+        crate::IMAGE_FILE_MAGIC => ImageLayer::new_for_path(path, file)?.dump(verbose, ctx)?,
+        crate::DELTA_FILE_MAGIC => DeltaLayer::new_for_path(path, file)?.dump(verbose, ctx)?,
         magic => bail!("unrecognized magic identifier: {:?}", magic),
     }
 
@@ -2966,7 +2984,7 @@ mod tests {
 
         // Branch the history, modify relation differently on the new timeline
         tenant
-            .branch_timeline(TIMELINE_ID, NEW_TIMELINE_ID, Some(Lsn(0x30)), &ctx)
+            .branch_timeline(&tline, NEW_TIMELINE_ID, Some(Lsn(0x30)), &ctx)
             .await?;
         let newtline = tenant
             .get_timeline(NEW_TIMELINE_ID, true)
@@ -3055,7 +3073,7 @@ mod tests {
 
         // try to branch at lsn 25, should fail because we already garbage collected the data
         match tenant
-            .branch_timeline(TIMELINE_ID, NEW_TIMELINE_ID, Some(Lsn(0x25)), &ctx)
+            .branch_timeline(&tline, NEW_TIMELINE_ID, Some(Lsn(0x25)), &ctx)
             .await
         {
             Ok(_) => panic!("branching should have failed"),
@@ -3078,12 +3096,13 @@ mod tests {
             TenantHarness::create("test_prohibit_branch_creation_on_pre_initdb_lsn")?
                 .load()
                 .await;
-        let tline =
-            tenant.create_empty_timeline(TIMELINE_ID, Lsn(0x50), DEFAULT_PG_VERSION, &ctx)?;
-        let _tline = tline.initialize(&ctx)?;
+
+        let tline = tenant
+            .create_empty_timeline(TIMELINE_ID, Lsn(0x50), DEFAULT_PG_VERSION, &ctx)?
+            .initialize(&ctx)?;
         // try to branch at lsn 0x25, should fail because initdb lsn is 0x50
         match tenant
-            .branch_timeline(TIMELINE_ID, NEW_TIMELINE_ID, Some(Lsn(0x25)), &ctx)
+            .branch_timeline(&tline, NEW_TIMELINE_ID, Some(Lsn(0x25)), &ctx)
             .await
         {
             Ok(_) => panic!("branching should have failed"),
@@ -3134,7 +3153,7 @@ mod tests {
         make_some_layers(tline.as_ref(), Lsn(0x20)).await?;
 
         tenant
-            .branch_timeline(TIMELINE_ID, NEW_TIMELINE_ID, Some(Lsn(0x40)), &ctx)
+            .branch_timeline(&tline, NEW_TIMELINE_ID, Some(Lsn(0x40)), &ctx)
             .await?;
         let newtline = tenant
             .get_timeline(NEW_TIMELINE_ID, true)
@@ -3158,7 +3177,7 @@ mod tests {
         make_some_layers(tline.as_ref(), Lsn(0x20)).await?;
 
         tenant
-            .branch_timeline(TIMELINE_ID, NEW_TIMELINE_ID, Some(Lsn(0x40)), &ctx)
+            .branch_timeline(&tline, NEW_TIMELINE_ID, Some(Lsn(0x40)), &ctx)
             .await?;
         let newtline = tenant
             .get_timeline(NEW_TIMELINE_ID, true)
@@ -3214,7 +3233,7 @@ mod tests {
             make_some_layers(tline.as_ref(), Lsn(0x20)).await?;
 
             tenant
-                .branch_timeline(TIMELINE_ID, NEW_TIMELINE_ID, Some(Lsn(0x40)), &ctx)
+                .branch_timeline(&tline, NEW_TIMELINE_ID, Some(Lsn(0x40)), &ctx)
                 .await?;
 
             let newtline = tenant
@@ -3499,16 +3518,14 @@ mod tests {
             keyspace.add_key(test_key);
         }
 
-        let mut tline_id = TIMELINE_ID;
         for _ in 0..50 {
             let new_tline_id = TimelineId::generate();
             tenant
-                .branch_timeline(tline_id, new_tline_id, Some(lsn), &ctx)
+                .branch_timeline(&tline, new_tline_id, Some(lsn), &ctx)
                 .await?;
             tline = tenant
                 .get_timeline(new_tline_id, true)
                 .expect("Should have the branched timeline");
-            tline_id = new_tline_id;
 
             for _ in 0..NUM_KEYS {
                 lsn = Lsn(lsn.0 + 0x10);
@@ -3565,18 +3582,16 @@ mod tests {
         let mut updated = [[Lsn(0); NUM_KEYS]; NUM_TLINES];
 
         let mut lsn = Lsn(0);
-        let mut tline_id = TIMELINE_ID;
 
         #[allow(clippy::needless_range_loop)]
         for idx in 0..NUM_TLINES {
             let new_tline_id = TimelineId::generate();
             tenant
-                .branch_timeline(tline_id, new_tline_id, Some(lsn), &ctx)
+                .branch_timeline(&tline, new_tline_id, Some(lsn), &ctx)
                 .await?;
             tline = tenant
                 .get_timeline(new_tline_id, true)
                 .expect("Should have the branched timeline");
-            tline_id = new_tline_id;
 
             for _ in 0..NUM_KEYS {
                 lsn = Lsn(lsn.0 + 0x10);

pageserver/src/tenant/layer_map.rs

@@ -46,6 +46,7 @@
 mod historic_layer_coverage;
 mod layer_coverage;
 
+use crate::context::RequestContext;
 use crate::keyspace::KeyPartitioning;
 use crate::metrics::NUM_ONDISK_LAYERS;
 use crate::repository::Key;
@@ -654,22 +655,22 @@ where
 
     /// debugging function to print out the contents of the layer map
     #[allow(unused)]
-    pub fn dump(&self, verbose: bool) -> Result<()> {
+    pub fn dump(&self, verbose: bool, ctx: &RequestContext) -> Result<()> {
         println!("Begin dump LayerMap");
 
         println!("open_layer:");
         if let Some(open_layer) = &self.open_layer {
-            open_layer.dump(verbose)?;
+            open_layer.dump(verbose, ctx)?;
         }
 
         println!("frozen_layers:");
         for frozen_layer in self.frozen_layers.iter() {
-            frozen_layer.dump(verbose)?;
+            frozen_layer.dump(verbose, ctx)?;
         }
 
         println!("historic_layers:");
         for layer in self.iter_historic_layers() {
-            layer.dump(verbose)?;
+            layer.dump(verbose, ctx)?;
         }
         println!("End dump LayerMap");
         Ok(())

pageserver/src/tenant/mgr.rs

@@ -291,10 +291,11 @@ pub async fn update_tenant_config(
     tenant_id: TenantId,
 ) -> anyhow::Result<()> {
     info!("configuring tenant {tenant_id}");
-    get_tenant(tenant_id, true)
-        .await?
-        .update_tenant_config(tenant_conf);
-    Tenant::persist_tenant_config(&conf.tenant_config_path(tenant_id), tenant_conf, false)?;
+    let tenant = get_tenant(tenant_id, true).await?;
+
+    tenant.update_tenant_config(tenant_conf);
+    let tenant_config_path = conf.tenant_config_path(tenant_id);
+    Tenant::persist_tenant_config(&tenant.tenant_id(), &tenant_config_path, tenant_conf, false)?;
     Ok(())
 }
 

pageserver/src/tenant/storage_layer.rs

@@ -1,11 +1,12 @@
 //! Common traits and structs for layers
 
-mod delta_layer;
+pub mod delta_layer;
 mod filename;
 mod image_layer;
 mod inmemory_layer;
 mod remote_layer;
 
+use crate::context::RequestContext;
 use crate::repository::{Key, Value};
 use crate::walrecord::NeonWalRecord;
 use anyhow::Result;
@@ -117,13 +118,14 @@ pub trait Layer: Send + Sync {
         key: Key,
         lsn_range: Range<Lsn>,
         reconstruct_data: &mut ValueReconstructState,
+        ctx: &RequestContext,
     ) -> Result<ValueReconstructResult>;
 
     /// A short ID string that uniquely identifies the given layer within a [`LayerMap`].
     fn short_id(&self) -> String;
 
     /// Dump summary of the contents of the layer to stdout
-    fn dump(&self, verbose: bool) -> Result<()>;
+    fn dump(&self, verbose: bool, ctx: &RequestContext) -> Result<()>;
 }
 
 /// Returned by [`Layer::iter`]
@@ -161,11 +163,11 @@ pub trait PersistentLayer: Layer {
     fn local_path(&self) -> Option<PathBuf>;
 
     /// Iterate through all keys and values stored in the layer
-    fn iter(&self) -> Result<LayerIter<'_>>;
+    fn iter(&self, ctx: &RequestContext) -> Result<LayerIter<'_>>;
 
     /// Iterate through all keys stored in the layer. Returns key, lsn and value size
     /// It is used only for compaction and so is currently implemented only for DeltaLayer
-    fn key_iter(&self) -> Result<LayerKeyIter<'_>> {
+    fn key_iter(&self, _ctx: &RequestContext) -> Result<LayerKeyIter<'_>> {
         panic!("Not implemented")
     }
 
@@ -231,6 +233,7 @@ impl Layer for LayerDescriptor {
         _key: Key,
         _lsn_range: Range<Lsn>,
         _reconstruct_data: &mut ValueReconstructState,
+        _ctx: &RequestContext,
     ) -> Result<ValueReconstructResult> {
         todo!("This method shouldn't be part of the Layer trait")
     }
@@ -239,7 +242,7 @@ impl Layer for LayerDescriptor {
         self.short_id.clone()
     }
 
-    fn dump(&self, _verbose: bool) -> Result<()> {
+    fn dump(&self, _verbose: bool, _ctx: &RequestContext) -> Result<()> {
         todo!()
     }
 }

pageserver/src/tenant/storage_layer/delta_layer.rs

@@ -24,6 +24,7 @@
 //! "values" part.
 //!
 use crate::config::PageServerConf;
+use crate::context::RequestContext;
 use crate::page_cache::{PageReadGuard, PAGE_SZ};
 use crate::repository::{Key, Value, KEY_SIZE};
 use crate::tenant::blob_io::{BlobCursor, BlobWriter, WriteBlobWriter};
@@ -62,7 +63,7 @@ use super::{DeltaFileName, Layer, LayerFileName, LayerIter, LayerKeyIter, PathOr
 /// the 'index' starts at the block indicated by 'index_start_blk'
 ///
 #[derive(Debug, Serialize, Deserialize, PartialEq, Eq)]
-struct Summary {
+pub struct Summary {
     /// Magic value to identify this as a neon delta file. Always DELTA_FILE_MAGIC.
     magic: u16,
     format_version: u16,
@@ -73,9 +74,9 @@ struct Summary {
     lsn_range: Range<Lsn>,
 
     /// Block number where the 'index' part of the file begins.
-    index_start_blk: u32,
+    pub index_start_blk: u32,
     /// Block within the 'index', where the B-tree root page is stored
-    index_root_blk: u32,
+    pub index_root_blk: u32,
 }
 
 impl From<&DeltaLayer> for Summary {
@@ -125,7 +126,7 @@ impl BlobRef {
     }
 }
 
-const DELTA_KEY_SIZE: usize = KEY_SIZE + 8;
+pub const DELTA_KEY_SIZE: usize = KEY_SIZE + 8;
 struct DeltaKey([u8; DELTA_KEY_SIZE]);
 
 ///
@@ -214,7 +215,7 @@ impl Layer for DeltaLayer {
         self.filename().file_name()
     }
     /// debugging function to print out the contents of the layer
-    fn dump(&self, verbose: bool) -> Result<()> {
+    fn dump(&self, verbose: bool, ctx: &RequestContext) -> Result<()> {
         println!(
             "----- delta layer for ten {} tli {} keys {}-{} lsn {}-{} ----",
             self.tenant_id,
@@ -229,7 +230,7 @@ impl Layer for DeltaLayer {
             return Ok(());
         }
 
-        let inner = self.load()?;
+        let inner = self.load(ctx)?;
 
         println!(
             "index_start_blk: {}, root {}",
@@ -293,6 +294,7 @@ impl Layer for DeltaLayer {
         key: Key,
         lsn_range: Range<Lsn>,
         reconstruct_state: &mut ValueReconstructState,
+        ctx: &RequestContext,
     ) -> anyhow::Result<ValueReconstructResult> {
         ensure!(lsn_range.start >= self.lsn_range.start);
         let mut need_image = true;
@@ -301,7 +303,7 @@ impl Layer for DeltaLayer {
 
         {
             // Open the file and lock the metadata in memory
-            let inner = self.load()?;
+            let inner = self.load(ctx)?;
 
             // Scan the page versions backwards, starting from `lsn`.
             let file = inner.file.as_ref().unwrap();
@@ -391,16 +393,16 @@ impl PersistentLayer for DeltaLayer {
         Some(self.path())
     }
 
-    fn iter(&self) -> Result<LayerIter<'_>> {
-        let inner = self.load().context("load delta layer")?;
+    fn iter(&self, ctx: &RequestContext) -> Result<LayerIter<'_>> {
+        let inner = self.load(ctx).context("load delta layer")?;
         Ok(match DeltaValueIter::new(inner) {
             Ok(iter) => Box::new(iter),
             Err(err) => Box::new(std::iter::once(Err(err))),
         })
     }
 
-    fn key_iter(&self) -> Result<LayerKeyIter<'_>> {
-        let inner = self.load()?;
+    fn key_iter(&self, ctx: &RequestContext) -> Result<LayerKeyIter<'_>> {
+        let inner = self.load(ctx)?;
         Ok(Box::new(
             DeltaKeyIter::new(inner).context("Layer index is corrupted")?,
         ))
@@ -459,7 +461,7 @@ impl DeltaLayer {
     /// Open the underlying file and read the metadata into memory, if it's
     /// not loaded already.
     ///
-    fn load(&self) -> Result<RwLockReadGuard<DeltaLayerInner>> {
+    fn load(&self, _ctx: &RequestContext) -> Result<RwLockReadGuard<DeltaLayerInner>> {
         loop {
             // Quick exit if already loaded
             let inner = self.inner.read().unwrap();

pageserver/src/tenant/storage_layer/image_layer.rs

@@ -20,6 +20,7 @@
 //! mapping from Key to an offset in the "values" part.  The
 //! actual page images are stored in the "values" part.
 use crate::config::PageServerConf;
+use crate::context::RequestContext;
 use crate::page_cache::PAGE_SZ;
 use crate::repository::{Key, KEY_SIZE};
 use crate::tenant::blob_io::{BlobCursor, BlobWriter, WriteBlobWriter};
@@ -143,7 +144,7 @@ impl Layer for ImageLayer {
     }
 
     /// debugging function to print out the contents of the layer
-    fn dump(&self, verbose: bool) -> Result<()> {
+    fn dump(&self, verbose: bool, ctx: &RequestContext) -> Result<()> {
         println!(
             "----- image layer for ten {} tli {} key {}-{} at {} ----",
             self.tenant_id, self.timeline_id, self.key_range.start, self.key_range.end, self.lsn
@@ -153,7 +154,7 @@ impl Layer for ImageLayer {
             return Ok(());
         }
 
-        let inner = self.load()?;
+        let inner = self.load(ctx)?;
         let file = inner.file.as_ref().unwrap();
         let tree_reader =
             DiskBtreeReader::<_, KEY_SIZE>::new(inner.index_start_blk, inner.index_root_blk, file);
@@ -174,12 +175,13 @@ impl Layer for ImageLayer {
         key: Key,
         lsn_range: Range<Lsn>,
         reconstruct_state: &mut ValueReconstructState,
+        ctx: &RequestContext,
     ) -> anyhow::Result<ValueReconstructResult> {
         assert!(self.key_range.contains(&key));
         assert!(lsn_range.start >= self.lsn);
         assert!(lsn_range.end >= self.lsn);
 
-        let inner = self.load()?;
+        let inner = self.load(ctx)?;
 
         let file = inner.file.as_ref().unwrap();
         let tree_reader = DiskBtreeReader::new(inner.index_start_blk, inner.index_root_blk, file);
@@ -220,7 +222,7 @@ impl PersistentLayer for ImageLayer {
     fn get_timeline_id(&self) -> TimelineId {
         self.timeline_id
     }
-    fn iter(&self) -> Result<LayerIter<'_>> {
+    fn iter(&self, _ctx: &RequestContext) -> Result<LayerIter<'_>> {
         unimplemented!();
     }
 
@@ -270,7 +272,7 @@ impl ImageLayer {
     /// Open the underlying file and read the metadata into memory, if it's
     /// not loaded already.
     ///
-    fn load(&self) -> Result<RwLockReadGuard<ImageLayerInner>> {
+    fn load(&self, _ctx: &RequestContext) -> Result<RwLockReadGuard<ImageLayerInner>> {
         loop {
             // Quick exit if already loaded
             let inner = self.inner.read().unwrap();

pageserver/src/tenant/storage_layer/inmemory_layer.rs

@@ -5,6 +5,7 @@
 //! its position in the file, is kept in memory, though.
 //!
 use crate::config::PageServerConf;
+use crate::context::RequestContext;
 use crate::repository::{Key, Value};
 use crate::tenant::blob_io::{BlobCursor, BlobWriter};
 use crate::tenant::block_io::BlockReader;
@@ -110,7 +111,7 @@ impl Layer for InMemoryLayer {
     }
 
     /// debugging function to print out the contents of the layer
-    fn dump(&self, verbose: bool) -> Result<()> {
+    fn dump(&self, verbose: bool, _ctx: &RequestContext) -> Result<()> {
         let inner = self.inner.read().unwrap();
 
         let end_str = inner
@@ -166,6 +167,7 @@ impl Layer for InMemoryLayer {
         key: Key,
         lsn_range: Range<Lsn>,
         reconstruct_state: &mut ValueReconstructState,
+        _ctx: &RequestContext,
     ) -> anyhow::Result<ValueReconstructResult> {
         ensure!(lsn_range.start >= self.start_lsn);
         let mut need_image = true;

pageserver/src/tenant/storage_layer/remote_layer.rs

@@ -2,6 +2,7 @@
 //! in remote storage.
 //!
 use crate::config::PageServerConf;
+use crate::context::RequestContext;
 use crate::repository::Key;
 use crate::tenant::remote_timeline_client::index::LayerFileMetadata;
 use crate::tenant::storage_layer::{Layer, ValueReconstructResult, ValueReconstructState};
@@ -51,6 +52,7 @@ impl Layer for RemoteLayer {
         _key: Key,
         _lsn_range: Range<Lsn>,
         _reconstruct_state: &mut ValueReconstructState,
+        _ctx: &RequestContext,
     ) -> Result<ValueReconstructResult> {
         bail!(
             "layer {} needs to be downloaded",
@@ -63,7 +65,7 @@ impl Layer for RemoteLayer {
     }
 
     /// debugging function to print out the contents of the layer
-    fn dump(&self, _verbose: bool) -> Result<()> {
+    fn dump(&self, _verbose: bool, _ctx: &RequestContext) -> Result<()> {
         println!(
             "----- remote layer for ten {} tli {} keys {}-{} lsn {}-{} ----",
             self.tenantid,
@@ -111,11 +113,11 @@ impl PersistentLayer for RemoteLayer {
         None
     }
 
-    fn iter(&self) -> Result<LayerIter<'_>> {
+    fn iter(&self, _ctx: &RequestContext) -> Result<LayerIter<'_>> {
         bail!("cannot iterate a remote layer");
     }
 
-    fn key_iter(&self) -> Result<LayerKeyIter<'_>> {
+    fn key_iter(&self, _ctx: &RequestContext) -> Result<LayerKeyIter<'_>> {
         bail!("cannot iterate a remote layer");
     }
 

pageserver/src/tenant/timeline.rs

@@ -1,5 +1,7 @@
 //!
 
+mod walreceiver;
+
 use anyhow::{anyhow, bail, ensure, Context};
 use bytes::Bytes;
 use fail::fail_point;
@@ -13,6 +15,7 @@ use pageserver_api::models::{
 use tokio::sync::{oneshot, watch, Semaphore, TryAcquireError};
 use tokio_util::sync::CancellationToken;
 use tracing::*;
+use utils::id::TenantTimelineId;
 
 use std::cmp::{max, min, Ordering};
 use std::collections::HashMap;
@@ -23,6 +26,7 @@ use std::sync::atomic::{AtomicI64, Ordering as AtomicOrdering};
 use std::sync::{Arc, Mutex, MutexGuard, RwLock, Weak};
 use std::time::{Duration, Instant, SystemTime};
 
+use crate::broker_client::is_broker_client_initialized;
 use crate::context::{DownloadBehavior, RequestContext};
 use crate::tenant::remote_timeline_client::{self, index::LayerFileMetadata};
 use crate::tenant::storage_layer::{
@@ -59,11 +63,11 @@ use crate::page_cache;
 use crate::repository::GcResult;
 use crate::repository::{Key, Value};
 use crate::task_mgr::TaskKind;
-use crate::walreceiver::{is_broker_client_initialized, spawn_connection_manager_task};
 use crate::walredo::WalRedoManager;
 use crate::METADATA_FILE_NAME;
 use crate::ZERO_PAGE;
 use crate::{is_temporary, task_mgr};
+use walreceiver::spawn_connection_manager_task;
 
 use super::remote_timeline_client::index::IndexPart;
 use super::remote_timeline_client::RemoteTimelineClient;
@@ -129,7 +133,6 @@ pub struct Timeline {
     ancestor_timeline: Option<Arc<Timeline>>,
     ancestor_lsn: Lsn,
 
-    // Metrics
     metrics: TimelineMetrics,
 
     /// Ensures layers aren't frozen by checkpointer between
@@ -379,7 +382,7 @@ pub enum PageReconstructError {
     Other(#[from] anyhow::Error), // source and Display delegate to anyhow::Error
 
     /// The operation would require downloading a layer that is missing locally.
-    NeedsDownload(Weak<Timeline>, Weak<RemoteLayer>),
+    NeedsDownload(TenantTimelineId, LayerFileName),
 
     /// The operation was cancelled
     Cancelled,
@@ -393,7 +396,14 @@ impl std::fmt::Debug for PageReconstructError {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> Result<(), std::fmt::Error> {
         match self {
             Self::Other(err) => err.fmt(f),
-            Self::NeedsDownload(_tli, _layer) => write!(f, "needs download"),
+            Self::NeedsDownload(tenant_timeline_id, layer_file_name) => {
+                write!(
+                    f,
+                    "layer {}/{} needs download",
+                    tenant_timeline_id,
+                    layer_file_name.file_name()
+                )
+            }
             Self::Cancelled => write!(f, "cancelled"),
             Self::WalRedo(err) => err.fmt(f),
         }
@@ -404,7 +414,14 @@ impl std::fmt::Display for PageReconstructError {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> Result<(), std::fmt::Error> {
         match self {
             Self::Other(err) => err.fmt(f),
-            Self::NeedsDownload(_tli, _layer) => write!(f, "needs download"),
+            Self::NeedsDownload(tenant_timeline_id, layer_file_name) => {
+                write!(
+                    f,
+                    "layer {}/{} needs download",
+                    tenant_timeline_id,
+                    layer_file_name.file_name()
+                )
+            }
             Self::Cancelled => write!(f, "cancelled"),
             Self::WalRedo(err) => err.fmt(f),
         }
@@ -679,7 +696,7 @@ impl Timeline {
 
                 // 3. Compact
                 let timer = self.metrics.compact_time_histo.start_timer();
-                self.compact_level0(target_file_size).await?;
+                self.compact_level0(target_file_size, ctx).await?;
                 timer.stop_and_record();
 
                 // If `create_image_layers' or `compact_level0` scheduled any
@@ -942,6 +959,10 @@ impl Timeline {
             };
             result.repartition_threshold = result.get_checkpoint_distance() / 10;
             result
+                .metrics
+                .last_record_gauge
+                .set(disk_consistent_lsn.0 as i64);
+            result
         })
     }
 
@@ -1746,6 +1767,7 @@ impl Timeline {
                                 key,
                                 lsn_floor..cont_lsn,
                                 reconstruct_state,
+                                ctx,
                             ) {
                                 Ok(result) => result,
                                 Err(e) => return Err(PageReconstructError::from(e)),
@@ -1771,6 +1793,7 @@ impl Timeline {
                                 key,
                                 lsn_floor..cont_lsn,
                                 reconstruct_state,
+                                ctx,
                             ) {
                                 Ok(result) => result,
                                 Err(e) => return Err(PageReconstructError::from(e)),
@@ -1804,6 +1827,7 @@ impl Timeline {
                                 key,
                                 lsn_floor..cont_lsn,
                                 reconstruct_state,
+                                ctx,
                             ) {
                                 Ok(result) => result,
                                 Err(e) => return Err(PageReconstructError::from(e)),
@@ -1872,8 +1896,8 @@ impl Timeline {
                     }
                     (DownloadBehavior::Error, false) => {
                         return Err(PageReconstructError::NeedsDownload(
-                            timeline.myself.clone(),
-                            Arc::downgrade(&remote_layer),
+                            TenantTimelineId::new(self.tenant_id, self.timeline_id),
+                            remote_layer.file_name.clone(),
                         ))
                     }
                 }
@@ -2442,6 +2466,7 @@ impl Timeline {
     async fn compact_level0_phase1(
         &self,
         target_file_size: u64,
+        ctx: &RequestContext,
     ) -> anyhow::Result<CompactLevel0Phase1Result> {
         let layers = self.layers.read().unwrap();
         let mut level0_deltas = layers.get_level0_deltas()?;
@@ -2501,8 +2526,9 @@ impl Timeline {
 
         // This iterator walks through all key-value pairs from all the layers
         // we're compacting, in key, LSN order.
-        let all_values_iter =
-            itertools::process_results(deltas_to_compact.iter().map(|l| l.iter()), |iter_iter| {
+        let all_values_iter = itertools::process_results(
+            deltas_to_compact.iter().map(|l| l.iter(ctx)),
+            |iter_iter| {
                 iter_iter.kmerge_by(|a, b| {
                     if let Ok((a_key, a_lsn, _)) = a {
                         if let Ok((b_key, b_lsn, _)) = b {
@@ -2518,11 +2544,12 @@ impl Timeline {
                         true
                     }
                 })
-            })?;
+            },
+        )?;
 
         // This iterator walks through all keys and is needed to calculate size used by each key
         let mut all_keys_iter = itertools::process_results(
-            deltas_to_compact.iter().map(|l| l.key_iter()),
+            deltas_to_compact.iter().map(|l| l.key_iter(ctx)),
             |iter_iter| {
                 iter_iter.kmerge_by(|a, b| {
                     let (a_key, a_lsn, _) = a;
@@ -2698,11 +2725,15 @@ impl Timeline {
     /// Collect a bunch of Level 0 layer files, and compact and reshuffle them as
     /// as Level 1 files.
     ///
-    async fn compact_level0(&self, target_file_size: u64) -> anyhow::Result<()> {
+    async fn compact_level0(
+        &self,
+        target_file_size: u64,
+        ctx: &RequestContext,
+    ) -> anyhow::Result<()> {
         let CompactLevel0Phase1Result {
             new_layers,
             deltas_to_compact,
-        } = self.compact_level0_phase1(target_file_size).await?;
+        } = self.compact_level0_phase1(target_file_size, ctx).await?;
 
         if new_layers.is_empty() && deltas_to_compact.is_empty() {
             // nothing to do
@@ -2872,6 +2903,8 @@ impl Timeline {
     /// obsolete.
     ///
     pub(super) async fn gc(&self) -> anyhow::Result<GcResult> {
+        let timer = self.metrics.garbage_collect_histo.start_timer();
+
         fail_point!("before-timeline-gc");
 
         let _layer_removal_cs = self.layer_removal_cs.lock().await;
@@ -2892,11 +2925,17 @@ impl Timeline {
 
         let new_gc_cutoff = Lsn::min(horizon_cutoff, pitr_cutoff);
 
-        self.gc_timeline(horizon_cutoff, pitr_cutoff, retain_lsns, new_gc_cutoff)
+        let res = self
+            .gc_timeline(horizon_cutoff, pitr_cutoff, retain_lsns, new_gc_cutoff)
             .instrument(
                 info_span!("gc_timeline", timeline = %self.timeline_id, cutoff = %new_gc_cutoff),
             )
-            .await
+            .await?;
+
+        // only record successes
+        timer.stop_and_record();
+
+        Ok(res)
     }
 
     async fn gc_timeline(

pageserver/src/tenant/timeline/walreceiver.rs

@@ -23,58 +23,15 @@
 mod connection_manager;
 mod walreceiver_connection;
 
-use crate::config::PageServerConf;
 use crate::task_mgr::WALRECEIVER_RUNTIME;
 
-use anyhow::Context;
-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;
 
-static BROKER_CLIENT: OnceCell<BrokerClientChannel> = OnceCell::new();
-
-///
-/// Initialize the broker client. This must be called once at page server startup.
-///
-pub async fn init_broker_client(conf: &'static PageServerConf) -> anyhow::Result<()> {
-    let broker_endpoint = conf.broker_endpoint.clone();
-
-    // Note: we do not attempt connecting here (but validate endpoints sanity).
-    let broker_client =
-        storage_broker::connect(broker_endpoint.clone(), conf.broker_keepalive_interval).context(
-            format!(
-                "Failed to create broker client to {}",
-                &conf.broker_endpoint
-            ),
-        )?;
-
-    if BROKER_CLIENT.set(broker_client).is_err() {
-        panic!("broker already initialized");
-    }
-
-    info!(
-        "Initialized broker client with endpoints: {}",
-        broker_endpoint
-    );
-    Ok(())
-}
-
-///
-/// Get a handle to the broker client
-///
-pub fn get_broker_client() -> &'static BrokerClientChannel {
-    BROKER_CLIENT.get().expect("broker client not initialized")
-}
-
-pub fn is_broker_client_initialized() -> bool {
-    BROKER_CLIENT.get().is_some()
-}
-
 /// 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 token that it can listen to for earlier interrupts.
@@ -95,7 +52,6 @@ pub enum TaskEvent<E> {
 
 #[derive(Debug, Clone)]
 pub enum TaskStateUpdate<E> {
-    Init,
     Started,
     Progress(E),
 }

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

@@ -11,11 +11,12 @@
 
 use std::{collections::HashMap, num::NonZeroU64, ops::ControlFlow, sync::Arc, time::Duration};
 
+use super::TaskStateUpdate;
+use crate::broker_client::get_broker_client;
 use crate::context::RequestContext;
-use crate::task_mgr::TaskKind;
 use crate::task_mgr::WALRECEIVER_RUNTIME;
+use crate::task_mgr::{self, TaskKind};
 use crate::tenant::Timeline;
-use crate::{task_mgr, walreceiver::TaskStateUpdate};
 use anyhow::Context;
 use chrono::{NaiveDateTime, Utc};
 use pageserver_api::models::TimelineState;
@@ -28,10 +29,7 @@ use storage_broker::Streaming;
 use tokio::{select, sync::watch};
 use tracing::*;
 
-use crate::{
-    exponential_backoff, walreceiver::get_broker_client, DEFAULT_BASE_BACKOFF_SECONDS,
-    DEFAULT_MAX_BACKOFF_SECONDS,
-};
+use crate::{exponential_backoff, DEFAULT_BASE_BACKOFF_SECONDS, DEFAULT_MAX_BACKOFF_SECONDS};
 use postgres_connection::{parse_host_port, PgConnectionConfig};
 use utils::{
     id::{NodeId, TenantTimelineId},
@@ -149,7 +147,7 @@ async fn connection_manager_loop_step(
                 let wal_connection = walreceiver_state.wal_connection.as_mut()
                     .expect("Should have a connection, as checked by the corresponding select! guard");
                 match wal_connection_update {
-                    TaskEvent::Update(TaskStateUpdate::Init | TaskStateUpdate::Started) => {},
+                    TaskEvent::Update(TaskStateUpdate::Started) => {},
                     TaskEvent::Update(TaskStateUpdate::Progress(new_status)) => {
                         if new_status.has_processed_wal {
                             // We have advanced last_record_lsn by processing the WAL received

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

@@ -22,8 +22,9 @@ use tokio_postgres::{replication::ReplicationStream, Client};
 use tokio_util::sync::CancellationToken;
 use tracing::{debug, error, info, trace, warn};
 
+use super::TaskStateUpdate;
 use crate::context::RequestContext;
-use crate::{metrics::LIVE_CONNECTIONS_COUNT, walreceiver::TaskStateUpdate};
+use crate::metrics::LIVE_CONNECTIONS_COUNT;
 use crate::{
     task_mgr,
     task_mgr::TaskKind,

pageserver/src/walredo.rs

@@ -22,16 +22,18 @@ use byteorder::{ByteOrder, LittleEndian};
 use bytes::{BufMut, Bytes, BytesMut};
 use nix::poll::*;
 use serde::Serialize;
+use std::collections::VecDeque;
 use std::fs::OpenOptions;
 use std::io::prelude::*;
 use std::io::{Error, ErrorKind};
 use std::ops::{Deref, DerefMut};
+use std::os::fd::RawFd;
 use std::os::unix::io::AsRawFd;
 use std::os::unix::prelude::CommandExt;
 use std::path::PathBuf;
 use std::process::Stdio;
 use std::process::{Child, ChildStderr, ChildStdin, ChildStdout, Command};
-use std::sync::Mutex;
+use std::sync::{Mutex, MutexGuard};
 use std::time::Duration;
 use std::time::Instant;
 use std::{fs, io};
@@ -90,6 +92,20 @@ pub trait WalRedoManager: Send + Sync {
     ) -> Result<Bytes, WalRedoError>;
 }
 
+struct ProcessInput {
+    child: NoLeakChild,
+    stdin: ChildStdin,
+    stderr_fd: RawFd,
+    stdout_fd: RawFd,
+    n_requests: usize,
+}
+
+struct ProcessOutput {
+    stdout: ChildStdout,
+    pending_responses: VecDeque<Option<Bytes>>,
+    n_processed_responses: usize,
+}
+
 ///
 /// This is the real implementation that uses a Postgres process to
 /// perform WAL replay. Only one thread can use the process at a time,
@@ -101,7 +117,9 @@ pub struct PostgresRedoManager {
     tenant_id: TenantId,
     conf: &'static PageServerConf,
 
-    process: Mutex<Option<PostgresRedoProcess>>,
+    stdout: Mutex<Option<ProcessOutput>>,
+    stdin: Mutex<Option<ProcessInput>>,
+    stderr: Mutex<Option<ChildStderr>>,
 }
 
 /// Can this request be served by neon redo functions
@@ -209,16 +227,17 @@ impl PostgresRedoManager {
         PostgresRedoManager {
             tenant_id,
             conf,
-            process: Mutex::new(None),
+            stdin: Mutex::new(None),
+            stdout: Mutex::new(None),
+            stderr: Mutex::new(None),
         }
     }
 
     /// Launch process pre-emptively. Should not be needed except for benchmarking.
-    pub fn launch_process(&mut self, pg_version: u32) -> anyhow::Result<()> {
-        let inner = self.process.get_mut().unwrap();
-        if inner.is_none() {
-            let p = PostgresRedoProcess::launch(self.conf, self.tenant_id, pg_version)?;
-            *inner = Some(p);
+    pub fn launch_process(&self, pg_version: u32) -> anyhow::Result<()> {
+        let mut proc = self.stdin.lock().unwrap();
+        if proc.is_none() {
+            self.launch(&mut proc, pg_version)?;
         }
         Ok(())
     }
@@ -241,22 +260,19 @@ impl PostgresRedoManager {
 
         let start_time = Instant::now();
 
-        let mut process_guard = self.process.lock().unwrap();
+        let mut proc = self.stdin.lock().unwrap();
         let lock_time = Instant::now();
 
         // launch the WAL redo process on first use
-        if process_guard.is_none() {
-            let p = PostgresRedoProcess::launch(self.conf, self.tenant_id, pg_version)?;
-            *process_guard = Some(p);
+        if proc.is_none() {
+            self.launch(&mut proc, pg_version)?;
         }
-        let process = process_guard.as_mut().unwrap();
-
         WAL_REDO_WAIT_TIME.observe(lock_time.duration_since(start_time).as_secs_f64());
 
         // Relational WAL records are applied using wal-redo-postgres
         let buf_tag = BufferTag { rel, blknum };
-        let result = process
-            .apply_wal_records(buf_tag, base_img, records, wal_redo_timeout)
+        let result = self
+            .apply_wal_records(proc, buf_tag, base_img, records, wal_redo_timeout)
             .map_err(WalRedoError::IoError);
 
         let end_time = Instant::now();
@@ -295,8 +311,22 @@ impl PostgresRedoManager {
 				base_img_lsn,
                 lsn
             );
-            let process = process_guard.take().unwrap();
-            process.kill();
+            // self.stdin only holds stdin & stderr as_raw_fd().
+            // Dropping it as part of take() doesn't close them.
+            // The owning objects (ChildStdout and ChildStderr) are stored in
+            // self.stdout and self.stderr, respsectively.
+            // We intentionally keep them open here to avoid a race between
+            // currently running `apply_wal_records()` and a `launch()` call
+            // after we return here.
+            // The currently running `apply_wal_records()` must not read from
+            // the newly launched process.
+            // By keeping self.stdout and self.stderr open here, `launch()` will
+            // get other file descriptors for the new child's stdout and stderr,
+            // and hence the current `apply_wal_records()` calls will observe
+            //  `output.stdout.as_raw_fd() != stdout_fd` .
+            if let Some(proc) = self.stdin.lock().unwrap().take() {
+                proc.child.kill_and_wait();
+            }
         }
         result
     }
@@ -595,32 +625,23 @@ impl<C: CommandExt> CloseFileDescriptors for C {
     }
 }
 
-///
-/// Handle to the Postgres WAL redo process
-///
-struct PostgresRedoProcess {
-    tenant_id: TenantId,
-    child: NoLeakChild,
-    stdin: ChildStdin,
-    stdout: ChildStdout,
-    stderr: ChildStderr,
-}
-
-impl PostgresRedoProcess {
+impl PostgresRedoManager {
     //
     // Start postgres binary in special WAL redo mode.
     //
-    #[instrument(skip_all,fields(tenant_id=%tenant_id, pg_version=pg_version))]
+    #[instrument(skip_all,fields(tenant_id=%self.tenant_id, pg_version=pg_version))]
     fn launch(
-        conf: &PageServerConf,
-        tenant_id: TenantId,
+        &self,
+        input: &mut MutexGuard<Option<ProcessInput>>,
         pg_version: u32,
-    ) -> Result<PostgresRedoProcess, Error> {
+    ) -> Result<(), Error> {
         // FIXME: We need a dummy Postgres cluster to run the process in. Currently, we
         // just create one with constant name. That fails if you try to launch more than
         // one WAL redo manager concurrently.
         let datadir = path_with_suffix_extension(
-            conf.tenant_path(&tenant_id).join("wal-redo-datadir"),
+            self.conf
+                .tenant_path(&self.tenant_id)
+                .join("wal-redo-datadir"),
             TEMP_FILE_SUFFIX,
         );
 
@@ -634,10 +655,12 @@ impl PostgresRedoProcess {
                 )
             })?;
         }
-        let pg_bin_dir_path = conf
+        let pg_bin_dir_path = self
+            .conf
             .pg_bin_dir(pg_version)
             .map_err(|e| Error::new(ErrorKind::Other, format!("incorrect pg_bin_dir path: {e}")))?;
-        let pg_lib_dir_path = conf
+        let pg_lib_dir_path = self
+            .conf
             .pg_lib_dir(pg_version)
             .map_err(|e| Error::new(ErrorKind::Other, format!("incorrect pg_lib_dir path: {e}")))?;
 
@@ -723,27 +746,31 @@ impl PostgresRedoProcess {
         // all fallible operations post-spawn are complete, so get rid of the guard
         let child = scopeguard::ScopeGuard::into_inner(child);
 
-        Ok(PostgresRedoProcess {
-            tenant_id,
+        **input = Some(ProcessInput {
             child,
+            stdout_fd: stdout.as_raw_fd(),
+            stderr_fd: stderr.as_raw_fd(),
             stdin,
+            n_requests: 0,
+        });
+
+        *self.stdout.lock().unwrap() = Some(ProcessOutput {
             stdout,
-            stderr,
-        })
+            pending_responses: VecDeque::new(),
+            n_processed_responses: 0,
+        });
+        *self.stderr.lock().unwrap() = Some(stderr);
+
+        Ok(())
     }
 
-    #[instrument(skip_all, fields(tenant_id=%self.tenant_id, pid=%self.child.id()))]
-    fn kill(self) {
-        self.child.kill_and_wait();
-    }
-
-    //
     // Apply given WAL records ('records') over an old page image. Returns
     // new page image.
     //
-    #[instrument(skip_all, fields(tenant_id=%self.tenant_id, pid=%self.child.id()))]
+    #[instrument(skip_all, fields(tenant_id=%self.tenant_id, pid=%input.as_ref().unwrap().child.id()))]
     fn apply_wal_records(
-        &mut self,
+        &self,
+        mut input: MutexGuard<Option<ProcessInput>>,
         tag: BufferTag,
         base_img: Option<Bytes>,
         records: &[(Lsn, NeonWalRecord)],
@@ -780,33 +807,23 @@ impl PostgresRedoProcess {
         build_get_page_msg(tag, &mut writebuf);
         WAL_REDO_RECORD_COUNTER.inc_by(records.len() as u64);
 
-        // The input is now in 'writebuf'. Do a blind write first, writing as much as
-        // we can, before calling poll(). That skips one call to poll() if the stdin is
-        // already available for writing, which it almost certainly is because the
-        // process is idle.
-        let mut nwrite = self.stdin.write(&writebuf)?;
-
-        // We expect the WAL redo process to respond with an 8k page image. We read it
-        // into this buffer.
-        let mut resultbuf = vec![0; BLCKSZ.into()];
-        let mut nresult: usize = 0; // # of bytes read into 'resultbuf' so far
+        let proc = input.as_mut().unwrap();
+        let mut nwrite = 0usize;
+        let stdout_fd = proc.stdout_fd;
 
         // Prepare for calling poll()
         let mut pollfds = [
-            PollFd::new(self.stdout.as_raw_fd(), PollFlags::POLLIN),
-            PollFd::new(self.stderr.as_raw_fd(), PollFlags::POLLIN),
-            PollFd::new(self.stdin.as_raw_fd(), PollFlags::POLLOUT),
+            PollFd::new(proc.stdin.as_raw_fd(), PollFlags::POLLOUT),
+            PollFd::new(proc.stderr_fd, PollFlags::POLLIN),
+            PollFd::new(stdout_fd, PollFlags::POLLIN),
         ];
 
-        // We do three things simultaneously: send the old base image and WAL records to
-        // the child process's stdin, read the result from child's stdout, and forward any logging
+        // We do two things simultaneously: send the old base image and WAL records to
+        // the child process's stdin and forward any logging
         // information that the child writes to its stderr to the page server's log.
-        while nresult < BLCKSZ.into() {
-            // If we have more data to write, wake up if 'stdin' becomes writeable or
-            // we have data to read. Otherwise only wake up if there's data to read.
-            let nfds = if nwrite < writebuf.len() { 3 } else { 2 };
+        while nwrite < writebuf.len() {
             let n = loop {
-                match nix::poll::poll(&mut pollfds[0..nfds], wal_redo_timeout.as_millis() as i32) {
+                match nix::poll::poll(&mut pollfds[0..2], wal_redo_timeout.as_millis() as i32) {
                     Err(e) if e == nix::errno::Errno::EINTR => continue,
                     res => break res,
                 }
@@ -820,14 +837,16 @@ impl PostgresRedoProcess {
             let err_revents = pollfds[1].revents().unwrap();
             if err_revents & (PollFlags::POLLERR | PollFlags::POLLIN) != PollFlags::empty() {
                 let mut errbuf: [u8; 16384] = [0; 16384];
-                let n = self.stderr.read(&mut errbuf)?;
+                let mut stderr_guard = self.stderr.lock().unwrap();
+                let stderr = stderr_guard.as_mut().unwrap();
+                let len = stderr.read(&mut errbuf)?;
 
                 // The message might not be split correctly into lines here. But this is
                 // good enough, the important thing is to get the message to the log.
-                if n > 0 {
+                if len > 0 {
                     error!(
                         "wal-redo-postgres: {}",
-                        String::from_utf8_lossy(&errbuf[0..n])
+                        String::from_utf8_lossy(&errbuf[0..len])
                     );
 
                     // To make sure we capture all log from the process if it fails, keep
@@ -841,33 +860,157 @@ impl PostgresRedoProcess {
                 ));
             }
 
-            // If we have more data to write and 'stdin' is writeable, do write.
-            if nwrite < writebuf.len() {
-                let in_revents = pollfds[2].revents().unwrap();
-                if in_revents & (PollFlags::POLLERR | PollFlags::POLLOUT) != PollFlags::empty() {
-                    nwrite += self.stdin.write(&writebuf[nwrite..])?;
-                } else if in_revents.contains(PollFlags::POLLHUP) {
-                    // We still have more data to write, but the process closed the pipe.
-                    return Err(Error::new(
-                        ErrorKind::BrokenPipe,
-                        "WAL redo process closed its stdin unexpectedly",
-                    ));
-                }
-            }
-
-            // If we have some data in stdout, read it to the result buffer.
-            let out_revents = pollfds[0].revents().unwrap();
-            if out_revents & (PollFlags::POLLERR | PollFlags::POLLIN) != PollFlags::empty() {
-                nresult += self.stdout.read(&mut resultbuf[nresult..])?;
-            } else if out_revents.contains(PollFlags::POLLHUP) {
+            // If 'stdin' is writeable, do write.
+            let in_revents = pollfds[0].revents().unwrap();
+            if in_revents & (PollFlags::POLLERR | PollFlags::POLLOUT) != PollFlags::empty() {
+                nwrite += proc.stdin.write(&writebuf[nwrite..])?;
+            } else if in_revents.contains(PollFlags::POLLHUP) {
+                // We still have more data to write, but the process closed the pipe.
                 return Err(Error::new(
                     ErrorKind::BrokenPipe,
-                    "WAL redo process closed its stdout unexpectedly",
+                    "WAL redo process closed its stdin unexpectedly",
                 ));
             }
         }
+        let request_no = proc.n_requests;
+        proc.n_requests += 1;
+        drop(input);
 
-        Ok(Bytes::from(resultbuf))
+        // To improve walredo performance we separate sending requests and receiving
+        // responses. Them are protected by different mutexes (output and input).
+        // If thread T1, T2, T3 send requests D1, D2, D3 to walredo process
+        // then there is not warranty that T1 will first granted output mutex lock.
+        // To address this issue we maintain number of sent requests, number of processed
+        // responses and ring buffer with pending responses. After sending response
+        // (under input mutex), threads remembers request number. Then it releases
+        // input mutex, locks output mutex and fetch in ring buffer all responses until
+        // its stored request number. The it takes correspondent element from
+        // pending responses ring buffer and truncate all empty elements from the front,
+        // advancing processed responses number.
+
+        let mut output_guard = self.stdout.lock().unwrap();
+        let output = output_guard.as_mut().unwrap();
+        if output.stdout.as_raw_fd() != stdout_fd {
+            // If stdout file descriptor is changed then it means that walredo process is crashed and restarted.
+            // As far as ProcessInput and ProcessOutout are protected by different mutexes,
+            // it can happen that we send request to one process and waiting response from another.
+            // To prevent such situation we compare stdout file descriptors.
+            // As far as old stdout pipe is destroyed only after new one is created,
+            // it can not reuse the same file descriptor, so this check is safe.
+            //
+            // Cross-read this with the comment in apply_batch_postgres if result.is_err().
+            // That's where we kill the child process.
+            return Err(Error::new(
+                ErrorKind::BrokenPipe,
+                "WAL redo process closed its stdout unexpectedly",
+            ));
+        }
+        let n_processed_responses = output.n_processed_responses;
+        while n_processed_responses + output.pending_responses.len() <= request_no {
+            // We expect the WAL redo process to respond with an 8k page image. We read it
+            // into this buffer.
+            let mut resultbuf = vec![0; BLCKSZ.into()];
+            let mut nresult: usize = 0; // # of bytes read into 'resultbuf' so far
+            while nresult < BLCKSZ.into() {
+                // We do two things simultaneously: reading response from stdout
+                // and forward any logging information that the child writes to its stderr to the page server's log.
+                let n = loop {
+                    match nix::poll::poll(&mut pollfds[1..3], wal_redo_timeout.as_millis() as i32) {
+                        Err(e) if e == nix::errno::Errno::EINTR => continue,
+                        res => break res,
+                    }
+                }?;
+
+                if n == 0 {
+                    return Err(Error::new(ErrorKind::Other, "WAL redo timed out"));
+                }
+
+                // If we have some messages in stderr, forward them to the log.
+                let err_revents = pollfds[1].revents().unwrap();
+                if err_revents & (PollFlags::POLLERR | PollFlags::POLLIN) != PollFlags::empty() {
+                    let mut errbuf: [u8; 16384] = [0; 16384];
+                    let mut stderr_guard = self.stderr.lock().unwrap();
+                    let stderr = stderr_guard.as_mut().unwrap();
+                    let len = stderr.read(&mut errbuf)?;
+
+                    // The message might not be split correctly into lines here. But this is
+                    // good enough, the important thing is to get the message to the log.
+                    if len > 0 {
+                        error!(
+                            "wal-redo-postgres: {}",
+                            String::from_utf8_lossy(&errbuf[0..len])
+                        );
+
+                        // To make sure we capture all log from the process if it fails, keep
+                        // reading from the stderr, before checking the stdout.
+                        continue;
+                    }
+                } else if err_revents.contains(PollFlags::POLLHUP) {
+                    return Err(Error::new(
+                        ErrorKind::BrokenPipe,
+                        "WAL redo process closed its stderr unexpectedly",
+                    ));
+                }
+
+                // If we have some data in stdout, read it to the result buffer.
+                let out_revents = pollfds[2].revents().unwrap();
+                if out_revents & (PollFlags::POLLERR | PollFlags::POLLIN) != PollFlags::empty() {
+                    nresult += output.stdout.read(&mut resultbuf[nresult..])?;
+                } else if out_revents.contains(PollFlags::POLLHUP) {
+                    return Err(Error::new(
+                        ErrorKind::BrokenPipe,
+                        "WAL redo process closed its stdout unexpectedly",
+                    ));
+                }
+            }
+            output
+                .pending_responses
+                .push_back(Some(Bytes::from(resultbuf)));
+        }
+        // Replace our request's response with None in `pending_responses`.
+        // Then make space in the ring buffer by clearing out any seqence of contiguous
+        // `None`'s from the front of `pending_responses`.
+        // NB: We can't pop_front() because other requests' responses because another
+        // requester might have grabbed the output mutex before us:
+        // T1: grab input mutex
+        // T1: send request_no 23
+        // T1: release input mutex
+        // T2: grab input mutex
+        // T2: send request_no 24
+        // T2: release input mutex
+        // T2: grab output mutex
+        // T2: n_processed_responses + output.pending_responses.len() <= request_no
+        //            23                                0                   24
+        // T2: enters poll loop that reads stdout
+        // T2: put response for 23 into pending_responses
+        // T2: put response for 24 into pending_resposnes
+        // pending_responses now looks like this: Front Some(response_23) Some(response_24) Back
+        // T2: takes its response_24
+        // pending_responses now looks like this: Front Some(response_23) None Back
+        // T2: does the while loop below
+        // pending_responses now looks like this: Front Some(response_23) None Back
+        // T2: releases output mutex
+        // T1: grabs output mutex
+        // T1: n_processed_responses + output.pending_responses.len() > request_no
+        //            23                                2                   23
+        // T1: skips poll loop that reads stdout
+        // T1: takes its response_23
+        // pending_responses now looks like this: Front None None Back
+        // T2: does the while loop below
+        // pending_responses now looks like this: Front Back
+        // n_processed_responses now has value 25
+        let res = output.pending_responses[request_no - n_processed_responses]
+            .take()
+            .expect("we own this request_no, nobody else is supposed to take it");
+        while let Some(front) = output.pending_responses.front() {
+            if front.is_none() {
+                output.pending_responses.pop_front();
+                output.n_processed_responses += 1;
+            } else {
+                break;
+            }
+        }
+        Ok(res)
     }
 }
 

proxy/Cargo.toml

@@ -6,58 +6,59 @@ license.workspace = true
 
 [dependencies]
 anyhow.workspace = true
+async-trait.workspace = true
 atty.workspace = true
 base64.workspace = true
 bstr.workspace = true
-bytes = {workspace = true, features = ['serde'] }
-clap.workspace = true
+bytes = { workspace = true, features = ["serde"] }
 chrono.workspace = true
+clap.workspace = true
 consumption_metrics.workspace = true
 futures.workspace = true
 git-version.workspace = true
 hashbrown.workspace = true
+hashlink.workspace = true
 hex.workspace = true
 hmac.workspace = true
-hyper.workspace = true
+hostname.workspace = true
+humantime.workspace = true
 hyper-tungstenite.workspace = true
+hyper.workspace = true
 itertools.workspace = true
 md5.workspace = true
+metrics.workspace = true
 once_cell.workspace = true
 parking_lot.workspace = true
 pin-project-lite.workspace = true
+pq_proto.workspace = true
+prometheus.workspace = true
 rand.workspace = true
 regex.workspace = true
-reqwest = { workspace = true, features = [ "json" ] }
+reqwest = { workspace = true, features = ["json"] }
 routerify.workspace = true
-rustls.workspace = true
 rustls-pemfile.workspace = true
+rustls.workspace = true
 scopeguard.workspace = true
 serde.workspace = true
 serde_json.workspace = true
 sha2.workspace = true
 socket2.workspace = true
 thiserror.workspace = true
-tokio.workspace = true
+tls-listener.workspace = true
 tokio-postgres.workspace = true
 tokio-rustls.workspace = true
-tls-listener.workspace = true
-tracing.workspace = true
+tokio.workspace = true
 tracing-subscriber.workspace = true
+tracing.workspace = true
 url.workspace = true
+utils.workspace = true
 uuid.workspace = true
 webpki-roots.workspace = true
 x509-parser.workspace = true
-metrics.workspace = true
-pq_proto.workspace = true
-utils.workspace = true
-prometheus.workspace = true
-humantime.workspace = true
-hostname.workspace = true
 
 workspace_hack.workspace = true
 
 [dev-dependencies]
-async-trait.workspace = true
 rcgen.workspace = true
 rstest.workspace = true
 tokio-postgres-rustls.workspace = true

proxy/src/auth.rs

@@ -1,7 +1,7 @@
 //! Client authentication mechanisms.
 
 pub mod backend;
-pub use backend::{BackendType, ConsoleReqExtra};
+pub use backend::BackendType;
 
 mod credentials;
 pub use credentials::ClientCredentials;
@@ -12,7 +12,7 @@ use password_hack::PasswordHackPayload;
 mod flow;
 pub use flow::*;
 
-use crate::error::UserFacingError;
+use crate::{console, error::UserFacingError};
 use std::io;
 use thiserror::Error;
 
@@ -26,10 +26,10 @@ pub enum AuthErrorImpl {
     Link(#[from] backend::LinkAuthError),
 
     #[error(transparent)]
-    GetAuthInfo(#[from] backend::GetAuthInfoError),
+    GetAuthInfo(#[from] console::errors::GetAuthInfoError),
 
     #[error(transparent)]
-    WakeCompute(#[from] backend::WakeComputeError),
+    WakeCompute(#[from] console::errors::WakeComputeError),
 
     /// SASL protocol errors (includes [SCRAM](crate::scram)).
     #[error(transparent)]

proxy/src/auth/backend.rs

@@ -1,48 +1,40 @@
-mod postgres;
+mod classic;
 
 mod link;
+use futures::TryFutureExt;
 pub use link::LinkAuthError;
 
-mod console;
-pub use console::{GetAuthInfoError, WakeComputeError};
-
 use crate::{
     auth::{self, AuthFlow, ClientCredentials},
-    compute,
-    console::messages::MetricsAuxInfo,
-    http, mgmt, stream, url,
-    waiters::{self, Waiter, Waiters},
+    console::{
+        self,
+        provider::{CachedNodeInfo, ConsoleReqExtra},
+        Api,
+    },
+    stream, url,
 };
-use once_cell::sync::Lazy;
 use std::borrow::Cow;
 use tokio::io::{AsyncRead, AsyncWrite};
 use tracing::{info, warn};
 
-static CPLANE_WAITERS: Lazy<Waiters<mgmt::ComputeReady>> = Lazy::new(Default::default);
-
-/// Give caller an opportunity to wait for the cloud's reply.
-pub async fn with_waiter<R, T, E>(
-    psql_session_id: impl Into<String>,
-    action: impl FnOnce(Waiter<'static, mgmt::ComputeReady>) -> R,
-) -> Result<T, E>
-where
-    R: std::future::Future<Output = Result<T, E>>,
-    E: From<waiters::RegisterError>,
-{
-    let waiter = CPLANE_WAITERS.register(psql_session_id.into())?;
-    action(waiter).await
+/// A product of successful authentication.
+pub struct AuthSuccess<T> {
+    /// Did we send [`pq_proto::BeMessage::AuthenticationOk`] to client?
+    pub reported_auth_ok: bool,
+    /// Something to be considered a positive result.
+    pub value: T,
 }
 
-pub fn notify(psql_session_id: &str, msg: mgmt::ComputeReady) -> Result<(), waiters::NotifyError> {
-    CPLANE_WAITERS.notify(psql_session_id, msg)
-}
-
-/// Extra query params we'd like to pass to the console.
-pub struct ConsoleReqExtra<'a> {
-    /// A unique identifier for a connection.
-    pub session_id: uuid::Uuid,
-    /// Name of client application, if set.
-    pub application_name: Option<&'a str>,
+impl<T> AuthSuccess<T> {
+    /// Very similar to [`std::option::Option::map`].
+    /// Maps [`AuthSuccess<T>`] to [`AuthSuccess<R>`] by applying
+    /// a function to a contained value.
+    pub fn map<R>(self, f: impl FnOnce(T) -> R) -> AuthSuccess<R> {
+        AuthSuccess {
+            reported_auth_ok: self.reported_auth_ok,
+            value: f(self.value),
+        }
+    }
 }
 
 /// This type serves two purposes:
@@ -53,12 +45,11 @@ pub struct ConsoleReqExtra<'a> {
 /// * However, when we substitute `T` with [`ClientCredentials`],
 ///   this helps us provide the credentials only to those auth
 ///   backends which require them for the authentication process.
-#[derive(Debug)]
 pub enum BackendType<'a, T> {
     /// Current Cloud API (V2).
-    Console(Cow<'a, http::Endpoint>, T),
+    Console(Cow<'a, console::provider::neon::Api>, T),
     /// Local mock of Cloud API (V2).
-    Postgres(Cow<'a, url::ApiUrl>, T),
+    Postgres(Cow<'a, console::provider::mock::Api>, T),
     /// Authentication via a web browser.
     Link(Cow<'a, url::ApiUrl>),
 }
@@ -67,14 +58,8 @@ impl std::fmt::Display for BackendType<'_, ()> {
     fn fmt(&self, fmt: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         use BackendType::*;
         match self {
-            Console(endpoint, _) => fmt
-                .debug_tuple("Console")
-                .field(&endpoint.url().as_str())
-                .finish(),
-            Postgres(endpoint, _) => fmt
-                .debug_tuple("Postgres")
-                .field(&endpoint.as_str())
-                .finish(),
+            Console(endpoint, _) => fmt.debug_tuple("Console").field(&endpoint.url()).finish(),
+            Postgres(endpoint, _) => fmt.debug_tuple("Postgres").field(&endpoint.url()).finish(),
             Link(url) => fmt.debug_tuple("Link").field(&url.as_str()).finish(),
         }
     }
@@ -120,30 +105,16 @@ impl<'a, T, E> BackendType<'a, Result<T, E>> {
     }
 }
 
-/// A product of successful authentication.
-pub struct AuthSuccess<T> {
-    /// Did we send [`pq_proto::BeMessage::AuthenticationOk`] to client?
-    pub reported_auth_ok: bool,
-    /// Something to be considered a positive result.
-    pub value: T,
-}
-
-/// Info for establishing a connection to a compute node.
-/// This is what we get after auth succeeded, but not before!
-pub struct NodeInfo {
-    /// Compute node connection params.
-    pub config: compute::ConnCfg,
-    /// Labels for proxy's metrics.
-    pub aux: MetricsAuxInfo,
-}
-
-impl BackendType<'_, ClientCredentials<'_>> {
+// TODO: get rid of explicit lifetimes in this block (there's a bug in rustc).
+// Read more: https://github.com/rust-lang/rust/issues/99190
+// Alleged fix: https://github.com/rust-lang/rust/pull/89056
+impl<'l> BackendType<'l, ClientCredentials<'_>> {
     /// Do something special if user didn't provide the `project` parameter.
-    async fn try_password_hack(
-        &mut self,
-        extra: &ConsoleReqExtra<'_>,
-        client: &mut stream::PqStream<impl AsyncRead + AsyncWrite + Unpin + Send>,
-    ) -> auth::Result<Option<AuthSuccess<NodeInfo>>> {
+    async fn try_password_hack<'a>(
+        &'a mut self,
+        extra: &'a ConsoleReqExtra<'a>,
+        client: &'a mut stream::PqStream<impl AsyncRead + AsyncWrite + Unpin>,
+    ) -> auth::Result<Option<AuthSuccess<CachedNodeInfo>>> {
         use BackendType::*;
 
         // If there's no project so far, that entails that client doesn't
@@ -179,33 +150,28 @@ impl BackendType<'_, ClientCredentials<'_>> {
 
         // 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() => {
+            Console(api, creds) if creds.project.is_none() => {
                 let payload = fetch_magic_payload(client).await?;
 
                 let mut creds = creds.as_ref();
                 creds.project = Some(payload.project.as_str().into());
-                let node = console::Api::new(endpoint, extra, &creds)
-                    .wake_compute()
-                    .await?;
+                let node = api.wake_compute(extra, &creds).await?;
 
                 (node, payload)
             }
-            Console(endpoint, creds) if creds.use_cleartext_password_flow => {
-                // This is a hack to allow cleartext password in secure connections (wss).
+            // This is a hack to allow cleartext password in secure connections (wss).
+            Console(api, creds) if creds.use_cleartext_password_flow => {
                 let payload = fetch_plaintext_password(client).await?;
-                let creds = creds.as_ref();
-                let node = console::Api::new(endpoint, extra, &creds)
-                    .wake_compute()
-                    .await?;
+                let node = api.wake_compute(extra, creds).await?;
 
                 (node, payload)
             }
-            Postgres(endpoint, creds) if creds.project.is_none() => {
+            Postgres(api, creds) if creds.project.is_none() => {
                 let payload = fetch_magic_payload(client).await?;
 
                 let mut creds = creds.as_ref();
                 creds.project = Some(payload.project.as_str().into());
-                let node = postgres::Api::new(endpoint, &creds).wake_compute().await?;
+                let node = api.wake_compute(extra, &creds).await?;
 
                 (node, payload)
             }
@@ -220,11 +186,11 @@ impl BackendType<'_, ClientCredentials<'_>> {
     }
 
     /// Authenticate the client via the requested backend, possibly using credentials.
-    pub async fn authenticate(
-        mut self,
-        extra: &ConsoleReqExtra<'_>,
-        client: &mut stream::PqStream<impl AsyncRead + AsyncWrite + Unpin + Send>,
-    ) -> auth::Result<AuthSuccess<NodeInfo>> {
+    pub async fn authenticate<'a>(
+        &mut self,
+        extra: &'a ConsoleReqExtra<'a>,
+        client: &'a mut stream::PqStream<impl AsyncRead + AsyncWrite + Unpin>,
+    ) -> auth::Result<AuthSuccess<CachedNodeInfo>> {
         use BackendType::*;
 
         // Handle cases when `project` is missing in `creds`.
@@ -235,7 +201,7 @@ impl BackendType<'_, ClientCredentials<'_>> {
         }
 
         let res = match self {
-            Console(endpoint, creds) => {
+            Console(api, creds) => {
                 info!(
                     user = creds.user,
                     project = creds.project(),
@@ -243,26 +209,40 @@ impl BackendType<'_, ClientCredentials<'_>> {
                 );
 
                 assert!(creds.project.is_some());
-                console::Api::new(&endpoint, extra, &creds)
-                    .handle_user(client)
-                    .await?
+                classic::handle_user(api.as_ref(), extra, creds, client).await?
             }
-            Postgres(endpoint, creds) => {
+            Postgres(api, creds) => {
                 info!("performing mock authentication using a local postgres instance");
 
                 assert!(creds.project.is_some());
-                postgres::Api::new(&endpoint, &creds)
-                    .handle_user(client)
-                    .await?
+                classic::handle_user(api.as_ref(), extra, creds, client).await?
             }
             // NOTE: this auth backend doesn't use client credentials.
             Link(url) => {
                 info!("performing link authentication");
-                link::handle_user(&url, client).await?
+
+                link::handle_user(url, client)
+                    .await?
+                    .map(CachedNodeInfo::new_uncached)
             }
         };
 
         info!("user successfully authenticated");
         Ok(res)
     }
+
+    /// When applicable, wake the compute node, gaining its connection info in the process.
+    /// The link auth flow doesn't support this, so we return [`None`] in that case.
+    pub async fn wake_compute<'a>(
+        &self,
+        extra: &'a ConsoleReqExtra<'a>,
+    ) -> Result<Option<CachedNodeInfo>, console::errors::WakeComputeError> {
+        use BackendType::*;
+
+        match self {
+            Console(api, creds) => api.wake_compute(extra, creds).map_ok(Some).await,
+            Postgres(api, creds) => api.wake_compute(extra, creds).map_ok(Some).await,
+            Link(_) => Ok(None),
+        }
+    }
 }

proxy/src/auth/backend/classic.rs

@@ -0,0 +1,61 @@
+use super::AuthSuccess;
+use crate::{
+    auth::{self, AuthFlow, ClientCredentials},
+    compute,
+    console::{self, AuthInfo, CachedNodeInfo, ConsoleReqExtra},
+    sasl, scram,
+    stream::PqStream,
+};
+use tokio::io::{AsyncRead, AsyncWrite};
+use tracing::info;
+
+pub(super) async fn handle_user(
+    api: &impl console::Api,
+    extra: &ConsoleReqExtra<'_>,
+    creds: &ClientCredentials<'_>,
+    client: &mut PqStream<impl AsyncRead + AsyncWrite + Unpin>,
+) -> auth::Result<AuthSuccess<CachedNodeInfo>> {
+    info!("fetching user's authentication info");
+    let info = api.get_auth_info(extra, creds).await?.unwrap_or_else(|| {
+        // If we don't have an authentication secret, we mock one to
+        // prevent malicious probing (possible due to missing protocol steps).
+        // This mocked secret will never lead to successful authentication.
+        info!("authentication info not found, mocking it");
+        AuthInfo::Scram(scram::ServerSecret::mock(creds.user, rand::random()))
+    });
+
+    let flow = AuthFlow::new(client);
+    let scram_keys = match info {
+        AuthInfo::Md5(_) => {
+            info!("auth endpoint chooses MD5");
+            return Err(auth::AuthError::bad_auth_method("MD5"));
+        }
+        AuthInfo::Scram(secret) => {
+            info!("auth endpoint chooses SCRAM");
+            let scram = auth::Scram(&secret);
+            let client_key = match flow.begin(scram).await?.authenticate().await? {
+                sasl::Outcome::Success(key) => key,
+                sasl::Outcome::Failure(reason) => {
+                    info!("auth backend failed with an error: {reason}");
+                    return Err(auth::AuthError::auth_failed(creds.user));
+                }
+            };
+
+            Some(compute::ScramKeys {
+                client_key: client_key.as_bytes(),
+                server_key: secret.server_key.as_bytes(),
+            })
+        }
+    };
+
+    let mut node = api.wake_compute(extra, creds).await?;
+    if let Some(keys) = scram_keys {
+        use tokio_postgres::config::AuthKeys;
+        node.config.auth_keys(AuthKeys::ScramSha256(keys));
+    }
+
+    Ok(AuthSuccess {
+        reported_auth_ok: false,
+        value: node,
+    })
+}

proxy/src/auth/backend/console.rs

@@ -1,365 +0,0 @@
-//! Cloud API V2.
-
-use super::{AuthSuccess, ConsoleReqExtra, NodeInfo};
-use crate::{
-    auth::{self, AuthFlow, ClientCredentials},
-    compute,
-    console::messages::{ConsoleError, GetRoleSecret, WakeCompute},
-    error::{io_error, UserFacingError},
-    http, sasl, scram,
-    stream::PqStream,
-};
-use futures::TryFutureExt;
-use reqwest::StatusCode as HttpStatusCode;
-use std::future::Future;
-use thiserror::Error;
-use tokio::io::{AsyncRead, AsyncWrite};
-use tracing::{error, info, info_span, warn, Instrument};
-
-/// A go-to error message which doesn't leak any detail.
-const REQUEST_FAILED: &str = "Console request failed";
-
-/// Common console API error.
-#[derive(Debug, Error)]
-pub enum ApiError {
-    /// Error returned by the console itself.
-    #[error("{REQUEST_FAILED} with {}: {}", .status, .text)]
-    Console {
-        status: HttpStatusCode,
-        text: Box<str>,
-    },
-
-    /// Various IO errors like broken pipe or malformed payload.
-    #[error("{REQUEST_FAILED}: {0}")]
-    Transport(#[from] std::io::Error),
-}
-
-impl ApiError {
-    /// Returns HTTP status code if it's the reason for failure.
-    fn http_status_code(&self) -> Option<HttpStatusCode> {
-        use ApiError::*;
-        match self {
-            Console { status, .. } => Some(*status),
-            _ => None,
-        }
-    }
-}
-
-impl UserFacingError for ApiError {
-    fn to_string_client(&self) -> String {
-        use ApiError::*;
-        match self {
-            // To minimize risks, only select errors are forwarded to users.
-            // Ask @neondatabase/control-plane for review before adding more.
-            Console { status, .. } => match *status {
-                HttpStatusCode::NOT_FOUND => {
-                    // Status 404: failed to get a project-related resource.
-                    format!("{REQUEST_FAILED}: endpoint cannot be found")
-                }
-                HttpStatusCode::NOT_ACCEPTABLE => {
-                    // Status 406: endpoint is disabled (we don't allow connections).
-                    format!("{REQUEST_FAILED}: endpoint is disabled")
-                }
-                HttpStatusCode::LOCKED => {
-                    // Status 423: project might be in maintenance mode (or bad state).
-                    format!("{REQUEST_FAILED}: endpoint is temporary unavailable")
-                }
-                _ => REQUEST_FAILED.to_owned(),
-            },
-            _ => REQUEST_FAILED.to_owned(),
-        }
-    }
-}
-
-// Helps eliminate graceless `.map_err` calls without introducing another ctor.
-impl From<reqwest::Error> for ApiError {
-    fn from(e: reqwest::Error) -> Self {
-        io_error(e).into()
-    }
-}
-
-#[derive(Debug, Error)]
-pub enum GetAuthInfoError {
-    // We shouldn't include the actual secret here.
-    #[error("Console responded with a malformed auth secret")]
-    BadSecret,
-
-    #[error(transparent)]
-    ApiError(ApiError),
-}
-
-// This allows more useful interactions than `#[from]`.
-impl<E: Into<ApiError>> From<E> for GetAuthInfoError {
-    fn from(e: E) -> Self {
-        Self::ApiError(e.into())
-    }
-}
-
-impl UserFacingError for GetAuthInfoError {
-    fn to_string_client(&self) -> String {
-        use GetAuthInfoError::*;
-        match self {
-            // We absolutely should not leak any secrets!
-            BadSecret => REQUEST_FAILED.to_owned(),
-            // However, API might return a meaningful error.
-            ApiError(e) => e.to_string_client(),
-        }
-    }
-}
-
-#[derive(Debug, Error)]
-pub enum WakeComputeError {
-    #[error("Console responded with a malformed compute address: {0}")]
-    BadComputeAddress(Box<str>),
-
-    #[error(transparent)]
-    ApiError(ApiError),
-}
-
-// This allows more useful interactions than `#[from]`.
-impl<E: Into<ApiError>> From<E> for WakeComputeError {
-    fn from(e: E) -> Self {
-        Self::ApiError(e.into())
-    }
-}
-
-impl UserFacingError for WakeComputeError {
-    fn to_string_client(&self) -> String {
-        use WakeComputeError::*;
-        match self {
-            // We shouldn't show user the address even if it's broken.
-            // Besides, user is unlikely to care about this detail.
-            BadComputeAddress(_) => REQUEST_FAILED.to_owned(),
-            // However, API might return a meaningful error.
-            ApiError(e) => e.to_string_client(),
-        }
-    }
-}
-
-/// Auth secret which is managed by the cloud.
-pub enum AuthInfo {
-    /// Md5 hash of user's password.
-    Md5([u8; 16]),
-
-    /// [SCRAM](crate::scram) authentication info.
-    Scram(scram::ServerSecret),
-}
-
-#[must_use]
-pub(super) struct Api<'a> {
-    endpoint: &'a http::Endpoint,
-    extra: &'a ConsoleReqExtra<'a>,
-    creds: &'a ClientCredentials<'a>,
-}
-
-impl<'a> AsRef<ClientCredentials<'a>> for Api<'a> {
-    fn as_ref(&self) -> &ClientCredentials<'a> {
-        self.creds
-    }
-}
-
-impl<'a> Api<'a> {
-    /// Construct an API object containing the auth parameters.
-    pub(super) fn new(
-        endpoint: &'a http::Endpoint,
-        extra: &'a ConsoleReqExtra<'a>,
-        creds: &'a ClientCredentials,
-    ) -> Self {
-        Self {
-            endpoint,
-            extra,
-            creds,
-        }
-    }
-
-    /// Authenticate the existing user or throw an error.
-    pub(super) async fn handle_user(
-        &'a self,
-        client: &mut PqStream<impl AsyncRead + AsyncWrite + Unpin + Send>,
-    ) -> auth::Result<AuthSuccess<NodeInfo>> {
-        handle_user(client, self, Self::get_auth_info, Self::wake_compute).await
-    }
-}
-
-impl Api<'_> {
-    async fn get_auth_info(&self) -> Result<Option<AuthInfo>, GetAuthInfoError> {
-        let request_id = uuid::Uuid::new_v4().to_string();
-        async {
-            let request = self
-                .endpoint
-                .get("proxy_get_role_secret")
-                .header("X-Request-ID", &request_id)
-                .query(&[("session_id", self.extra.session_id)])
-                .query(&[
-                    ("application_name", self.extra.application_name),
-                    ("project", Some(self.creds.project().expect("impossible"))),
-                    ("role", Some(self.creds.user)),
-                ])
-                .build()?;
-
-            info!(url = request.url().as_str(), "sending http request");
-            let response = self.endpoint.execute(request).await?;
-            let body = match parse_body::<GetRoleSecret>(response).await {
-                Ok(body) => body,
-                // Error 404 is special: it's ok not to have a secret.
-                Err(e) => match e.http_status_code() {
-                    Some(HttpStatusCode::NOT_FOUND) => return Ok(None),
-                    _otherwise => return Err(e.into()),
-                },
-            };
-
-            let secret = scram::ServerSecret::parse(&body.role_secret)
-                .map(AuthInfo::Scram)
-                .ok_or(GetAuthInfoError::BadSecret)?;
-
-            Ok(Some(secret))
-        }
-        .map_err(crate::error::log_error)
-        .instrument(info_span!("get_auth_info", id = request_id))
-        .await
-    }
-
-    /// Wake up the compute node and return the corresponding connection info.
-    pub async fn wake_compute(&self) -> Result<NodeInfo, WakeComputeError> {
-        let request_id = uuid::Uuid::new_v4().to_string();
-        async {
-            let request = self
-                .endpoint
-                .get("proxy_wake_compute")
-                .header("X-Request-ID", &request_id)
-                .query(&[("session_id", self.extra.session_id)])
-                .query(&[
-                    ("application_name", self.extra.application_name),
-                    ("project", Some(self.creds.project().expect("impossible"))),
-                ])
-                .build()?;
-
-            info!(url = request.url().as_str(), "sending http request");
-            let response = self.endpoint.execute(request).await?;
-            let body = parse_body::<WakeCompute>(response).await?;
-
-            // Unfortunately, ownership won't let us use `Option::ok_or` here.
-            let (host, port) = match parse_host_port(&body.address) {
-                None => return Err(WakeComputeError::BadComputeAddress(body.address)),
-                Some(x) => x,
-            };
-
-            let mut config = compute::ConnCfg::new();
-            config
-                .host(host)
-                .port(port)
-                .dbname(self.creds.dbname)
-                .user(self.creds.user);
-
-            Ok(NodeInfo {
-                config,
-                aux: body.aux,
-            })
-        }
-        .map_err(crate::error::log_error)
-        .instrument(info_span!("wake_compute", id = request_id))
-        .await
-    }
-}
-
-/// Common logic for user handling in API V2.
-/// We reuse this for a mock API implementation in [`super::postgres`].
-pub(super) async fn handle_user<'a, Endpoint, GetAuthInfo, WakeCompute>(
-    client: &mut PqStream<impl AsyncRead + AsyncWrite + Unpin>,
-    endpoint: &'a Endpoint,
-    get_auth_info: impl FnOnce(&'a Endpoint) -> GetAuthInfo,
-    wake_compute: impl FnOnce(&'a Endpoint) -> WakeCompute,
-) -> auth::Result<AuthSuccess<NodeInfo>>
-where
-    Endpoint: AsRef<ClientCredentials<'a>>,
-    GetAuthInfo: Future<Output = Result<Option<AuthInfo>, GetAuthInfoError>>,
-    WakeCompute: Future<Output = Result<NodeInfo, WakeComputeError>>,
-{
-    let creds = endpoint.as_ref();
-
-    info!("fetching user's authentication info");
-    let info = get_auth_info(endpoint).await?.unwrap_or_else(|| {
-        // If we don't have an authentication secret, we mock one to
-        // prevent malicious probing (possible due to missing protocol steps).
-        // This mocked secret will never lead to successful authentication.
-        info!("authentication info not found, mocking it");
-        AuthInfo::Scram(scram::ServerSecret::mock(creds.user, rand::random()))
-    });
-
-    let flow = AuthFlow::new(client);
-    let scram_keys = match info {
-        AuthInfo::Md5(_) => {
-            info!("auth endpoint chooses MD5");
-            return Err(auth::AuthError::bad_auth_method("MD5"));
-        }
-        AuthInfo::Scram(secret) => {
-            info!("auth endpoint chooses SCRAM");
-            let scram = auth::Scram(&secret);
-            let client_key = match flow.begin(scram).await?.authenticate().await? {
-                sasl::Outcome::Success(key) => key,
-                sasl::Outcome::Failure(reason) => {
-                    info!("auth backend failed with an error: {reason}");
-                    return Err(auth::AuthError::auth_failed(creds.user));
-                }
-            };
-
-            Some(compute::ScramKeys {
-                client_key: client_key.as_bytes(),
-                server_key: secret.server_key.as_bytes(),
-            })
-        }
-    };
-
-    let mut node = wake_compute(endpoint).await?;
-    if let Some(keys) = scram_keys {
-        use tokio_postgres::config::AuthKeys;
-        node.config.auth_keys(AuthKeys::ScramSha256(keys));
-    }
-
-    Ok(AuthSuccess {
-        reported_auth_ok: false,
-        value: node,
-    })
-}
-
-/// Parse http response body, taking status code into account.
-async fn parse_body<T: for<'a> serde::Deserialize<'a>>(
-    response: reqwest::Response,
-) -> Result<T, ApiError> {
-    let status = response.status();
-    if status.is_success() {
-        // We shouldn't log raw body because it may contain secrets.
-        info!("request succeeded, processing the body");
-        return Ok(response.json().await?);
-    }
-
-    // Don't throw an error here because it's not as important
-    // as the fact that the request itself has failed.
-    let body = response.json().await.unwrap_or_else(|e| {
-        warn!("failed to parse error body: {e}");
-        ConsoleError {
-            error: "reason unclear (malformed error message)".into(),
-        }
-    });
-
-    let text = body.error;
-    error!("console responded with an error ({status}): {text}");
-    Err(ApiError::Console { status, text })
-}
-
-fn parse_host_port(input: &str) -> Option<(&str, u16)> {
-    let (host, port) = input.split_once(':')?;
-    Some((host, port.parse().ok()?))
-}
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-
-    #[test]
-    fn test_parse_host_port() {
-        let (host, port) = parse_host_port("127.0.0.1:5432").expect("failed to parse");
-        assert_eq!(host, "127.0.0.1");
-        assert_eq!(port, 5432);
-    }
-}

proxy/src/auth/backend/link.rs

@@ -1,5 +1,11 @@
-use super::{AuthSuccess, NodeInfo};
-use crate::{auth, compute, error::UserFacingError, stream::PqStream, waiters};
+use super::AuthSuccess;
+use crate::{
+    auth, compute,
+    console::{self, provider::NodeInfo},
+    error::UserFacingError,
+    stream::PqStream,
+    waiters,
+};
 use pq_proto::BeMessage as Be;
 use thiserror::Error;
 use tokio::io::{AsyncRead, AsyncWrite};
@@ -47,7 +53,7 @@ pub fn new_psql_session_id() -> String {
     hex::encode(rand::random::<[u8; 8]>())
 }
 
-pub async fn handle_user(
+pub(super) async fn handle_user(
     link_uri: &reqwest::Url,
     client: &mut PqStream<impl AsyncRead + AsyncWrite + Unpin>,
 ) -> auth::Result<AuthSuccess<NodeInfo>> {
@@ -55,7 +61,7 @@ pub async fn handle_user(
     let span = info_span!("link", psql_session_id = &psql_session_id);
     let greeting = hello_message(link_uri, &psql_session_id);
 
-    let db_info = super::with_waiter(psql_session_id, |waiter| async {
+    let db_info = console::mgmt::with_waiter(psql_session_id, |waiter| async {
         // Give user a URL to spawn a new database.
         info!(parent: &span, "sending the auth URL to the user");
         client
@@ -80,14 +86,14 @@ pub async fn handle_user(
         .user(&db_info.user);
 
     if let Some(password) = db_info.password {
-        config.password(password);
+        config.password(password.as_ref());
     }
 
     Ok(AuthSuccess {
         reported_auth_ok: true,
         value: NodeInfo {
             config,
-            aux: db_info.aux,
+            aux: db_info.aux.into(),
         },
     })
 }

proxy/src/auth/credentials.rs

@@ -33,6 +33,7 @@ impl UserFacingError for ClientCredsParseError {}
 pub struct ClientCredentials<'a> {
     pub user: &'a str,
     pub dbname: &'a str,
+    // TODO: this is a severe misnomer! We should think of a new name ASAP.
     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.

proxy/src/cache.rs

@@ -0,0 +1,304 @@
+use std::{
+    borrow::Borrow,
+    hash::Hash,
+    ops::{Deref, DerefMut},
+    time::{Duration, Instant},
+};
+use tracing::debug;
+
+// This seems to make more sense than `lru` or `cached`:
+//
+// * `near/nearcore` ditched `cached` in favor of `lru`
+//   (https://github.com/near/nearcore/issues?q=is%3Aissue+lru+is%3Aclosed).
+//
+// * `lru` methods use an obscure `KeyRef` type in their contraints (which is deliberately excluded from docs).
+//   This severely hinders its usage both in terms of creating wrappers and supported key types.
+//
+// On the other hand, `hashlink` has good download stats and appears to be maintained.
+use hashlink::{linked_hash_map::RawEntryMut, LruCache};
+
+/// A generic trait which exposes types of cache's key and value,
+/// as well as the notion of cache entry invalidation.
+/// This is useful for [`timed_lru::Cached`].
+pub trait Cache {
+    /// Entry's key.
+    type Key;
+
+    /// Entry's value.
+    type Value;
+
+    /// Used for entry invalidation.
+    type LookupInfo<Key>;
+
+    /// Invalidate an entry using a lookup info.
+    /// We don't have an empty default impl because it's error-prone.
+    fn invalidate(&self, _: &Self::LookupInfo<Self::Key>);
+}
+
+impl<C: Cache> Cache for &C {
+    type Key = C::Key;
+    type Value = C::Value;
+    type LookupInfo<Key> = C::LookupInfo<Key>;
+
+    fn invalidate(&self, info: &Self::LookupInfo<Self::Key>) {
+        C::invalidate(self, info)
+    }
+}
+
+pub use timed_lru::TimedLru;
+pub mod timed_lru {
+    use super::*;
+
+    /// An implementation of timed LRU cache with fixed capacity.
+    /// Key properties:
+    ///
+    /// * Whenever a new entry is inserted, the least recently accessed one is evicted.
+    ///   The cache also keeps track of entry's insertion time (`created_at`) and TTL (`expires_at`).
+    ///
+    /// * When the entry is about to be retrieved, we check its expiration timestamp.
+    ///   If the entry has expired, we remove it from the cache; Otherwise we bump the
+    ///   expiration timestamp (e.g. +5mins) and change its place in LRU list to prolong
+    ///   its existence.
+    ///
+    /// * There's an API for immediate invalidation (removal) of a cache entry;
+    ///   It's useful in case we know for sure that the entry is no longer correct.
+    ///   See [`timed_lru::LookupInfo`] & [`timed_lru::Cached`] for more information.
+    ///
+    /// * Expired entries are kept in the cache, until they are evicted by the LRU policy,
+    ///   or by a successful lookup (i.e. the entry hasn't expired yet).
+    ///   There is no background job to reap the expired records.
+    ///
+    /// * It's possible for an entry that has not yet expired entry to be evicted
+    ///   before expired items. That's a bit wasteful, but probably fine in practice.
+    pub struct TimedLru<K, V> {
+        /// Cache's name for tracing.
+        name: &'static str,
+
+        /// The underlying cache implementation.
+        cache: parking_lot::Mutex<LruCache<K, Entry<V>>>,
+
+        /// Default time-to-live of a single entry.
+        ttl: Duration,
+    }
+
+    impl<K: Hash + Eq, V> Cache for TimedLru<K, V> {
+        type Key = K;
+        type Value = V;
+        type LookupInfo<Key> = LookupInfo<Key>;
+
+        fn invalidate(&self, info: &Self::LookupInfo<K>) {
+            self.invalidate_raw(info)
+        }
+    }
+
+    struct Entry<T> {
+        created_at: Instant,
+        expires_at: Instant,
+        value: T,
+    }
+
+    impl<K: Hash + Eq, V> TimedLru<K, V> {
+        /// Construct a new LRU cache with timed entries.
+        pub fn new(name: &'static str, capacity: usize, ttl: Duration) -> Self {
+            Self {
+                name,
+                cache: LruCache::new(capacity).into(),
+                ttl,
+            }
+        }
+
+        /// Drop an entry from the cache if it's outdated.
+        #[tracing::instrument(level = "debug", fields(cache = self.name), skip_all)]
+        fn invalidate_raw(&self, info: &LookupInfo<K>) {
+            let now = Instant::now();
+
+            // Do costly things before taking the lock.
+            let mut cache = self.cache.lock();
+            let raw_entry = match cache.raw_entry_mut().from_key(&info.key) {
+                RawEntryMut::Vacant(_) => return,
+                RawEntryMut::Occupied(x) => x,
+            };
+
+            // Remove the entry if it was created prior to lookup timestamp.
+            let entry = raw_entry.get();
+            let (created_at, expires_at) = (entry.created_at, entry.expires_at);
+            let should_remove = created_at <= info.created_at || expires_at <= now;
+
+            if should_remove {
+                raw_entry.remove();
+            }
+
+            drop(cache); // drop lock before logging
+            debug!(
+                created_at = format_args!("{created_at:?}"),
+                expires_at = format_args!("{expires_at:?}"),
+                entry_removed = should_remove,
+                "processed a cache entry invalidation event"
+            );
+        }
+
+        /// Try retrieving an entry by its key, then execute `extract` if it exists.
+        #[tracing::instrument(level = "debug", fields(cache = self.name), skip_all)]
+        fn get_raw<Q, R>(&self, key: &Q, extract: impl FnOnce(&K, &Entry<V>) -> R) -> Option<R>
+        where
+            K: Borrow<Q>,
+            Q: Hash + Eq + ?Sized,
+        {
+            let now = Instant::now();
+            let deadline = now.checked_add(self.ttl).expect("time overflow");
+
+            // Do costly things before taking the lock.
+            let mut cache = self.cache.lock();
+            let mut raw_entry = match cache.raw_entry_mut().from_key(key) {
+                RawEntryMut::Vacant(_) => return None,
+                RawEntryMut::Occupied(x) => x,
+            };
+
+            // Immeditely drop the entry if it has expired.
+            let entry = raw_entry.get();
+            if entry.expires_at <= now {
+                raw_entry.remove();
+                return None;
+            }
+
+            let value = extract(raw_entry.key(), entry);
+            let (created_at, expires_at) = (entry.created_at, entry.expires_at);
+
+            // Update the deadline and the entry's position in the LRU list.
+            raw_entry.get_mut().expires_at = deadline;
+            raw_entry.to_back();
+
+            drop(cache); // drop lock before logging
+            debug!(
+                created_at = format_args!("{created_at:?}"),
+                old_expires_at = format_args!("{expires_at:?}"),
+                new_expires_at = format_args!("{deadline:?}"),
+                "accessed a cache entry"
+            );
+
+            Some(value)
+        }
+
+        /// Insert an entry to the cache. If an entry with the same key already
+        /// existed, return the previous value and its creation timestamp.
+        #[tracing::instrument(level = "debug", fields(cache = self.name), skip_all)]
+        fn insert_raw(&self, key: K, value: V) -> (Instant, Option<V>) {
+            let created_at = Instant::now();
+            let expires_at = created_at.checked_add(self.ttl).expect("time overflow");
+
+            let entry = Entry {
+                created_at,
+                expires_at,
+                value,
+            };
+
+            // Do costly things before taking the lock.
+            let old = self
+                .cache
+                .lock()
+                .insert(key, entry)
+                .map(|entry| entry.value);
+
+            debug!(
+                created_at = format_args!("{created_at:?}"),
+                expires_at = format_args!("{expires_at:?}"),
+                replaced = old.is_some(),
+                "created a cache entry"
+            );
+
+            (created_at, old)
+        }
+    }
+
+    impl<K: Hash + Eq + Clone, V: Clone> TimedLru<K, V> {
+        pub fn insert(&self, key: K, value: V) -> (Option<V>, Cached<&Self>) {
+            let (created_at, old) = self.insert_raw(key.clone(), value.clone());
+
+            let cached = Cached {
+                token: Some((self, LookupInfo { created_at, key })),
+                value,
+            };
+
+            (old, cached)
+        }
+    }
+
+    impl<K: Hash + Eq, V: Clone> TimedLru<K, V> {
+        /// Retrieve a cached entry in convenient wrapper.
+        pub fn get<Q>(&self, key: &Q) -> Option<timed_lru::Cached<&Self>>
+        where
+            K: Borrow<Q> + Clone,
+            Q: Hash + Eq + ?Sized,
+        {
+            self.get_raw(key, |key, entry| {
+                let info = LookupInfo {
+                    created_at: entry.created_at,
+                    key: key.clone(),
+                };
+
+                Cached {
+                    token: Some((self, info)),
+                    value: entry.value.clone(),
+                }
+            })
+        }
+    }
+
+    /// Lookup information for key invalidation.
+    pub struct LookupInfo<K> {
+        /// Time of creation of a cache [`Entry`].
+        /// We use this during invalidation lookups to prevent eviction of a newer
+        /// entry sharing the same key (it might've been inserted by a different
+        /// task after we got the entry we're trying to invalidate now).
+        created_at: Instant,
+
+        /// Search by this key.
+        key: K,
+    }
+
+    /// Wrapper for convenient entry invalidation.
+    pub struct Cached<C: Cache> {
+        /// Cache + lookup info.
+        token: Option<(C, C::LookupInfo<C::Key>)>,
+
+        /// The value itself.
+        pub value: C::Value,
+    }
+
+    impl<C: Cache> Cached<C> {
+        /// Place any entry into this wrapper; invalidation will be a no-op.
+        /// Unfortunately, rust doesn't let us implement [`From`] or [`Into`].
+        pub fn new_uncached(value: impl Into<C::Value>) -> Self {
+            Self {
+                token: None,
+                value: value.into(),
+            }
+        }
+
+        /// Drop this entry from a cache if it's still there.
+        pub fn invalidate(&self) {
+            if let Some((cache, info)) = &self.token {
+                cache.invalidate(info);
+            }
+        }
+
+        /// Tell if this entry is actually cached.
+        pub fn cached(&self) -> bool {
+            self.token.is_some()
+        }
+    }
+
+    impl<C: Cache> Deref for Cached<C> {
+        type Target = C::Value;
+
+        fn deref(&self) -> &Self::Target {
+            &self.value
+        }
+    }
+
+    impl<C: Cache> DerefMut for Cached<C> {
+        fn deref_mut(&mut self) -> &mut Self::Target {
+            &mut self.value
+        }
+    }
+}

proxy/src/cancellation.rs

@@ -1,6 +1,5 @@
 use anyhow::{anyhow, Context};
 use hashbrown::HashMap;
-use parking_lot::Mutex;
 use pq_proto::CancelKeyData;
 use std::net::SocketAddr;
 use tokio::net::TcpStream;
@@ -9,14 +8,15 @@ use tracing::info;
 
 /// Enables serving `CancelRequest`s.
 #[derive(Default)]
-pub struct CancelMap(Mutex<HashMap<CancelKeyData, Option<CancelClosure>>>);
+pub struct CancelMap(parking_lot::RwLock<HashMap<CancelKeyData, Option<CancelClosure>>>);
 
 impl CancelMap {
     /// Cancel a running query for the corresponding connection.
     pub async fn cancel_session(&self, key: CancelKeyData) -> anyhow::Result<()> {
+        // NB: we should immediately release the lock after cloning the token.
         let cancel_closure = self
             .0
-            .lock()
+            .read()
             .get(&key)
             .and_then(|x| x.clone())
             .with_context(|| format!("query cancellation key not found: {key}"))?;
@@ -41,14 +41,14 @@ impl CancelMap {
         // Random key collisions are unlikely to happen here, but they're still possible,
         // which is why we have to take care not to rewrite an existing key.
         self.0
-            .lock()
+            .write()
             .try_insert(key, None)
             .map_err(|_| anyhow!("query cancellation key already exists: {key}"))?;
 
         // This will guarantee that the session gets dropped
         // as soon as the future is finished.
         scopeguard::defer! {
-            self.0.lock().remove(&key);
+            self.0.write().remove(&key);
             info!("dropped query cancellation key {key}");
         }
 
@@ -59,12 +59,12 @@ impl CancelMap {
 
     #[cfg(test)]
     fn contains(&self, session: &Session) -> bool {
-        self.0.lock().contains_key(&session.key)
+        self.0.read().contains_key(&session.key)
     }
 
     #[cfg(test)]
     fn is_empty(&self) -> bool {
-        self.0.lock().is_empty()
+        self.0.read().is_empty()
     }
 }
 
@@ -115,7 +115,7 @@ impl Session<'_> {
         info!("enabling query cancellation for this session");
         self.cancel_map
             .0
-            .lock()
+            .write()
             .insert(self.key, Some(cancel_closure));
 
         self.key

proxy/src/compute.rs

@@ -42,14 +42,65 @@ pub type ScramKeys = tokio_postgres::config::ScramKeys<32>;
 /// A config for establishing a connection to compute node.
 /// Eventually, `tokio_postgres` will be replaced with something better.
 /// Newtype allows us to implement methods on top of it.
+#[derive(Clone)]
 #[repr(transparent)]
 pub struct ConnCfg(Box<tokio_postgres::Config>);
 
+/// Creation and initialization routines.
 impl ConnCfg {
-    /// Construct a new connection config.
     pub fn new() -> Self {
         Self(Default::default())
     }
+
+    /// Reuse password or auth keys from the other config.
+    pub fn reuse_password(&mut self, other: &Self) {
+        if let Some(password) = other.get_password() {
+            self.password(password);
+        }
+
+        if let Some(keys) = other.get_auth_keys() {
+            self.auth_keys(keys);
+        }
+    }
+
+    /// Apply startup message params to the connection config.
+    pub fn set_startup_params(&mut self, params: &StartupMessageParams) {
+        if let Some(options) = params.options_raw() {
+            // We must drop all proxy-specific parameters.
+            #[allow(unstable_name_collisions)]
+            let options: String = options
+                .filter(|opt| !opt.starts_with("project="))
+                .intersperse(" ") // TODO: use impl from std once it's stabilized
+                .collect();
+
+            self.options(&options);
+        }
+
+        if let Some(app_name) = params.get("application_name") {
+            self.application_name(app_name);
+        }
+
+        // TODO: This is especially ugly...
+        if let Some(replication) = params.get("replication") {
+            use tokio_postgres::config::ReplicationMode;
+            match replication {
+                "true" | "on" | "yes" | "1" => {
+                    self.replication_mode(ReplicationMode::Physical);
+                }
+                "database" => {
+                    self.replication_mode(ReplicationMode::Logical);
+                }
+                _other => {}
+            }
+        }
+
+        // TODO: extend the list of the forwarded startup parameters.
+        // Currently, tokio-postgres doesn't allow us to pass
+        // arbitrary parameters, but the ones above are a good start.
+        //
+        // This and the reverse params problem can be better addressed
+        // in a bespoke connection machinery (a new library for that sake).
+    }
 }
 
 impl std::ops::Deref for ConnCfg {
@@ -132,50 +183,13 @@ pub struct PostgresConnection {
     pub stream: TcpStream,
     /// PostgreSQL connection parameters.
     pub params: std::collections::HashMap<String, String>,
+    /// Query cancellation token.
+    pub cancel_closure: CancelClosure,
 }
 
 impl ConnCfg {
     /// Connect to a corresponding compute node.
-    pub async fn connect(
-        mut self,
-        params: &StartupMessageParams,
-    ) -> Result<(PostgresConnection, CancelClosure), ConnectionError> {
-        if let Some(options) = params.options_raw() {
-            // We must drop all proxy-specific parameters.
-            #[allow(unstable_name_collisions)]
-            let options: String = options
-                .filter(|opt| !opt.starts_with("project="))
-                .intersperse(" ") // TODO: use impl from std once it's stabilized
-                .collect();
-
-            self.0.options(&options);
-        }
-
-        if let Some(app_name) = params.get("application_name") {
-            self.0.application_name(app_name);
-        }
-
-        // TODO: This is especially ugly...
-        if let Some(replication) = params.get("replication") {
-            use tokio_postgres::config::ReplicationMode;
-            match replication {
-                "true" | "on" | "yes" | "1" => {
-                    self.0.replication_mode(ReplicationMode::Physical);
-                }
-                "database" => {
-                    self.0.replication_mode(ReplicationMode::Logical);
-                }
-                _other => {}
-            }
-        }
-
-        // TODO: extend the list of the forwarded startup parameters.
-        // Currently, tokio-postgres doesn't allow us to pass
-        // arbitrary parameters, but the ones above are a good start.
-        //
-        // This and the reverse params problem can be better addressed
-        // in a bespoke connection machinery (a new library for that sake).
-
+    pub async fn connect(&self) -> Result<PostgresConnection, ConnectionError> {
         // TODO: establish a secure connection to the DB.
         let (socket_addr, mut stream) = self.connect_raw().await?;
         let (client, connection) = self.0.connect_raw(&mut stream, NoTls).await?;
@@ -189,8 +203,13 @@ impl ConnCfg {
         // NB: CancelToken is supposed to hold socket_addr, but we use connect_raw.
         // Yet another reason to rework the connection establishing code.
         let cancel_closure = CancelClosure::new(socket_addr, client.cancel_token());
-        let db = PostgresConnection { stream, params };
 
-        Ok((db, cancel_closure))
+        let connection = PostgresConnection {
+            stream,
+            params,
+            cancel_closure,
+        };
+
+        Ok(connection)
     }
 }

proxy/src/config.rs

@@ -1,16 +1,16 @@
 use crate::auth;
-use anyhow::{ensure, Context};
-use std::sync::Arc;
+use anyhow::{bail, ensure, Context};
+use std::{str::FromStr, sync::Arc, time::Duration};
 
 pub struct ProxyConfig {
     pub tls_config: Option<TlsConfig>,
     pub auth_backend: auth::BackendType<'static, ()>,
-    pub metric_collection_config: Option<MetricCollectionConfig>,
+    pub metric_collection: Option<MetricCollectionConfig>,
 }
 
 pub struct MetricCollectionConfig {
     pub endpoint: reqwest::Url,
-    pub interval: std::time::Duration,
+    pub interval: Duration,
 }
 
 pub struct TlsConfig {
@@ -37,6 +37,7 @@ pub fn configure_tls(key_path: &str, cert_path: &str) -> anyhow::Result<TlsConfi
 
     let cert_chain_bytes = std::fs::read(cert_path)
         .context(format!("Failed to read TLS cert file at '{cert_path}.'"))?;
+
     let cert_chain = {
         rustls_pemfile::certs(&mut &cert_chain_bytes[..])
             .context(format!(
@@ -73,3 +74,80 @@ pub fn configure_tls(key_path: &str, cert_path: &str) -> anyhow::Result<TlsConfi
         common_name,
     })
 }
+
+/// Helper for cmdline cache options parsing.
+pub struct CacheOptions {
+    /// Max number of entries.
+    pub size: usize,
+    /// Entry's time-to-live.
+    pub ttl: Duration,
+}
+
+impl CacheOptions {
+    /// Default options for [`crate::auth::caches::NodeInfoCache`].
+    pub const DEFAULT_OPTIONS_NODE_INFO: &str = "size=4000,ttl=5m";
+
+    /// Parse cache options passed via cmdline.
+    /// Example: [`Self::DEFAULT_OPTIONS_NODE_INFO`].
+    fn parse(options: &str) -> anyhow::Result<Self> {
+        let mut size = None;
+        let mut ttl = None;
+
+        for option in options.split(',') {
+            let (key, value) = option
+                .split_once('=')
+                .with_context(|| format!("bad key-value pair: {option}"))?;
+
+            match key {
+                "size" => size = Some(value.parse()?),
+                "ttl" => ttl = Some(humantime::parse_duration(value)?),
+                unknown => bail!("unknown key: {unknown}"),
+            }
+        }
+
+        // TTL doesn't matter if cache is always empty.
+        if let Some(0) = size {
+            ttl.get_or_insert(Duration::default());
+        }
+
+        Ok(Self {
+            size: size.context("missing `size`")?,
+            ttl: ttl.context("missing `ttl`")?,
+        })
+    }
+}
+
+impl FromStr for CacheOptions {
+    type Err = anyhow::Error;
+
+    fn from_str(options: &str) -> Result<Self, Self::Err> {
+        let error = || format!("failed to parse cache options '{options}'");
+        Self::parse(options).with_context(error)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_parse_cache_options() -> anyhow::Result<()> {
+        let CacheOptions { size, ttl } = "size=4096,ttl=5min".parse()?;
+        assert_eq!(size, 4096);
+        assert_eq!(ttl, Duration::from_secs(5 * 60));
+
+        let CacheOptions { size, ttl } = "ttl=4m,size=2".parse()?;
+        assert_eq!(size, 2);
+        assert_eq!(ttl, Duration::from_secs(4 * 60));
+
+        let CacheOptions { size, ttl } = "size=0,ttl=1s".parse()?;
+        assert_eq!(size, 0);
+        assert_eq!(ttl, Duration::from_secs(1));
+
+        let CacheOptions { size, ttl } = "size=0".parse()?;
+        assert_eq!(size, 0);
+        assert_eq!(ttl, Duration::default());
+
+        Ok(())
+    }
+}

proxy/src/console.rs

@@ -3,3 +3,15 @@
 
 /// Payloads used in the console's APIs.
 pub mod messages;
+
+/// Wrappers for console APIs and their mocks.
+pub mod provider;
+pub use provider::{errors, Api, AuthInfo, CachedNodeInfo, ConsoleReqExtra, NodeInfo};
+
+/// Various cache-related types.
+pub mod caches {
+    pub use super::provider::{ApiCaches, NodeInfoCache};
+}
+
+/// Console's management API.
+pub mod mgmt;

proxy/src/console/messages.rs

@@ -63,13 +63,13 @@ impl KickSession<'_> {
 /// Compute node connection params.
 #[derive(Deserialize)]
 pub struct DatabaseInfo {
-    pub host: String,
+    pub host: Box<str>,
     pub port: u16,
-    pub dbname: String,
-    pub user: String,
+    pub dbname: Box<str>,
+    pub user: Box<str>,
     /// Console always provides a password, but it might
     /// be inconvenient for debug with local PG instance.
-    pub password: Option<String>,
+    pub password: Option<Box<str>>,
     pub aux: MetricsAuxInfo,
 }
 

proxy/src/console/mgmt.rs

@@ -1,8 +1,9 @@
 use crate::{
-    auth,
     console::messages::{DatabaseInfo, KickSession},
+    waiters::{self, Waiter, Waiters},
 };
 use anyhow::Context;
+use once_cell::sync::Lazy;
 use pq_proto::{BeMessage, SINGLE_COL_ROWDESC};
 use std::{
     net::{TcpListener, TcpStream},
@@ -14,6 +15,25 @@ use utils::{
     postgres_backend_async::QueryError,
 };
 
+static CPLANE_WAITERS: Lazy<Waiters<ComputeReady>> = Lazy::new(Default::default);
+
+/// Give caller an opportunity to wait for the cloud's reply.
+pub async fn with_waiter<R, T, E>(
+    psql_session_id: impl Into<String>,
+    action: impl FnOnce(Waiter<'static, ComputeReady>) -> R,
+) -> Result<T, E>
+where
+    R: std::future::Future<Output = Result<T, E>>,
+    E: From<waiters::RegisterError>,
+{
+    let waiter = CPLANE_WAITERS.register(psql_session_id.into())?;
+    action(waiter).await
+}
+
+pub fn notify(psql_session_id: &str, msg: ComputeReady) -> Result<(), waiters::NotifyError> {
+    CPLANE_WAITERS.notify(psql_session_id, msg)
+}
+
 /// Console management API listener thread.
 /// It spawns console response handlers needed for the link auth.
 pub fn thread_main(listener: TcpListener) -> anyhow::Result<()> {
@@ -76,7 +96,7 @@ fn try_process_query(pgb: &mut PostgresBackend, query: &str) -> Result<(), Query
     let _enter = span.enter();
     info!("got response: {:?}", resp.result);
 
-    match auth::backend::notify(resp.session_id, Ok(resp.result)) {
+    match notify(resp.session_id, Ok(resp.result)) {
         Ok(()) => {
             pgb.write_message_noflush(&SINGLE_COL_ROWDESC)?
                 .write_message_noflush(&BeMessage::DataRow(&[Some(b"ok")]))?

proxy/src/console/provider.rs

@@ -0,0 +1,194 @@
+pub mod mock;
+pub mod neon;
+
+use super::messages::MetricsAuxInfo;
+use crate::{
+    auth::ClientCredentials,
+    cache::{timed_lru, TimedLru},
+    compute, scram,
+};
+use async_trait::async_trait;
+use std::sync::Arc;
+
+pub mod errors {
+    use crate::error::{io_error, UserFacingError};
+    use reqwest::StatusCode as HttpStatusCode;
+    use thiserror::Error;
+
+    /// A go-to error message which doesn't leak any detail.
+    const REQUEST_FAILED: &str = "Console request failed";
+
+    /// Common console API error.
+    #[derive(Debug, Error)]
+    pub enum ApiError {
+        /// Error returned by the console itself.
+        #[error("{REQUEST_FAILED} with {}: {}", .status, .text)]
+        Console {
+            status: HttpStatusCode,
+            text: Box<str>,
+        },
+
+        /// Various IO errors like broken pipe or malformed payload.
+        #[error("{REQUEST_FAILED}: {0}")]
+        Transport(#[from] std::io::Error),
+    }
+
+    impl ApiError {
+        /// Returns HTTP status code if it's the reason for failure.
+        pub fn http_status_code(&self) -> Option<HttpStatusCode> {
+            use ApiError::*;
+            match self {
+                Console { status, .. } => Some(*status),
+                _ => None,
+            }
+        }
+    }
+
+    impl UserFacingError for ApiError {
+        fn to_string_client(&self) -> String {
+            use ApiError::*;
+            match self {
+                // To minimize risks, only select errors are forwarded to users.
+                // Ask @neondatabase/control-plane for review before adding more.
+                Console { status, .. } => match *status {
+                    HttpStatusCode::NOT_FOUND => {
+                        // Status 404: failed to get a project-related resource.
+                        format!("{REQUEST_FAILED}: endpoint cannot be found")
+                    }
+                    HttpStatusCode::NOT_ACCEPTABLE => {
+                        // Status 406: endpoint is disabled (we don't allow connections).
+                        format!("{REQUEST_FAILED}: endpoint is disabled")
+                    }
+                    HttpStatusCode::LOCKED => {
+                        // Status 423: project might be in maintenance mode (or bad state).
+                        format!("{REQUEST_FAILED}: endpoint is temporary unavailable")
+                    }
+                    _ => REQUEST_FAILED.to_owned(),
+                },
+                _ => REQUEST_FAILED.to_owned(),
+            }
+        }
+    }
+
+    // Helps eliminate graceless `.map_err` calls without introducing another ctor.
+    impl From<reqwest::Error> for ApiError {
+        fn from(e: reqwest::Error) -> Self {
+            io_error(e).into()
+        }
+    }
+
+    #[derive(Debug, Error)]
+    pub enum GetAuthInfoError {
+        // We shouldn't include the actual secret here.
+        #[error("Console responded with a malformed auth secret")]
+        BadSecret,
+
+        #[error(transparent)]
+        ApiError(ApiError),
+    }
+
+    // This allows more useful interactions than `#[from]`.
+    impl<E: Into<ApiError>> From<E> for GetAuthInfoError {
+        fn from(e: E) -> Self {
+            Self::ApiError(e.into())
+        }
+    }
+
+    impl UserFacingError for GetAuthInfoError {
+        fn to_string_client(&self) -> String {
+            use GetAuthInfoError::*;
+            match self {
+                // We absolutely should not leak any secrets!
+                BadSecret => REQUEST_FAILED.to_owned(),
+                // However, API might return a meaningful error.
+                ApiError(e) => e.to_string_client(),
+            }
+        }
+    }
+    #[derive(Debug, Error)]
+    pub enum WakeComputeError {
+        #[error("Console responded with a malformed compute address: {0}")]
+        BadComputeAddress(Box<str>),
+
+        #[error(transparent)]
+        ApiError(ApiError),
+    }
+
+    // This allows more useful interactions than `#[from]`.
+    impl<E: Into<ApiError>> From<E> for WakeComputeError {
+        fn from(e: E) -> Self {
+            Self::ApiError(e.into())
+        }
+    }
+
+    impl UserFacingError for WakeComputeError {
+        fn to_string_client(&self) -> String {
+            use WakeComputeError::*;
+            match self {
+                // We shouldn't show user the address even if it's broken.
+                // Besides, user is unlikely to care about this detail.
+                BadComputeAddress(_) => REQUEST_FAILED.to_owned(),
+                // However, API might return a meaningful error.
+                ApiError(e) => e.to_string_client(),
+            }
+        }
+    }
+}
+
+/// Extra query params we'd like to pass to the console.
+pub struct ConsoleReqExtra<'a> {
+    /// A unique identifier for a connection.
+    pub session_id: uuid::Uuid,
+    /// Name of client application, if set.
+    pub application_name: Option<&'a str>,
+}
+
+/// Auth secret which is managed by the cloud.
+pub enum AuthInfo {
+    /// Md5 hash of user's password.
+    Md5([u8; 16]),
+
+    /// [SCRAM](crate::scram) authentication info.
+    Scram(scram::ServerSecret),
+}
+
+/// Info for establishing a connection to a compute node.
+/// This is what we get after auth succeeded, but not before!
+#[derive(Clone)]
+pub struct NodeInfo {
+    /// Compute node connection params.
+    /// It's sad that we have to clone this, but this will improve
+    /// once we migrate to a bespoke connection logic.
+    pub config: compute::ConnCfg,
+
+    /// Labels for proxy's metrics.
+    pub aux: Arc<MetricsAuxInfo>,
+}
+
+pub type NodeInfoCache = TimedLru<Arc<str>, NodeInfo>;
+pub type CachedNodeInfo = timed_lru::Cached<&'static NodeInfoCache>;
+
+/// This will allocate per each call, but the http requests alone
+/// already require a few allocations, so it should be fine.
+#[async_trait]
+pub trait Api {
+    /// Get the client's auth secret for authentication.
+    async fn get_auth_info(
+        &self,
+        extra: &ConsoleReqExtra<'_>,
+        creds: &ClientCredentials<'_>,
+    ) -> Result<Option<AuthInfo>, errors::GetAuthInfoError>;
+
+    /// Wake up the compute node and return the corresponding connection info.
+    async fn wake_compute(
+        &self,
+        extra: &ConsoleReqExtra<'_>,
+        creds: &ClientCredentials<'_>,
+    ) -> Result<CachedNodeInfo, errors::WakeComputeError>;
+}
+
+/// Various caches for [`console`].
+pub struct ApiCaches {
+    /// Cache for the `wake_compute` API method.
+    pub node_info: NodeInfoCache,
+}

proxy/src/console/provider/mock.rs

@@ -1,21 +1,14 @@
-//! Local mock of Cloud API V2.
+//! Mock console backend which relies on a user-provided postgres instance.
 
 use super::{
-    console::{self, AuthInfo, GetAuthInfoError, WakeComputeError},
-    AuthSuccess, NodeInfo,
-};
-use crate::{
-    auth::{self, ClientCredentials},
-    compute,
-    error::io_error,
-    scram,
-    stream::PqStream,
-    url::ApiUrl,
+    errors::{ApiError, GetAuthInfoError, WakeComputeError},
+    AuthInfo, CachedNodeInfo, ConsoleReqExtra, NodeInfo,
 };
+use crate::{auth::ClientCredentials, compute, error::io_error, scram, url::ApiUrl};
+use async_trait::async_trait;
 use futures::TryFutureExt;
 use thiserror::Error;
-use tokio::io::{AsyncRead, AsyncWrite};
-use tracing::{info, info_span, warn, Instrument};
+use tracing::{error, info, info_span, warn, Instrument};
 
 #[derive(Debug, Error)]
 enum MockApiError {
@@ -23,49 +16,36 @@ enum MockApiError {
     PasswordNotSet(tokio_postgres::Error),
 }
 
-impl From<MockApiError> for console::ApiError {
+impl From<MockApiError> for ApiError {
     fn from(e: MockApiError) -> Self {
         io_error(e).into()
     }
 }
 
-impl From<tokio_postgres::Error> for console::ApiError {
+impl From<tokio_postgres::Error> for ApiError {
     fn from(e: tokio_postgres::Error) -> Self {
         io_error(e).into()
     }
 }
 
-#[must_use]
-pub(super) struct Api<'a> {
-    endpoint: &'a ApiUrl,
-    creds: &'a ClientCredentials<'a>,
+#[derive(Clone)]
+pub struct Api {
+    endpoint: ApiUrl,
 }
 
-impl<'a> AsRef<ClientCredentials<'a>> for Api<'a> {
-    fn as_ref(&self) -> &ClientCredentials<'a> {
-        self.creds
-    }
-}
-
-impl<'a> Api<'a> {
-    /// Construct an API object containing the auth parameters.
-    pub(super) fn new(endpoint: &'a ApiUrl, creds: &'a ClientCredentials) -> Self {
-        Self { endpoint, creds }
+impl Api {
+    pub fn new(endpoint: ApiUrl) -> Self {
+        Self { endpoint }
     }
 
-    /// Authenticate the existing user or throw an error.
-    pub(super) async fn handle_user(
-        &'a self,
-        client: &mut PqStream<impl AsyncRead + AsyncWrite + Unpin + Send>,
-    ) -> auth::Result<AuthSuccess<NodeInfo>> {
-        // We reuse user handling logic from a production module.
-        console::handle_user(client, self, Self::get_auth_info, Self::wake_compute).await
+    pub fn url(&self) -> &str {
+        self.endpoint.as_str()
     }
-}
 
-impl Api<'_> {
-    /// This implementation fetches the auth info from a local postgres instance.
-    async fn get_auth_info(&self) -> Result<Option<AuthInfo>, GetAuthInfoError> {
+    async fn do_get_auth_info(
+        &self,
+        creds: &ClientCredentials<'_>,
+    ) -> Result<Option<AuthInfo>, GetAuthInfoError> {
         async {
             // Perhaps we could persist this connection, but then we'd have to
             // write more code for reopening it if it got closed, which doesn't
@@ -75,7 +55,7 @@ impl Api<'_> {
 
             tokio::spawn(connection);
             let query = "select rolpassword from pg_catalog.pg_authid where rolname = $1";
-            let rows = client.query(query, &[&self.creds.user]).await?;
+            let rows = client.query(query, &[&creds.user]).await?;
 
             // We can get at most one row, because `rolname` is unique.
             let row = match rows.get(0) {
@@ -84,7 +64,7 @@ impl Api<'_> {
                 // However, this is still a *valid* outcome which is very similar
                 // to getting `404 Not found` from the Neon console.
                 None => {
-                    warn!("user '{}' does not exist", self.creds.user);
+                    warn!("user '{}' does not exist", creds.user);
                     return Ok(None);
                 }
             };
@@ -98,23 +78,50 @@ impl Api<'_> {
             Ok(secret.or_else(|| parse_md5(entry).map(AuthInfo::Md5)))
         }
         .map_err(crate::error::log_error)
-        .instrument(info_span!("get_auth_info", mock = self.endpoint.as_str()))
+        .instrument(info_span!("postgres", url = self.endpoint.as_str()))
         .await
     }
 
-    /// We don't need to wake anything locally, so we just return the connection info.
-    pub async fn wake_compute(&self) -> Result<NodeInfo, WakeComputeError> {
+    async fn do_wake_compute(
+        &self,
+        creds: &ClientCredentials<'_>,
+    ) -> Result<NodeInfo, WakeComputeError> {
         let mut config = compute::ConnCfg::new();
         config
             .host(self.endpoint.host_str().unwrap_or("localhost"))
             .port(self.endpoint.port().unwrap_or(5432))
-            .dbname(self.creds.dbname)
-            .user(self.creds.user);
+            .dbname(creds.dbname)
+            .user(creds.user);
 
-        Ok(NodeInfo {
+        let node = NodeInfo {
             config,
             aux: Default::default(),
-        })
+        };
+
+        Ok(node)
+    }
+}
+
+#[async_trait]
+impl super::Api for Api {
+    #[tracing::instrument(skip_all)]
+    async fn get_auth_info(
+        &self,
+        _extra: &ConsoleReqExtra<'_>,
+        creds: &ClientCredentials<'_>,
+    ) -> Result<Option<AuthInfo>, GetAuthInfoError> {
+        self.do_get_auth_info(creds).await
+    }
+
+    #[tracing::instrument(skip_all)]
+    async fn wake_compute(
+        &self,
+        _extra: &ConsoleReqExtra<'_>,
+        creds: &ClientCredentials<'_>,
+    ) -> Result<CachedNodeInfo, WakeComputeError> {
+        self.do_wake_compute(creds)
+            .map_ok(CachedNodeInfo::new_uncached)
+            .await
     }
 }
 

proxy/src/console/provider/neon.rs

@@ -0,0 +1,196 @@
+//! Production console backend.
+
+use super::{
+    super::messages::{ConsoleError, GetRoleSecret, WakeCompute},
+    errors::{ApiError, GetAuthInfoError, WakeComputeError},
+    ApiCaches, AuthInfo, CachedNodeInfo, ConsoleReqExtra, NodeInfo,
+};
+use crate::{auth::ClientCredentials, compute, http, scram};
+use async_trait::async_trait;
+use futures::TryFutureExt;
+use reqwest::StatusCode as HttpStatusCode;
+use tracing::{error, info, info_span, warn, Instrument};
+
+#[derive(Clone)]
+pub struct Api {
+    endpoint: http::Endpoint,
+    caches: &'static ApiCaches,
+}
+
+impl Api {
+    /// Construct an API object containing the auth parameters.
+    pub fn new(endpoint: http::Endpoint, caches: &'static ApiCaches) -> Self {
+        Self { endpoint, caches }
+    }
+
+    pub fn url(&self) -> &str {
+        self.endpoint.url().as_str()
+    }
+
+    async fn do_get_auth_info(
+        &self,
+        extra: &ConsoleReqExtra<'_>,
+        creds: &ClientCredentials<'_>,
+    ) -> Result<Option<AuthInfo>, GetAuthInfoError> {
+        let request_id = uuid::Uuid::new_v4().to_string();
+        async {
+            let request = self
+                .endpoint
+                .get("proxy_get_role_secret")
+                .header("X-Request-ID", &request_id)
+                .query(&[("session_id", extra.session_id)])
+                .query(&[
+                    ("application_name", extra.application_name),
+                    ("project", Some(creds.project().expect("impossible"))),
+                    ("role", Some(creds.user)),
+                ])
+                .build()?;
+
+            info!(url = request.url().as_str(), "sending http request");
+            let response = self.endpoint.execute(request).await?;
+            let body = match parse_body::<GetRoleSecret>(response).await {
+                Ok(body) => body,
+                // Error 404 is special: it's ok not to have a secret.
+                Err(e) => match e.http_status_code() {
+                    Some(HttpStatusCode::NOT_FOUND) => return Ok(None),
+                    _otherwise => return Err(e.into()),
+                },
+            };
+
+            let secret = scram::ServerSecret::parse(&body.role_secret)
+                .map(AuthInfo::Scram)
+                .ok_or(GetAuthInfoError::BadSecret)?;
+
+            Ok(Some(secret))
+        }
+        .map_err(crate::error::log_error)
+        .instrument(info_span!("http", id = request_id))
+        .await
+    }
+
+    async fn do_wake_compute(
+        &self,
+        extra: &ConsoleReqExtra<'_>,
+        creds: &ClientCredentials<'_>,
+    ) -> Result<NodeInfo, WakeComputeError> {
+        let project = creds.project().expect("impossible");
+        let request_id = uuid::Uuid::new_v4().to_string();
+        async {
+            let request = self
+                .endpoint
+                .get("proxy_wake_compute")
+                .header("X-Request-ID", &request_id)
+                .query(&[("session_id", extra.session_id)])
+                .query(&[
+                    ("application_name", extra.application_name),
+                    ("project", Some(project)),
+                ])
+                .build()?;
+
+            info!(url = request.url().as_str(), "sending http request");
+            let response = self.endpoint.execute(request).await?;
+            let body = parse_body::<WakeCompute>(response).await?;
+
+            // Unfortunately, ownership won't let us use `Option::ok_or` here.
+            let (host, port) = match parse_host_port(&body.address) {
+                None => return Err(WakeComputeError::BadComputeAddress(body.address)),
+                Some(x) => x,
+            };
+
+            let mut config = compute::ConnCfg::new();
+            config
+                .host(host)
+                .port(port)
+                .dbname(creds.dbname)
+                .user(creds.user);
+
+            let node = NodeInfo {
+                config,
+                aux: body.aux.into(),
+            };
+
+            Ok(node)
+        }
+        .map_err(crate::error::log_error)
+        .instrument(info_span!("http", id = request_id))
+        .await
+    }
+}
+
+#[async_trait]
+impl super::Api for Api {
+    #[tracing::instrument(skip_all)]
+    async fn get_auth_info(
+        &self,
+        extra: &ConsoleReqExtra<'_>,
+        creds: &ClientCredentials<'_>,
+    ) -> Result<Option<AuthInfo>, GetAuthInfoError> {
+        self.do_get_auth_info(extra, creds).await
+    }
+
+    #[tracing::instrument(skip_all)]
+    async fn wake_compute(
+        &self,
+        extra: &ConsoleReqExtra<'_>,
+        creds: &ClientCredentials<'_>,
+    ) -> Result<CachedNodeInfo, WakeComputeError> {
+        let key = creds.project().expect("impossible");
+
+        // Every time we do a wakeup http request, the compute node will stay up
+        // for some time (highly depends on the console's scale-to-zero policy);
+        // The connection info remains the same during that period of time,
+        // which means that we might cache it to reduce the load and latency.
+        if let Some(cached) = self.caches.node_info.get(key) {
+            info!(key = key, "found cached compute node info");
+            return Ok(cached);
+        }
+
+        let node = self.do_wake_compute(extra, creds).await?;
+        let (_, cached) = self.caches.node_info.insert(key.into(), node);
+        info!(key = key, "created a cache entry for compute node info");
+
+        Ok(cached)
+    }
+}
+
+/// Parse http response body, taking status code into account.
+async fn parse_body<T: for<'a> serde::Deserialize<'a>>(
+    response: reqwest::Response,
+) -> Result<T, ApiError> {
+    let status = response.status();
+    if status.is_success() {
+        // We shouldn't log raw body because it may contain secrets.
+        info!("request succeeded, processing the body");
+        return Ok(response.json().await?);
+    }
+
+    // Don't throw an error here because it's not as important
+    // as the fact that the request itself has failed.
+    let body = response.json().await.unwrap_or_else(|e| {
+        warn!("failed to parse error body: {e}");
+        ConsoleError {
+            error: "reason unclear (malformed error message)".into(),
+        }
+    });
+
+    let text = body.error;
+    error!("console responded with an error ({status}): {text}");
+    Err(ApiError::Console { status, text })
+}
+
+fn parse_host_port(input: &str) -> Option<(&str, u16)> {
+    let (host, port) = input.split_once(':')?;
+    Some((host, port.parse().ok()?))
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_parse_host_port() {
+        let (host, port) = parse_host_port("127.0.0.1:5432").expect("failed to parse");
+        assert_eq!(host, "127.0.0.1");
+        assert_eq!(port, 5432);
+    }
+}

proxy/src/http/server.rs

@@ -9,8 +9,7 @@ async fn status_handler(_: Request<Body>) -> Result<Response<Body>, ApiError> {
 }
 
 fn make_router() -> RouterBuilder<hyper::Body, ApiError> {
-    let router = endpoint::make_router();
-    router.get("/v1/status", status_handler)
+    endpoint::make_router().get("/v1/status", status_handler)
 }
 
 pub async fn task_main(http_listener: TcpListener) -> anyhow::Result<()> {

proxy/src/http/websocket.rs

@@ -1,6 +1,6 @@
 use bytes::{Buf, Bytes};
 use futures::{Sink, Stream, StreamExt};
-use hyper::server::accept::{self};
+use hyper::server::accept;
 use hyper::server::conn::AddrIncoming;
 use hyper::upgrade::Upgraded;
 use hyper::{Body, Request, Response, StatusCode};
@@ -161,7 +161,7 @@ impl AsyncBufRead for WebSocketRW {
 
 async fn serve_websocket(
     websocket: HyperWebsocket,
-    config: &ProxyConfig,
+    config: &'static ProxyConfig,
     cancel_map: &CancelMap,
     session_id: uuid::Uuid,
     hostname: Option<String>,

proxy/src/main.rs

@@ -5,6 +5,7 @@
 //! in somewhat transparent manner (again via communication with control plane API).
 
 mod auth;
+mod cache;
 mod cancellation;
 mod compute;
 mod config;
@@ -12,7 +13,6 @@ mod console;
 mod error;
 mod http;
 mod metrics;
-mod mgmt;
 mod parse;
 mod proxy;
 mod sasl;
@@ -21,7 +21,6 @@ mod stream;
 mod url;
 mod waiters;
 
-use ::metrics::set_build_info_metric;
 use anyhow::{bail, Context};
 use clap::{self, Arg};
 use config::ProxyConfig;
@@ -29,8 +28,7 @@ use futures::FutureExt;
 use std::{borrow::Cow, future::Future, net::SocketAddr};
 use tokio::{net::TcpListener, task::JoinError};
 use tracing::{info, info_span, Instrument};
-use utils::project_git_version;
-use utils::sentry_init::init_sentry;
+use utils::{project_git_version, sentry_init::init_sentry};
 
 project_git_version!(GIT_VERSION);
 
@@ -51,124 +49,133 @@ async fn main() -> anyhow::Result<()> {
     // initialize sentry if SENTRY_DSN is provided
     let _sentry_guard = init_sentry(Some(GIT_VERSION.into()), &[]);
 
-    let arg_matches = cli().get_matches();
-
-    let tls_config = match (
-        arg_matches.get_one::<String>("tls-key"),
-        arg_matches.get_one::<String>("tls-cert"),
-    ) {
-        (Some(key_path), Some(cert_path)) => Some(config::configure_tls(key_path, cert_path)?),
-        (None, None) => None,
-        _ => bail!("either both or neither tls-key and tls-cert must be specified"),
-    };
-
-    let proxy_address: SocketAddr = arg_matches.get_one::<String>("proxy").unwrap().parse()?;
-    let mgmt_address: SocketAddr = arg_matches.get_one::<String>("mgmt").unwrap().parse()?;
-    let http_address: SocketAddr = arg_matches.get_one::<String>("http").unwrap().parse()?;
-
-    let metric_collection_config = match
-    (
-        arg_matches.get_one::<String>("metric-collection-endpoint"),
-        arg_matches.get_one::<String>("metric-collection-interval"),
-    ) {
-
-        (Some(endpoint), Some(interval)) => {
-            Some(config::MetricCollectionConfig {
-                endpoint: endpoint.parse()?,
-                interval: humantime::parse_duration(interval)?,
-            })
-        }
-        (None, None) => None,
-        _ => bail!("either both or neither metric-collection-endpoint and metric-collection-interval must be specified"),
-    };
-
-    let auth_backend = match arg_matches
-        .get_one::<String>("auth-backend")
-        .unwrap()
-        .as_str()
-    {
-        "console" => {
-            let url = arg_matches
-                .get_one::<String>("auth-endpoint")
-                .unwrap()
-                .parse()?;
-            let endpoint = http::Endpoint::new(url, reqwest::Client::new());
-            auth::BackendType::Console(Cow::Owned(endpoint), ())
-        }
-        "postgres" => {
-            let url = arg_matches
-                .get_one::<String>("auth-endpoint")
-                .unwrap()
-                .parse()?;
-            auth::BackendType::Postgres(Cow::Owned(url), ())
-        }
-        "link" => {
-            let url = arg_matches.get_one::<String>("uri").unwrap().parse()?;
-            auth::BackendType::Link(Cow::Owned(url))
-        }
-        other => bail!("unsupported auth backend: {other}"),
-    };
-
-    let config: &ProxyConfig = Box::leak(Box::new(ProxyConfig {
-        tls_config,
-        auth_backend,
-        metric_collection_config,
-    }));
-
     info!("Version: {GIT_VERSION}");
+    ::metrics::set_build_info_metric(GIT_VERSION);
+
+    let args = cli().get_matches();
+    let config = build_config(&args)?;
+
     info!("Authentication backend: {}", config.auth_backend);
 
     // Check that we can bind to address before further initialization
+    let http_address: SocketAddr = args.get_one::<String>("http").unwrap().parse()?;
     info!("Starting http on {http_address}");
     let http_listener = TcpListener::bind(http_address).await?.into_std()?;
 
+    let mgmt_address: SocketAddr = args.get_one::<String>("mgmt").unwrap().parse()?;
     info!("Starting mgmt on {mgmt_address}");
     let mgmt_listener = TcpListener::bind(mgmt_address).await?.into_std()?;
 
+    let proxy_address: SocketAddr = args.get_one::<String>("proxy").unwrap().parse()?;
     info!("Starting proxy on {proxy_address}");
     let proxy_listener = TcpListener::bind(proxy_address).await?;
 
     let 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)),
+        tokio::task::spawn_blocking(move || console::mgmt::thread_main(mgmt_listener)),
     ];
 
-    if let Some(wss_address) = arg_matches.get_one::<String>("wss") {
+    if let Some(wss_address) = args.get_one::<String>("wss") {
         let wss_address: SocketAddr = wss_address.parse()?;
-        info!("Starting wss on {}", wss_address);
+        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,
         )));
     }
 
-    if let Some(metric_collection_config) = &config.metric_collection_config {
+    // TODO: refactor.
+    if let Some(metric_collection) = &config.metric_collection {
         let hostname = hostname::get()?
             .into_string()
             .map_err(|e| anyhow::anyhow!("failed to get hostname {e:?}"))?;
 
         tasks.push(tokio::spawn(
             metrics::collect_metrics(
-                &metric_collection_config.endpoint,
-                metric_collection_config.interval,
+                &metric_collection.endpoint,
+                metric_collection.interval,
                 hostname,
             )
             .instrument(info_span!("collect_metrics")),
         ));
     }
 
-    let tasks = tasks.into_iter().map(flatten_err);
-
-    set_build_info_metric(GIT_VERSION);
     // This will block until all tasks have completed.
     // Furthermore, the first one to fail will cancel the rest.
+    let tasks = tasks.into_iter().map(flatten_err);
     let _: Vec<()> = futures::future::try_join_all(tasks).await?;
 
     Ok(())
 }
 
+/// ProxyConfig is created at proxy startup, and lives forever.
+fn build_config(args: &clap::ArgMatches) -> anyhow::Result<&'static ProxyConfig> {
+    let tls_config = match (
+        args.get_one::<String>("tls-key"),
+        args.get_one::<String>("tls-cert"),
+    ) {
+        (Some(key_path), Some(cert_path)) => Some(config::configure_tls(key_path, cert_path)?),
+        (None, None) => None,
+        _ => bail!("either both or neither tls-key and tls-cert must be specified"),
+    };
+
+    let metric_collection = match (
+        args.get_one::<String>("metric-collection-endpoint"),
+        args.get_one::<String>("metric-collection-interval"),
+    ) {
+        (Some(endpoint), Some(interval)) => Some(config::MetricCollectionConfig {
+            endpoint: endpoint.parse()?,
+            interval: humantime::parse_duration(interval)?,
+        }),
+        (None, None) => None,
+        _ => bail!(
+            "either both or neither metric-collection-endpoint \
+             and metric-collection-interval must be specified"
+        ),
+    };
+
+    let auth_backend = match args.get_one::<String>("auth-backend").unwrap().as_str() {
+        "console" => {
+            let config::CacheOptions { size, ttl } = args
+                .get_one::<String>("wake-compute-cache")
+                .unwrap()
+                .parse()?;
+
+            info!("Using NodeInfoCache (wake_compute) with size={size} ttl={ttl:?}");
+            let caches = Box::leak(Box::new(console::caches::ApiCaches {
+                node_info: console::caches::NodeInfoCache::new("node_info_cache", size, ttl),
+            }));
+
+            let url = args.get_one::<String>("auth-endpoint").unwrap().parse()?;
+            let endpoint = http::Endpoint::new(url, reqwest::Client::new());
+
+            let api = console::provider::neon::Api::new(endpoint, caches);
+            auth::BackendType::Console(Cow::Owned(api), ())
+        }
+        "postgres" => {
+            let url = args.get_one::<String>("auth-endpoint").unwrap().parse()?;
+            let api = console::provider::mock::Api::new(url);
+            auth::BackendType::Postgres(Cow::Owned(api), ())
+        }
+        "link" => {
+            let url = args.get_one::<String>("uri").unwrap().parse()?;
+            auth::BackendType::Link(Cow::Owned(url))
+        }
+        other => bail!("unsupported auth backend: {other}"),
+    };
+
+    let config = Box::leak(Box::new(ProxyConfig {
+        tls_config,
+        auth_backend,
+        metric_collection,
+    }));
+
+    Ok(config)
+}
+
 fn cli() -> clap::Command {
     clap::Command::new("Neon proxy/router")
         .disable_help_flag(true)
@@ -235,16 +242,27 @@ fn cli() -> clap::Command {
         .arg(
             Arg::new("metric-collection-endpoint")
                 .long("metric-collection-endpoint")
-                .help("metric collection HTTP endpoint"),
+                .help("http endpoint to receive periodic metric updates"),
         )
         .arg(
             Arg::new("metric-collection-interval")
                 .long("metric-collection-interval")
-                .help("metric collection interval"),
+                .help("how often metrics should be sent to a collection endpoint"),
+        )
+        .arg(
+            Arg::new("wake-compute-cache")
+                .long("wake-compute-cache")
+                .help("cache for `wake_compute` api method (use `size=0` to disable)")
+                .default_value(config::CacheOptions::DEFAULT_OPTIONS_NODE_INFO),
         )
 }
 
-#[test]
-fn verify_cli() {
-    cli().debug_assert();
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn verify_cli() {
+        cli().debug_assert();
+    }
 }

proxy/src/proxy.rs

@@ -2,9 +2,12 @@
 mod tests;
 
 use crate::{
-    auth,
+    auth::{self, backend::AuthSuccess},
     cancellation::{self, CancelMap},
+    compute::{self, PostgresConnection},
     config::{ProxyConfig, TlsConfig},
+    console::{self, messages::MetricsAuxInfo},
+    error::io_error,
     stream::{MeasuredStream, PqStream, Stream},
 };
 use anyhow::{bail, Context};
@@ -14,7 +17,10 @@ use once_cell::sync::Lazy;
 use pq_proto::{BeMessage as Be, FeStartupPacket, StartupMessageParams};
 use std::sync::Arc;
 use tokio::io::{AsyncRead, AsyncWrite};
-use tracing::{error, info, info_span, Instrument};
+use tracing::{error, info, info_span, warn, Instrument};
+
+/// Number of times we should retry the `/proxy_wake_compute` http request.
+const NUM_RETRIES_WAKE_COMPUTE: usize = 1;
 
 const ERR_INSECURE_CONNECTION: &str = "connection is insecure (try using `sslmode=require`)";
 const ERR_PROTO_VIOLATION: &str = "protocol violation";
@@ -35,6 +41,15 @@ static NUM_CONNECTIONS_CLOSED_COUNTER: Lazy<IntCounter> = Lazy::new(|| {
     .unwrap()
 });
 
+static NUM_CONNECTION_FAILURES: Lazy<IntCounterVec> = Lazy::new(|| {
+    register_int_counter_vec!(
+        "proxy_connection_failures_total",
+        "Number of connection failures (per kind).",
+        &["kind"],
+    )
+    .unwrap()
+});
+
 static NUM_BYTES_PROXIED_COUNTER: Lazy<IntCounterVec> = Lazy::new(|| {
     register_int_counter_vec!(
         "proxy_io_bytes_per_client",
@@ -82,11 +97,12 @@ pub async fn task_main(
     }
 }
 
+// TODO(tech debt): unite this with its twin below.
 pub async fn handle_ws_client(
-    config: &ProxyConfig,
+    config: &'static ProxyConfig,
     cancel_map: &CancelMap,
     session_id: uuid::Uuid,
-    stream: impl AsyncRead + AsyncWrite + Unpin + Send,
+    stream: impl AsyncRead + AsyncWrite + Unpin,
     hostname: Option<String>,
 ) -> anyhow::Result<()> {
     // The `closed` counter will increase when this future is destroyed.
@@ -99,7 +115,7 @@ pub async fn handle_ws_client(
     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 do_handshake = handshake(stream, None, cancel_map);
     let (mut stream, params) = match do_handshake.await? {
         Some(x) => x,
         None => return Ok(()), // it's a cancellation request
@@ -124,10 +140,10 @@ pub async fn handle_ws_client(
 }
 
 async fn handle_client(
-    config: &ProxyConfig,
+    config: &'static ProxyConfig,
     cancel_map: &CancelMap,
     session_id: uuid::Uuid,
-    stream: impl AsyncRead + AsyncWrite + Unpin + Send,
+    stream: impl AsyncRead + AsyncWrite + Unpin,
 ) -> anyhow::Result<()> {
     // The `closed` counter will increase when this future is destroyed.
     NUM_CONNECTIONS_ACCEPTED_COUNTER.inc();
@@ -136,7 +152,7 @@ async fn handle_client(
     }
 
     let tls = config.tls_config.as_ref();
-    let do_handshake = handshake(stream, tls, cancel_map).instrument(info_span!("handshake"));
+    let do_handshake = handshake(stream, tls, cancel_map);
     let (mut stream, params) = match do_handshake.await? {
         Some(x) => x,
         None => return Ok(()), // it's a cancellation request
@@ -165,6 +181,7 @@ async fn handle_client(
 /// For better testing experience, `stream` can be any object satisfying the traits.
 /// It's easier to work with owned `stream` here as we need to upgrade it to TLS;
 /// we also take an extra care of propagating only the select handshake errors to client.
+#[tracing::instrument(skip_all)]
 async fn handshake<S: AsyncRead + AsyncWrite + Unpin>(
     stream: S,
     mut tls: Option<&TlsConfig>,
@@ -226,6 +243,133 @@ async fn handshake<S: AsyncRead + AsyncWrite + Unpin>(
     }
 }
 
+/// Try to connect to the compute node once.
+#[tracing::instrument(name = "connect_once", skip_all)]
+async fn connect_to_compute_once(
+    node_info: &console::CachedNodeInfo,
+) -> Result<PostgresConnection, compute::ConnectionError> {
+    // If we couldn't connect, a cached connection info might be to blame
+    // (e.g. the compute node's address might've changed at the wrong time).
+    // Invalidate the cache entry (if any) to prevent subsequent errors.
+    let invalidate_cache = |_: &compute::ConnectionError| {
+        let is_cached = node_info.cached();
+        if is_cached {
+            warn!("invalidating stalled compute node info cache entry");
+            node_info.invalidate();
+        }
+
+        let label = match is_cached {
+            true => "compute_cached",
+            false => "compute_uncached",
+        };
+        NUM_CONNECTION_FAILURES.with_label_values(&[label]).inc();
+    };
+
+    node_info
+        .config
+        .connect()
+        .inspect_err(invalidate_cache)
+        .await
+}
+
+/// Try to connect to the compute node, retrying if necessary.
+/// This function might update `node_info`, so we take it by `&mut`.
+#[tracing::instrument(skip_all)]
+async fn connect_to_compute(
+    node_info: &mut console::CachedNodeInfo,
+    params: &StartupMessageParams,
+    extra: &console::ConsoleReqExtra<'_>,
+    creds: &auth::BackendType<'_, auth::ClientCredentials<'_>>,
+) -> Result<PostgresConnection, compute::ConnectionError> {
+    let mut num_retries: usize = NUM_RETRIES_WAKE_COMPUTE;
+    loop {
+        // Apply startup params to the (possibly, cached) compute node info.
+        node_info.config.set_startup_params(params);
+        match connect_to_compute_once(node_info).await {
+            Err(e) if num_retries > 0 => {
+                info!("compute node's state has changed; requesting a wake-up");
+                match creds.wake_compute(extra).map_err(io_error).await? {
+                    // Update `node_info` and try one more time.
+                    Some(mut new) => {
+                        new.config.reuse_password(&node_info.config);
+                        *node_info = new;
+                    }
+                    // Link auth doesn't work that way, so we just exit.
+                    None => return Err(e),
+                }
+            }
+            other => return other,
+        }
+
+        num_retries -= 1;
+        info!("retrying after wake-up ({num_retries} attempts left)");
+    }
+}
+
+/// Finish client connection initialization: confirm auth success, send params, etc.
+#[tracing::instrument(skip_all)]
+async fn prepare_client_connection(
+    node: &compute::PostgresConnection,
+    reported_auth_ok: bool,
+    session: cancellation::Session<'_>,
+    stream: &mut PqStream<impl AsyncRead + AsyncWrite + Unpin>,
+) -> anyhow::Result<()> {
+    // Register compute's query cancellation token and produce a new, unique one.
+    // The new token (cancel_key_data) will be sent to the client.
+    let cancel_key_data = session.enable_query_cancellation(node.cancel_closure.clone());
+
+    // Report authentication success if we haven't done this already.
+    // Note that we do this only (for the most part) after we've connected
+    // to a compute (see above) which performs its own authentication.
+    if !reported_auth_ok {
+        stream.write_message_noflush(&Be::AuthenticationOk)?;
+    }
+
+    // Forward all postgres connection params to the client.
+    // Right now the implementation is very hacky and inefficent (ideally,
+    // we don't need an intermediate hashmap), but at least it should be correct.
+    for (name, value) in &node.params {
+        // TODO: Theoretically, this could result in a big pile of params...
+        stream.write_message_noflush(&Be::ParameterStatus {
+            name: name.as_bytes(),
+            value: value.as_bytes(),
+        })?;
+    }
+
+    stream
+        .write_message_noflush(&Be::BackendKeyData(cancel_key_data))?
+        .write_message(&Be::ReadyForQuery)
+        .await?;
+
+    Ok(())
+}
+
+/// Forward bytes in both directions (client <-> compute).
+#[tracing::instrument(skip_all)]
+async fn proxy_pass(
+    client: impl AsyncRead + AsyncWrite + Unpin,
+    compute: impl AsyncRead + AsyncWrite + Unpin,
+    aux: &MetricsAuxInfo,
+) -> anyhow::Result<()> {
+    let m_sent = NUM_BYTES_PROXIED_COUNTER.with_label_values(&aux.traffic_labels("tx"));
+    let mut client = MeasuredStream::new(client, |cnt| {
+        // Number of bytes we sent to the client (outbound).
+        m_sent.inc_by(cnt as u64);
+    });
+
+    let m_recv = NUM_BYTES_PROXIED_COUNTER.with_label_values(&aux.traffic_labels("rx"));
+    let mut compute = MeasuredStream::new(compute, |cnt| {
+        // Number of bytes the client sent to the compute node (inbound).
+        m_recv.inc_by(cnt as u64);
+    });
+
+    // Starting from here we only proxy the client's traffic.
+    info!("performing the proxy pass...");
+    let _ = tokio::io::copy_bidirectional(&mut client, &mut compute).await?;
+
+    Ok(())
+}
+
 /// Thin connection context.
 struct Client<'a, S> {
     /// The underlying libpq protocol stream.
@@ -255,17 +399,17 @@ impl<'a, S> Client<'a, S> {
     }
 }
 
-impl<S: AsyncRead + AsyncWrite + Unpin + Send> Client<'_, S> {
+impl<S: AsyncRead + AsyncWrite + Unpin> Client<'_, S> {
     /// Let the client authenticate and connect to the designated compute node.
     async fn connect_to_db(self, session: cancellation::Session<'_>) -> anyhow::Result<()> {
         let Self {
             mut stream,
-            creds,
+            mut creds,
             params,
             session_id,
         } = self;
 
-        let extra = auth::ConsoleReqExtra {
+        let extra = console::ConsoleReqExtra {
             session_id, // aka this connection's id
             application_name: params.get("application_name"),
         };
@@ -278,54 +422,16 @@ impl<S: AsyncRead + AsyncWrite + Unpin + Send> Client<'_, S> {
         .instrument(info_span!("auth"))
         .await?;
 
-        let node = auth_result.value;
-        let (db, cancel_closure) = node
-            .config
-            .connect(params)
+        let AuthSuccess {
+            reported_auth_ok,
+            value: mut node_info,
+        } = auth_result;
+
+        let node = connect_to_compute(&mut node_info, params, &extra, &creds)
             .or_else(|e| stream.throw_error(e))
             .await?;
 
-        let cancel_key_data = session.enable_query_cancellation(cancel_closure);
-
-        // Report authentication success if we haven't done this already.
-        // Note that we do this only (for the most part) after we've connected
-        // to a compute (see above) which performs its own authentication.
-        if !auth_result.reported_auth_ok {
-            stream.write_message_noflush(&Be::AuthenticationOk)?;
-        }
-
-        // Forward all postgres connection params to the client.
-        // Right now the implementation is very hacky and inefficent (ideally,
-        // we don't need an intermediate hashmap), but at least it should be correct.
-        for (name, value) in &db.params {
-            // TODO: Theoretically, this could result in a big pile of params...
-            stream.write_message_noflush(&Be::ParameterStatus {
-                name: name.as_bytes(),
-                value: value.as_bytes(),
-            })?;
-        }
-
-        stream
-            .write_message_noflush(&Be::BackendKeyData(cancel_key_data))?
-            .write_message(&Be::ReadyForQuery)
-            .await?;
-
-        let m_sent = NUM_BYTES_PROXIED_COUNTER.with_label_values(&node.aux.traffic_labels("tx"));
-        let mut client = MeasuredStream::new(stream.into_inner(), |cnt| {
-            // Number of bytes we sent to the client (outbound).
-            m_sent.inc_by(cnt as u64);
-        });
-
-        let m_recv = NUM_BYTES_PROXIED_COUNTER.with_label_values(&node.aux.traffic_labels("rx"));
-        let mut db = MeasuredStream::new(db.stream, |cnt| {
-            // Number of bytes the client sent to the compute node (inbound).
-            m_recv.inc_by(cnt as u64);
-        });
-
-        // Starting from here we only proxy the client's traffic.
-        info!("performing the proxy pass...");
-        let _ = tokio::io::copy_bidirectional(&mut client, &mut db).await?;
-
-        Ok(())
+        prepare_client_connection(&node, reported_auth_ok, session, &mut stream).await?;
+        proxy_pass(stream.into_inner(), node.stream, &node_info.aux).await
     }
 }

pyproject.toml

@@ -33,7 +33,7 @@ psutil = "^5.9.4"
 types-psutil = "^5.9.5.4"
 types-toml = "^0.10.8"
 pytest-httpserver = "^1.0.6"
-aiohttp = "3.7"
+aiohttp = "3.7.4"
 
 [tool.poetry.dev-dependencies]
 flake8 = "^5.0.4"

scripts/export_import_between_pageservers.py

@@ -293,7 +293,7 @@ class NeonPageserverHttpClient(requests.Session):
 
     def timeline_detail(self, tenant_id: uuid.UUID, timeline_id: uuid.UUID) -> Dict[Any, Any]:
         res = self.get(
-            f"http://localhost:{self.port}/v1/tenant/{tenant_id.hex}/timeline/{timeline_id.hex}?include-non-incremental-logical-size=1"
+            f"http://localhost:{self.port}/v1/tenant/{tenant_id.hex}/timeline/{timeline_id.hex}?include-non-incremental-logical-size=true"
         )
         self.verbose_error(res)
         res_json = res.json()

scripts/pysync

@@ -4,6 +4,10 @@
 # It is intended to be a primary endpoint for all the people who want to
 # just setup test environment without going into details of python package management
 
+# on ubuntu 22.04 based `six` fails to install without this, because something imports keyring.
+# null keyring is fine, which means no secrets should be accessed.
+export PYTHON_KEYRING_BACKEND=keyring.backends.null.Keyring
+
 poetry config --list
 
 if [ -z "${CI}" ]; then

test_runner/fixtures/metrics.py

@@ -46,6 +46,12 @@ PAGESERVER_PER_TENANT_REMOTE_TIMELINE_CLIENT_METRICS: Tuple[str, ...] = (
     "pageserver_remote_physical_size",
 )
 
+PAGESERVER_GLOBAL_METRICS: Tuple[str, ...] = (
+    "pageserver_storage_operations_seconds_global_count",
+    "pageserver_storage_operations_seconds_global_sum",
+    "pageserver_storage_operations_seconds_global_bucket",
+)
+
 PAGESERVER_PER_TENANT_METRICS: Tuple[str, ...] = (
     "pageserver_current_logical_size",
     "pageserver_resident_physical_size",
@@ -61,13 +67,13 @@ PAGESERVER_PER_TENANT_METRICS: Tuple[str, ...] = (
     "pageserver_smgr_query_seconds_bucket",
     "pageserver_smgr_query_seconds_count",
     "pageserver_smgr_query_seconds_sum",
-    "pageserver_storage_operations_seconds_bucket",
-    "pageserver_storage_operations_seconds_count",
-    "pageserver_storage_operations_seconds_sum",
+    "pageserver_storage_operations_seconds_count_total",
+    "pageserver_storage_operations_seconds_sum_total",
     "pageserver_wait_lsn_seconds_bucket",
     "pageserver_wait_lsn_seconds_count",
     "pageserver_wait_lsn_seconds_sum",
     "pageserver_created_persistent_files_total",
     "pageserver_written_persistent_bytes_total",
+    "pageserver_tenant_states_count",
     *PAGESERVER_PER_TENANT_REMOTE_TIMELINE_CLIENT_METRICS,
 )

test_runner/fixtures/neon_fixtures.py

@@ -1232,9 +1232,9 @@ class PageserverHttpClient(requests.Session):
 
         params = {}
         if include_non_incremental_logical_size:
-            params["include-non-incremental-logical-size"] = "yes"
+            params["include-non-incremental-logical-size"] = "true"
         if include_timeline_dir_layer_file_size_sum:
-            params["include-timeline-dir-layer-file-size-sum"] = "yes"
+            params["include-timeline-dir-layer-file-size-sum"] = "true"
 
         res = self.get(
             f"http://localhost:{self.port}/v1/tenant/{tenant_id}/timeline", params=params
@@ -1276,9 +1276,9 @@ class PageserverHttpClient(requests.Session):
     ) -> Dict[Any, Any]:
         params = {}
         if include_non_incremental_logical_size:
-            params["include-non-incremental-logical-size"] = "yes"
+            params["include-non-incremental-logical-size"] = "true"
         if include_timeline_dir_layer_file_size_sum:
-            params["include-timeline-dir-layer-file-size-sum"] = "yes"
+            params["include-timeline-dir-layer-file-size-sum"] = "true"
 
         res = self.get(
             f"http://localhost:{self.port}/v1/tenant/{tenant_id}/timeline/{timeline_id}",

test_runner/regress/test_tenant_conf.py

@@ -2,7 +2,15 @@ from contextlib import closing
 
 import psycopg2.extras
 from fixtures.log_helper import log
-from fixtures.neon_fixtures import NeonEnvBuilder
+from fixtures.neon_fixtures import (
+    LocalFsStorage,
+    NeonEnvBuilder,
+    RemoteStorageKind,
+    assert_tenant_status,
+    wait_for_upload,
+)
+from fixtures.types import Lsn
+from fixtures.utils import wait_until
 
 
 def test_tenant_config(neon_env_builder: NeonEnvBuilder):
@@ -158,3 +166,46 @@ tenant_config={checkpoint_distance = 10000, compaction_target_size = 1048576}"""
                     "pitr_interval": 60,
                 }.items()
             )
+
+
+def test_creating_tenant_conf_after_attach(neon_env_builder: NeonEnvBuilder):
+    neon_env_builder.enable_remote_storage(
+        remote_storage_kind=RemoteStorageKind.LOCAL_FS,
+        test_name="test_creating_tenant_conf_after_attach",
+    )
+
+    env = neon_env_builder.init_start()
+    assert isinstance(env.remote_storage, LocalFsStorage)
+
+    # tenant is created with defaults, as in without config file
+    (tenant_id, timeline_id) = env.neon_cli.create_tenant()
+    config_path = env.repo_dir / "tenants" / str(tenant_id) / "config"
+    assert config_path.exists(), "config file is always initially created"
+
+    http_client = env.pageserver.http_client()
+
+    detail = http_client.timeline_detail(tenant_id, timeline_id)
+    last_record_lsn = Lsn(detail["last_record_lsn"])
+    assert last_record_lsn.lsn_int != 0, "initdb must have executed"
+
+    wait_for_upload(http_client, tenant_id, timeline_id, last_record_lsn)
+
+    http_client.tenant_detach(tenant_id)
+
+    assert not config_path.exists(), "detach did not remove config file"
+
+    http_client.tenant_attach(tenant_id)
+    wait_until(
+        number_of_iterations=5,
+        interval=1,
+        func=lambda: assert_tenant_status(http_client, tenant_id, "Active"),
+    )
+
+    env.neon_cli.config_tenant(tenant_id, {"gc_horizon": "1000000"})
+    contents_first = config_path.read_text()
+    env.neon_cli.config_tenant(tenant_id, {"gc_horizon": "0"})
+    contents_later = config_path.read_text()
+
+    # dont test applying the setting here, we have that another test case to show it
+    # we just care about being able to create the file
+    assert len(contents_first) > len(contents_later)

test_runner/regress/test_tenant_detach.py

@@ -6,6 +6,7 @@ from threading import Thread
 import asyncpg
 import pytest
 from fixtures.log_helper import log
+from fixtures.metrics import parse_metrics
 from fixtures.neon_fixtures import (
     NeonEnv,
     NeonEnvBuilder,
@@ -59,11 +60,11 @@ def test_tenant_reattach(
     # create new nenant
     tenant_id, timeline_id = env.neon_cli.create_tenant()
 
-    pg = env.postgres.create_start("main", tenant_id=tenant_id)
-    with pg.cursor() as cur:
-        cur.execute("CREATE TABLE t(key int primary key, value text)")
-        cur.execute("INSERT INTO t SELECT generate_series(1,100000), 'payload'")
-        current_lsn = Lsn(query_scalar(cur, "SELECT pg_current_wal_flush_lsn()"))
+    with env.postgres.create_start("main", tenant_id=tenant_id) as pg:
+        with pg.cursor() as cur:
+            cur.execute("CREATE TABLE t(key int primary key, value text)")
+            cur.execute("INSERT INTO t SELECT generate_series(1,100000), 'payload'")
+            current_lsn = Lsn(query_scalar(cur, "SELECT pg_current_wal_flush_lsn()"))
 
     # Wait for the all data to be processed by the pageserver and uploaded in remote storage
     wait_for_last_record_lsn(pageserver_http, tenant_id, timeline_id, current_lsn)
@@ -78,15 +79,34 @@ def test_tenant_reattach(
         ".*failed to perform remote task UploadMetadata.*, will retry.*"
     )
 
+    ps_metrics = parse_metrics(pageserver_http.get_metrics(), "pageserver")
+    tenant_metric_filter = {
+        "tenant_id": str(tenant_id),
+        "timeline_id": str(timeline_id),
+    }
+    pageserver_last_record_lsn_before_detach = int(
+        ps_metrics.query_one("pageserver_last_record_lsn", filter=tenant_metric_filter).value
+    )
+
     pageserver_http.tenant_detach(tenant_id)
     pageserver_http.tenant_attach(tenant_id)
 
-    with pg.cursor() as cur:
-        assert query_scalar(cur, "SELECT count(*) FROM t") == 100000
+    time.sleep(1)  # for metrics propagation
 
-    # Check that we had to retry the downloads
-    assert env.pageserver.log_contains(".*list prefixes.*failed, will retry.*")
-    assert env.pageserver.log_contains(".*download.*failed, will retry.*")
+    ps_metrics = parse_metrics(pageserver_http.get_metrics(), "pageserver")
+    pageserver_last_record_lsn = int(
+        ps_metrics.query_one("pageserver_last_record_lsn", filter=tenant_metric_filter).value
+    )
+
+    assert pageserver_last_record_lsn_before_detach == pageserver_last_record_lsn
+
+    with env.postgres.create_start("main", tenant_id=tenant_id) as pg:
+        with pg.cursor() as cur:
+            assert query_scalar(cur, "SELECT count(*) FROM t") == 100000
+
+        # Check that we had to retry the downloads
+        assert env.pageserver.log_contains(".*list prefixes.*failed, will retry.*")
+        assert env.pageserver.log_contains(".*download.*failed, will retry.*")
 
 
 num_connections = 10

test_runner/regress/test_tenants.py

@@ -1,5 +1,6 @@
 import os
 import shutil
+import time
 from contextlib import closing
 from datetime import datetime
 from pathlib import Path
@@ -8,6 +9,7 @@ from typing import List
 import pytest
 from fixtures.log_helper import log
 from fixtures.metrics import (
+    PAGESERVER_GLOBAL_METRICS,
     PAGESERVER_PER_TENANT_METRICS,
     PAGESERVER_PER_TENANT_REMOTE_TIMELINE_CLIENT_METRICS,
     parse_metrics,
@@ -160,6 +162,14 @@ def test_metrics_normal_work(neon_env_builder: NeonEnvBuilder):
             f"process_start_time_seconds (UTC): {datetime.fromtimestamp(metrics.query_one('process_start_time_seconds').value)}"
         )
 
+    # Test (a subset of) pageserver global metrics
+    for metric in PAGESERVER_GLOBAL_METRICS:
+        ps_samples = ps_metrics.query_all(metric, {})
+        assert len(ps_samples) > 0
+        for sample in ps_samples:
+            labels = ",".join([f'{key}="{value}"' for key, value in sample.labels.items()])
+            log.info(f"{sample.name}{{{labels}}} {sample.value}")
+
 
 @pytest.mark.parametrize(
     "remote_storage_kind",
@@ -259,7 +269,7 @@ def test_pageserver_with_empty_tenants(
         files_in_timelines_dir == 0
     ), f"Tenant {tenant_with_empty_timelines_dir} should have an empty timelines/ directory"
 
-    # Trigger timeline reinitialization after pageserver restart
+    # Trigger timeline re-initialization after pageserver restart
     env.postgres.stop_all()
     env.pageserver.stop()
 
@@ -278,7 +288,51 @@ def test_pageserver_with_empty_tenants(
         broken_tenant["state"] == "Broken"
     ), f"Tenant {tenant_without_timelines_dir} without timelines dir should be broken"
 
+    broken_tenant_status = client.tenant_status(tenant_without_timelines_dir)
+    assert (
+        broken_tenant_status["state"] == "Broken"
+    ), f"Tenant {tenant_without_timelines_dir} without timelines dir should be broken"
+
+    assert env.pageserver.log_contains(".*Setting tenant as Broken state, reason:.*")
+
     [loaded_tenant] = [t for t in tenants if t["id"] == str(tenant_with_empty_timelines_dir)]
     assert (
         loaded_tenant["state"] == "Active"
     ), "Tenant {tenant_with_empty_timelines_dir} with empty timelines dir should be active and ready for timeline creation"
+
+    loaded_tenant_status = client.tenant_status(tenant_with_empty_timelines_dir)
+    assert (
+        loaded_tenant_status["state"] == "Active"
+    ), f"Tenant {tenant_with_empty_timelines_dir} without timelines dir should be active"
+
+    time.sleep(1)  # to allow metrics propagation
+
+    ps_metrics = parse_metrics(client.get_metrics(), "pageserver")
+    broken_tenants_metric_filter = {
+        "tenant_id": str(tenant_without_timelines_dir),
+        "state": "broken",
+    }
+    active_tenants_metric_filter = {
+        "tenant_id": str(tenant_with_empty_timelines_dir),
+        "state": "active",
+    }
+
+    tenant_active_count = int(
+        ps_metrics.query_one(
+            "pageserver_tenant_states_count", filter=active_tenants_metric_filter
+        ).value
+    )
+
+    assert (
+        tenant_active_count == 1
+    ), f"Tenant {tenant_with_empty_timelines_dir} should have metric as active"
+
+    tenant_broken_count = int(
+        ps_metrics.query_one(
+            "pageserver_tenant_states_count", filter=broken_tenants_metric_filter
+        ).value
+    )
+
+    assert (
+        tenant_broken_count == 1
+    ), f"Tenant {tenant_without_timelines_dir} should have metric as broken"

workspace_hack/Cargo.toml

@@ -25,6 +25,7 @@ futures-channel = { version = "0.3", features = ["sink"] }
 futures-executor = { version = "0.3" }
 futures-task = { version = "0.3", default-features = false, features = ["std"] }
 futures-util = { version = "0.3", features = ["channel", "io", "sink"] }
+hashbrown = { version = "0.12", features = ["raw"] }
 indexmap = { version = "1", default-features = false, features = ["std"] }
 itertools = { version = "0.10" }
 libc = { version = "0.2", features = ["extra_traits"] }
@@ -58,6 +59,7 @@ url = { version = "2", features = ["serde"] }
 anyhow = { version = "1", features = ["backtrace"] }
 bytes = { version = "1", features = ["serde"] }
 either = { version = "1" }
+hashbrown = { version = "0.12", features = ["raw"] }
 indexmap = { version = "1", default-features = false, features = ["std"] }
 itertools = { version = "0.10" }
 libc = { version = "0.2", features = ["extra_traits"] }