need less ports

## Describe your changes
Add Error enum for tenant state response to allow better error handling
in mgmt api
## Issue ticket number and link
#2238
## Checklist before requesting a review
- [x] I have performed a self-review of my code.
- [ ] If it is a core feature, I have added thorough tests.
- [ ] Do we need to implement analytics? if so did you add the relevant
metrics to the dashboard?
- [ ] If this PR requires public announcement, mark it with
/release-notes label and add several sentences in this section.

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

@@ -19,7 +19,7 @@ inputs:
   run_in_parallel:
     description: 'Whether to run tests in parallel'
     required: false
-    default: 'true'
+    default: 'false'
   save_perf_report:
     description: 'Whether to upload the performance report, if true PERF_TEST_RESULT_CONNSTR env variable should be set'
     required: false
@@ -171,7 +171,7 @@ runs:
           --junitxml=$TEST_OUTPUT/junit.xml \
           --alluredir=$TEST_OUTPUT/allure/results \
           --tb=short \
-          --verbose \
+          --verbose -k "test_forward or test_create_snapsh" -x \
           -rA $TEST_SELECTION $EXTRA_PARAMS
 
         if [[ "${{ inputs.save_perf_report }}" == "true" ]]; then

.github/ansible/deploy.yaml

@@ -91,6 +91,15 @@
       tags:
       - pageserver
 
+    # used in `pageserver.service` template
+    - name: learn current availability_zone
+      shell:
+        cmd: "curl -s http://169.254.169.254/latest/meta-data/placement/availability-zone"
+      register: ec2_availability_zone
+
+    - set_fact: 
+        ec2_availability_zone={{ ec2_availability_zone.stdout }}
+
     - name: upload systemd service definition
       ansible.builtin.template:
         src: systemd/pageserver.service
@@ -118,7 +127,7 @@
         cmd: |
           INSTANCE_ID=$(curl -s http://169.254.169.254/latest/meta-data/instance-id)
           curl -sfS -H "Authorization: Bearer {{ CONSOLE_API_TOKEN }}" {{ console_mgmt_base_url }}/management/api/v2/pageservers/$INSTANCE_ID | jq '.version = {{ current_version }}' > /tmp/new_version
-          curl -sfS -H "Authorization: Bearer {{ CONSOLE_API_TOKEN }}" -X POST -d@/tmp/new_version {{ console_mgmt_base_url }}/management/api/v2/pageservers
+          curl -sfS -H "Authorization: Bearer {{ CONSOLE_API_TOKEN }}" -H "Content-Type: application/json" -X POST -d@/tmp/new_version {{ console_mgmt_base_url }}/management/api/v2/pageservers
       tags:
       - pageserver
 
@@ -153,6 +162,15 @@
       tags:
       - safekeeper
 
+    # used in `safekeeper.service` template
+    - name: learn current availability_zone
+      shell:
+        cmd: "curl -s http://169.254.169.254/latest/meta-data/placement/availability-zone"
+      register: ec2_availability_zone
+
+    - set_fact: 
+        ec2_availability_zone={{ ec2_availability_zone.stdout }}
+
     # in the future safekeepers should discover pageservers byself
     # but currently use first pageserver that was discovered
     - name: set first pageserver var for safekeepers
@@ -188,6 +206,6 @@
         cmd: |
           INSTANCE_ID=$(curl -s http://169.254.169.254/latest/meta-data/instance-id)
           curl -sfS -H "Authorization: Bearer {{ CONSOLE_API_TOKEN }}" {{ console_mgmt_base_url }}/management/api/v2/safekeepers/$INSTANCE_ID | jq '.version = {{ current_version }}' > /tmp/new_version
-          curl -sfS -H "Authorization: Bearer {{ CONSOLE_API_TOKEN }}" -X POST -d@/tmp/new_version {{ console_mgmt_base_url }}/management/api/v2/safekeepers
+          curl -sfS -H "Authorization: Bearer {{ CONSOLE_API_TOKEN }}" -H "Content-Type: application/json" -X POST -d@/tmp/new_version {{ console_mgmt_base_url }}/management/api/v2/safekeepers
       tags:
       - safekeeper

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

@@ -27,6 +27,8 @@ storage:
           ansible_host:  i-0cd8d316ecbb715be
         pageserver-1.eu-central-1.aws.neon.tech:
           ansible_host:  i-090044ed3d383fef0
+        pageserver-2.eu-central-1.aws.neon.tech:
+          ansible_host:  i-033584edf3f4b6742
 
     safekeepers:
       hosts:

.github/ansible/scripts/init_pageserver.sh

@@ -26,7 +26,7 @@ EOF
 if ! curl -sf -H "Authorization: Bearer {{ CONSOLE_API_TOKEN }}" {{ console_mgmt_base_url }}/management/api/v2/pageservers/${INSTANCE_ID} -o /dev/null; then
 
     # not registered, so register it now
-    ID=$(curl -sf -X POST -H "Authorization: Bearer {{ CONSOLE_API_TOKEN }}" {{ console_mgmt_base_url }}/management/api/v2/pageservers -d@/tmp/payload | jq -r '.id')
+    ID=$(curl -sf -X POST -H "Authorization: Bearer {{ CONSOLE_API_TOKEN }}" -H "Content-Type: application/json" {{ console_mgmt_base_url }}/management/api/v2/pageservers -d@/tmp/payload | jq -r '.id')
 
     # init pageserver
     sudo -u pageserver /usr/local/bin/pageserver -c "id=${ID}" -c "pg_distrib_dir='/usr/local'" --init -D /storage/pageserver/data

.github/ansible/scripts/init_safekeeper.sh

@@ -25,7 +25,7 @@ EOF
 if ! curl -sf -H "Authorization: Bearer {{ CONSOLE_API_TOKEN }}" {{ console_mgmt_base_url }}/management/api/v2/safekeepers/${INSTANCE_ID} -o /dev/null; then
 
     # not registered, so register it now
-    ID=$(curl -sf -X POST -H "Authorization: Bearer {{ CONSOLE_API_TOKEN }}" {{ console_mgmt_base_url }}/management/api/v2/safekeepers -d@/tmp/payload | jq -r '.id')
+    ID=$(curl -sf -X POST -H "Authorization: Bearer {{ CONSOLE_API_TOKEN }}" -H "Content-Type: application/json" {{ console_mgmt_base_url }}/management/api/v2/safekeepers -d@/tmp/payload | jq -r '.id')
     # init safekeeper
     sudo -u safekeeper /usr/local/bin/safekeeper --id ${ID} --init -D /storage/safekeeper/data
 fi

.github/ansible/systemd/pageserver.service

@@ -6,7 +6,7 @@ After=network.target auditd.service
 Type=simple
 User=pageserver
 Environment=RUST_BACKTRACE=1 NEON_REPO_DIR=/storage/pageserver LD_LIBRARY_PATH=/usr/local/v14/lib SENTRY_DSN={{ SENTRY_URL_PAGESERVER }} SENTRY_ENVIRONMENT={{ sentry_environment }}
-ExecStart=/usr/local/bin/pageserver -c "pg_distrib_dir='/usr/local'" -c "listen_pg_addr='0.0.0.0:6400'" -c "listen_http_addr='0.0.0.0:9898'" -c "broker_endpoint='{{ broker_endpoint }}'" -D /storage/pageserver/data
+ExecStart=/usr/local/bin/pageserver -c "pg_distrib_dir='/usr/local'" -c "listen_pg_addr='0.0.0.0:6400'" -c "listen_http_addr='0.0.0.0:9898'" -c "broker_endpoint='{{ broker_endpoint }}'" -c "availability_zone='{{ ec2_availability_zone }}'" -D /storage/pageserver/data
 ExecReload=/bin/kill -HUP $MAINPID
 KillMode=mixed
 KillSignal=SIGINT

.github/ansible/systemd/safekeeper.service

@@ -6,7 +6,7 @@ After=network.target auditd.service
 Type=simple
 User=safekeeper
 Environment=RUST_BACKTRACE=1 NEON_REPO_DIR=/storage/safekeeper/data LD_LIBRARY_PATH=/usr/local/v14/lib SENTRY_DSN={{ SENTRY_URL_SAFEKEEPER }} SENTRY_ENVIRONMENT={{ sentry_environment }}
-ExecStart=/usr/local/bin/safekeeper -l {{ inventory_hostname }}{{ hostname_suffix }}:6500 --listen-http {{ inventory_hostname }}{{ hostname_suffix }}:7676 -D /storage/safekeeper/data --broker-endpoint={{ broker_endpoint }} --remote-storage='{bucket_name="{{bucket_name}}", bucket_region="{{bucket_region}}", prefix_in_bucket="{{ safekeeper_s3_prefix }}"}'
+ExecStart=/usr/local/bin/safekeeper -l {{ inventory_hostname }}{{ hostname_suffix }}:6500 --listen-http {{ inventory_hostname }}{{ hostname_suffix }}:7676 -D /storage/safekeeper/data --broker-endpoint={{ broker_endpoint }} --remote-storage='{bucket_name="{{bucket_name}}", bucket_region="{{bucket_region}}", prefix_in_bucket="{{ safekeeper_s3_prefix }}"}' --availability-zone={{ ec2_availability_zone }}
 ExecReload=/bin/kill -HUP $MAINPID
 KillMode=mixed
 KillSignal=SIGINT

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

@@ -1,6 +1,22 @@
 # Helm chart values for neon-proxy-scram.
 # This is a YAML-formatted file.
 
+deploymentStrategy:
+  type: RollingUpdate
+  rollingUpdate:
+    maxSurge: 100%
+    maxUnavailable: 50%
+
+# Delay the kill signal by 7 days (7 * 24 * 60 * 60)
+# The pod(s) will stay in Terminating, keeps the existing connections
+# but doesn't receive new ones
+containerLifecycle:
+  preStop:
+    exec:
+      command: ["/bin/sh", "-c", "sleep 604800"]
+terminationGracePeriodSeconds: 604800
+
+
 image:
   repository: neondatabase/neon
 

.github/workflows/build_and_test.yml

@@ -5,6 +5,7 @@ on:
     branches:
       - main
       - release
+      - tmp-repro
   pull_request:
 
 defaults:
@@ -74,15 +75,12 @@ jobs:
       - name: Install Python deps
         run: ./scripts/pysync
 
-      - name: Run isort to ensure code format
-        run: poetry run isort --diff --check .
+      - name: Run ruff to ensure code format
+        run: poetry run ruff .
 
       - name: Run black to ensure code format
         run: poetry run black --diff --check .
 
-      - name: Run flake8 to ensure code format
-        run: poetry run flake8 .
-
       - name: Run mypy to check types
         run: poetry run mypy .
 
@@ -551,6 +549,48 @@ jobs:
       - name: Cleanup ECR folder
         run: rm -rf ~/.ecr
 
+
+  neon-image-depot:
+    # For testing this will run side-by-side for a few merges.
+    # This action is not really optimized yet, but gets the job done
+    runs-on: [ self-hosted, gen3, small ]
+    needs: [ tag ]
+    container: 369495373322.dkr.ecr.eu-central-1.amazonaws.com/base:pinned
+    permissions:
+      contents: read
+      id-token: write
+
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v3
+        with:
+          submodules: true
+          fetch-depth: 0
+
+      - name: Setup go
+        uses: actions/setup-go@v3
+        with:
+          go-version: '1.19'
+
+      - name: Set up Depot CLI
+        uses: depot/setup-action@v1
+
+      - name: Install Crane & ECR helper
+        run: 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: Build and push
+        uses: depot/build-push-action@v1
+        with:
+          # if no depot.json file is at the root of your repo, you must specify the project id
+          project: nrdv0s4kcs
+          push: true
+          tags: 369495373322.dkr.ecr.eu-central-1.amazonaws.com/neon:depot-${{needs.tag.outputs.build-tag}}
+
   compute-tools-image:
     runs-on: [ self-hosted, gen3, large ]
     needs: [ tag ]

.github/workflows/release.yml

@@ -31,3 +31,4 @@ jobs:
         head: releases/${{ steps.date.outputs.date }}
         base: release
         title: Release ${{ steps.date.outputs.date }}
+        team_reviewers: release

Cargo.lock

@@ -851,6 +851,7 @@ dependencies = [
  "futures",
  "hyper",
  "notify",
+ "num_cpus",
  "opentelemetry",
  "postgres",
  "regex",
@@ -3323,6 +3324,7 @@ dependencies = [
  "async-trait",
  "byteorder",
  "bytes",
+ "chrono",
  "clap 4.1.4",
  "const_format",
  "crc32c",
@@ -3332,7 +3334,6 @@ dependencies = [
  "humantime",
  "hyper",
  "metrics",
- "nix",
  "once_cell",
  "parking_lot",
  "postgres",
@@ -4531,9 +4532,9 @@ dependencies = [
  "metrics",
  "nix",
  "once_cell",
+ "pin-project-lite",
  "rand",
  "routerify",
- "rustls",
  "sentry",
  "serde",
  "serde_json",
@@ -4544,7 +4545,6 @@ dependencies = [
  "tempfile",
  "thiserror",
  "tokio",
- "tokio-rustls",
  "tracing",
  "tracing-subscriber",
  "url",
@@ -4885,6 +4885,7 @@ dependencies = [
  "socket2",
  "syn",
  "tokio",
+ "tokio-rustls",
  "tokio-util",
  "tonic",
  "tower",

Cargo.toml

@@ -64,6 +64,7 @@ md5 = "0.7.0"
 memoffset = "0.8"
 nix = "0.26"
 notify = "5.0.0"
+num_cpus = "1.15"
 num-traits = "0.2.15"
 once_cell = "1.13"
 opentelemetry = "0.18.0"

Dockerfile

@@ -39,7 +39,7 @@ ARG CACHEPOT_BUCKET=neon-github-dev
 
 COPY --from=pg-build /home/nonroot/pg_install/v14/include/postgresql/server pg_install/v14/include/postgresql/server
 COPY --from=pg-build /home/nonroot/pg_install/v15/include/postgresql/server pg_install/v15/include/postgresql/server
-COPY . .
+COPY --chown=nonroot . .
 
 # Show build caching stats to check if it was used in the end.
 # Has to be the part of the same RUN since cachepot daemon is killed in the end of this RUN, losing the compilation stats.

Dockerfile.compute-node

@@ -225,6 +225,81 @@ RUN wget https://github.com/iCyberon/pg_hashids/archive/refs/tags/v1.2.1.tar.gz
     make -j $(getconf _NPROCESSORS_ONLN) install PG_CONFIG=/usr/local/pgsql/bin/pg_config USE_PGXS=1 && \
     echo 'trusted = true' >> /usr/local/pgsql/share/extension/pg_hashids.control
 
+#########################################################################################
+#
+# Layer "rum-pg-build"
+# compile rum extension
+#
+#########################################################################################
+FROM build-deps AS rum-pg-build
+COPY --from=pg-build /usr/local/pgsql/ /usr/local/pgsql/
+
+RUN wget https://github.com/postgrespro/rum/archive/refs/tags/1.3.13.tar.gz -O rum.tar.gz && \
+    mkdir rum-src && cd rum-src && tar xvzf ../rum.tar.gz --strip-components=1 -C . && \
+    make -j $(getconf _NPROCESSORS_ONLN) PG_CONFIG=/usr/local/pgsql/bin/pg_config USE_PGXS=1 && \
+    make -j $(getconf _NPROCESSORS_ONLN) install PG_CONFIG=/usr/local/pgsql/bin/pg_config USE_PGXS=1 && \
+    echo 'trusted = true' >> /usr/local/pgsql/share/extension/rum.control
+
+#########################################################################################
+#
+# Layer "pgtap-pg-build"
+# compile pgTAP extension
+#
+#########################################################################################
+FROM build-deps AS pgtap-pg-build
+COPY --from=pg-build /usr/local/pgsql/ /usr/local/pgsql/
+
+RUN wget https://github.com/theory/pgtap/archive/refs/tags/v1.2.0.tar.gz -O pgtap.tar.gz && \
+    mkdir pgtap-src && cd pgtap-src && tar xvzf ../pgtap.tar.gz --strip-components=1 -C . && \
+    make -j $(getconf _NPROCESSORS_ONLN) PG_CONFIG=/usr/local/pgsql/bin/pg_config && \
+    make -j $(getconf _NPROCESSORS_ONLN) install PG_CONFIG=/usr/local/pgsql/bin/pg_config && \
+    echo 'trusted = true' >> /usr/local/pgsql/share/extension/pgtap.control
+
+#########################################################################################
+#
+# Layer "prefix-pg-build"
+# compile Prefix extension
+#
+#########################################################################################
+FROM build-deps AS prefix-pg-build
+COPY --from=pg-build /usr/local/pgsql/ /usr/local/pgsql/
+
+RUN wget https://github.com/dimitri/prefix/archive/refs/tags/v1.2.9.tar.gz -O prefix.tar.gz && \
+    mkdir prefix-src && cd prefix-src && tar xvzf ../prefix.tar.gz --strip-components=1 -C . && \
+    make -j $(getconf _NPROCESSORS_ONLN) PG_CONFIG=/usr/local/pgsql/bin/pg_config && \
+    make -j $(getconf _NPROCESSORS_ONLN) install PG_CONFIG=/usr/local/pgsql/bin/pg_config && \
+    echo 'trusted = true' >> /usr/local/pgsql/share/extension/prefix.control
+
+#########################################################################################
+#
+# Layer "hll-pg-build"
+# compile hll extension
+#
+#########################################################################################
+FROM build-deps AS hll-pg-build
+COPY --from=pg-build /usr/local/pgsql/ /usr/local/pgsql/
+
+RUN wget https://github.com/citusdata/postgresql-hll/archive/refs/tags/v2.17.tar.gz -O hll.tar.gz && \
+    mkdir hll-src && cd hll-src && tar xvzf ../hll.tar.gz --strip-components=1 -C . && \
+    make -j $(getconf _NPROCESSORS_ONLN) PG_CONFIG=/usr/local/pgsql/bin/pg_config && \
+    make -j $(getconf _NPROCESSORS_ONLN) install PG_CONFIG=/usr/local/pgsql/bin/pg_config && \
+    echo 'trusted = true' >> /usr/local/pgsql/share/extension/hll.control
+
+#########################################################################################
+#
+# Layer "plpgsql-check-pg-build"
+# compile plpgsql_check extension
+#
+#########################################################################################
+FROM build-deps AS plpgsql-check-pg-build
+COPY --from=pg-build /usr/local/pgsql/ /usr/local/pgsql/
+
+RUN wget https://github.com/okbob/plpgsql_check/archive/refs/tags/v2.3.2.tar.gz -O plpgsql_check.tar.gz && \
+    mkdir plpgsql_check-src && cd plpgsql_check-src && tar xvzf ../plpgsql_check.tar.gz --strip-components=1 -C . && \
+    make -j $(getconf _NPROCESSORS_ONLN) PG_CONFIG=/usr/local/pgsql/bin/pg_config USE_PGXS=1 && \
+    make -j $(getconf _NPROCESSORS_ONLN) install PG_CONFIG=/usr/local/pgsql/bin/pg_config USE_PGXS=1 && \
+    echo 'trusted = true' >> /usr/local/pgsql/share/extension/plpgsql_check.control
+
 #########################################################################################
 # 
 # Layer "rust extensions"
@@ -248,7 +323,7 @@ RUN curl -sSO https://static.rust-lang.org/rustup/dist/$(uname -m)-unknown-linux
     chmod +x rustup-init && \
     ./rustup-init -y --no-modify-path --profile minimal --default-toolchain stable && \
     rm rustup-init && \
-    cargo install --git https://github.com/vadim2404/pgx --branch neon_abi_v0.6.1 --locked cargo-pgx && \
+    cargo install --locked --version 0.7.3 cargo-pgx && \
     /bin/bash -c 'cargo pgx init --pg${PG_VERSION:1}=/usr/local/pgsql/bin/pg_config'
 
 USER root
@@ -262,11 +337,11 @@ USER root
 
 FROM rust-extensions-build AS pg-jsonschema-pg-build
 
-RUN git clone --depth=1 --single-branch --branch neon_abi_v0.1.4 https://github.com/vadim2404/pg_jsonschema/ && \
-    cd pg_jsonschema && \
+# there is no release tag yet, but we need it due to the superuser fix in the control file
+RUN wget https://github.com/supabase/pg_jsonschema/archive/caeab60d70b2fd3ae421ec66466a3abbb37b7ee6.tar.gz -O pg_jsonschema.tar.gz && \
+    mkdir pg_jsonschema-src && cd pg_jsonschema-src && tar xvzf ../pg_jsonschema.tar.gz --strip-components=1 -C . && \
+    sed -i 's/pgx = "0.7.1"/pgx = { version = "0.7.3", features = [ "unsafe-postgres" ] }/g' Cargo.toml && \
     cargo pgx install --release && \
-    # it's needed to enable extension because it uses untrusted C language
-    sed -i 's/superuser = false/superuser = true/g' /usr/local/pgsql/share/extension/pg_jsonschema.control && \
     echo "trusted = true" >> /usr/local/pgsql/share/extension/pg_jsonschema.control
 
 #########################################################################################
@@ -278,13 +353,32 @@ RUN git clone --depth=1 --single-branch --branch neon_abi_v0.1.4 https://github.
 
 FROM rust-extensions-build AS pg-graphql-pg-build
 
-RUN git clone --depth=1 --single-branch --branch neon_abi_v1.1.0 https://github.com/vadim2404/pg_graphql && \
-    cd pg_graphql && \  
+# Currently pgx version bump to >= 0.7.2  causes "call to unsafe function" compliation errors in
+# pgx-contrib-spiext. There is a branch that removes that dependency, so use it. It is on the
+# same 1.1 version we've used before.
+RUN git clone -b remove-pgx-contrib-spiext --single-branch https://github.com/yrashk/pg_graphql && \
+    cd pg_graphql && \
+    sed -i 's/pgx = "~0.7.1"/pgx = { version = "0.7.3", features = [ "unsafe-postgres" ] }/g' Cargo.toml && \
+    sed -i 's/pgx-tests = "~0.7.1"/pgx-tests = "0.7.3"/g' Cargo.toml && \
     cargo pgx install --release && \
     # it's needed to enable extension because it uses untrusted C language
     sed -i 's/superuser = false/superuser = true/g' /usr/local/pgsql/share/extension/pg_graphql.control && \
     echo "trusted = true" >> /usr/local/pgsql/share/extension/pg_graphql.control
 
+#########################################################################################
+#
+# Layer "pg-tiktoken-build"
+# Compile "pg_tiktoken" extension
+#
+#########################################################################################
+
+FROM rust-extensions-build AS pg-tiktoken-pg-build
+
+RUN git clone --depth=1 --single-branch https://github.com/kelvich/pg_tiktoken && \
+    cd pg_tiktoken && \
+    cargo pgx install --release && \
+    echo "trusted = true" >> /usr/local/pgsql/share/extension/pg_tiktoken.control
+
 #########################################################################################
 #
 # Layer "neon-pg-ext-build"
@@ -302,13 +396,23 @@ COPY --from=vector-pg-build /usr/local/pgsql/ /usr/local/pgsql/
 COPY --from=pgjwt-pg-build /usr/local/pgsql/ /usr/local/pgsql/
 COPY --from=pg-jsonschema-pg-build /usr/local/pgsql/ /usr/local/pgsql/
 COPY --from=pg-graphql-pg-build /usr/local/pgsql/ /usr/local/pgsql/
+COPY --from=pg-tiktoken-pg-build /usr/local/pgsql/ /usr/local/pgsql/
 COPY --from=hypopg-pg-build /usr/local/pgsql/ /usr/local/pgsql/
 COPY --from=pg-hashids-pg-build /usr/local/pgsql/ /usr/local/pgsql/
+COPY --from=rum-pg-build /usr/local/pgsql/ /usr/local/pgsql/
+COPY --from=pgtap-pg-build /usr/local/pgsql/ /usr/local/pgsql/
+COPY --from=prefix-pg-build /usr/local/pgsql/ /usr/local/pgsql/
+COPY --from=hll-pg-build /usr/local/pgsql/ /usr/local/pgsql/
+COPY --from=plpgsql-check-pg-build /usr/local/pgsql/ /usr/local/pgsql/
 COPY pgxn/ pgxn/
 
 RUN make -j $(getconf _NPROCESSORS_ONLN) \
         PG_CONFIG=/usr/local/pgsql/bin/pg_config \
         -C pgxn/neon \
+        -s install && \
+    make -j $(getconf _NPROCESSORS_ONLN) \
+        PG_CONFIG=/usr/local/pgsql/bin/pg_config \
+        -C pgxn/neon_utils \
         -s install
 
 #########################################################################################
@@ -363,7 +467,7 @@ COPY --from=compute-tools --chown=postgres /home/nonroot/target/release-line-deb
 
 # Install:
 # libreadline8 for psql
-# libicu67, locales for collations (including ICU)
+# libicu67, locales for collations (including ICU and plpgsql_check)
 # libossp-uuid16 for extension ossp-uuid
 # libgeos, libgdal, libsfcgal1, libproj and libprotobuf-c1 for PostGIS
 # libxml2, libxslt1.1 for xml2

Dockerfile.vm-compute-node

@@ -1,32 +1,70 @@
 # Note: this file *mostly* just builds on Dockerfile.compute-node
 
 ARG SRC_IMAGE
-ARG VM_INFORMANT_VERSION=v0.1.6
+ARG VM_INFORMANT_VERSION=v0.1.14
+# on libcgroup update, make sure to check bootstrap.sh for changes
+ARG LIBCGROUP_VERSION=v2.0.3
 
-# Pull VM informant and set up inittab
+# Pull VM informant, to copy from later
 FROM neondatabase/vm-informant:$VM_INFORMANT_VERSION as informant
 
+# Build cgroup-tools
+#
+# At time of writing (2023-03-14), debian bullseye has a version of cgroup-tools (technically
+# libcgroup) that doesn't support cgroup v2 (version 0.41-11). Unfortunately, the vm-informant
+# requires cgroup v2, so we'll build cgroup-tools ourselves.
+FROM debian:bullseye-slim as libcgroup-builder
+ARG LIBCGROUP_VERSION
+
+RUN set -exu \
+	&& apt update \
+	&& apt install --no-install-recommends -y \
+		git \
+		ca-certificates \
+		automake \
+		cmake \
+		make \
+		gcc \
+		byacc \
+		flex \
+		libtool \
+		libpam0g-dev \
+	&& git clone --depth 1 -b $LIBCGROUP_VERSION https://github.com/libcgroup/libcgroup \
+	&& INSTALL_DIR="/libcgroup-install" \
+	&& mkdir -p "$INSTALL_DIR/bin" "$INSTALL_DIR/include" \
+	&& cd libcgroup \
+	# extracted from bootstrap.sh, with modified flags:
+	&& (test -d m4 || mkdir m4) \
+	&& autoreconf -fi \
+	&& rm -rf autom4te.cache \
+	&& CFLAGS="-O3" ./configure --prefix="$INSTALL_DIR" --sysconfdir=/etc --localstatedir=/var --enable-opaque-hierarchy="name=systemd" \
+	# actually build the thing...
+	&& make install
+
+# Combine, starting from non-VM compute node image.
+FROM $SRC_IMAGE as base
+
+# Temporarily set user back to root so we can run adduser, set inittab
+USER root
+RUN adduser vm-informant --disabled-password --no-create-home
+
 RUN set -e \
 	&& rm -f /etc/inittab \
 	&& touch /etc/inittab
 
 RUN set -e \
 	&& echo "::sysinit:cgconfigparser -l /etc/cgconfig.conf -s 1664" >> /etc/inittab \
-	&& echo "::respawn:su vm-informant -c '/usr/local/bin/vm-informant --auto-restart --cgroup=neon-postgres'" >> /etc/inittab
+	&& CONNSTR="dbname=neondb user=cloud_admin sslmode=disable" \
+	&& ARGS="--auto-restart --cgroup=neon-postgres --pgconnstr=\"$CONNSTR\"" \
+	&& echo "::respawn:su vm-informant -c '/usr/local/bin/vm-informant $ARGS'" >> /etc/inittab
 
-# Combine, starting from non-VM compute node image.
-FROM $SRC_IMAGE as base
-
-# Temporarily set user back to root so we can run apt update and adduser
-USER root
-RUN apt update && \
-	apt install --no-install-recommends -y \
-        cgroup-tools
-RUN adduser vm-informant --disabled-password --no-create-home
 USER postgres
 
 ADD vm-cgconfig.conf /etc/cgconfig.conf
-COPY --from=informant /etc/inittab /etc/inittab
 COPY --from=informant /usr/bin/vm-informant /usr/local/bin/vm-informant
 
+COPY --from=libcgroup-builder /libcgroup-install/bin/* /usr/bin/
+COPY --from=libcgroup-builder /libcgroup-install/lib/* /usr/lib/
+COPY --from=libcgroup-builder /libcgroup-install/sbin/* /usr/sbin/
+
 ENTRYPOINT ["/usr/sbin/cgexec", "-g", "*:neon-postgres", "/usr/local/bin/compute_ctl"]

Makefile

@@ -133,6 +133,11 @@ neon-pg-ext-%: postgres-%
 	$(MAKE) PG_CONFIG=$(POSTGRES_INSTALL_DIR)/$*/bin/pg_config CFLAGS='$(PG_CFLAGS) $(COPT)' \
 		-C $(POSTGRES_INSTALL_DIR)/build/neon-test-utils-$* \
 		-f $(ROOT_PROJECT_DIR)/pgxn/neon_test_utils/Makefile install
+	+@echo "Compiling neon_utils $*"
+	mkdir -p $(POSTGRES_INSTALL_DIR)/build/neon-utils-$*
+	$(MAKE) PG_CONFIG=$(POSTGRES_INSTALL_DIR)/$*/bin/pg_config CFLAGS='$(PG_CFLAGS) $(COPT)' \
+		-C $(POSTGRES_INSTALL_DIR)/build/neon-utils-$* \
+		-f $(ROOT_PROJECT_DIR)/pgxn/neon_utils/Makefile install
 
 .PHONY: neon-pg-ext-clean-%
 neon-pg-ext-clean-%:
@@ -145,6 +150,9 @@ neon-pg-ext-clean-%:
 	$(MAKE) PG_CONFIG=$(POSTGRES_INSTALL_DIR)/$*/bin/pg_config \
 	-C $(POSTGRES_INSTALL_DIR)/build/neon-test-utils-$* \
 	-f $(ROOT_PROJECT_DIR)/pgxn/neon_test_utils/Makefile clean
+	$(MAKE) PG_CONFIG=$(POSTGRES_INSTALL_DIR)/$*/bin/pg_config \
+	-C $(POSTGRES_INSTALL_DIR)/build/neon-utils-$* \
+	-f $(ROOT_PROJECT_DIR)/pgxn/neon_utils/Makefile clean
 
 .PHONY: neon-pg-ext
 neon-pg-ext: \

compute_tools/Cargo.toml

@@ -11,6 +11,7 @@ clap.workspace = true
 futures.workspace = true
 hyper = { workspace = true, features = ["full"] }
 notify.workspace = true
+num_cpus.workspace = true
 opentelemetry.workspace = true
 postgres.workspace = true
 regex.workspace = true

compute_tools/src/compute.rs

@@ -25,6 +25,7 @@ use anyhow::{Context, Result};
 use chrono::{DateTime, Utc};
 use postgres::{Client, NoTls};
 use serde::{Serialize, Serializer};
+use tokio_postgres;
 use tracing::{info, instrument, warn};
 
 use crate::checker::create_writability_check_data;
@@ -284,6 +285,7 @@ impl ComputeNode {
         handle_role_deletions(self, &mut client)?;
         handle_grants(self, &mut client)?;
         create_writability_check_data(&mut client)?;
+        handle_extensions(&self.spec, &mut client)?;
 
         // 'Close' connection
         drop(client);
@@ -400,4 +402,43 @@ impl ComputeNode {
 
         Ok(())
     }
+
+    /// Select `pg_stat_statements` data and return it as a stringified JSON
+    pub async fn collect_insights(&self) -> String {
+        let mut result_rows: Vec<String> = Vec::new();
+        let connect_result = tokio_postgres::connect(self.connstr.as_str(), NoTls).await;
+        let (client, connection) = connect_result.unwrap();
+        tokio::spawn(async move {
+            if let Err(e) = connection.await {
+                eprintln!("connection error: {}", e);
+            }
+        });
+        let result = client
+            .simple_query(
+                "SELECT
+    row_to_json(pg_stat_statements)
+FROM
+    pg_stat_statements
+WHERE
+    userid != 'cloud_admin'::regrole::oid
+ORDER BY
+    (mean_exec_time + mean_plan_time) DESC
+LIMIT 100",
+            )
+            .await;
+
+        if let Ok(raw_rows) = result {
+            for message in raw_rows.iter() {
+                if let postgres::SimpleQueryMessage::Row(row) = message {
+                    if let Some(json) = row.get(0) {
+                        result_rows.push(json.to_string());
+                    }
+                }
+            }
+
+            format!("{{\"pg_stat_statements\": [{}]}}", result_rows.join(","))
+        } else {
+            "{{\"pg_stat_statements\": []}}".to_string()
+        }
+    }
 }

compute_tools/src/http/api.rs

@@ -7,6 +7,7 @@ use crate::compute::ComputeNode;
 use anyhow::Result;
 use hyper::service::{make_service_fn, service_fn};
 use hyper::{Body, Method, Request, Response, Server, StatusCode};
+use num_cpus;
 use serde_json;
 use tracing::{error, info};
 use tracing_utils::http::OtelName;
@@ -33,6 +34,13 @@ async fn routes(req: Request<Body>, compute: &Arc<ComputeNode>) -> Response<Body
             Response::new(Body::from(serde_json::to_string(&compute.metrics).unwrap()))
         }
 
+        // Collect Postgres current usage insights
+        (&Method::GET, "/insights") => {
+            info!("serving /insights GET request");
+            let insights = compute.collect_insights().await;
+            Response::new(Body::from(insights))
+        }
+
         (&Method::POST, "/check_writability") => {
             info!("serving /check_writability POST request");
             let res = crate::checker::check_writability(compute).await;
@@ -42,6 +50,17 @@ async fn routes(req: Request<Body>, compute: &Arc<ComputeNode>) -> Response<Body
             }
         }
 
+        (&Method::GET, "/info") => {
+            let num_cpus = num_cpus::get_physical();
+            info!("serving /info GET request. num_cpus: {}", num_cpus);
+            Response::new(Body::from(
+                serde_json::json!({
+                    "num_cpus": num_cpus,
+                })
+                .to_string(),
+            ))
+        }
+
         // Return the `404 Not Found` for any other routes.
         _ => {
             let mut not_found = Response::new(Body::from("404 Not Found"));

compute_tools/src/http/openapi_spec.yaml

@@ -10,12 +10,12 @@ paths:
   /status:
     get:
       tags:
-      - "info"
+      - Info
       summary: Get compute node internal status
       description: ""
       operationId: getComputeStatus
       responses:
-        "200":
+        200:
           description: ComputeState
           content:
             application/json:
@@ -25,27 +25,58 @@ paths:
   /metrics.json:
     get:
       tags:
-      - "info"
+      - Info
       summary: Get compute node startup metrics in JSON format
       description: ""
       operationId: getComputeMetricsJSON
       responses:
-        "200":
+        200:
           description: ComputeMetrics
           content:
             application/json:
               schema:
                 $ref: "#/components/schemas/ComputeMetrics"
 
+  /insights:
+    get:
+      tags:
+      - Info
+      summary: Get current compute insights in JSON format
+      description: |
+        Note, that this doesn't include any historical data
+      operationId: getComputeInsights
+      responses:
+        200:
+          description: Compute insights
+          content:
+            application/json:
+              schema:
+                $ref: "#/components/schemas/ComputeInsights"
+
+  /info:
+    get:
+      tags:
+      - "info"
+      summary: Get info about the compute Pod/VM
+      description: ""
+      operationId: getInfo
+      responses:
+        "200":
+          description: Info
+          content:
+            application/json:
+              schema:
+                $ref: "#/components/schemas/Info"
+
   /check_writability:
     post:
       tags:
-      - "check"
+      - Check
       summary: Check that we can write new data on this compute
       description: ""
       operationId: checkComputeWritability
       responses:
-        "200":
+        200:
           description: Check result
           content:
             text/plain:
@@ -80,6 +111,15 @@ components:
         total_startup_ms:
           type: integer
 
+    Info:
+      type: object
+      description: Information about VM/Pod
+      required:
+        - num_cpus
+      properties:
+        num_cpus:
+          type: integer
+
     ComputeState:
       type: object
       required:
@@ -96,6 +136,15 @@ components:
           type: string
           description: Text of the error during compute startup, if any
 
+    ComputeInsights:
+      type: object
+      properties:
+        pg_stat_statements:
+          description: Contains raw output from pg_stat_statements in JSON format
+          type: array
+          items:
+            type: object
+
     ComputeStatus:
       type: string
       enum:

compute_tools/src/pg_helpers.rs

@@ -47,12 +47,23 @@ pub struct GenericOption {
 /// declare a `trait` on it.
 pub type GenericOptions = Option<Vec<GenericOption>>;
 
+/// Escape a string for including it in a SQL literal
+fn escape_literal(s: &str) -> String {
+    s.replace('\'', "''").replace('\\', "\\\\")
+}
+
+/// Escape a string so that it can be used in postgresql.conf.
+/// Same as escape_literal, currently.
+fn escape_conf_value(s: &str) -> String {
+    s.replace('\'', "''").replace('\\', "\\\\")
+}
+
 impl GenericOption {
     /// Represent `GenericOption` as SQL statement parameter.
     pub fn to_pg_option(&self) -> String {
         if let Some(val) = &self.value {
             match self.vartype.as_ref() {
-                "string" => format!("{} '{}'", self.name, val),
+                "string" => format!("{} '{}'", self.name, escape_literal(val)),
                 _ => format!("{} {}", self.name, val),
             }
         } else {
@@ -63,6 +74,8 @@ impl GenericOption {
     /// Represent `GenericOption` as configuration option.
     pub fn to_pg_setting(&self) -> String {
         if let Some(val) = &self.value {
+            // TODO: check in the console DB that we don't have these settings
+            // set for any non-deleted project and drop this override.
             let name = match self.name.as_str() {
                 "safekeepers" => "neon.safekeepers",
                 "wal_acceptor_reconnect" => "neon.safekeeper_reconnect_timeout",
@@ -71,7 +84,7 @@ impl GenericOption {
             };
 
             match self.vartype.as_ref() {
-                "string" => format!("{} = '{}'", name, val),
+                "string" => format!("{} = '{}'", name, escape_conf_value(val)),
                 _ => format!("{} = {}", name, val),
             }
         } else {
@@ -107,6 +120,7 @@ impl PgOptionsSerialize for GenericOptions {
                 .map(|op| op.to_pg_setting())
                 .collect::<Vec<String>>()
                 .join("\n")
+                + "\n" // newline after last setting
         } else {
             "".to_string()
         }

compute_tools/src/spec.rs

@@ -515,3 +515,18 @@ pub fn handle_grants(node: &ComputeNode, client: &mut Client) -> Result<()> {
 
     Ok(())
 }
+
+/// Create required system extensions
+#[instrument(skip_all)]
+pub fn handle_extensions(spec: &ComputeSpec, client: &mut Client) -> Result<()> {
+    if let Some(libs) = spec.cluster.settings.find("shared_preload_libraries") {
+        if libs.contains("pg_stat_statements") {
+            // Create extension only if this compute really needs it
+            let query = "CREATE EXTENSION IF NOT EXISTS pg_stat_statements";
+            info!("creating system extensions with query: {}", query);
+            client.simple_query(query)?;
+        }
+    }
+
+    Ok(())
+}

compute_tools/tests/cluster_spec.json

@@ -178,6 +178,11 @@
                 "name": "neon.pageserver_connstring",
                 "value": "host=127.0.0.1 port=6400",
                 "vartype": "string"
+            },
+            {
+                "name": "test.escaping",
+                "value": "here's a backslash \\ and a quote ' and a double-quote \" hooray",
+                "vartype": "string"
             }
         ]
     },

compute_tools/tests/pg_helpers_tests.rs

@@ -28,7 +28,30 @@ mod pg_helpers_tests {
 
         assert_eq!(
             spec.cluster.settings.as_pg_settings(),
-            "fsync = off\nwal_level = replica\nhot_standby = on\nneon.safekeepers = '127.0.0.1:6502,127.0.0.1:6503,127.0.0.1:6501'\nwal_log_hints = on\nlog_connections = on\nshared_buffers = 32768\nport = 55432\nmax_connections = 100\nmax_wal_senders = 10\nlisten_addresses = '0.0.0.0'\nwal_sender_timeout = 0\npassword_encryption = md5\nmaintenance_work_mem = 65536\nmax_parallel_workers = 8\nmax_worker_processes = 8\nneon.tenant_id = 'b0554b632bd4d547a63b86c3630317e8'\nmax_replication_slots = 10\nneon.timeline_id = '2414a61ffc94e428f14b5758fe308e13'\nshared_preload_libraries = 'neon'\nsynchronous_standby_names = 'walproposer'\nneon.pageserver_connstring = 'host=127.0.0.1 port=6400'"
+            r#"fsync = off
+wal_level = replica
+hot_standby = on
+neon.safekeepers = '127.0.0.1:6502,127.0.0.1:6503,127.0.0.1:6501'
+wal_log_hints = on
+log_connections = on
+shared_buffers = 32768
+port = 55432
+max_connections = 100
+max_wal_senders = 10
+listen_addresses = '0.0.0.0'
+wal_sender_timeout = 0
+password_encryption = md5
+maintenance_work_mem = 65536
+max_parallel_workers = 8
+max_worker_processes = 8
+neon.tenant_id = 'b0554b632bd4d547a63b86c3630317e8'
+max_replication_slots = 10
+neon.timeline_id = '2414a61ffc94e428f14b5758fe308e13'
+shared_preload_libraries = 'neon'
+synchronous_standby_names = 'walproposer'
+neon.pageserver_connstring = 'host=127.0.0.1 port=6400'
+test.escaping = 'here''s a backslash \\ and a quote '' and a double-quote " hooray'
+"#
         );
     }
 

control_plane/safekeepers.conf

@@ -2,7 +2,8 @@
 [pageserver]
 listen_pg_addr = '127.0.0.1:64000'
 listen_http_addr = '127.0.0.1:9898'
-auth_type = 'Trust'
+pg_auth_type = 'Trust'
+http_auth_type = 'Trust'
 
 [[safekeepers]]
 id = 1

control_plane/simple.conf

@@ -3,7 +3,8 @@
 [pageserver]
 listen_pg_addr = '127.0.0.1:64000'
 listen_http_addr = '127.0.0.1:9898'
-auth_type = 'Trust'
+pg_auth_type = 'Trust'
+http_auth_type = 'Trust'
 
 [[safekeepers]]
 id = 1

control_plane/src/bin/neon_local.rs

@@ -53,14 +53,15 @@ listen_addr = '{DEFAULT_BROKER_ADDR}'
 id = {DEFAULT_PAGESERVER_ID}
 listen_pg_addr = '{DEFAULT_PAGESERVER_PG_ADDR}'
 listen_http_addr = '{DEFAULT_PAGESERVER_HTTP_ADDR}'
-auth_type = '{pageserver_auth_type}'
+pg_auth_type = '{trust_auth}'
+http_auth_type = '{trust_auth}'
 
 [[safekeepers]]
 id = {DEFAULT_SAFEKEEPER_ID}
 pg_port = {DEFAULT_SAFEKEEPER_PG_PORT}
 http_port = {DEFAULT_SAFEKEEPER_HTTP_PORT}
 "#,
-        pageserver_auth_type = AuthType::Trust,
+        trust_auth = AuthType::Trust,
     )
 }
 
@@ -627,7 +628,7 @@ fn handle_pg(pg_match: &ArgMatches, env: &local_env::LocalEnv) -> Result<()> {
 
             let node = cplane.nodes.get(&(tenant_id, node_name.to_string()));
 
-            let auth_token = if matches!(env.pageserver.auth_type, AuthType::NeonJWT) {
+            let auth_token = if matches!(env.pageserver.pg_auth_type, AuthType::NeonJWT) {
                 let claims = Claims::new(Some(tenant_id), Scope::Tenant);
 
                 Some(env.generate_auth_token(&claims)?)

control_plane/src/compute.rs

@@ -97,7 +97,7 @@ impl ComputeControlPlane {
         });
 
         node.create_pgdata()?;
-        node.setup_pg_conf(self.env.pageserver.auth_type)?;
+        node.setup_pg_conf(self.env.pageserver.pg_auth_type)?;
 
         self.nodes
             .insert((tenant_id, node.name.clone()), Arc::clone(&node));

control_plane/src/local_env.rs

@@ -110,12 +110,14 @@ impl NeonBroker {
 pub struct PageServerConf {
     // node id
     pub id: NodeId,
+
     // Pageserver connection settings
     pub listen_pg_addr: String,
     pub listen_http_addr: String,
 
-    // used to determine which auth type is used
-    pub auth_type: AuthType,
+    // auth type used for the PG and HTTP ports
+    pub pg_auth_type: AuthType,
+    pub http_auth_type: AuthType,
 
     // jwt auth token used for communication with pageserver
     pub auth_token: String,
@@ -127,7 +129,8 @@ impl Default for PageServerConf {
             id: NodeId(0),
             listen_pg_addr: String::new(),
             listen_http_addr: String::new(),
-            auth_type: AuthType::Trust,
+            pg_auth_type: AuthType::Trust,
+            http_auth_type: AuthType::Trust,
             auth_token: String::new(),
         }
     }

control_plane/src/pageserver.rs

@@ -82,7 +82,7 @@ impl PageServerNode {
         let (host, port) = parse_host_port(&env.pageserver.listen_pg_addr)
             .expect("Unable to parse listen_pg_addr");
         let port = port.unwrap_or(5432);
-        let password = if env.pageserver.auth_type == AuthType::NeonJWT {
+        let password = if env.pageserver.pg_auth_type == AuthType::NeonJWT {
             Some(env.pageserver.auth_token.clone())
         } else {
             None
@@ -106,25 +106,32 @@ impl PageServerNode {
             self.env.pg_distrib_dir_raw().display()
         );
 
-        let authg_type_param = format!("auth_type='{}'", self.env.pageserver.auth_type);
+        let http_auth_type_param =
+            format!("http_auth_type='{}'", self.env.pageserver.http_auth_type);
         let listen_http_addr_param = format!(
             "listen_http_addr='{}'",
             self.env.pageserver.listen_http_addr
         );
+
+        let pg_auth_type_param = format!("pg_auth_type='{}'", self.env.pageserver.pg_auth_type);
         let listen_pg_addr_param =
             format!("listen_pg_addr='{}'", self.env.pageserver.listen_pg_addr);
+
         let broker_endpoint_param = format!("broker_endpoint='{}'", self.env.broker.client_url());
 
         let mut overrides = vec![
             id,
             pg_distrib_dir_param,
-            authg_type_param,
+            http_auth_type_param,
+            pg_auth_type_param,
             listen_http_addr_param,
             listen_pg_addr_param,
             broker_endpoint_param,
         ];
 
-        if self.env.pageserver.auth_type != AuthType::Trust {
+        if self.env.pageserver.http_auth_type != AuthType::Trust
+            || self.env.pageserver.pg_auth_type != AuthType::Trust
+        {
             overrides.push("auth_validation_public_key_path='auth_public_key.pem'".to_owned());
         }
         overrides
@@ -247,7 +254,10 @@ impl PageServerNode {
     }
 
     fn pageserver_env_variables(&self) -> anyhow::Result<Vec<(String, String)>> {
-        Ok(if self.env.pageserver.auth_type != AuthType::Trust {
+        // FIXME: why is this tied to pageserver's auth type? Whether or not the safekeeper
+        // needs a token, and how to generate that token, seems independent to whether
+        // the pageserver requires a token in incoming requests.
+        Ok(if self.env.pageserver.http_auth_type != AuthType::Trust {
             // Generate a token to connect from the pageserver to a safekeeper
             let token = self
                 .env
@@ -283,7 +293,7 @@ impl PageServerNode {
 
     fn http_request<U: IntoUrl>(&self, method: Method, url: U) -> RequestBuilder {
         let mut builder = self.http_client.request(method, url);
-        if self.env.pageserver.auth_type == AuthType::NeonJWT {
+        if self.env.pageserver.http_auth_type == AuthType::NeonJWT {
             builder = builder.bearer_auth(&self.env.pageserver.auth_token)
         }
         builder

control_plane/src/safekeeper.rs

@@ -115,6 +115,10 @@ impl SafekeeperNode {
         let datadir = self.datadir_path();
 
         let id_string = id.to_string();
+        // TODO: add availability_zone to the config.
+        // Right now we just specify any value here and use it to check metrics in tests.
+        let availability_zone = format!("sk-{}", id_string);
+
         let mut args = vec![
             "-D",
             datadir.to_str().with_context(|| {
@@ -126,6 +130,8 @@ impl SafekeeperNode {
             &listen_pg,
             "--listen-http",
             &listen_http,
+            "--availability-zone",
+            &availability_zone,
         ];
         if !self.conf.sync {
             args.push("--no-sync");

docs/authentication.md

@@ -29,6 +29,41 @@ These components should not have access to the private key and may only get toke
 The key pair is generated once for an installation of compute/pageserver/safekeeper, e.g. by `neon_local init`.
 There is currently no way to rotate the key without bringing down all components.
 
+### Token format
+
+The JWT tokens in Neon use RSA as the algorithm. Example:
+
+Header:
+
+```
+{
+  "alg": "RS512",     # RS256, RS384, or RS512
+  "typ": "JWT"
+}
+```
+
+Payload:
+
+```
+{
+  "scope": "tenant",  # "tenant", "pageserverapi", or "safekeeperdata"
+  "tenant_id": "5204921ff44f09de8094a1390a6a50f6",
+}
+```
+
+
+Meanings of scope:
+
+"tenant": Provides access to all data for a specific tenant
+
+"pageserverapi": Provides blanket access to all tenants on the pageserver plus pageserver-wide APIs.
+Should only be used e.g. for status check/tenant creation/list.
+
+"safekeeperdata": Provides blanket access to all data on the safekeeper plus safekeeper-wide APIs.
+Should only be used e.g. for status check.
+Currently also used for connection from any pageserver to any safekeeper.
+
+
 ### CLI
 CLI generates a key pair during call to `neon_local init` with the following commands:
 
@@ -102,10 +137,12 @@ Each compute should present a token valid for the timeline's tenant.
 Pageserver also has HTTP API: some parts are per-tenant,
 some parts are server-wide, these are different scopes.
 
-The `auth_type` configuration variable in Pageserver's config may have
-either of three values:
+Authentication can be enabled separately for the HTTP mgmt API, and
+for the libpq connections from compute. The `http_auth_type` and
+`pg_auth_type` configuration variables in Pageserver's config may
+have one of these values:
 
-* `Trust` removes all authentication. The outdated `MD5` value does likewise
+* `Trust` removes all authentication.
 * `NeonJWT` enables JWT validation.
    Tokens are validated using the public key which lies in a PEM file
    specified in the `auth_validation_public_key_path` config.

docs/sourcetree.md

@@ -129,13 +129,12 @@ Run `poetry shell` to activate the virtual environment.
 Alternatively, use `poetry run` to run a single command in the venv, e.g. `poetry run pytest`.
 
 ### Obligatory checks
-We force code formatting via `black`, `isort` and type hints via `mypy`.
+We force code formatting via `black`, `ruff`, and type hints via `mypy`.
 Run the following commands in the repository's root (next to `pyproject.toml`):
 
 ```bash
-poetry run isort .  # Imports are reformatted
 poetry run black .  # All code is reformatted
-poetry run flake8 .  # Python linter
+poetry run ruff .  # Python linter
 poetry run mypy .  # Ensure there are no typing errors
 ```
 

libs/postgres_backend/src/lib.rs

@@ -13,9 +13,8 @@ use std::sync::Arc;
 use std::task::{ready, Poll};
 use std::{fmt, io};
 use std::{future::Future, str::FromStr};
-use tokio::io::{AsyncRead, AsyncWrite, ReadHalf, WriteHalf};
+use tokio::io::{AsyncRead, AsyncWrite};
 use tokio_rustls::TlsAcceptor;
-
 use tracing::{debug, error, info, trace};
 
 use pq_proto::framed::{ConnectionError, Framed, FramedReader, FramedWriter};
@@ -60,14 +59,14 @@ pub fn is_expected_io_error(e: &io::Error) -> bool {
 }
 
 #[async_trait::async_trait]
-pub trait Handler {
+pub trait Handler<IO> {
     /// Handle single query.
     /// postgres_backend will issue ReadyForQuery after calling this (this
     /// might be not what we want after CopyData streaming, but currently we don't
     /// care). It will also flush out the output buffer.
     async fn process_query(
         &mut self,
-        pgb: &mut PostgresBackend,
+        pgb: &mut PostgresBackend<IO>,
         query_string: &str,
     ) -> Result<(), QueryError>;
 
@@ -78,7 +77,7 @@ pub trait Handler {
     /// to override whole init logic in implementations.
     fn startup(
         &mut self,
-        _pgb: &mut PostgresBackend,
+        _pgb: &mut PostgresBackend<IO>,
         _sm: &FeStartupPacket,
     ) -> Result<(), QueryError> {
         Ok(())
@@ -87,7 +86,7 @@ pub trait Handler {
     /// Check auth jwt
     fn check_auth_jwt(
         &mut self,
-        _pgb: &mut PostgresBackend,
+        _pgb: &mut PostgresBackend<IO>,
         _jwt_response: &[u8],
     ) -> Result<(), QueryError> {
         Err(QueryError::Other(anyhow::anyhow!("JWT auth failed")))
@@ -116,12 +115,12 @@ pub enum ProcessMsgResult {
 }
 
 /// Either plain TCP stream or encrypted one, implementing AsyncRead + AsyncWrite.
-pub enum MaybeTlsStream {
-    Unencrypted(tokio::net::TcpStream),
-    Tls(Box<tokio_rustls::server::TlsStream<tokio::net::TcpStream>>),
+pub enum MaybeTlsStream<IO> {
+    Unencrypted(IO),
+    Tls(Box<tokio_rustls::server::TlsStream<IO>>),
 }
 
-impl AsyncWrite for MaybeTlsStream {
+impl<IO: AsyncRead + AsyncWrite + Unpin> AsyncWrite for MaybeTlsStream<IO> {
     fn poll_write(
         self: Pin<&mut Self>,
         cx: &mut std::task::Context<'_>,
@@ -148,7 +147,7 @@ impl AsyncWrite for MaybeTlsStream {
         }
     }
 }
-impl AsyncRead for MaybeTlsStream {
+impl<IO: AsyncRead + AsyncWrite + Unpin> AsyncRead for MaybeTlsStream<IO> {
     fn poll_read(
         self: Pin<&mut Self>,
         cx: &mut std::task::Context<'_>,
@@ -193,13 +192,13 @@ impl fmt::Display for AuthType {
 /// PostgresBackend after call to `split`. In principle we could always store a
 /// pair of splitted handles, but that would force to to pay splitting price
 /// (Arc and kinda mutex inside polling) for all uses (e.g. pageserver).
-enum MaybeWriteOnly {
-    Full(Framed<MaybeTlsStream>),
-    WriteOnly(FramedWriter<WriteHalf<MaybeTlsStream>>),
+enum MaybeWriteOnly<IO> {
+    Full(Framed<MaybeTlsStream<IO>>),
+    WriteOnly(FramedWriter<MaybeTlsStream<IO>>),
     Broken, // temporary value palmed off during the split
 }
 
-impl MaybeWriteOnly {
+impl<IO: AsyncRead + AsyncWrite + Unpin> MaybeWriteOnly<IO> {
     async fn read_startup_message(&mut self) -> Result<Option<FeStartupPacket>, ConnectionError> {
         match self {
             MaybeWriteOnly::Full(framed) => framed.read_startup_message().await,
@@ -245,8 +244,8 @@ impl MaybeWriteOnly {
     }
 }
 
-pub struct PostgresBackend {
-    framed: MaybeWriteOnly,
+pub struct PostgresBackend<IO> {
+    framed: MaybeWriteOnly<IO>,
 
     pub state: ProtoState,
 
@@ -256,6 +255,8 @@ pub struct PostgresBackend {
     pub tls_config: Option<Arc<rustls::ServerConfig>>,
 }
 
+pub type PostgresBackendTCP = PostgresBackend<tokio::net::TcpStream>;
+
 pub fn query_from_cstring(query_string: Bytes) -> Vec<u8> {
     let mut query_string = query_string.to_vec();
     if let Some(ch) = query_string.last() {
@@ -272,7 +273,7 @@ fn cstr_to_str(bytes: &[u8]) -> anyhow::Result<&str> {
     std::str::from_utf8(without_null).map_err(|e| e.into())
 }
 
-impl PostgresBackend {
+impl PostgresBackend<tokio::net::TcpStream> {
     pub fn new(
         socket: tokio::net::TcpStream,
         auth_type: AuthType,
@@ -289,6 +290,25 @@ impl PostgresBackend {
             peer_addr,
         })
     }
+}
+
+impl<IO: AsyncRead + AsyncWrite + Unpin> PostgresBackend<IO> {
+    pub fn new_from_io(
+        socket: IO,
+        peer_addr: SocketAddr,
+        auth_type: AuthType,
+        tls_config: Option<Arc<rustls::ServerConfig>>,
+    ) -> io::Result<Self> {
+        let stream = MaybeTlsStream::Unencrypted(socket);
+
+        Ok(Self {
+            framed: MaybeWriteOnly::Full(Framed::new(stream)),
+            state: ProtoState::Initialization,
+            auth_type,
+            tls_config,
+            peer_addr,
+        })
+    }
 
     pub fn get_peer_addr(&self) -> &SocketAddr {
         &self.peer_addr
@@ -347,14 +367,14 @@ impl PostgresBackend {
     /// to it in CopyData messages, and writes them to the connection
     ///
     /// The caller is responsible for sending CopyOutResponse and CopyDone messages.
-    pub fn copyout_writer(&mut self) -> CopyDataWriter {
+    pub fn copyout_writer(&mut self) -> CopyDataWriter<IO> {
         CopyDataWriter { pgb: self }
     }
 
     /// Wrapper for run_message_loop() that shuts down socket when we are done
     pub async fn run<F, S>(
         mut self,
-        handler: &mut impl Handler,
+        handler: &mut impl Handler<IO>,
         shutdown_watcher: F,
     ) -> Result<(), QueryError>
     where
@@ -370,7 +390,7 @@ impl PostgresBackend {
 
     async fn run_message_loop<F, S>(
         &mut self,
-        handler: &mut impl Handler,
+        handler: &mut impl Handler<IO>,
         shutdown_watcher: F,
     ) -> Result<(), QueryError>
     where
@@ -427,9 +447,9 @@ impl PostgresBackend {
 
     /// Try to upgrade MaybeTlsStream into actual TLS one, performing handshake.
     async fn tls_upgrade(
-        src: MaybeTlsStream,
+        src: MaybeTlsStream<IO>,
         tls_config: Arc<rustls::ServerConfig>,
-    ) -> anyhow::Result<MaybeTlsStream> {
+    ) -> anyhow::Result<MaybeTlsStream<IO>> {
         match src {
             MaybeTlsStream::Unencrypted(s) => {
                 let acceptor = TlsAcceptor::from(tls_config);
@@ -467,7 +487,7 @@ impl PostgresBackend {
 
     /// Split off owned read part from which messages can be read in different
     /// task/thread.
-    pub fn split(&mut self) -> anyhow::Result<PostgresBackendReader> {
+    pub fn split(&mut self) -> anyhow::Result<PostgresBackendReader<IO>> {
         // temporary replace stream with fake to cook split one, Indiana Jones style
         match std::mem::replace(&mut self.framed, MaybeWriteOnly::Broken) {
             MaybeWriteOnly::Full(framed) => {
@@ -483,7 +503,7 @@ impl PostgresBackend {
     }
 
     /// Join read part back.
-    pub fn unsplit(&mut self, reader: PostgresBackendReader) -> anyhow::Result<()> {
+    pub fn unsplit(&mut self, reader: PostgresBackendReader<IO>) -> anyhow::Result<()> {
         // temporary replace stream with fake to cook joined one, Indiana Jones style
         match std::mem::replace(&mut self.framed, MaybeWriteOnly::Broken) {
             MaybeWriteOnly::Full(_) => {
@@ -500,7 +520,7 @@ impl PostgresBackend {
 
     /// Perform handshake with the client, transitioning to Established.
     /// In case of EOF during handshake logs this, sets state to Closed and returns Ok(()).
-    async fn handshake(&mut self, handler: &mut impl Handler) -> Result<(), QueryError> {
+    async fn handshake(&mut self, handler: &mut impl Handler<IO>) -> Result<(), QueryError> {
         while self.state < ProtoState::Authentication {
             match self.framed.read_startup_message().await? {
                 Some(msg) => {
@@ -566,7 +586,7 @@ impl PostgresBackend {
     ///   actual startup packet.
     async fn process_startup_message(
         &mut self,
-        handler: &mut impl Handler,
+        handler: &mut impl Handler<IO>,
         msg: FeStartupPacket,
     ) -> Result<(), QueryError> {
         assert!(self.state < ProtoState::Authentication);
@@ -630,7 +650,7 @@ impl PostgresBackend {
 
     async fn process_message(
         &mut self,
-        handler: &mut impl Handler,
+        handler: &mut impl Handler<IO>,
         msg: FeMessage,
         unnamed_query_string: &mut Bytes,
     ) -> Result<ProcessMsgResult, QueryError> {
@@ -777,9 +797,9 @@ impl PostgresBackend {
     }
 }
 
-pub struct PostgresBackendReader(FramedReader<ReadHalf<MaybeTlsStream>>);
+pub struct PostgresBackendReader<IO>(FramedReader<MaybeTlsStream<IO>>);
 
-impl PostgresBackendReader {
+impl<IO: AsyncRead + AsyncWrite + Unpin> PostgresBackendReader<IO> {
     /// Read full message or return None if connection is cleanly closed with no
     /// unprocessed data.
     pub async fn read_message(&mut self) -> Result<Option<FeMessage>, ConnectionError> {
@@ -813,11 +833,11 @@ impl PostgresBackendReader {
 /// messages.
 ///
 
-pub struct CopyDataWriter<'a> {
-    pgb: &'a mut PostgresBackend,
+pub struct CopyDataWriter<'a, IO> {
+    pgb: &'a mut PostgresBackend<IO>,
 }
 
-impl<'a> AsyncWrite for CopyDataWriter<'a> {
+impl<'a, IO: AsyncRead + AsyncWrite + Unpin> AsyncWrite for CopyDataWriter<'a, IO> {
     fn poll_write(
         self: Pin<&mut Self>,
         cx: &mut std::task::Context<'_>,

libs/postgres_backend/tests/simple_select.rs

@@ -4,6 +4,7 @@ use postgres_backend::{AuthType, Handler, PostgresBackend, QueryError};
 use pq_proto::{BeMessage, RowDescriptor};
 use std::io::Cursor;
 use std::{future, sync::Arc};
+use tokio::io::{AsyncRead, AsyncWrite};
 use tokio::net::{TcpListener, TcpStream};
 use tokio_postgres::config::SslMode;
 use tokio_postgres::tls::MakeTlsConnect;
@@ -22,11 +23,11 @@ async fn make_tcp_pair() -> (TcpStream, TcpStream) {
 struct TestHandler {}
 
 #[async_trait::async_trait]
-impl Handler for TestHandler {
+impl<IO: AsyncRead + AsyncWrite + Unpin + Send> Handler<IO> for TestHandler {
     // return single col 'hey' for any query
     async fn process_query(
         &mut self,
-        pgb: &mut PostgresBackend,
+        pgb: &mut PostgresBackend<IO>,
         _query_string: &str,
     ) -> Result<(), QueryError> {
         pgb.write_message_noflush(&BeMessage::RowDescription(&[RowDescriptor::text_col(

libs/postgres_ffi/build.rs

@@ -63,10 +63,7 @@ fn main() -> anyhow::Result<()> {
             pg_install_dir_versioned = cwd.join("..").join("..").join(pg_install_dir_versioned);
         }
 
-        let pg_config_bin = pg_install_dir_versioned
-            .join(pg_version)
-            .join("bin")
-            .join("pg_config");
+        let pg_config_bin = pg_install_dir_versioned.join("bin").join("pg_config");
         let inc_server_path: String = if pg_config_bin.exists() {
             let output = Command::new(pg_config_bin)
                 .arg("--includedir-server")

libs/pq_proto/src/framed.rs

@@ -63,9 +63,9 @@ impl<S> Framed<S> {
         &self.stream
     }
 
-    /// Extract the underlying stream.
-    pub fn into_inner(self) -> S {
-        self.stream
+    /// Deconstruct into the underlying stream and read buffer.
+    pub fn into_inner(self) -> (S, BytesMut) {
+        (self.stream, self.read_buf)
     }
 
     /// Return new Framed with stream type transformed by async f, for TLS
@@ -118,7 +118,7 @@ impl<S: AsyncRead + AsyncWrite + Unpin> Framed<S> {
     /// Split into owned read and write parts. Beware of potential issues with
     /// using halves in different tasks on TLS stream:
     /// https://github.com/tokio-rs/tls/issues/40
-    pub fn split(self) -> (FramedReader<ReadHalf<S>>, FramedWriter<WriteHalf<S>>) {
+    pub fn split(self) -> (FramedReader<S>, FramedWriter<S>) {
         let (read_half, write_half) = tokio::io::split(self.stream);
         let reader = FramedReader {
             stream: read_half,
@@ -132,7 +132,7 @@ impl<S: AsyncRead + AsyncWrite + Unpin> Framed<S> {
     }
 
     /// Join read and write parts back.
-    pub fn unsplit(reader: FramedReader<ReadHalf<S>>, writer: FramedWriter<WriteHalf<S>>) -> Self {
+    pub fn unsplit(reader: FramedReader<S>, writer: FramedWriter<S>) -> Self {
         Self {
             stream: reader.stream.unsplit(writer.stream),
             read_buf: reader.read_buf,
@@ -143,7 +143,7 @@ impl<S: AsyncRead + AsyncWrite + Unpin> Framed<S> {
 
 /// Read-only version of `Framed`.
 pub struct FramedReader<S> {
-    stream: S,
+    stream: ReadHalf<S>,
     read_buf: BytesMut,
 }
 
@@ -155,17 +155,10 @@ impl<S: AsyncRead + Unpin> FramedReader<S> {
 
 /// Write-only version of `Framed`.
 pub struct FramedWriter<S> {
-    stream: S,
+    stream: WriteHalf<S>,
     write_buf: BytesMut,
 }
 
-impl<S> FramedWriter<S> {
-    /// Get a mut reference to the underlying stream.
-    pub fn get_mut(&mut self) -> &mut S {
-        &mut self.stream
-    }
-}
-
 impl<S: AsyncWrite + Unpin> FramedWriter<S> {
     /// Write next message to the output buffer; doesn't flush.
     pub fn write_message_noflush(&mut self, msg: &BeMessage<'_>) -> Result<(), ProtocolError> {

libs/pq_proto/src/lib.rs

@@ -265,11 +265,9 @@ impl FeMessage {
             b'c' => Ok(Some(FeMessage::CopyDone)),
             b'f' => Ok(Some(FeMessage::CopyFail)),
             b'p' => Ok(Some(FeMessage::PasswordMessage(msg))),
-            tag => {
-                return Err(ProtocolError::Protocol(format!(
-                    "unknown message tag: {tag},'{msg:?}'"
-                )))
-            }
+            tag => Err(ProtocolError::Protocol(format!(
+                "unknown message tag: {tag},'{msg:?}'"
+            ))),
         }
     }
 }
@@ -818,7 +816,7 @@ impl<'a> BeMessage<'a> {
             BeMessage::ErrorResponse(error_msg, pg_error_code) => {
                 // 'E' signalizes ErrorResponse messages
                 buf.put_u8(b'E');
-                write_body(buf, |buf| -> Result<(), ProtocolError> {
+                write_body(buf, |buf| {
                     buf.put_u8(b'S'); // severity
                     buf.put_slice(b"ERROR\0");
 
@@ -843,7 +841,7 @@ impl<'a> BeMessage<'a> {
 
                 // 'N' signalizes NoticeResponse messages
                 buf.put_u8(b'N');
-                write_body(buf, |buf| -> Result<(), ProtocolError> {
+                write_body(buf, |buf| {
                     buf.put_u8(b'S'); // severity
                     buf.put_slice(b"NOTICE\0");
 
@@ -898,7 +896,7 @@ impl<'a> BeMessage<'a> {
 
             BeMessage::RowDescription(rows) => {
                 buf.put_u8(b'T');
-                write_body(buf, |buf| -> Result<(), ProtocolError> {
+                write_body(buf, |buf| {
                     buf.put_i16(rows.len() as i16); // # of fields
                     for row in rows.iter() {
                         write_cstr(row.name, buf)?;

libs/utils/Cargo.toml

@@ -18,17 +18,16 @@ futures = { workspace = true}
 jsonwebtoken.workspace = true
 nix.workspace = true
 once_cell.workspace = true
+pin-project-lite.workspace = true
 routerify.workspace = true
 serde.workspace = true
 serde_json.workspace = true
 signal-hook.workspace = true
 thiserror.workspace = true
 tokio.workspace = true
-tokio-rustls.workspace = true
 tracing.workspace = true
 tracing-subscriber = { workspace = true, features = ["json"] }
 rand.workspace = true
-rustls.workspace = true
 serde_with.workspace = true
 strum.workspace = true
 strum_macros.workspace = true

libs/utils/src/auth.rs

@@ -9,16 +9,28 @@ use std::path::Path;
 
 use anyhow::Result;
 use jsonwebtoken::{
-    decode, encode, Algorithm, DecodingKey, EncodingKey, Header, TokenData, Validation,
+    decode, encode, Algorithm, Algorithm::*, DecodingKey, EncodingKey, Header, TokenData,
+    Validation,
 };
 use serde::{Deserialize, Serialize};
 use serde_with::{serde_as, DisplayFromStr};
 
 use crate::id::TenantId;
 
-const JWT_ALGORITHM: Algorithm = Algorithm::RS256;
+/// Algorithms accepted during validation.
+///
+/// Accept all RSA-based algorithms. We pass this list to jsonwebtoken::decode,
+/// which checks that the algorithm in the token is one of these.
+///
+/// XXX: It also fails the validation if there are any algorithms in this list that belong
+/// to different family than the token's algorithm. In other words, we can *not* list any
+/// non-RSA algorithms here, or the validation always fails with InvalidAlgorithm error.
+const ACCEPTED_ALGORITHMS: &[Algorithm] = &[RS256, RS384, RS512];
 
-#[derive(Debug, Serialize, Deserialize, Clone)]
+/// Algorithm to use when generating a new token in [`encode_from_key_file`]
+const ENCODE_ALGORITHM: Algorithm = Algorithm::RS256;
+
+#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
 #[serde(rename_all = "lowercase")]
 pub enum Scope {
     // Provides access to all data for a specific tenant (specified in `struct Claims` below)
@@ -33,8 +45,9 @@ pub enum Scope {
     SafekeeperData,
 }
 
+/// JWT payload. See docs/authentication.md for the format
 #[serde_as]
-#[derive(Debug, Serialize, Deserialize, Clone)]
+#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
 pub struct Claims {
     #[serde(default)]
     #[serde_as(as = "Option<DisplayFromStr>")]
@@ -55,7 +68,8 @@ pub struct JwtAuth {
 
 impl JwtAuth {
     pub fn new(decoding_key: DecodingKey) -> Self {
-        let mut validation = Validation::new(JWT_ALGORITHM);
+        let mut validation = Validation::default();
+        validation.algorithms = ACCEPTED_ALGORITHMS.into();
         // The default 'required_spec_claims' is 'exp'. But we don't want to require
         // expiration.
         validation.required_spec_claims = [].into();
@@ -86,5 +100,113 @@ impl std::fmt::Debug for JwtAuth {
 // this function is used only for testing purposes in CLI e g generate tokens during init
 pub fn encode_from_key_file(claims: &Claims, key_data: &[u8]) -> Result<String> {
     let key = EncodingKey::from_rsa_pem(key_data)?;
-    Ok(encode(&Header::new(JWT_ALGORITHM), claims, &key)?)
+    Ok(encode(&Header::new(ENCODE_ALGORITHM), claims, &key)?)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::str::FromStr;
+
+    // generated with:
+    //
+    // openssl genpkey -algorithm rsa -out storage-auth-priv.pem
+    // openssl pkey -in storage-auth-priv.pem -pubout -out storage-auth-pub.pem
+    const TEST_PUB_KEY_RSA: &[u8] = br#"
+-----BEGIN PUBLIC KEY-----
+MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEAy6OZ+/kQXcueVJA/KTzO
+v4ljxylc/Kcb0sXWuXg1GB8k3nDA1gK66LFYToH0aTnqrnqG32Vu6wrhwuvqsZA7
+jQvP0ZePAbWhpEqho7EpNunDPcxZ/XDy5TQlB1P58F9I3lkJXDC+DsHYLuuzwhAv
+vo2MtWRdYlVHblCVLyZtANHhUMp2HUhgjHnJh5UrLIKOl4doCBxkM3rK0wjKsNCt
+M92PCR6S9rvYzldfeAYFNppBkEQrXt2CgUqZ4KaS4LXtjTRUJxljijA4HWffhxsr
+euRu3ufq8kVqie7fum0rdZZSkONmce0V0LesQ4aE2jB+2Sn48h6jb4dLXGWdq8TV
+wQIDAQAB
+-----END PUBLIC KEY-----
+"#;
+    const TEST_PRIV_KEY_RSA: &[u8] = br#"
+-----BEGIN PRIVATE KEY-----
+MIIEvQIBADANBgkqhkiG9w0BAQEFAASCBKcwggSjAgEAAoIBAQDLo5n7+RBdy55U
+kD8pPM6/iWPHKVz8pxvSxda5eDUYHyTecMDWArrosVhOgfRpOequeobfZW7rCuHC
+6+qxkDuNC8/Rl48BtaGkSqGjsSk26cM9zFn9cPLlNCUHU/nwX0jeWQlcML4Owdgu
+67PCEC++jYy1ZF1iVUduUJUvJm0A0eFQynYdSGCMecmHlSssgo6Xh2gIHGQzesrT
+CMqw0K0z3Y8JHpL2u9jOV194BgU2mkGQRCte3YKBSpngppLgte2NNFQnGWOKMDgd
+Z9+HGyt65G7e5+ryRWqJ7t+6bSt1llKQ42Zx7RXQt6xDhoTaMH7ZKfjyHqNvh0tc
+ZZ2rxNXBAgMBAAECggEAVz3u4Wlx3o02dsoZlSQs+xf0PEX3RXKeU+1YMbtTG9Nz
+6yxpIQaoZrpbt76rJE2gwkFR+PEu1NmjoOuLb6j4KlQuI4AHz1auOoGSwFtM6e66
+K4aZ4x95oEJ3vqz2fkmEIWYJwYpMUmwvnuJx76kZm0xvROMLsu4QHS2+zCVtO5Tr
+hvS05IMVuZ2TdQBZw0+JaFdwXbgDjQnQGY5n9MoTWSx1a4s/FF4Eby65BbDutcpn
+Vt3jQAOmO1X2kbPeWSGuPJRzyUs7Kg8qfeglBIR3ppGP3vPYAdWX+ho00bmsVkSp
+Q8vjul6C3WiM+kjwDxotHSDgbl/xldAl7OqPh0bfAQKBgQDnycXuq14Vg8nZvyn9
+rTnvucO8RBz5P6G+FZ+44cAS2x79+85onARmMnm+9MKYLSMo8fOvsK034NDI68XM
+04QQ/vlfouvFklMTGJIurgEImTZbGCmlMYCvFyIxaEWixon8OpeI4rFe4Hmbiijh
+PxhxWg221AwvBS2sco8J/ylEkQKBgQDg6Rh2QYb/j0Wou1rJPbuy3NhHofd5Rq35
+4YV3f2lfVYcPrgRhwe3T9SVII7Dx8LfwzsX5TAlf48ESlI3Dzv40uOCDM+xdtBRI
+r96SfSm+jup6gsXU3AsdNkrRK3HoOG9Z/TkrUp213QAIlVnvIx65l4ckFMlpnPJ0
+lo1LDXZWMQKBgFArzjZ7N5OhfdO+9zszC3MLgdRAivT7OWqR+CjujIz5FYMr8Xzl
+WfAvTUTrS9Nu6VZkObFvHrrRG+YjBsuN7YQjbQXTSFGSBwH34bgbn2fl9pMTjHQC
+50uoaL9GHa/rlBaV/YvvPQJgCi/uXa1rMX0jdNLkDULGO8IF7cu7Yf7BAoGBAIUU
+J29BkpmAst0GDs/ogTlyR18LTR0rXyHt+UUd1MGeH859TwZw80JpWWf4BmkB4DTS
+hH3gKePdJY7S65ci0XNsuRupC4DeXuorde0DtkGU2tUmr9wlX0Ynq9lcdYfMbMa4
+eK1TsxG69JwfkxlWlIWITWRiEFM3lJa7xlrUWmLhAoGAFpKWF/hn4zYg3seU9gai
+EYHKSbhxA4mRb+F0/9IlCBPMCqFrL5yftUsYIh2XFKn8+QhO97Nmk8wJSK6TzQ5t
+ZaSRmgySrUUhx4nZ/MgqWCFv8VUbLM5MBzwxPKhXkSTfR4z2vLYLJwVY7Tb4kZtp
+8ismApXVGHpOCstzikV9W7k=
+-----END PRIVATE KEY-----
+"#;
+
+    #[test]
+    fn test_decode() -> Result<(), anyhow::Error> {
+        let expected_claims = Claims {
+            tenant_id: Some(TenantId::from_str("3d1f7595b468230304e0b73cecbcb081")?),
+            scope: Scope::Tenant,
+        };
+
+        // Here are tokens containing the following payload, signed using TEST_PRIV_KEY_RSA
+        // using RS512, RS384 and RS256 algorithms:
+        //
+        // ```
+        // {
+        //   "scope": "tenant",
+        //   "tenant_id": "3d1f7595b468230304e0b73cecbcb081",
+        //   "iss": "neon.controlplane",
+        //   "exp": 1709200879,
+        //   "iat": 1678442479
+        // }
+        // ```
+        //
+        // These were encoded with the online debugger at https://jwt.io
+        //
+        let encoded_rs512 = "eyJhbGciOiJSUzUxMiIsInR5cCI6IkpXVCJ9.eyJzY29wZSI6InRlbmFudCIsInRlbmFudF9pZCI6IjNkMWY3NTk1YjQ2ODIzMDMwNGUwYjczY2VjYmNiMDgxIiwiaXNzIjoibmVvbi5jb250cm9scGxhbmUiLCJleHAiOjE3MDkyMDA4NzksImlhdCI6MTY3ODQ0MjQ3OX0.QmqfteDQmDGoxQ5EFkasbt35Lx0W0Nh63muQnYZvFq93DSh4ZbOG9Mc4yaiXZoiS5HgeKtFKv3mbWkDqjz3En06aY17hWwguBtAsGASX48lYeCPADYGlGAuaWnOnVRwe3iiOC7tvPFvwX_45S84X73sNUXyUiXv6nLdcDqVXudtNrGST_DnZDnjuUJX11w7sebtKqQQ8l9-iGHiXOl5yevpMCoB1OcTWcT6DfDtffoNuMHDC3fyhmEGG5oKAt1qBybqAIiyC9-UBAowRZXhdfxrzUl-I9jzKWvk85c5ulhVRwbPeP6TTTlPKwFzBNHg1i2U-1GONew5osQ3aoptwsA";
+
+        let encoded_rs384 = "eyJhbGciOiJSUzM4NCIsInR5cCI6IkpXVCJ9.eyJzY29wZSI6InRlbmFudCIsInRlbmFudF9pZCI6IjNkMWY3NTk1YjQ2ODIzMDMwNGUwYjczY2VjYmNiMDgxIiwiaXNzIjoibmVvbi5jb250cm9scGxhbmUiLCJleHAiOjE3MDkyMDA4NzksImlhdCI6MTY3ODQ0MjQ3OX0.qqk4nkxKzOJP38c_g57_w_SfdQVmCsDT_bsLmdFj_N6LIB22gr6U6_P_5mvk3pIAsp0VCTDwPrCU908TxqjibEkwvQoJwbogHamSGHpD7eJBxGblSnA-Nr3MlEMxpFtec8QokSm6C5mH7DoBYjB2xzeOlxAmpR2GAzInKiMkU4kZ_OcqqrmVcMXY_6VnbxZWMekuw56zE1-PP_qNF1HvYOH-P08ONP8qdo5UPtBG7QBEFlCqZXJZCFihQaI4Vzil9rDuZGCm3I7xQJ8-yh1PX3BTbGo8EzqLdRyBeTpr08UTuRbp_MJDWevHpP3afvJetAItqZXIoZQrbJjcByHqKw";
+
+        let encoded_rs256 = "eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.eyJzY29wZSI6InRlbmFudCIsInRlbmFudF9pZCI6IjNkMWY3NTk1YjQ2ODIzMDMwNGUwYjczY2VjYmNiMDgxIiwiaXNzIjoibmVvbi5jb250cm9scGxhbmUiLCJleHAiOjE3MDkyMDA4NzksImlhdCI6MTY3ODQ0MjQ3OX0.dF2N9KXG8ftFKHYbd5jQtXMQqv0Ej8FISGp1b_dmqOCotXj5S1y2AWjwyB_EXHM77JXfbEoJPAPrFFBNfd8cWtkCSTvpxWoHaecGzegDFGv5ZSc5AECFV1Daahc3PI3jii9wEiGkFOiwiBNfZ5INomOAsV--XXxlqIwKbTcgSYI7lrOTfecXAbAHiMKQlQYiIBSGnytRCgafhRkyGzPAL8ismthFJ9RHfeejyskht-9GbVHURw02bUyijuHEulpf9eEY3ZiB28de6jnCdU7ftIYaUMaYWt0nZQGkzxKPSfSLZNy14DTOYLDS04DVstWQPqnCUW_ojg0wJETOOfo9Zw";
+
+        // Check that RS512, RS384 and RS256 tokens can all be validated
+        let auth = JwtAuth::new(DecodingKey::from_rsa_pem(TEST_PUB_KEY_RSA)?);
+
+        for encoded in [encoded_rs512, encoded_rs384, encoded_rs256] {
+            let claims_from_token = auth.decode(encoded)?.claims;
+            assert_eq!(claims_from_token, expected_claims);
+        }
+        Ok(())
+    }
+
+    #[test]
+    fn test_encode() -> Result<(), anyhow::Error> {
+        let claims = Claims {
+            tenant_id: Some(TenantId::from_str("3d1f7595b468230304e0b73cecbcb081")?),
+            scope: Scope::Tenant,
+        };
+
+        let encoded = encode_from_key_file(&claims, TEST_PRIV_KEY_RSA)?;
+
+        // decode it back
+        let auth = JwtAuth::new(DecodingKey::from_rsa_pem(TEST_PUB_KEY_RSA)?);
+        let decoded = auth.decode(&encoded)?;
+
+        assert_eq!(decoded.claims, claims);
+
+        Ok(())
+    }
 }

libs/utils/src/fs_ext.rs

@@ -11,7 +11,7 @@ where
     P: AsRef<Path>,
 {
     fn is_empty_dir(&self) -> io::Result<bool> {
-        Ok(fs::read_dir(self)?.into_iter().next().is_none())
+        Ok(fs::read_dir(self)?.next().is_none())
     }
 }
 

libs/utils/src/http/endpoint.rs

@@ -3,14 +3,14 @@ use crate::http::error;
 use anyhow::{anyhow, Context};
 use hyper::header::{HeaderName, AUTHORIZATION};
 use hyper::http::HeaderValue;
+use hyper::Method;
 use hyper::{header::CONTENT_TYPE, Body, Request, Response, Server};
-use hyper::{Method, StatusCode};
 use metrics::{register_int_counter, Encoder, IntCounter, TextEncoder};
 use once_cell::sync::Lazy;
 use routerify::ext::RequestExt;
 use routerify::{Middleware, RequestInfo, Router, RouterBuilder, RouterService};
 use tokio::task::JoinError;
-use tracing;
+use tracing::{self, debug, info, info_span, warn, Instrument};
 
 use std::future::Future;
 use std::net::TcpListener;
@@ -32,31 +32,77 @@ static X_REQUEST_ID_HEADER: HeaderName = HeaderName::from_static(X_REQUEST_ID_HE
 #[derive(Debug, Default, Clone)]
 struct RequestId(String);
 
-async fn logger(res: Response<Body>, info: RequestInfo) -> Result<Response<Body>, ApiError> {
-    let request_id = info.context::<RequestId>().unwrap_or_default().0;
+/// Adds a tracing info_span! instrumentation around the handler events,
+/// logs the request start and end events for non-GET requests and non-200 responses.
+///
+/// Use this to distinguish between logs of different HTTP requests: every request handler wrapped
+/// in this type will get request info logged in the wrapping span, including the unique request ID.
+///
+/// There could be other ways to implement similar functionality:
+///
+/// * procmacros placed on top of all handler methods
+/// With all the drawbacks of procmacros, brings no difference implementation-wise,
+/// and little code reduction compared to the existing approach.
+///
+/// * Another `TraitExt` with e.g. the `get_with_span`, `post_with_span` methods to do similar logic,
+/// implemented for [`RouterBuilder`].
+/// Could be simpler, but we don't want to depend on [`routerify`] more, targeting to use other library later.
+///
+/// * In theory, a span guard could've been created in a pre-request middleware and placed into a global collection, to be dropped
+/// later, in a post-response middleware.
+/// Due to suspendable nature of the futures, would give contradictive results which is exactly the opposite of what `tracing-futures`
+/// tries to achive with its `.instrument` used in the current approach.
+///
+/// If needed, a declarative macro to substitute the |r| ... closure boilerplate could be introduced.
+pub struct RequestSpan<E, R, H>(pub H)
+where
+    E: Into<Box<dyn std::error::Error + Send + Sync>> + 'static,
+    R: Future<Output = Result<Response<Body>, E>> + Send + 'static,
+    H: Fn(Request<Body>) -> R + Send + Sync + 'static;
 
-    // cannot factor out the Level to avoid the repetition
-    // because tracing can only work with const Level
-    // which is not the case here
+impl<E, R, H> RequestSpan<E, R, H>
+where
+    E: Into<Box<dyn std::error::Error + Send + Sync>> + 'static,
+    R: Future<Output = Result<Response<Body>, E>> + Send + 'static,
+    H: Fn(Request<Body>) -> R + Send + Sync + 'static,
+{
+    /// Creates a tracing span around inner request handler and executes the request handler in the contex of that span.
+    /// Use as `|r| RequestSpan(my_handler).handle(r)` instead of `my_handler` as the request handler to get the span enabled.
+    pub async fn handle(self, request: Request<Body>) -> Result<Response<Body>, E> {
+        let request_id = request.context::<RequestId>().unwrap_or_default().0;
+        let method = request.method();
+        let path = request.uri().path();
+        let request_span = info_span!("request", %method, %path, %request_id);
 
-    if info.method() == Method::GET && res.status() == StatusCode::OK {
-        tracing::debug!(
-            "{} {} {} {}",
-            info.method(),
-            info.uri().path(),
-            request_id,
-            res.status()
-        );
-    } else {
-        tracing::info!(
-            "{} {} {} {}",
-            info.method(),
-            info.uri().path(),
-            request_id,
-            res.status()
-        );
+        let log_quietly = method == Method::GET;
+        async move {
+            if log_quietly {
+                debug!("Handling request");
+            } else {
+                info!("Handling request");
+            }
+
+            // Note that we reuse `error::handler` here and not returning and error at all,
+            // yet cannot use `!` directly in the method signature due to `routerify::RouterBuilder` limitation.
+            // Usage of the error handler also means that we expect only the `ApiError` errors to be raised in this call.
+            //
+            // Panics are not handled separately, there's a `tracing_panic_hook` from another module to do that globally.
+            match (self.0)(request).await {
+                Ok(response) => {
+                    let response_status = response.status();
+                    if log_quietly && response_status.is_success() {
+                        debug!("Request handled, status: {response_status}");
+                    } else {
+                        info!("Request handled, status: {response_status}");
+                    }
+                    Ok(response)
+                }
+                Err(e) => Ok(error::handler(e.into()).await),
+            }
+        }
+        .instrument(request_span)
+        .await
     }
-    Ok(res)
 }
 
 async fn prometheus_metrics_handler(_req: Request<Body>) -> Result<Response<Body>, ApiError> {
@@ -96,12 +142,6 @@ pub fn add_request_id_middleware<B: hyper::body::HttpBody + Send + Sync + 'stati
                 request_id.to_string()
             }
         };
-
-        if req.method() == Method::GET {
-            tracing::debug!("{} {} {}", req.method(), req.uri().path(), request_id);
-        } else {
-            tracing::info!("{} {} {}", req.method(), req.uri().path(), request_id);
-        }
         req.set_context(RequestId(request_id));
 
         Ok(req)
@@ -125,11 +165,12 @@ async fn add_request_id_header_to_response(
 pub fn make_router() -> RouterBuilder<hyper::Body, ApiError> {
     Router::builder()
         .middleware(add_request_id_middleware())
-        .middleware(Middleware::post_with_info(logger))
         .middleware(Middleware::post_with_info(
             add_request_id_header_to_response,
         ))
-        .get("/metrics", prometheus_metrics_handler)
+        .get("/metrics", |r| {
+            RequestSpan(prometheus_metrics_handler).handle(r)
+        })
         .err_handler(error::handler)
 }
 
@@ -139,40 +180,43 @@ pub fn attach_openapi_ui(
     spec_mount_path: &'static str,
     ui_mount_path: &'static str,
 ) -> RouterBuilder<hyper::Body, ApiError> {
-    router_builder.get(spec_mount_path, move |_| async move {
-        Ok(Response::builder().body(Body::from(spec)).unwrap())
-    }).get(ui_mount_path, move |_| async move {
-        Ok(Response::builder().body(Body::from(format!(r#"
-            <!DOCTYPE html>
-            <html lang="en">
-            <head>
-            <title>rweb</title>
-            <link href="https://cdn.jsdelivr.net/npm/swagger-ui-dist@3/swagger-ui.css" rel="stylesheet">
-            </head>
-            <body>
-                <div id="swagger-ui"></div>
-                <script src="https://cdn.jsdelivr.net/npm/swagger-ui-dist@3/swagger-ui-bundle.js" charset="UTF-8"> </script>
-                <script>
-                    window.onload = function() {{
-                    const ui = SwaggerUIBundle({{
-                        "dom_id": "\#swagger-ui",
-                        presets: [
-                        SwaggerUIBundle.presets.apis,
-                        SwaggerUIBundle.SwaggerUIStandalonePreset
-                        ],
-                        layout: "BaseLayout",
-                        deepLinking: true,
-                        showExtensions: true,
-                        showCommonExtensions: true,
-                        url: "{}",
-                    }})
-                    window.ui = ui;
-                }};
-            </script>
-            </body>
-            </html>
-        "#, spec_mount_path))).unwrap())
-    })
+    router_builder
+        .get(spec_mount_path, move |r| {
+            RequestSpan(move |_| async move { Ok(Response::builder().body(Body::from(spec)).unwrap()) })
+                .handle(r)
+        })
+        .get(ui_mount_path, move |r| RequestSpan( move |_| async move {
+            Ok(Response::builder().body(Body::from(format!(r#"
+                <!DOCTYPE html>
+                <html lang="en">
+                <head>
+                <title>rweb</title>
+                <link href="https://cdn.jsdelivr.net/npm/swagger-ui-dist@3/swagger-ui.css" rel="stylesheet">
+                </head>
+                <body>
+                    <div id="swagger-ui"></div>
+                    <script src="https://cdn.jsdelivr.net/npm/swagger-ui-dist@3/swagger-ui-bundle.js" charset="UTF-8"> </script>
+                    <script>
+                        window.onload = function() {{
+                        const ui = SwaggerUIBundle({{
+                            "dom_id": "\#swagger-ui",
+                            presets: [
+                            SwaggerUIBundle.presets.apis,
+                            SwaggerUIBundle.SwaggerUIStandalonePreset
+                            ],
+                            layout: "BaseLayout",
+                            deepLinking: true,
+                            showExtensions: true,
+                            showCommonExtensions: true,
+                            url: "{}",
+                        }})
+                        window.ui = ui;
+                    }};
+                </script>
+                </body>
+                </html>
+            "#, spec_mount_path))).unwrap())
+        }).handle(r))
 }
 
 fn parse_token(header_value: &str) -> Result<&str, ApiError> {
@@ -234,7 +278,7 @@ where
             async move {
                 let headers = response.headers_mut();
                 if headers.contains_key(&name) {
-                    tracing::warn!(
+                    warn!(
                         "{} response already contains header {:?}",
                         request_info.uri(),
                         &name,
@@ -274,7 +318,7 @@ pub fn serve_thread_main<S>(
 where
     S: Future<Output = ()> + Send + Sync,
 {
-    tracing::info!("Starting an HTTP endpoint at {}", listener.local_addr()?);
+    info!("Starting an HTTP endpoint at {}", listener.local_addr()?);
 
     // Create a Service from the router above to handle incoming requests.
     let service = RouterService::new(router_builder.build().map_err(|err| anyhow!(err))?).unwrap();

libs/utils/src/http/json.rs

@@ -1,7 +1,9 @@
+use std::fmt::Display;
+
 use anyhow::Context;
 use bytes::Buf;
 use hyper::{header, Body, Request, Response, StatusCode};
-use serde::{Deserialize, Serialize};
+use serde::{Deserialize, Serialize, Serializer};
 
 use super::error::ApiError;
 
@@ -31,3 +33,12 @@ pub fn json_response<T: Serialize>(
         .map_err(|e| ApiError::InternalServerError(e.into()))?;
     Ok(response)
 }
+
+/// Serialize through Display trait.
+pub fn display_serialize<S, F>(z: &F, s: S) -> Result<S::Ok, S::Error>
+where
+    S: Serializer,
+    F: Display,
+{
+    s.serialize_str(&format!("{}", z))
+}

libs/utils/src/lib.rs

@@ -49,6 +49,8 @@ pub mod fs_ext;
 
 pub mod history_buffer;
 
+pub mod measured_stream;
+
 /// use with fail::cfg("$name", "return(2000)")
 #[macro_export]
 macro_rules! failpoint_sleep_millis_async {

libs/utils/src/lsn.rs

@@ -6,6 +6,7 @@ use std::ops::{Add, AddAssign};
 use std::path::Path;
 use std::str::FromStr;
 use std::sync::atomic::{AtomicU64, Ordering};
+use tracing::info;
 
 use crate::seqwait::MonotonicCounter;
 
@@ -239,6 +240,7 @@ impl MonotonicCounter<Lsn> for RecordLsn {
         let new_prev = self.last;
         self.last = lsn;
         self.prev = new_prev;
+        info!("advanced record lsn to {}/{}", self.last, self.prev);
     }
     fn cnt_value(&self) -> Lsn {
         self.last

libs/utils/src/measured_stream.rs

@@ -0,0 +1,77 @@
+use pin_project_lite::pin_project;
+use std::pin::Pin;
+use std::{io, task};
+use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
+
+pin_project! {
+    /// This stream tracks all writes and calls user provided
+    /// callback when the underlying stream is flushed.
+    pub struct MeasuredStream<S, R, W> {
+        #[pin]
+        stream: S,
+        write_count: usize,
+        inc_read_count: R,
+        inc_write_count: W,
+    }
+}
+
+impl<S, R, W> MeasuredStream<S, R, W> {
+    pub fn new(stream: S, inc_read_count: R, inc_write_count: W) -> Self {
+        Self {
+            stream,
+            write_count: 0,
+            inc_read_count,
+            inc_write_count,
+        }
+    }
+}
+
+impl<S: AsyncRead + Unpin, R: FnMut(usize), W> AsyncRead for MeasuredStream<S, R, W> {
+    fn poll_read(
+        self: Pin<&mut Self>,
+        context: &mut task::Context<'_>,
+        buf: &mut ReadBuf<'_>,
+    ) -> task::Poll<io::Result<()>> {
+        let this = self.project();
+        let filled = buf.filled().len();
+        this.stream.poll_read(context, buf).map_ok(|()| {
+            let cnt = buf.filled().len() - filled;
+            // Increment the read count.
+            (this.inc_read_count)(cnt);
+        })
+    }
+}
+
+impl<S: AsyncWrite + Unpin, R, W: FnMut(usize)> AsyncWrite for MeasuredStream<S, R, W> {
+    fn poll_write(
+        self: Pin<&mut Self>,
+        context: &mut task::Context<'_>,
+        buf: &[u8],
+    ) -> task::Poll<io::Result<usize>> {
+        let this = self.project();
+        this.stream.poll_write(context, buf).map_ok(|cnt| {
+            // Increment the write count.
+            *this.write_count += cnt;
+            cnt
+        })
+    }
+
+    fn poll_flush(
+        self: Pin<&mut Self>,
+        context: &mut task::Context<'_>,
+    ) -> task::Poll<io::Result<()>> {
+        let this = self.project();
+        this.stream.poll_flush(context).map_ok(|()| {
+            // Call the user provided callback and reset the write count.
+            (this.inc_write_count)(*this.write_count);
+            *this.write_count = 0;
+        })
+    }
+
+    fn poll_shutdown(
+        self: Pin<&mut Self>,
+        context: &mut task::Context<'_>,
+    ) -> task::Poll<io::Result<()>> {
+        self.project().stream.poll_shutdown(context)
+    }
+}

pageserver/src/bin/pageserver.rs

@@ -270,43 +270,43 @@ fn start_pageserver(
     WALRECEIVER_RUNTIME.block_on(pageserver::broker_client::init_broker_client(conf))?;
 
     // Initialize authentication for incoming connections
-    let auth = match &conf.auth_type {
-        AuthType::Trust => None,
-        AuthType::NeonJWT => {
-            // unwrap is ok because check is performed when creating config, so path is set and file exists
-            let key_path = conf.auth_validation_public_key_path.as_ref().unwrap();
-            Some(JwtAuth::from_key_path(key_path)?.into())
-        }
-    };
-    info!("Using auth: {:#?}", conf.auth_type);
+    let http_auth;
+    let pg_auth;
+    if conf.http_auth_type == AuthType::NeonJWT || conf.pg_auth_type == AuthType::NeonJWT {
+        // unwrap is ok because check is performed when creating config, so path is set and file exists
+        let key_path = conf.auth_validation_public_key_path.as_ref().unwrap();
+        info!(
+            "Loading public key for verifying JWT tokens from {:#?}",
+            key_path
+        );
+        let auth: Arc<JwtAuth> = Arc::new(JwtAuth::from_key_path(key_path)?);
 
-    // TODO: remove ZENITH_AUTH_TOKEN once it's not used anywhere in development/staging/prod configuration.
-    match (var("ZENITH_AUTH_TOKEN"), var("NEON_AUTH_TOKEN")) {
-        (old, Ok(v)) => {
+        http_auth = match &conf.http_auth_type {
+            AuthType::Trust => None,
+            AuthType::NeonJWT => Some(auth.clone()),
+        };
+        pg_auth = match &conf.pg_auth_type {
+            AuthType::Trust => None,
+            AuthType::NeonJWT => Some(auth),
+        };
+    } else {
+        http_auth = None;
+        pg_auth = None;
+    }
+    info!("Using auth for http API: {:#?}", conf.http_auth_type);
+    info!("Using auth for pg connections: {:#?}", conf.pg_auth_type);
+
+    match var("NEON_AUTH_TOKEN") {
+        Ok(v) => {
             info!("Loaded JWT token for authentication with Safekeeper");
-            if let Ok(v_old) = old {
-                warn!(
-                    "JWT token for Safekeeper is specified twice, ZENITH_AUTH_TOKEN is deprecated"
-                );
-                if v_old != v {
-                    warn!("JWT token for Safekeeper has two different values, choosing NEON_AUTH_TOKEN");
-                }
-            }
             pageserver::config::SAFEKEEPER_AUTH_TOKEN
                 .set(Arc::new(v))
                 .map_err(|_| anyhow!("Could not initialize SAFEKEEPER_AUTH_TOKEN"))?;
         }
-        (Ok(v), _) => {
-            info!("Loaded JWT token for authentication with Safekeeper");
-            warn!("Please update pageserver configuration: the JWT token should be NEON_AUTH_TOKEN, not ZENITH_AUTH_TOKEN");
-            pageserver::config::SAFEKEEPER_AUTH_TOKEN
-                .set(Arc::new(v))
-                .map_err(|_| anyhow!("Could not initialize SAFEKEEPER_AUTH_TOKEN"))?;
-        }
-        (_, Err(VarError::NotPresent)) => {
+        Err(VarError::NotPresent) => {
             info!("No JWT token for authentication with Safekeeper detected");
         }
-        (_, Err(e)) => {
+        Err(e) => {
             return Err(e).with_context(|| {
                 "Failed to either load to detect non-present NEON_AUTH_TOKEN environment variable"
             })
@@ -324,7 +324,7 @@ fn start_pageserver(
     {
         let _rt_guard = MGMT_REQUEST_RUNTIME.enter();
 
-        let router = http::make_router(conf, launch_ts, auth.clone(), remote_storage)?
+        let router = http::make_router(conf, launch_ts, http_auth, remote_storage)?
             .build()
             .map_err(|err| anyhow!(err))?;
         let service = utils::http::RouterService::new(router).unwrap();
@@ -398,9 +398,9 @@ fn start_pageserver(
             async move {
                 page_service::libpq_listener_main(
                     conf,
-                    auth,
+                    pg_auth,
                     pageserver_listener,
-                    conf.auth_type,
+                    conf.pg_auth_type,
                     libpq_ctx,
                 )
                 .await

pageserver/src/config.rs

@@ -118,6 +118,9 @@ pub struct PageServerConf {
     /// Example (default): 127.0.0.1:9898
     pub listen_http_addr: String,
 
+    /// Current availability zone. Used for traffic metrics.
+    pub availability_zone: Option<String>,
+
     // Timeout when waiting for WAL receiver to catch up to an LSN given in a GetPage@LSN call.
     pub wait_lsn_timeout: Duration,
     // How long to wait for WAL redo to complete.
@@ -138,9 +141,15 @@ pub struct PageServerConf {
 
     pub pg_distrib_dir: PathBuf,
 
-    pub auth_type: AuthType,
-
+    // Authentication
+    /// authentication method for the HTTP mgmt API
+    pub http_auth_type: AuthType,
+    /// authentication method for libpq connections from compute
+    pub pg_auth_type: AuthType,
+    /// Path to a file containing public key for verifying JWT tokens.
+    /// Used for both mgmt and compute auth, if enabled.
     pub auth_validation_public_key_path: Option<PathBuf>,
+
     pub remote_storage_config: Option<RemoteStorageConfig>,
 
     pub default_tenant_conf: TenantConf,
@@ -196,6 +205,8 @@ struct PageServerConfigBuilder {
 
     listen_http_addr: BuilderValue<String>,
 
+    availability_zone: BuilderValue<Option<String>>,
+
     wait_lsn_timeout: BuilderValue<Duration>,
     wal_redo_timeout: BuilderValue<Duration>,
 
@@ -208,7 +219,8 @@ struct PageServerConfigBuilder {
 
     pg_distrib_dir: BuilderValue<PathBuf>,
 
-    auth_type: BuilderValue<AuthType>,
+    http_auth_type: BuilderValue<AuthType>,
+    pg_auth_type: BuilderValue<AuthType>,
 
     //
     auth_validation_public_key_path: BuilderValue<Option<PathBuf>>,
@@ -240,6 +252,7 @@ impl Default for PageServerConfigBuilder {
         Self {
             listen_pg_addr: Set(DEFAULT_PG_LISTEN_ADDR.to_string()),
             listen_http_addr: Set(DEFAULT_HTTP_LISTEN_ADDR.to_string()),
+            availability_zone: Set(None),
             wait_lsn_timeout: Set(humantime::parse_duration(DEFAULT_WAIT_LSN_TIMEOUT)
                 .expect("cannot parse default wait lsn timeout")),
             wal_redo_timeout: Set(humantime::parse_duration(DEFAULT_WAL_REDO_TIMEOUT)
@@ -251,7 +264,8 @@ impl Default for PageServerConfigBuilder {
             pg_distrib_dir: Set(env::current_dir()
                 .expect("cannot access current directory")
                 .join("pg_install")),
-            auth_type: Set(AuthType::Trust),
+            http_auth_type: Set(AuthType::Trust),
+            pg_auth_type: Set(AuthType::Trust),
             auth_validation_public_key_path: Set(None),
             remote_storage_config: Set(None),
             id: NotSet,
@@ -295,6 +309,10 @@ impl PageServerConfigBuilder {
         self.listen_http_addr = BuilderValue::Set(listen_http_addr)
     }
 
+    pub fn availability_zone(&mut self, availability_zone: Option<String>) {
+        self.availability_zone = BuilderValue::Set(availability_zone)
+    }
+
     pub fn wait_lsn_timeout(&mut self, wait_lsn_timeout: Duration) {
         self.wait_lsn_timeout = BuilderValue::Set(wait_lsn_timeout)
     }
@@ -323,8 +341,12 @@ impl PageServerConfigBuilder {
         self.pg_distrib_dir = BuilderValue::Set(pg_distrib_dir)
     }
 
-    pub fn auth_type(&mut self, auth_type: AuthType) {
-        self.auth_type = BuilderValue::Set(auth_type)
+    pub fn http_auth_type(&mut self, auth_type: AuthType) {
+        self.http_auth_type = BuilderValue::Set(auth_type)
+    }
+
+    pub fn pg_auth_type(&mut self, auth_type: AuthType) {
+        self.pg_auth_type = BuilderValue::Set(auth_type)
     }
 
     pub fn auth_validation_public_key_path(
@@ -402,6 +424,9 @@ impl PageServerConfigBuilder {
             listen_http_addr: self
                 .listen_http_addr
                 .ok_or(anyhow!("missing listen_http_addr"))?,
+            availability_zone: self
+                .availability_zone
+                .ok_or(anyhow!("missing availability_zone"))?,
             wait_lsn_timeout: self
                 .wait_lsn_timeout
                 .ok_or(anyhow!("missing wait_lsn_timeout"))?,
@@ -419,7 +444,10 @@ impl PageServerConfigBuilder {
             pg_distrib_dir: self
                 .pg_distrib_dir
                 .ok_or(anyhow!("missing pg_distrib_dir"))?,
-            auth_type: self.auth_type.ok_or(anyhow!("missing auth_type"))?,
+            http_auth_type: self
+                .http_auth_type
+                .ok_or(anyhow!("missing http_auth_type"))?,
+            pg_auth_type: self.pg_auth_type.ok_or(anyhow!("missing pg_auth_type"))?,
             auth_validation_public_key_path: self
                 .auth_validation_public_key_path
                 .ok_or(anyhow!("missing auth_validation_public_key_path"))?,
@@ -599,6 +627,7 @@ impl PageServerConf {
             match key {
                 "listen_pg_addr" => builder.listen_pg_addr(parse_toml_string(key, item)?),
                 "listen_http_addr" => builder.listen_http_addr(parse_toml_string(key, item)?),
+                "availability_zone" => builder.availability_zone(Some(parse_toml_string(key, item)?)),
                 "wait_lsn_timeout" => builder.wait_lsn_timeout(parse_toml_duration(key, item)?),
                 "wal_redo_timeout" => builder.wal_redo_timeout(parse_toml_duration(key, item)?),
                 "initial_superuser_name" => builder.superuser(parse_toml_string(key, item)?),
@@ -612,7 +641,8 @@ impl PageServerConf {
                 "auth_validation_public_key_path" => builder.auth_validation_public_key_path(Some(
                     PathBuf::from(parse_toml_string(key, item)?),
                 )),
-                "auth_type" => builder.auth_type(parse_toml_from_str(key, item)?),
+                "http_auth_type" => builder.http_auth_type(parse_toml_from_str(key, item)?),
+                "pg_auth_type" => builder.pg_auth_type(parse_toml_from_str(key, item)?),
                 "remote_storage" => {
                     builder.remote_storage_config(RemoteStorageConfig::from_toml(item)?)
                 }
@@ -647,7 +677,7 @@ impl PageServerConf {
 
         let mut conf = builder.build().context("invalid config")?;
 
-        if conf.auth_type == AuthType::NeonJWT {
+        if conf.http_auth_type == AuthType::NeonJWT || conf.pg_auth_type == AuthType::NeonJWT {
             let auth_validation_public_key_path = conf
                 .auth_validation_public_key_path
                 .get_or_insert_with(|| workdir.join("auth_public_key.pem"));
@@ -698,6 +728,12 @@ impl PageServerConf {
                 Some(parse_toml_u64("compaction_threshold", compaction_threshold)?.try_into()?);
         }
 
+        if let Some(image_creation_threshold) = item.get("image_creation_threshold") {
+            t_conf.image_creation_threshold = Some(
+                parse_toml_u64("image_creation_threshold", image_creation_threshold)?.try_into()?,
+            );
+        }
+
         if let Some(gc_horizon) = item.get("gc_horizon") {
             t_conf.gc_horizon = Some(parse_toml_u64("gc_horizon", gc_horizon)?);
         }
@@ -757,10 +793,12 @@ impl PageServerConf {
             max_file_descriptors: defaults::DEFAULT_MAX_FILE_DESCRIPTORS,
             listen_pg_addr: defaults::DEFAULT_PG_LISTEN_ADDR.to_string(),
             listen_http_addr: defaults::DEFAULT_HTTP_LISTEN_ADDR.to_string(),
+            availability_zone: None,
             superuser: "cloud_admin".to_string(),
             workdir: repo_dir,
             pg_distrib_dir,
-            auth_type: AuthType::Trust,
+            http_auth_type: AuthType::Trust,
+            pg_auth_type: AuthType::Trust,
             auth_validation_public_key_path: None,
             remote_storage_config: None,
             default_tenant_conf: TenantConf::default(),
@@ -938,6 +976,7 @@ log_format = 'json'
                 id: NodeId(10),
                 listen_pg_addr: defaults::DEFAULT_PG_LISTEN_ADDR.to_string(),
                 listen_http_addr: defaults::DEFAULT_HTTP_LISTEN_ADDR.to_string(),
+                availability_zone: None,
                 wait_lsn_timeout: humantime::parse_duration(defaults::DEFAULT_WAIT_LSN_TIMEOUT)?,
                 wal_redo_timeout: humantime::parse_duration(defaults::DEFAULT_WAL_REDO_TIMEOUT)?,
                 superuser: defaults::DEFAULT_SUPERUSER.to_string(),
@@ -945,7 +984,8 @@ log_format = 'json'
                 max_file_descriptors: defaults::DEFAULT_MAX_FILE_DESCRIPTORS,
                 workdir,
                 pg_distrib_dir,
-                auth_type: AuthType::Trust,
+                http_auth_type: AuthType::Trust,
+                pg_auth_type: AuthType::Trust,
                 auth_validation_public_key_path: None,
                 remote_storage_config: None,
                 default_tenant_conf: TenantConf::default(),
@@ -995,6 +1035,7 @@ log_format = 'json'
                 id: NodeId(10),
                 listen_pg_addr: "127.0.0.1:64000".to_string(),
                 listen_http_addr: "127.0.0.1:9898".to_string(),
+                availability_zone: None,
                 wait_lsn_timeout: Duration::from_secs(111),
                 wal_redo_timeout: Duration::from_secs(111),
                 superuser: "zzzz".to_string(),
@@ -1002,7 +1043,8 @@ log_format = 'json'
                 max_file_descriptors: 333,
                 workdir,
                 pg_distrib_dir,
-                auth_type: AuthType::Trust,
+                http_auth_type: AuthType::Trust,
+                pg_auth_type: AuthType::Trust,
                 auth_validation_public_key_path: None,
                 remote_storage_config: None,
                 default_tenant_conf: TenantConf::default(),

pageserver/src/http/openapi_spec.yml

@@ -245,6 +245,53 @@ paths:
             application/json:
               schema:
                 $ref: "#/components/schemas/Error"
+  /v1/tenant/{tenant_id}/timeline/{timeline_id}/do_gc:
+    parameters:
+      - name: tenant_id
+        in: path
+        required: true
+        schema:
+          type: string
+          format: hex
+      - name: timeline_id
+        in: path
+        required: true
+        schema:
+          type: string
+          format: hex
+    put:
+      description: Garbage collect given timeline
+      responses:
+        "200":
+          description: OK
+          content:
+            application/json:
+              schema:
+                type: string
+        "400":
+          description: Error when no tenant id found in path, no timeline id or invalid timestamp
+          content:
+            application/json:
+              schema:
+                $ref: "#/components/schemas/Error"
+        "401":
+          description: Unauthorized Error
+          content:
+            application/json:
+              schema:
+                $ref: "#/components/schemas/UnauthorizedError"
+        "403":
+          description: Forbidden Error
+          content:
+            application/json:
+              schema:
+                $ref: "#/components/schemas/ForbiddenError"
+        "500":
+          description: Generic operation error
+          content:
+            application/json:
+              schema:
+                $ref: "#/components/schemas/Error"
   /v1/tenant/{tenant_id}/attach:
     parameters:
       - name: tenant_id

pageserver/src/http/routes.rs

@@ -10,6 +10,7 @@ use remote_storage::GenericRemoteStorage;
 use tenant_size_model::{SizeResult, StorageModel};
 use tokio_util::sync::CancellationToken;
 use tracing::*;
+use utils::http::endpoint::RequestSpan;
 use utils::http::request::{get_request_param, must_get_query_param, parse_query_param};
 
 use super::models::{
@@ -20,7 +21,7 @@ use crate::context::{DownloadBehavior, RequestContext};
 use crate::pgdatadir_mapping::LsnForTimestamp;
 use crate::task_mgr::TaskKind;
 use crate::tenant::config::TenantConfOpt;
-use crate::tenant::mgr::TenantMapInsertError;
+use crate::tenant::mgr::{TenantMapInsertError, TenantStateError};
 use crate::tenant::size::ModelInputs;
 use crate::tenant::storage_layer::LayerAccessStatsReset;
 use crate::tenant::{PageReconstructError, Timeline};
@@ -81,38 +82,52 @@ fn get_config(request: &Request<Body>) -> &'static PageServerConf {
     get_state(request).conf
 }
 
+/// Check that the requester is authorized to operate on given tenant
 fn check_permission(request: &Request<Body>, tenant_id: Option<TenantId>) -> Result<(), ApiError> {
     check_permission_with(request, |claims| {
         crate::auth::check_permission(claims, tenant_id)
     })
 }
 
-fn apierror_from_prerror(err: PageReconstructError) -> ApiError {
-    match err {
-        PageReconstructError::Other(err) => ApiError::InternalServerError(err),
-        PageReconstructError::NeedsDownload(_, _) => {
-            // This shouldn't happen, because we use a RequestContext that requests to
-            // download any missing layer files on-demand.
-            ApiError::InternalServerError(anyhow::anyhow!("need to download remote layer file"))
-        }
-        PageReconstructError::Cancelled => {
-            ApiError::InternalServerError(anyhow::anyhow!("request was cancelled"))
-        }
-        PageReconstructError::WalRedo(err) => {
-            ApiError::InternalServerError(anyhow::Error::new(err))
+impl From<PageReconstructError> for ApiError {
+    fn from(pre: PageReconstructError) -> ApiError {
+        match pre {
+            PageReconstructError::Other(pre) => ApiError::InternalServerError(pre),
+            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!("need to download remote layer file"))
+            }
+            PageReconstructError::Cancelled => {
+                ApiError::InternalServerError(anyhow::anyhow!("request was cancelled"))
+            }
+            PageReconstructError::WalRedo(pre) => {
+                ApiError::InternalServerError(anyhow::Error::new(pre))
+            }
         }
     }
 }
 
-fn apierror_from_tenant_map_insert_error(e: TenantMapInsertError) -> ApiError {
-    match e {
-        TenantMapInsertError::StillInitializing | TenantMapInsertError::ShuttingDown => {
-            ApiError::InternalServerError(anyhow::Error::new(e))
+impl From<TenantMapInsertError> for ApiError {
+    fn from(tmie: TenantMapInsertError) -> ApiError {
+        match tmie {
+            TenantMapInsertError::StillInitializing | TenantMapInsertError::ShuttingDown => {
+                ApiError::InternalServerError(anyhow::Error::new(tmie))
+            }
+            TenantMapInsertError::TenantAlreadyExists(id, state) => {
+                ApiError::Conflict(format!("tenant {id} already exists, state: {state:?}"))
+            }
+            TenantMapInsertError::Closure(e) => ApiError::InternalServerError(e),
         }
-        TenantMapInsertError::TenantAlreadyExists(id, state) => {
-            ApiError::Conflict(format!("tenant {id} already exists, state: {state:?}"))
+    }
+}
+
+impl From<TenantStateError> for ApiError {
+    fn from(tse: TenantStateError) -> ApiError {
+        match tse {
+            TenantStateError::NotFound(tid) => ApiError::NotFound(anyhow!("tenant {}", tid)),
+            _ => ApiError::InternalServerError(anyhow::Error::new(tse)),
         }
-        TenantMapInsertError::Closure(e) => ApiError::InternalServerError(e),
     }
 }
 
@@ -216,9 +231,7 @@ async fn timeline_create_handler(mut request: Request<Body>) -> Result<Response<
 
     let ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Error);
 
-    let tenant = mgr::get_tenant(tenant_id, true)
-        .await
-        .map_err(ApiError::NotFound)?;
+    let tenant = mgr::get_tenant(tenant_id, true).await?;
     match tenant.create_timeline(
         new_timeline_id,
         request_data.ancestor_timeline_id.map(TimelineId::from),
@@ -248,9 +261,7 @@ async fn timeline_list_handler(request: Request<Body>) -> Result<Response<Body>,
     let ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Download);
 
     let response_data = async {
-        let tenant = mgr::get_tenant(tenant_id, true)
-            .await
-            .map_err(ApiError::NotFound)?;
+        let tenant = mgr::get_tenant(tenant_id, true).await?;
         let timelines = tenant.list_timelines();
 
         let mut response_data = Vec::with_capacity(timelines.len());
@@ -266,7 +277,7 @@ async fn timeline_list_handler(request: Request<Body>) -> Result<Response<Body>,
 
             response_data.push(timeline_info);
         }
-        Ok(response_data)
+        Ok::<Vec<TimelineInfo>, ApiError>(response_data)
     }
     .instrument(info_span!("timeline_list", tenant = %tenant_id))
     .await?;
@@ -285,9 +296,7 @@ async fn timeline_detail_handler(request: Request<Body>) -> Result<Response<Body
     let ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Download);
 
     let timeline_info = async {
-        let tenant = mgr::get_tenant(tenant_id, true)
-            .await
-            .map_err(ApiError::NotFound)?;
+        let tenant = mgr::get_tenant(tenant_id, true).await?;
 
         let timeline = tenant
             .get_timeline(timeline_id, false)
@@ -323,10 +332,7 @@ async fn get_lsn_by_timestamp_handler(request: Request<Body>) -> Result<Response
 
     let ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Download);
     let timeline = active_timeline_of_active_tenant(tenant_id, timeline_id).await?;
-    let result = timeline
-        .find_lsn_for_timestamp(timestamp_pg, &ctx)
-        .await
-        .map_err(apierror_from_prerror)?;
+    let result = timeline.find_lsn_for_timestamp(timestamp_pg, &ctx).await?;
 
     let result = match result {
         LsnForTimestamp::Present(lsn) => format!("{lsn}"),
@@ -351,8 +357,7 @@ async fn tenant_attach_handler(request: Request<Body>) -> Result<Response<Body>,
     if let Some(remote_storage) = &state.remote_storage {
         mgr::attach_tenant(state.conf, tenant_id, remote_storage.clone(), &ctx)
             .instrument(info_span!("tenant_attach", tenant = %tenant_id))
-            .await
-            .map_err(apierror_from_tenant_map_insert_error)?;
+            .await?;
     } else {
         return Err(ApiError::BadRequest(anyhow!(
             "attach_tenant is not possible because pageserver was configured without remote storage"
@@ -371,11 +376,7 @@ async fn timeline_delete_handler(request: Request<Body>) -> Result<Response<Body
 
     mgr::delete_timeline(tenant_id, timeline_id, &ctx)
         .instrument(info_span!("timeline_delete", tenant = %tenant_id, timeline = %timeline_id))
-        .await
-        // FIXME: Errors from `delete_timeline` can occur for a number of reasons, incuding both
-        // user and internal errors. Replace this with better handling once the error type permits
-        // it.
-        .map_err(ApiError::InternalServerError)?;
+        .await?;
 
     json_response(StatusCode::OK, ())
 }
@@ -388,10 +389,7 @@ async fn tenant_detach_handler(request: Request<Body>) -> Result<Response<Body>,
     let conf = state.conf;
     mgr::detach_tenant(conf, tenant_id)
         .instrument(info_span!("tenant_detach", tenant = %tenant_id))
-        .await
-        // FIXME: Errors from `detach_tenant` can be caused by both both user and internal errors.
-        // Replace this with better handling once the error type permits it.
-        .map_err(ApiError::InternalServerError)?;
+        .await?;
 
     json_response(StatusCode::OK, ())
 }
@@ -405,8 +403,7 @@ async fn tenant_load_handler(request: Request<Body>) -> Result<Response<Body>, A
     let state = get_state(&request);
     mgr::load_tenant(state.conf, tenant_id, state.remote_storage.clone(), &ctx)
         .instrument(info_span!("load", tenant = %tenant_id))
-        .await
-        .map_err(apierror_from_tenant_map_insert_error)?;
+        .await?;
 
     json_response(StatusCode::ACCEPTED, ())
 }
@@ -419,10 +416,7 @@ async fn tenant_ignore_handler(request: Request<Body>) -> Result<Response<Body>,
     let conf = state.conf;
     mgr::ignore_tenant(conf, tenant_id)
         .instrument(info_span!("ignore_tenant", tenant = %tenant_id))
-        .await
-        // FIXME: Errors from `ignore_tenant` can be caused by both both user and internal errors.
-        // Replace this with better handling once the error type permits it.
-        .map_err(ApiError::InternalServerError)?;
+        .await?;
 
     json_response(StatusCode::OK, ())
 }
@@ -496,9 +490,7 @@ async fn tenant_size_handler(request: Request<Body>) -> Result<Response<Body>, A
     let headers = request.headers();
 
     let ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Download);
-    let tenant = mgr::get_tenant(tenant_id, true)
-        .await
-        .map_err(ApiError::InternalServerError)?;
+    let tenant = mgr::get_tenant(tenant_id, true).await?;
 
     // this can be long operation
     let inputs = tenant
@@ -761,8 +753,7 @@ async fn tenant_create_handler(mut request: Request<Body>) -> Result<Response<Bo
         &ctx,
     )
     .instrument(info_span!("tenant_create", tenant = ?target_tenant_id))
-    .await
-    .map_err(apierror_from_tenant_map_insert_error)?;
+    .await?;
 
     // We created the tenant. Existing API semantics are that the tenant
     // is Active when this function returns.
@@ -786,9 +777,7 @@ async fn get_tenant_config_handler(request: Request<Body>) -> Result<Response<Bo
     let tenant_id: TenantId = parse_request_param(&request, "tenant_id")?;
     check_permission(&request, Some(tenant_id))?;
 
-    let tenant = mgr::get_tenant(tenant_id, false)
-        .await
-        .map_err(ApiError::NotFound)?;
+    let tenant = mgr::get_tenant(tenant_id, false).await?;
 
     let response = HashMap::from([
         (
@@ -883,10 +872,7 @@ async fn update_tenant_config_handler(
     let state = get_state(&request);
     mgr::set_new_tenant_config(state.conf, tenant_conf, tenant_id)
         .instrument(info_span!("tenant_config", tenant = ?tenant_id))
-        .await
-        // FIXME: `update_tenant_config` can fail because of both user and internal errors.
-        // Replace this `map_err` with better error handling once the type permits it
-        .map_err(ApiError::InternalServerError)?;
+        .await?;
 
     json_response(StatusCode::OK, ())
 }
@@ -971,19 +957,22 @@ async fn timeline_checkpoint_handler(request: Request<Body>) -> Result<Response<
     let tenant_id: TenantId = parse_request_param(&request, "tenant_id")?;
     let timeline_id: TimelineId = parse_request_param(&request, "timeline_id")?;
     check_permission(&request, Some(tenant_id))?;
+    async {
+        let ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Download);
+        let timeline = active_timeline_of_active_tenant(tenant_id, timeline_id).await?;
+        timeline
+            .freeze_and_flush()
+            .await
+            .map_err(ApiError::InternalServerError)?;
+        timeline
+            .compact(&ctx)
+            .await
+            .map_err(ApiError::InternalServerError)?;
 
-    let ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Download);
-    let timeline = active_timeline_of_active_tenant(tenant_id, timeline_id).await?;
-    timeline
-        .freeze_and_flush()
-        .await
-        .map_err(ApiError::InternalServerError)?;
-    timeline
-        .compact(&ctx)
-        .await
-        .map_err(ApiError::InternalServerError)?;
-
-    json_response(StatusCode::OK, ())
+        json_response(StatusCode::OK, ())
+    }
+    .instrument(info_span!("manual_checkpoint", tenant_id = %tenant_id, timeline_id = %timeline_id))
+    .await
 }
 
 async fn timeline_download_remote_layers_handler_post(
@@ -1020,9 +1009,7 @@ async fn active_timeline_of_active_tenant(
     tenant_id: TenantId,
     timeline_id: TimelineId,
 ) -> Result<Arc<Timeline>, ApiError> {
-    let tenant = mgr::get_tenant(tenant_id, true)
-        .await
-        .map_err(ApiError::NotFound)?;
+    let tenant = mgr::get_tenant(tenant_id, true).await?;
     tenant
         .get_timeline(timeline_id, true)
         .map_err(ApiError::NotFound)
@@ -1088,7 +1075,8 @@ pub fn make_router(
             let handler = $handler;
             #[cfg(not(feature = "testing"))]
             let handler = cfg_disabled;
-            handler
+
+            move |r| RequestSpan(handler).handle(r)
         }};
     }
 
@@ -1096,35 +1084,55 @@ pub fn make_router(
         .data(Arc::new(
             State::new(conf, auth, remote_storage).context("Failed to initialize router state")?,
         ))
-        .get("/v1/status", status_handler)
+        .get("/v1/status", |r| RequestSpan(status_handler).handle(r))
         .put(
             "/v1/failpoints",
             testing_api!("manage failpoints", failpoints_handler),
         )
-        .get("/v1/tenant", tenant_list_handler)
-        .post("/v1/tenant", tenant_create_handler)
-        .get("/v1/tenant/:tenant_id", tenant_status)
-        .get("/v1/tenant/:tenant_id/synthetic_size", tenant_size_handler)
-        .put("/v1/tenant/config", update_tenant_config_handler)
-        .get("/v1/tenant/:tenant_id/config", get_tenant_config_handler)
-        .get("/v1/tenant/:tenant_id/timeline", timeline_list_handler)
-        .post("/v1/tenant/:tenant_id/timeline", timeline_create_handler)
-        .post("/v1/tenant/:tenant_id/attach", tenant_attach_handler)
-        .post("/v1/tenant/:tenant_id/detach", tenant_detach_handler)
-        .post("/v1/tenant/:tenant_id/load", tenant_load_handler)
-        .post("/v1/tenant/:tenant_id/ignore", tenant_ignore_handler)
-        .get(
-            "/v1/tenant/:tenant_id/timeline/:timeline_id",
-            timeline_detail_handler,
-        )
+        .get("/v1/tenant", |r| RequestSpan(tenant_list_handler).handle(r))
+        .post("/v1/tenant", |r| {
+            RequestSpan(tenant_create_handler).handle(r)
+        })
+        .get("/v1/tenant/:tenant_id", |r| {
+            RequestSpan(tenant_status).handle(r)
+        })
+        .get("/v1/tenant/:tenant_id/synthetic_size", |r| {
+            RequestSpan(tenant_size_handler).handle(r)
+        })
+        .put("/v1/tenant/config", |r| {
+            RequestSpan(update_tenant_config_handler).handle(r)
+        })
+        .get("/v1/tenant/:tenant_id/config", |r| {
+            RequestSpan(get_tenant_config_handler).handle(r)
+        })
+        .get("/v1/tenant/:tenant_id/timeline", |r| {
+            RequestSpan(timeline_list_handler).handle(r)
+        })
+        .post("/v1/tenant/:tenant_id/timeline", |r| {
+            RequestSpan(timeline_create_handler).handle(r)
+        })
+        .post("/v1/tenant/:tenant_id/attach", |r| {
+            RequestSpan(tenant_attach_handler).handle(r)
+        })
+        .post("/v1/tenant/:tenant_id/detach", |r| {
+            RequestSpan(tenant_detach_handler).handle(r)
+        })
+        .post("/v1/tenant/:tenant_id/load", |r| {
+            RequestSpan(tenant_load_handler).handle(r)
+        })
+        .post("/v1/tenant/:tenant_id/ignore", |r| {
+            RequestSpan(tenant_ignore_handler).handle(r)
+        })
+        .get("/v1/tenant/:tenant_id/timeline/:timeline_id", |r| {
+            RequestSpan(timeline_detail_handler).handle(r)
+        })
         .get(
             "/v1/tenant/:tenant_id/timeline/:timeline_id/get_lsn_by_timestamp",
-            get_lsn_by_timestamp_handler,
-        )
-        .put(
-            "/v1/tenant/:tenant_id/timeline/:timeline_id/do_gc",
-            timeline_gc_handler,
+            |r| RequestSpan(get_lsn_by_timestamp_handler).handle(r),
         )
+        .put("/v1/tenant/:tenant_id/timeline/:timeline_id/do_gc", |r| {
+            RequestSpan(timeline_gc_handler).handle(r)
+        })
         .put(
             "/v1/tenant/:tenant_id/timeline/:timeline_id/compact",
             testing_api!("run timeline compaction", timeline_compact_handler),
@@ -1135,28 +1143,26 @@ pub fn make_router(
         )
         .post(
             "/v1/tenant/:tenant_id/timeline/:timeline_id/download_remote_layers",
-            timeline_download_remote_layers_handler_post,
+            |r| RequestSpan(timeline_download_remote_layers_handler_post).handle(r),
         )
         .get(
             "/v1/tenant/:tenant_id/timeline/:timeline_id/download_remote_layers",
-            timeline_download_remote_layers_handler_get,
-        )
-        .delete(
-            "/v1/tenant/:tenant_id/timeline/:timeline_id",
-            timeline_delete_handler,
-        )
-        .get(
-            "/v1/tenant/:tenant_id/timeline/:timeline_id/layer",
-            layer_map_info_handler,
+            |r| RequestSpan(timeline_download_remote_layers_handler_get).handle(r),
         )
+        .delete("/v1/tenant/:tenant_id/timeline/:timeline_id", |r| {
+            RequestSpan(timeline_delete_handler).handle(r)
+        })
+        .get("/v1/tenant/:tenant_id/timeline/:timeline_id/layer", |r| {
+            RequestSpan(layer_map_info_handler).handle(r)
+        })
         .get(
             "/v1/tenant/:tenant_id/timeline/:timeline_id/layer/:layer_file_name",
-            layer_download_handler,
+            |r| RequestSpan(layer_download_handler).handle(r),
         )
         .delete(
             "/v1/tenant/:tenant_id/timeline/:timeline_id/layer/:layer_file_name",
-            evict_timeline_layer_handler,
+            |r| RequestSpan(evict_timeline_layer_handler).handle(r),
         )
-        .get("/v1/panic", always_panic_handler)
+        .get("/v1/panic", |r| RequestSpan(always_panic_handler).handle(r))
         .any(handler_404))
 }

pageserver/src/metrics.rs

@@ -123,6 +123,22 @@ static REMOTE_PHYSICAL_SIZE: Lazy<UIntGaugeVec> = Lazy::new(|| {
     .expect("failed to define a metric")
 });
 
+pub static REMOTE_ONDEMAND_DOWNLOADED_LAYERS: Lazy<IntCounter> = Lazy::new(|| {
+    register_int_counter!(
+        "pageserver_remote_ondemand_downloaded_layers_total",
+        "Total on-demand downloaded layers"
+    )
+    .unwrap()
+});
+
+pub static REMOTE_ONDEMAND_DOWNLOADED_BYTES: Lazy<IntCounter> = Lazy::new(|| {
+    register_int_counter!(
+        "pageserver_remote_ondemand_downloaded_bytes_total",
+        "Total bytes of layers on-demand downloaded",
+    )
+    .unwrap()
+});
+
 static CURRENT_LOGICAL_SIZE: Lazy<UIntGaugeVec> = Lazy::new(|| {
     register_uint_gauge_vec!(
         "pageserver_current_logical_size",

pageserver/src/page_service.rs

@@ -12,7 +12,7 @@
 use anyhow::Context;
 use bytes::Buf;
 use bytes::Bytes;
-use futures::{Stream, StreamExt};
+use futures::Stream;
 use pageserver_api::models::TenantState;
 use pageserver_api::models::{
     PagestreamBeMessage, PagestreamDbSizeRequest, PagestreamDbSizeResponse,
@@ -20,6 +20,7 @@ use pageserver_api::models::{
     PagestreamFeMessage, PagestreamGetPageRequest, PagestreamGetPageResponse,
     PagestreamNblocksRequest, PagestreamNblocksResponse,
 };
+use postgres_backend::PostgresBackendTCP;
 use postgres_backend::{self, is_expected_io_error, AuthType, PostgresBackend, QueryError};
 use pq_proto::framed::ConnectionError;
 use pq_proto::FeStartupPacket;
@@ -30,6 +31,7 @@ use std::str;
 use std::str::FromStr;
 use std::sync::Arc;
 use std::time::Duration;
+use tokio_util::io::StreamReader;
 use tracing::*;
 use utils::id::ConnectionId;
 use utils::{
@@ -54,7 +56,7 @@ use crate::trace::Tracer;
 use postgres_ffi::pg_constants::DEFAULTTABLESPACE_OID;
 use postgres_ffi::BLCKSZ;
 
-fn copyin_stream(pgb: &mut PostgresBackend) -> impl Stream<Item = io::Result<Bytes>> + '_ {
+fn copyin_stream(pgb: &mut PostgresBackendTCP) -> impl Stream<Item = io::Result<Bytes>> + '_ {
     async_stream::try_stream! {
         loop {
             let msg = tokio::select! {
@@ -114,6 +116,49 @@ fn copyin_stream(pgb: &mut PostgresBackend) -> impl Stream<Item = io::Result<Byt
     }
 }
 
+/// Read the end of a tar archive.
+///
+/// A tar archive normally ends with two consecutive blocks of zeros, 512 bytes each.
+/// `tokio_tar` already read the first such block. Read the second all-zeros block,
+/// and check that there is no more data after the EOF marker.
+///
+/// XXX: Currently, any trailing data after the EOF marker prints a warning.
+/// Perhaps it should be a hard error?
+async fn read_tar_eof(mut reader: (impl tokio::io::AsyncRead + Unpin)) -> anyhow::Result<()> {
+    use tokio::io::AsyncReadExt;
+    let mut buf = [0u8; 512];
+
+    // Read the all-zeros block, and verify it
+    let mut total_bytes = 0;
+    while total_bytes < 512 {
+        let nbytes = reader.read(&mut buf[total_bytes..]).await?;
+        total_bytes += nbytes;
+        if nbytes == 0 {
+            break;
+        }
+    }
+    if total_bytes < 512 {
+        anyhow::bail!("incomplete or invalid tar EOF marker");
+    }
+    if !buf.iter().all(|&x| x == 0) {
+        anyhow::bail!("invalid tar EOF marker");
+    }
+
+    // Drain any data after the EOF marker
+    let mut trailing_bytes = 0;
+    loop {
+        let nbytes = reader.read(&mut buf).await?;
+        trailing_bytes += nbytes;
+        if nbytes == 0 {
+            break;
+        }
+    }
+    if trailing_bytes > 0 {
+        warn!("ignored {trailing_bytes} unexpected bytes after the tar archive");
+    }
+    Ok(())
+}
+
 ///////////////////////////////////////////////////////////////////////////////
 
 ///
@@ -288,7 +333,7 @@ impl PageServerHandler {
     #[instrument(skip(self, pgb, ctx))]
     async fn handle_pagerequests(
         &self,
-        pgb: &mut PostgresBackend,
+        pgb: &mut PostgresBackendTCP,
         tenant_id: TenantId,
         timeline_id: TimelineId,
         ctx: RequestContext,
@@ -392,7 +437,7 @@ impl PageServerHandler {
     #[instrument(skip(self, pgb, ctx))]
     async fn handle_import_basebackup(
         &self,
-        pgb: &mut PostgresBackend,
+        pgb: &mut PostgresBackendTCP,
         tenant_id: TenantId,
         timeline_id: TimelineId,
         base_lsn: Lsn,
@@ -421,19 +466,14 @@ impl PageServerHandler {
         pgb.write_message_noflush(&BeMessage::CopyInResponse)?;
         pgb.flush().await?;
 
-        let mut copyin_stream = Box::pin(copyin_stream(pgb));
+        let copyin_reader = StreamReader::new(copyin_stream(pgb));
+        tokio::pin!(copyin_reader);
         timeline
-            .import_basebackup_from_tar(&mut copyin_stream, base_lsn, &ctx)
+            .import_basebackup_from_tar(&mut copyin_reader, base_lsn, &ctx)
             .await?;
 
-        // Drain the rest of the Copy data
-        let mut bytes_after_tar = 0;
-        while let Some(bytes) = copyin_stream.next().await {
-            bytes_after_tar += bytes?.len();
-        }
-        if bytes_after_tar > 0 {
-            warn!("ignored {bytes_after_tar} unexpected bytes after the tar archive");
-        }
+        // Read the end of the tar archive.
+        read_tar_eof(copyin_reader).await?;
 
         // TODO check checksum
         // Meanwhile you can verify client-side by taking fullbackup
@@ -448,7 +488,7 @@ impl PageServerHandler {
     #[instrument(skip(self, pgb, ctx))]
     async fn handle_import_wal(
         &self,
-        pgb: &mut PostgresBackend,
+        pgb: &mut PostgresBackendTCP,
         tenant_id: TenantId,
         timeline_id: TimelineId,
         start_lsn: Lsn,
@@ -472,19 +512,13 @@ impl PageServerHandler {
         info!("importing wal");
         pgb.write_message_noflush(&BeMessage::CopyInResponse)?;
         pgb.flush().await?;
-        let mut copyin_stream = Box::pin(copyin_stream(pgb));
-        let mut reader = tokio_util::io::StreamReader::new(&mut copyin_stream);
-        import_wal_from_tar(&timeline, &mut reader, start_lsn, end_lsn, &ctx).await?;
+        let copyin_reader = StreamReader::new(copyin_stream(pgb));
+        tokio::pin!(copyin_reader);
+        import_wal_from_tar(&timeline, &mut copyin_reader, start_lsn, end_lsn, &ctx).await?;
         info!("wal import complete");
 
-        // Drain the rest of the Copy data
-        let mut bytes_after_tar = 0;
-        while let Some(bytes) = copyin_stream.next().await {
-            bytes_after_tar += bytes?.len();
-        }
-        if bytes_after_tar > 0 {
-            warn!("ignored {bytes_after_tar} unexpected bytes after the tar archive");
-        }
+        // Read the end of the tar archive.
+        read_tar_eof(copyin_reader).await?;
 
         // TODO Does it make sense to overshoot?
         if timeline.get_last_record_lsn() < end_lsn {
@@ -659,7 +693,7 @@ impl PageServerHandler {
     #[instrument(skip(self, pgb, ctx))]
     async fn handle_basebackup_request(
         &mut self,
-        pgb: &mut PostgresBackend,
+        pgb: &mut PostgresBackendTCP,
         tenant_id: TenantId,
         timeline_id: TimelineId,
         lsn: Option<Lsn>,
@@ -723,10 +757,10 @@ impl PageServerHandler {
 }
 
 #[async_trait::async_trait]
-impl postgres_backend::Handler for PageServerHandler {
+impl postgres_backend::Handler<tokio::net::TcpStream> for PageServerHandler {
     fn check_auth_jwt(
         &mut self,
-        _pgb: &mut PostgresBackend,
+        _pgb: &mut PostgresBackendTCP,
         jwt_response: &[u8],
     ) -> Result<(), QueryError> {
         // this unwrap is never triggered, because check_auth_jwt only called when auth_type is NeonJWT
@@ -754,7 +788,7 @@ impl postgres_backend::Handler for PageServerHandler {
 
     fn startup(
         &mut self,
-        _pgb: &mut PostgresBackend,
+        _pgb: &mut PostgresBackendTCP,
         _sm: &FeStartupPacket,
     ) -> Result<(), QueryError> {
         Ok(())
@@ -762,7 +796,7 @@ impl postgres_backend::Handler for PageServerHandler {
 
     async fn process_query(
         &mut self,
-        pgb: &mut PostgresBackend,
+        pgb: &mut PostgresBackendTCP,
         query_string: &str,
     ) -> Result<(), QueryError> {
         let ctx = self.connection_ctx.attached_child();
@@ -1073,7 +1107,10 @@ async fn get_active_tenant_with_timeout(
     tenant_id: TenantId,
     _ctx: &RequestContext, /* require get a context to support cancellation in the future */
 ) -> Result<Arc<Tenant>, GetActiveTenantError> {
-    let tenant = mgr::get_tenant(tenant_id, false).await?;
+    let tenant = match mgr::get_tenant(tenant_id, false).await {
+        Ok(tenant) => tenant,
+        Err(e) => return Err(GetActiveTenantError::Other(e.into())),
+    };
     let wait_time = Duration::from_secs(30);
     match tokio::time::timeout(wait_time, tenant.wait_to_become_active()).await {
         Ok(Ok(())) => Ok(tenant),

pageserver/src/tenant.rs

@@ -12,9 +12,7 @@
 //!
 
 use anyhow::{bail, Context};
-use bytes::Bytes;
 use futures::FutureExt;
-use futures::Stream;
 use pageserver_api::models::TimelineState;
 use remote_storage::DownloadError;
 use remote_storage::GenericRemoteStorage;
@@ -239,14 +237,13 @@ impl UninitializedTimeline<'_> {
     /// Prepares timeline data by loading it from the basebackup archive.
     pub async fn import_basebackup_from_tar(
         self,
-        copyin_stream: &mut (impl Stream<Item = io::Result<Bytes>> + Sync + Send + Unpin),
+        copyin_read: &mut (impl tokio::io::AsyncRead + Send + Sync + Unpin),
         base_lsn: Lsn,
         ctx: &RequestContext,
     ) -> anyhow::Result<Arc<Timeline>> {
         let raw_timeline = self.raw_timeline()?;
 
-        let mut reader = tokio_util::io::StreamReader::new(copyin_stream);
-        import_datadir::import_basebackup_from_tar(raw_timeline, &mut reader, base_lsn, ctx)
+        import_datadir::import_basebackup_from_tar(raw_timeline, copyin_read, base_lsn, ctx)
             .await
             .context("Failed to import basebackup")?;
 
@@ -1243,11 +1240,8 @@ impl Tenant {
             "Cannot run GC iteration on inactive tenant"
         );
 
-        let gc_result = self
-            .gc_iteration_internal(target_timeline_id, horizon, pitr, ctx)
-            .await;
-
-        gc_result
+        self.gc_iteration_internal(target_timeline_id, horizon, pitr, ctx)
+            .await
     }
 
     /// Perform one compaction iteration.
@@ -3176,6 +3170,44 @@ mod tests {
     }
      */
 
+    #[tokio::test]
+    async fn test_get_branchpoints_from_an_inactive_timeline() -> anyhow::Result<()> {
+        let (tenant, ctx) =
+            TenantHarness::create("test_get_branchpoints_from_an_inactive_timeline")?
+                .load()
+                .await;
+        let tline = tenant
+            .create_empty_timeline(TIMELINE_ID, Lsn(0), DEFAULT_PG_VERSION, &ctx)?
+            .initialize(&ctx)?;
+        make_some_layers(tline.as_ref(), Lsn(0x20)).await?;
+
+        tenant
+            .branch_timeline(&tline, NEW_TIMELINE_ID, Some(Lsn(0x40)), &ctx)
+            .await?;
+        let newtline = tenant
+            .get_timeline(NEW_TIMELINE_ID, true)
+            .expect("Should have a local timeline");
+
+        make_some_layers(newtline.as_ref(), Lsn(0x60)).await?;
+
+        tline.set_state(TimelineState::Broken);
+
+        tenant
+            .gc_iteration(Some(TIMELINE_ID), 0x10, Duration::ZERO, &ctx)
+            .await?;
+
+        assert_eq!(
+            newtline.get(*TEST_KEY, Lsn(0x50), &ctx).await?,
+            TEST_IMG(&format!("foo at {}", Lsn(0x40)))
+        );
+
+        let branchpoints = &tline.gc_info.read().unwrap().retain_lsns;
+        assert_eq!(branchpoints.len(), 1);
+        assert_eq!(branchpoints[0], Lsn(0x40));
+
+        Ok(())
+    }
+
     #[tokio::test]
     async fn test_retain_data_in_parent_which_is_needed_for_child() -> anyhow::Result<()> {
         let (tenant, ctx) =

pageserver/src/tenant/block_io.rs

@@ -51,9 +51,6 @@ where
 ///
 /// A "cursor" for efficiently reading multiple pages from a BlockReader
 ///
-/// A cursor caches the last accessed page, allowing for faster access if the
-/// same block is accessed repeatedly.
-///
 /// You can access the last page with `*cursor`. 'read_blk' returns 'self', so
 /// that in many cases you can use a BlockCursor as a drop-in replacement for
 /// the underlying BlockReader. For example:
@@ -73,8 +70,6 @@ where
     R: BlockReader,
 {
     reader: R,
-    /// last accessed page
-    cache: Option<(u32, R::BlockLease)>,
 }
 
 impl<R> BlockCursor<R>
@@ -82,40 +77,13 @@ where
     R: BlockReader,
 {
     pub fn new(reader: R) -> Self {
-        BlockCursor {
-            reader,
-            cache: None,
-        }
+        BlockCursor { reader }
     }
 
-    pub fn read_blk(&mut self, blknum: u32) -> Result<&Self, std::io::Error> {
-        // Fast return if this is the same block as before
-        if let Some((cached_blk, _buf)) = &self.cache {
-            if *cached_blk == blknum {
-                return Ok(self);
-            }
-        }
-
-        // Read the block from the underlying reader, and cache it
-        self.cache = None;
-        let buf = self.reader.read_blk(blknum)?;
-        self.cache = Some((blknum, buf));
-
-        Ok(self)
+    pub fn read_blk(&mut self, blknum: u32) -> Result<R::BlockLease, std::io::Error> {
+        self.reader.read_blk(blknum)
     }
 }
-
-impl<R> Deref for BlockCursor<R>
-where
-    R: BlockReader,
-{
-    type Target = [u8; PAGE_SZ];
-
-    fn deref(&self) -> &<Self as Deref>::Target {
-        &self.cache.as_ref().unwrap().1
-    }
-}
-
 static NEXT_ID: AtomicU64 = AtomicU64::new(1);
 
 /// An adapter for reading a (virtual) file using the page cache.

pageserver/src/tenant/config.rs

@@ -103,6 +103,7 @@ pub struct TenantConfOpt {
     pub checkpoint_distance: Option<u64>,
 
     #[serde(skip_serializing_if = "Option::is_none")]
+    #[serde(with = "humantime_serde")]
     #[serde(default)]
     pub checkpoint_timeout: Option<Duration>,
 
@@ -221,48 +222,6 @@ impl TenantConfOpt {
             eviction_policy: self.eviction_policy.unwrap_or(global_conf.eviction_policy),
         }
     }
-
-    pub fn update(&mut self, other: &TenantConfOpt) {
-        if let Some(checkpoint_distance) = other.checkpoint_distance {
-            self.checkpoint_distance = Some(checkpoint_distance);
-        }
-        if let Some(checkpoint_timeout) = other.checkpoint_timeout {
-            self.checkpoint_timeout = Some(checkpoint_timeout);
-        }
-        if let Some(compaction_target_size) = other.compaction_target_size {
-            self.compaction_target_size = Some(compaction_target_size);
-        }
-        if let Some(compaction_period) = other.compaction_period {
-            self.compaction_period = Some(compaction_period);
-        }
-        if let Some(compaction_threshold) = other.compaction_threshold {
-            self.compaction_threshold = Some(compaction_threshold);
-        }
-        if let Some(gc_horizon) = other.gc_horizon {
-            self.gc_horizon = Some(gc_horizon);
-        }
-        if let Some(gc_period) = other.gc_period {
-            self.gc_period = Some(gc_period);
-        }
-        if let Some(image_creation_threshold) = other.image_creation_threshold {
-            self.image_creation_threshold = Some(image_creation_threshold);
-        }
-        if let Some(pitr_interval) = other.pitr_interval {
-            self.pitr_interval = Some(pitr_interval);
-        }
-        if let Some(walreceiver_connect_timeout) = other.walreceiver_connect_timeout {
-            self.walreceiver_connect_timeout = Some(walreceiver_connect_timeout);
-        }
-        if let Some(lagging_wal_timeout) = other.lagging_wal_timeout {
-            self.lagging_wal_timeout = Some(lagging_wal_timeout);
-        }
-        if let Some(max_lsn_wal_lag) = other.max_lsn_wal_lag {
-            self.max_lsn_wal_lag = Some(max_lsn_wal_lag);
-        }
-        if let Some(trace_read_requests) = other.trace_read_requests {
-            self.trace_read_requests = Some(trace_read_requests);
-        }
-    }
 }
 
 impl Default for TenantConf {

pageserver/src/tenant/ephemeral_file.rs

@@ -2,9 +2,7 @@
 //! used to keep in-memory layers spilled on disk.
 
 use crate::config::PageServerConf;
-use crate::page_cache;
-use crate::page_cache::PAGE_SZ;
-use crate::page_cache::{ReadBufResult, WriteBufResult};
+use crate::page_cache::{self, ReadBufResult, WriteBufResult, PAGE_SZ};
 use crate::tenant::blob_io::BlobWriter;
 use crate::tenant::block_io::BlockReader;
 use crate::virtual_file::VirtualFile;
@@ -427,7 +425,6 @@ mod tests {
             let actual = cursor.read_blob(pos)?;
             assert_eq!(actual, expected);
         }
-        drop(cursor);
 
         // Test a large blob that spans multiple pages
         let mut large_data = Vec::new();

pageserver/src/tenant/layer_map.rs

@@ -154,11 +154,7 @@ where
         expected: &Arc<L>,
         new: Arc<L>,
     ) -> anyhow::Result<Replacement<Arc<L>>> {
-        fail::fail_point!("layermap-replace-notfound", |_| Ok(
-            // this is not what happens if an L0 layer was not found a anyhow error but perhaps
-            // that should be changed. this is good enough to show a replacement failure.
-            Replacement::NotFound
-        ));
+        fail::fail_point!("layermap-replace-notfound", |_| Ok(Replacement::NotFound));
 
         self.layer_map.replace_historic_noflush(expected, new)
     }
@@ -340,12 +336,15 @@ where
 
         let l0_index = if expected_l0 {
             // find the index in case replace worked, we need to replace that as well
-            Some(
-                self.l0_delta_layers
-                    .iter()
-                    .position(|slot| Self::compare_arced_layers(slot, expected))
-                    .ok_or_else(|| anyhow::anyhow!("existing l0 delta layer was not found"))?,
-            )
+            let pos = self
+                .l0_delta_layers
+                .iter()
+                .position(|slot| Self::compare_arced_layers(slot, expected));
+
+            if pos.is_none() {
+                return Ok(Replacement::NotFound);
+            }
+            pos
         } else {
             None
         };
@@ -804,6 +803,26 @@ mod tests {
             )
         }
 
+        #[test]
+        fn replacing_missing_l0_is_notfound() {
+            // original impl had an oversight, and L0 was an anyhow::Error. anyhow::Error should
+            // however only happen for precondition failures.
+
+            let layer = "000000000000000000000000000000000000-FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF__0000000053423C21-0000000053424D69";
+            let layer = LayerFileName::from_str(layer).unwrap();
+            let layer = LayerDescriptor::from(layer);
+
+            // same skeletan construction; see scenario below
+            let not_found: Arc<dyn Layer> = Arc::new(layer.clone());
+            let new_version: Arc<dyn Layer> = Arc::new(layer);
+
+            let mut map = LayerMap::default();
+
+            let res = map.batch_update().replace_historic(&not_found, new_version);
+
+            assert!(matches!(res, Ok(Replacement::NotFound)), "{res:?}");
+        }
+
         fn l0_delta_layers_updated_scenario(layer_name: &str, expected_l0: bool) {
             let name = LayerFileName::from_str(layer_name).unwrap();
             let skeleton = LayerDescriptor::from(name);
@@ -813,7 +832,8 @@ mod tests {
 
             let mut map = LayerMap::default();
 
-            // two disjoint Arcs in different lifecycle phases.
+            // two disjoint Arcs in different lifecycle phases. even if it seems they must be the
+            // same layer, we use LayerMap::compare_arced_layers as the identity of layers.
             assert!(!LayerMap::compare_arced_layers(&remote, &downloaded));
 
             let expected_in_counts = (1, usize::from(expected_l0));

pageserver/src/tenant/mgr.rs

@@ -289,7 +289,7 @@ pub async fn set_new_tenant_config(
     conf: &'static PageServerConf,
     new_tenant_conf: TenantConfOpt,
     tenant_id: TenantId,
-) -> anyhow::Result<()> {
+) -> Result<(), TenantStateError> {
     info!("configuring tenant {tenant_id}");
     let tenant = get_tenant(tenant_id, true).await?;
 
@@ -306,16 +306,20 @@ pub async fn set_new_tenant_config(
 
 /// Gets the tenant from the in-memory data, erroring if it's absent or is not fitting to the query.
 /// `active_only = true` allows to query only tenants that are ready for operations, erroring on other kinds of tenants.
-pub async fn get_tenant(tenant_id: TenantId, active_only: bool) -> anyhow::Result<Arc<Tenant>> {
+pub async fn get_tenant(
+    tenant_id: TenantId,
+    active_only: bool,
+) -> Result<Arc<Tenant>, TenantStateError> {
     let m = TENANTS.read().await;
     let tenant = m
         .get(&tenant_id)
-        .with_context(|| format!("Tenant {tenant_id} not found in the local state"))?;
+        .ok_or(TenantStateError::NotFound(tenant_id))?;
     if active_only && !tenant.is_active() {
-        anyhow::bail!(
+        tracing::warn!(
             "Tenant {tenant_id} is not active. Current state: {:?}",
             tenant.current_state()
-        )
+        );
+        Err(TenantStateError::NotActive(tenant_id))
     } else {
         Ok(Arc::clone(tenant))
     }
@@ -325,21 +329,28 @@ pub async fn delete_timeline(
     tenant_id: TenantId,
     timeline_id: TimelineId,
     ctx: &RequestContext,
-) -> anyhow::Result<()> {
-    match get_tenant(tenant_id, true).await {
-        Ok(tenant) => {
-            tenant.delete_timeline(timeline_id, ctx).await?;
-        }
-        Err(e) => anyhow::bail!("Cannot access tenant {tenant_id} in local tenant state: {e:?}"),
-    }
-
+) -> Result<(), TenantStateError> {
+    let tenant = get_tenant(tenant_id, true).await?;
+    tenant.delete_timeline(timeline_id, ctx).await?;
     Ok(())
 }
 
+#[derive(Debug, thiserror::Error)]
+pub enum TenantStateError {
+    #[error("Tenant {0} not found")]
+    NotFound(TenantId),
+    #[error("Tenant {0} is stopping")]
+    IsStopping(TenantId),
+    #[error("Tenant {0} is not active")]
+    NotActive(TenantId),
+    #[error(transparent)]
+    Other(#[from] anyhow::Error),
+}
+
 pub async fn detach_tenant(
     conf: &'static PageServerConf,
     tenant_id: TenantId,
-) -> anyhow::Result<()> {
+) -> Result<(), TenantStateError> {
     remove_tenant_from_memory(tenant_id, async {
         let local_tenant_directory = conf.tenant_path(&tenant_id);
         fs::remove_dir_all(&local_tenant_directory)
@@ -379,7 +390,7 @@ pub async fn load_tenant(
 pub async fn ignore_tenant(
     conf: &'static PageServerConf,
     tenant_id: TenantId,
-) -> anyhow::Result<()> {
+) -> Result<(), TenantStateError> {
     remove_tenant_from_memory(tenant_id, async {
         let ignore_mark_file = conf.tenant_ignore_mark_file_path(tenant_id);
         fs::File::create(&ignore_mark_file)
@@ -489,7 +500,7 @@ where
 async fn remove_tenant_from_memory<V, F>(
     tenant_id: TenantId,
     tenant_cleanup: F,
-) -> anyhow::Result<V>
+) -> Result<V, TenantStateError>
 where
     F: std::future::Future<Output = anyhow::Result<V>>,
 {
@@ -505,11 +516,9 @@ where
                 | TenantState::Loading
                 | TenantState::Broken
                 | TenantState::Active => tenant.set_stopping(),
-                TenantState::Stopping => {
-                    anyhow::bail!("Tenant {tenant_id} is stopping already")
-                }
+                TenantState::Stopping => return Err(TenantStateError::IsStopping(tenant_id)),
             },
-            None => anyhow::bail!("Tenant not found for id {tenant_id}"),
+            None => return Err(TenantStateError::NotFound(tenant_id)),
         }
     }
 
@@ -532,10 +541,15 @@ where
         Err(e) => {
             let tenants_accessor = TENANTS.read().await;
             match tenants_accessor.get(&tenant_id) {
-                Some(tenant) => tenant.set_broken(&e.to_string()),
-                None => warn!("Tenant {tenant_id} got removed from memory"),
+                Some(tenant) => {
+                    tenant.set_broken(&e.to_string());
+                }
+                None => {
+                    warn!("Tenant {tenant_id} got removed from memory");
+                    return Err(TenantStateError::NotFound(tenant_id));
+                }
             }
-            Err(e)
+            Err(TenantStateError::Other(e))
         }
     }
 }
@@ -555,7 +569,7 @@ pub async fn immediate_gc(
     let tenant = guard
         .get(&tenant_id)
         .map(Arc::clone)
-        .with_context(|| format!("Tenant {tenant_id} not found"))
+        .with_context(|| format!("tenant {tenant_id}"))
         .map_err(ApiError::NotFound)?;
 
     let gc_horizon = gc_req.gc_horizon.unwrap_or_else(|| tenant.get_gc_horizon());
@@ -605,7 +619,7 @@ pub async fn immediate_compact(
     let tenant = guard
         .get(&tenant_id)
         .map(Arc::clone)
-        .with_context(|| format!("Tenant {tenant_id} not found"))
+        .with_context(|| format!("tenant {tenant_id}"))
         .map_err(ApiError::NotFound)?;
 
     let timeline = tenant

pageserver/src/tenant/remote_timeline_client.rs

@@ -218,9 +218,10 @@ use tracing::{debug, info, warn};
 use tracing::{info_span, Instrument};
 use utils::lsn::Lsn;
 
-use crate::metrics::RemoteOpFileKind;
-use crate::metrics::RemoteOpKind;
-use crate::metrics::{MeasureRemoteOp, RemoteTimelineClientMetrics};
+use crate::metrics::{
+    MeasureRemoteOp, RemoteOpFileKind, RemoteOpKind, RemoteTimelineClientMetrics,
+    REMOTE_ONDEMAND_DOWNLOADED_BYTES, REMOTE_ONDEMAND_DOWNLOADED_LAYERS,
+};
 use crate::tenant::remote_timeline_client::index::LayerFileMetadata;
 use crate::{
     config::PageServerConf,
@@ -446,6 +447,10 @@ impl RemoteTimelineClient {
                 );
             }
         }
+
+        REMOTE_ONDEMAND_DOWNLOADED_LAYERS.inc();
+        REMOTE_ONDEMAND_DOWNLOADED_BYTES.inc_by(downloaded_size);
+
         Ok(downloaded_size)
     }
 

pageserver/src/tenant/remote_timeline_client/download.rs

@@ -6,11 +6,13 @@
 use std::collections::HashSet;
 use std::future::Future;
 use std::path::Path;
+use std::time::Duration;
 
 use anyhow::{anyhow, Context};
 use tokio::fs;
 use tokio::io::AsyncWriteExt;
-use tracing::{error, info, warn};
+
+use tracing::{info, warn};
 
 use crate::config::PageServerConf;
 use crate::tenant::storage_layer::LayerFileName;
@@ -26,6 +28,8 @@ async fn fsync_path(path: impl AsRef<std::path::Path>) -> Result<(), std::io::Er
     fs::File::open(path).await?.sync_all().await
 }
 
+static MAX_DOWNLOAD_DURATION: Duration = Duration::from_secs(120);
+
 ///
 /// If 'metadata' is given, we will validate that the downloaded file's size matches that
 /// in the metadata. (In the future, we might do more cross-checks, like CRC validation)
@@ -64,22 +68,28 @@ pub async fn download_layer_file<'a>(
             // TODO: this doesn't use the cached fd for some reason?
             let mut destination_file = fs::File::create(&temp_file_path).await.with_context(|| {
                 format!(
-                    "Failed to create a destination file for layer '{}'",
+                    "create a destination file for layer '{}'",
                     temp_file_path.display()
                 )
             })
             .map_err(DownloadError::Other)?;
             let mut download = storage.download(&remote_path).await.with_context(|| {
                 format!(
-                    "Failed to open a download stream for layer with remote storage path '{remote_path:?}'"
+                    "open a download stream for layer with remote storage path '{remote_path:?}'"
                 )
             })
             .map_err(DownloadError::Other)?;
-            let bytes_amount = tokio::io::copy(&mut download.download_stream, &mut destination_file).await.with_context(|| {
-                format!("Failed to download layer with remote storage path '{remote_path:?}' into file {temp_file_path:?}")
-            })
-            .map_err(DownloadError::Other)?;
+
+            let bytes_amount = tokio::time::timeout(MAX_DOWNLOAD_DURATION, tokio::io::copy(&mut download.download_stream, &mut destination_file))
+                .await
+                .map_err(|e| DownloadError::Other(anyhow::anyhow!("Timed out  {:?}", e)))?
+                .with_context(|| {
+                    format!("Failed to download layer with remote storage path '{remote_path:?}' into file {temp_file_path:?}")
+                })
+                .map_err(DownloadError::Other)?;
+
             Ok((destination_file, bytes_amount))
+
         },
         &format!("download {remote_path:?}"),
     ).await?;
@@ -300,7 +310,7 @@ where
             }
             Err(DownloadError::Other(ref err)) => {
                 // Operation failed FAILED_DOWNLOAD_RETRIES times. Time to give up.
-                error!("{description} still failed after {attempts} retries, giving up: {err:?}");
+                warn!("{description} still failed after {attempts} retries, giving up: {err:?}");
                 return result;
             }
         }

pageserver/src/tenant/storage_layer.rs

@@ -364,7 +364,7 @@ pub trait PersistentLayer: Layer {
     }
 
     /// Permanently remove this layer from disk.
-    fn delete(&self) -> Result<()>;
+    fn delete_resident_layer_file(&self) -> Result<()>;
 
     fn downcast_remote_layer(self: Arc<Self>) -> Option<std::sync::Arc<RemoteLayer>> {
         None

pageserver/src/tenant/storage_layer/delta_layer.rs

@@ -438,7 +438,7 @@ impl PersistentLayer for DeltaLayer {
         ))
     }
 
-    fn delete(&self) -> Result<()> {
+    fn delete_resident_layer_file(&self) -> Result<()> {
         // delete underlying file
         fs::remove_file(self.path())?;
         Ok(())

pageserver/src/tenant/storage_layer/image_layer.rs

@@ -252,7 +252,7 @@ impl PersistentLayer for ImageLayer {
         unimplemented!();
     }
 
-    fn delete(&self) -> Result<()> {
+    fn delete_resident_layer_file(&self) -> Result<()> {
         // delete underlying file
         fs::remove_file(self.path())?;
         Ok(())

pageserver/src/tenant/storage_layer/remote_layer.rs

@@ -155,8 +155,8 @@ impl PersistentLayer for RemoteLayer {
         bail!("cannot iterate a remote layer");
     }
 
-    fn delete(&self) -> Result<()> {
-        Ok(())
+    fn delete_resident_layer_file(&self) -> Result<()> {
+        bail!("remote layer has no layer file");
     }
 
     fn downcast_remote_layer<'a>(self: Arc<Self>) -> Option<std::sync::Arc<RemoteLayer>> {

pageserver/src/tenant/timeline.rs

@@ -662,8 +662,8 @@ impl Timeline {
             // update the index file on next flush iteration too. But it
             // could take a while until that happens.
             //
-            // Additionally, only do this on the terminal round before sleeping.
-            if last_round {
+            // Additionally, only do this once before we return from this function.
+            if last_round || res.is_ok() {
                 if let Some(remote_client) = &self.remote_client {
                     remote_client.schedule_index_upload_for_file_changes()?;
                 }
@@ -1047,11 +1047,12 @@ impl Timeline {
             return Ok(false);
         }
 
-        let layer_metadata = LayerFileMetadata::new(
-            local_layer
-                .file_size()
-                .expect("Local layer should have a file size"),
-        );
+        let layer_file_size = local_layer
+            .file_size()
+            .expect("Local layer should have a file size");
+
+        let layer_metadata = LayerFileMetadata::new(layer_file_size);
+
         let new_remote_layer = Arc::new(match local_layer.filename() {
             LayerFileName::Image(image_name) => RemoteLayer::new_img(
                 self.tenant_id,
@@ -1075,15 +1076,22 @@ impl Timeline {
 
         let replaced = match batch_updates.replace_historic(local_layer, new_remote_layer)? {
             Replacement::Replaced { .. } => {
-                let layer_size = local_layer.file_size();
-
-                if let Err(e) = local_layer.delete() {
+                if let Err(e) = local_layer.delete_resident_layer_file() {
                     error!("failed to remove layer file on evict after replacement: {e:#?}");
                 }
-
-                if let Some(layer_size) = layer_size {
-                    self.metrics.resident_physical_size_gauge.sub(layer_size);
-                }
+                // Always decrement the physical size gauge, even if we failed to delete the file.
+                // Rationale: we already replaced the layer with a remote layer in the layer map,
+                // and any subsequent download_remote_layer will
+                // 1. overwrite the file on disk and
+                // 2. add the downloaded size to the resident size gauge.
+                //
+                // If there is no re-download, and we restart the pageserver, then load_layer_map
+                // will treat the file as a local layer again, count it towards resident size,
+                // and it'll be like the layer removal never happened.
+                // The bump in resident size is perhaps unexpected but overall a robust behavior.
+                self.metrics
+                    .resident_physical_size_gauge
+                    .sub(layer_file_size);
 
                 true
             }
@@ -1238,6 +1246,11 @@ impl Timeline {
 
                 state,
             };
+            info!(
+                "initialized lsrlsn to {}/{}",
+                disk_consistent_lsn,
+                metadata.prev_record_lsn().unwrap_or(Lsn(0))
+            );
             result.repartition_threshold = result.get_checkpoint_distance() / 10;
             result
                 .metrics
@@ -1327,6 +1340,7 @@ impl Timeline {
             lagging_wal_timeout,
             max_lsn_wal_lag,
             crate::config::SAFEKEEPER_AUTH_TOKEN.get().cloned(),
+            self.conf.availability_zone.clone(),
             background_ctx,
         );
     }
@@ -1942,11 +1956,14 @@ impl Timeline {
         layer: Arc<dyn PersistentLayer>,
         updates: &mut BatchedUpdates<'_, dyn PersistentLayer>,
     ) -> anyhow::Result<()> {
-        let layer_size = layer.file_size();
-
-        layer.delete()?;
-        if let Some(layer_size) = layer_size {
-            self.metrics.resident_physical_size_gauge.sub(layer_size);
+        if !layer.is_remote_layer() {
+            layer.delete_resident_layer_file()?;
+            let layer_file_size = layer
+                .file_size()
+                .expect("Local layer should have a file size");
+            self.metrics
+                .resident_physical_size_gauge
+                .sub(layer_file_size);
         }
 
         // TODO Removing from the bottom of the layer map is expensive.
@@ -2704,10 +2721,22 @@ impl Timeline {
     ) -> Result<HashMap<LayerFileName, LayerFileMetadata>, PageReconstructError> {
         let timer = self.metrics.create_images_time_histo.start_timer();
         let mut image_layers: Vec<ImageLayer> = Vec::new();
+
+        // We need to avoid holes between generated image layers.
+        // Otherwise LayerMap::image_layer_exists will return false if key range of some layer is covered by more than one
+        // image layer with hole between them. In this case such layer can not be utilized by GC.
+        //
+        // How such hole between partitions can appear?
+        // if we have relation with relid=1 and size 100 and relation with relid=2 with size 200 then result of
+        // KeySpace::partition may contain partitions <100000000..100000099> and <200000000..200000199>.
+        // If there is delta layer <100000000..300000000> then it never be garbage collected because
+        // image layers  <100000000..100000099> and <200000000..200000199> are not completely covering it.
+        let mut start = Key::MIN;
+
         for partition in partitioning.parts.iter() {
+            let img_range = start..partition.ranges.last().unwrap().end;
+            start = img_range.end;
             if force || self.time_for_new_image_layer(partition, lsn)? {
-                let img_range =
-                    partition.ranges.first().unwrap().start..partition.ranges.last().unwrap().end;
                 let mut image_layer_writer = ImageLayerWriter::new(
                     self.conf,
                     self.timeline_id,
@@ -2721,7 +2750,6 @@ impl Timeline {
                         "failpoint image-layer-writer-fail-before-finish"
                     )))
                 });
-
                 for range in &partition.ranges {
                     let mut key = range.start;
                     while key < range.end {
@@ -3136,9 +3164,7 @@ impl Timeline {
             }
 
             fail_point!("delta-layer-writer-fail-before-finish", |_| {
-                return Err(
-                    anyhow::anyhow!("failpoint delta-layer-writer-fail-before-finish").into(),
-                );
+                Err(anyhow::anyhow!("failpoint delta-layer-writer-fail-before-finish").into())
             });
 
             writer.as_mut().unwrap().put_value(key, lsn, value)?;
@@ -3808,7 +3834,7 @@ impl Timeline {
                     remote_layer.ongoing_download.close();
                 } else {
                     // Keep semaphore open. We'll drop the permit at the end of the function.
-                    info!("on-demand download failed: {:?}", result.as_ref().unwrap_err());
+                    error!("on-demand download failed: {:?}", result.as_ref().unwrap_err());
                 }
 
                 // Don't treat it as an error if the task that triggered the download

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

@@ -45,6 +45,7 @@ pub fn spawn_connection_manager_task(
     lagging_wal_timeout: Duration,
     max_lsn_wal_lag: NonZeroU64,
     auth_token: Option<Arc<String>>,
+    availability_zone: Option<String>,
     ctx: RequestContext,
 ) {
     let mut broker_client = get_broker_client().clone();
@@ -67,6 +68,7 @@ pub fn spawn_connection_manager_task(
                 lagging_wal_timeout,
                 max_lsn_wal_lag,
                 auth_token,
+                availability_zone,
             );
             loop {
                 select! {
@@ -334,6 +336,7 @@ struct WalreceiverState {
     /// Data about all timelines, available for connection, fetched from storage broker, grouped by their corresponding safekeeper node id.
     wal_stream_candidates: HashMap<NodeId, BrokerSkTimeline>,
     auth_token: Option<Arc<String>>,
+    availability_zone: Option<String>,
 }
 
 /// Current connection data.
@@ -381,6 +384,7 @@ impl WalreceiverState {
         lagging_wal_timeout: Duration,
         max_lsn_wal_lag: NonZeroU64,
         auth_token: Option<Arc<String>>,
+        availability_zone: Option<String>,
     ) -> Self {
         let id = TenantTimelineId {
             tenant_id: timeline.tenant_id,
@@ -396,6 +400,7 @@ impl WalreceiverState {
             wal_stream_candidates: HashMap::new(),
             wal_connection_retries: HashMap::new(),
             auth_token,
+            availability_zone,
         }
     }
 
@@ -740,6 +745,7 @@ impl WalreceiverState {
                         None => None,
                         Some(x) => Some(x),
                     },
+                    self.availability_zone.as_deref(),
                 ) {
                     Ok(connstr) => Some((*sk_id, info, connstr)),
                     Err(e) => {
@@ -824,17 +830,24 @@ fn wal_stream_connection_config(
     }: TenantTimelineId,
     listen_pg_addr_str: &str,
     auth_token: Option<&str>,
+    availability_zone: Option<&str>,
 ) -> anyhow::Result<PgConnectionConfig> {
     let (host, port) =
         parse_host_port(listen_pg_addr_str).context("Unable to parse listen_pg_addr_str")?;
     let port = port.unwrap_or(5432);
-    Ok(PgConnectionConfig::new_host_port(host, port)
+    let mut connstr = PgConnectionConfig::new_host_port(host, port)
         .extend_options([
             "-c".to_owned(),
             format!("timeline_id={}", timeline_id),
             format!("tenant_id={}", tenant_id),
         ])
-        .set_password(auth_token.map(|s| s.to_owned())))
+        .set_password(auth_token.map(|s| s.to_owned()));
+
+    if let Some(availability_zone) = availability_zone {
+        connstr = connstr.extend_options([format!("availability_zone={}", availability_zone)]);
+    }
+
+    Ok(connstr)
 }
 
 #[cfg(test)]
@@ -1273,6 +1286,7 @@ mod tests {
             wal_stream_candidates: HashMap::new(),
             wal_connection_retries: HashMap::new(),
             auth_token: None,
+            availability_zone: None,
         }
     }
 }

pageserver/src/walredo.rs

@@ -23,13 +23,11 @@ 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::unix::io::{AsRawFd, RawFd};
 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, MutexGuard};
@@ -256,52 +254,53 @@ impl PostgresRedoManager {
         pg_version: u32,
     ) -> Result<Bytes, WalRedoError> {
         let (rel, blknum) = key_to_rel_block(key).or(Err(WalRedoError::InvalidRecord))?;
-
+        const MAX_RETRY_ATTEMPTS: u32 = 1;
         let start_time = Instant::now();
+        let mut n_attempts = 0u32;
+        loop {
+            let mut proc = self.stdin.lock().unwrap();
+            let lock_time = Instant::now();
 
-        let mut proc = self.stdin.lock().unwrap();
-        let lock_time = Instant::now();
+            // launch the WAL redo process on first use
+            if proc.is_none() {
+                self.launch(&mut proc, pg_version)?;
+            }
+            WAL_REDO_WAIT_TIME.observe(lock_time.duration_since(start_time).as_secs_f64());
 
-        // launch the WAL redo process on first use
-        if proc.is_none() {
-            self.launch(&mut proc, pg_version)?;
-        }
-        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 = self
+                .apply_wal_records(proc, buf_tag, &base_img, records, wal_redo_timeout)
+                .map_err(WalRedoError::IoError);
 
-        // Relational WAL records are applied using wal-redo-postgres
-        let buf_tag = BufferTag { rel, blknum };
-        let result = self
-            .apply_wal_records(proc, buf_tag, base_img, records, wal_redo_timeout)
-            .map_err(WalRedoError::IoError);
+            let end_time = Instant::now();
+            let duration = end_time.duration_since(lock_time);
 
-        let end_time = Instant::now();
-        let duration = end_time.duration_since(lock_time);
+            let len = records.len();
+            let nbytes = records.iter().fold(0, |acumulator, record| {
+                acumulator
+                    + match &record.1 {
+                        NeonWalRecord::Postgres { rec, .. } => rec.len(),
+                        _ => unreachable!("Only PostgreSQL records are accepted in this batch"),
+                    }
+            });
 
-        let len = records.len();
-        let nbytes = records.iter().fold(0, |acumulator, record| {
-            acumulator
-                + match &record.1 {
-                    NeonWalRecord::Postgres { rec, .. } => rec.len(),
-                    _ => unreachable!("Only PostgreSQL records are accepted in this batch"),
-                }
-        });
+            WAL_REDO_TIME.observe(duration.as_secs_f64());
+            WAL_REDO_RECORDS_HISTOGRAM.observe(len as f64);
+            WAL_REDO_BYTES_HISTOGRAM.observe(nbytes as f64);
 
-        WAL_REDO_TIME.observe(duration.as_secs_f64());
-        WAL_REDO_RECORDS_HISTOGRAM.observe(len as f64);
-        WAL_REDO_BYTES_HISTOGRAM.observe(nbytes as f64);
+            debug!(
+				"postgres applied {} WAL records ({} bytes) in {} us to reconstruct page image at LSN {}",
+				len,
+				nbytes,
+				duration.as_micros(),
+				lsn
+			);
 
-        debug!(
-            "postgres applied {} WAL records ({} bytes) in {} us to reconstruct page image at LSN {}",
-            len,
-            nbytes,
-            duration.as_micros(),
-            lsn
-        );
-
-        // If something went wrong, don't try to reuse the process. Kill it, and
-        // next request will launch a new one.
-        if result.is_err() {
-            error!(
+            // If something went wrong, don't try to reuse the process. Kill it, and
+            // next request will launch a new one.
+            if result.is_err() {
+                error!(
                 "error applying {} WAL records {}..{} ({} bytes) to base image with LSN {} to reconstruct page image at LSN {}",
                 records.len(),
 				records.first().map(|p| p.0).unwrap_or(Lsn(0)),
@@ -310,24 +309,28 @@ impl PostgresRedoManager {
 				base_img_lsn,
                 lsn
             );
-            // 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();
+                // 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();
+                }
+            }
+            n_attempts += 1;
+            if n_attempts > MAX_RETRY_ATTEMPTS || result.is_ok() {
+                return result;
             }
         }
-        result
     }
 
     ///
@@ -634,26 +637,26 @@ impl PostgresRedoManager {
         input: &mut MutexGuard<Option<ProcessInput>>,
         pg_version: u32,
     ) -> 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(
+        // Previous versions of wal-redo required data directory and that directories
+        // occupied some space on disk. Remove it if we face it.
+        //
+        // This code could be dropped after one release cycle.
+        let legacy_datadir = path_with_suffix_extension(
             self.conf
                 .tenant_path(&self.tenant_id)
                 .join("wal-redo-datadir"),
             TEMP_FILE_SUFFIX,
         );
-
-        // Create empty data directory for wal-redo postgres, deleting old one first.
-        if datadir.exists() {
-            info!("old temporary datadir {datadir:?} exists, removing");
-            fs::remove_dir_all(&datadir).map_err(|e| {
+        if legacy_datadir.exists() {
+            info!("legacy wal-redo datadir {legacy_datadir:?} exists, removing");
+            fs::remove_dir_all(&legacy_datadir).map_err(|e| {
                 Error::new(
                     e.kind(),
-                    format!("Old temporary dir {datadir:?} removal failure: {e}"),
+                    format!("legacy wal-redo datadir {legacy_datadir:?} removal failure: {e}"),
                 )
             })?;
         }
+
         let pg_bin_dir_path = self
             .conf
             .pg_bin_dir(pg_version)
@@ -663,35 +666,6 @@ impl PostgresRedoManager {
             .pg_lib_dir(pg_version)
             .map_err(|e| Error::new(ErrorKind::Other, format!("incorrect pg_lib_dir path: {e}")))?;
 
-        info!("running initdb in {}", datadir.display());
-        let initdb = Command::new(pg_bin_dir_path.join("initdb"))
-            .args(["-D", &datadir.to_string_lossy()])
-            .arg("-N")
-            .env_clear()
-            .env("LD_LIBRARY_PATH", &pg_lib_dir_path)
-            .env("DYLD_LIBRARY_PATH", &pg_lib_dir_path) // macOS
-            .close_fds()
-            .output()
-            .map_err(|e| Error::new(e.kind(), format!("failed to execute initdb: {e}")))?;
-
-        if !initdb.status.success() {
-            return Err(Error::new(
-                ErrorKind::Other,
-                format!(
-                    "initdb failed\nstdout: {}\nstderr:\n{}",
-                    String::from_utf8_lossy(&initdb.stdout),
-                    String::from_utf8_lossy(&initdb.stderr)
-                ),
-            ));
-        } else {
-            // Limit shared cache for wal-redo-postgres
-            let mut config = OpenOptions::new()
-                .append(true)
-                .open(PathBuf::from(&datadir).join("postgresql.conf"))?;
-            config.write_all(b"shared_buffers=128kB\n")?;
-            config.write_all(b"fsync=off\n")?;
-        }
-
         // Start postgres itself
         let child = Command::new(pg_bin_dir_path.join("postgres"))
             .arg("--wal-redo")
@@ -701,7 +675,6 @@ impl PostgresRedoManager {
             .env_clear()
             .env("LD_LIBRARY_PATH", &pg_lib_dir_path)
             .env("DYLD_LIBRARY_PATH", &pg_lib_dir_path)
-            .env("PGDATA", &datadir)
             // The redo process is not trusted, and runs in seccomp mode that
             // doesn't allow it to open any files. We have to also make sure it
             // doesn't inherit any file descriptors from the pageserver, that
@@ -771,7 +744,7 @@ impl PostgresRedoManager {
         &self,
         mut input: MutexGuard<Option<ProcessInput>>,
         tag: BufferTag,
-        base_img: Option<Bytes>,
+        base_img: &Option<Bytes>,
         records: &[(Lsn, NeonWalRecord)],
         wal_redo_timeout: Duration,
     ) -> Result<Bytes, std::io::Error> {
@@ -787,7 +760,7 @@ impl PostgresRedoManager {
         let mut writebuf: Vec<u8> = Vec::with_capacity((BLCKSZ as usize) * 3);
         build_begin_redo_for_block_msg(tag, &mut writebuf);
         if let Some(img) = base_img {
-            build_push_page_msg(tag, &img, &mut writebuf);
+            build_push_page_msg(tag, img, &mut writebuf);
         }
         for (lsn, rec) in records.iter() {
             if let NeonWalRecord::Postgres {

pgxn/neon/libpagestore.c

@@ -32,6 +32,9 @@
 
 #define PageStoreTrace DEBUG5
 
+#define MAX_RECONNECT_ATTEMPTS 5
+#define RECONNECT_INTERVAL_USEC 1000000
+
 bool		connected = false;
 PGconn	   *pageserver_conn = NULL;
 
@@ -52,8 +55,8 @@ int			readahead_buffer_size = 128;
 
 static void pageserver_flush(void);
 
-static void
-pageserver_connect()
+static bool
+pageserver_connect(int elevel)
 {
 	char	   *query;
 	int			ret;
@@ -69,10 +72,11 @@ pageserver_connect()
 		PQfinish(pageserver_conn);
 		pageserver_conn = NULL;
 
-		ereport(ERROR,
+		ereport(elevel,
 				(errcode(ERRCODE_SQLCLIENT_UNABLE_TO_ESTABLISH_SQLCONNECTION),
 				 errmsg(NEON_TAG "could not establish connection to pageserver"),
 				 errdetail_internal("%s", msg)));
+		return false;
 	}
 
 	query = psprintf("pagestream %s %s", neon_tenant, neon_timeline);
@@ -81,7 +85,8 @@ pageserver_connect()
 	{
 		PQfinish(pageserver_conn);
 		pageserver_conn = NULL;
-		neon_log(ERROR, "could not send pagestream command to pageserver");
+		neon_log(elevel, "could not send pagestream command to pageserver");
+		return false;
 	}
 
 	pageserver_conn_wes = CreateWaitEventSet(TopMemoryContext, 3);
@@ -113,8 +118,9 @@ pageserver_connect()
 				FreeWaitEventSet(pageserver_conn_wes);
 				pageserver_conn_wes = NULL;
 
-				neon_log(ERROR, "could not complete handshake with pageserver: %s",
+				neon_log(elevel, "could not complete handshake with pageserver: %s",
 						 msg);
+				return false;
 			}
 		}
 	}
@@ -122,6 +128,7 @@ pageserver_connect()
 	neon_log(LOG, "libpagestore: connected to '%s'", page_server_connstring_raw);
 
 	connected = true;
+	return true;
 }
 
 /*
@@ -149,8 +156,11 @@ retry:
 		if (event.events & WL_SOCKET_READABLE)
 		{
 			if (!PQconsumeInput(pageserver_conn))
-				neon_log(ERROR, "could not get response from pageserver: %s",
+			{
+				neon_log(LOG, "could not get response from pageserver: %s",
 						 PQerrorMessage(pageserver_conn));
+				return -1;
+			}
 		}
 
 		goto retry;
@@ -190,31 +200,62 @@ static void
 pageserver_send(NeonRequest * request)
 {
 	StringInfoData req_buff;
+	int n_reconnect_attempts = 0;
 
 	/* If the connection was lost for some reason, reconnect */
 	if (connected && PQstatus(pageserver_conn) == CONNECTION_BAD)
 		pageserver_disconnect();
 
-	if (!connected)
-		pageserver_connect();
 
 	req_buff = nm_pack_request(request);
 
 	/*
-	 * Send request.
-	 *
-	 * In principle, this could block if the output buffer is full, and we
-	 * should use async mode and check for interrupts while waiting. In
-	 * practice, our requests are small enough to always fit in the output and
-	 * TCP buffer.
+	 * If pageserver is stopped, the connections from compute node are broken.
+	 * The compute node doesn't notice that immediately, but it will cause the next request to fail, usually on the next query.
+	 * That causes user-visible errors if pageserver is restarted, or the tenant is moved from one pageserver to another.
+	 * See https://github.com/neondatabase/neon/issues/1138
+	 * So try to reestablish connection in case of failure.
 	 */
-	if (PQputCopyData(pageserver_conn, req_buff.data, req_buff.len) <= 0)
+	while (true)
 	{
-		char	   *msg = pchomp(PQerrorMessage(pageserver_conn));
+		if (!connected)
+		{
+			if (!pageserver_connect(n_reconnect_attempts < MAX_RECONNECT_ATTEMPTS ? LOG : ERROR))
+			{
+				n_reconnect_attempts += 1;
+				pg_usleep(RECONNECT_INTERVAL_USEC);
+				continue;
+			}
+		}
 
-		pageserver_disconnect();
-		neon_log(ERROR, "failed to send page request: %s", msg);
+		/*
+		 * Send request.
+		 *
+		 * In principle, this could block if the output buffer is full, and we
+		 * should use async mode and check for interrupts while waiting. In
+		 * practice, our requests are small enough to always fit in the output and
+		 * TCP buffer.
+		 */
+		if (PQputCopyData(pageserver_conn, req_buff.data, req_buff.len) <= 0)
+		{
+			char	   *msg = pchomp(PQerrorMessage(pageserver_conn));
+			if (n_reconnect_attempts < MAX_RECONNECT_ATTEMPTS)
+			{
+				neon_log(LOG, "failed to send page request (try to reconnect): %s", msg);
+				if (n_reconnect_attempts != 0) /* do not sleep before first reconnect attempt, assuming that pageserver is already restarted */
+					pg_usleep(RECONNECT_INTERVAL_USEC);
+				n_reconnect_attempts += 1;
+				continue;
+			}
+			else
+			{
+				pageserver_disconnect();
+				neon_log(ERROR, "failed to send page request: %s", msg);
+			}
+		}
+		break;
 	}
+
 	pfree(req_buff.data);
 
 	n_unflushed_requests++;

pgxn/neon_utils/Makefile

@@ -0,0 +1,15 @@
+# pgxs/neon_utils/Makefile
+
+
+MODULE_big = neon_utils
+OBJS = \
+	$(WIN32RES) \
+	neon_utils.o
+
+EXTENSION = neon_utils
+DATA = neon_utils--1.0.sql
+PGFILEDESC = "neon_utils - small useful functions"
+
+PG_CONFIG = pg_config
+PGXS := $(shell $(PG_CONFIG) --pgxs)
+include $(PGXS)

pgxn/neon_utils/neon_utils--1.0.sql

@@ -0,0 +1,6 @@
+CREATE FUNCTION num_cpus()
+RETURNS int
+AS 'MODULE_PATHNAME', 'num_cpus'
+LANGUAGE C STRICT
+PARALLEL UNSAFE
+VOLATILE;

pgxn/neon_utils/neon_utils.c

@@ -0,0 +1,35 @@
+/*-------------------------------------------------------------------------
+ *
+ * neon_utils.c
+ *	  neon_utils - small useful functions
+ *
+ * IDENTIFICATION
+ *	 contrib/neon_utils/neon_utils.c
+ *
+ *-------------------------------------------------------------------------
+ */
+#ifdef _WIN32
+#include <windows.h>
+#else
+#include <unistd.h>
+#endif
+
+#include "postgres.h"
+#include "fmgr.h"
+
+PG_MODULE_MAGIC;
+
+PG_FUNCTION_INFO_V1(num_cpus);
+
+Datum
+num_cpus(PG_FUNCTION_ARGS)
+{
+#ifdef _WIN32
+	SYSTEM_INFO sysinfo;
+	GetSystemInfo(&sysinfo);
+	uint32 num_cpus = (uint32) sysinfo.dwNumberOfProcessors;
+#else
+	uint32 num_cpus = (uint32) sysconf(_SC_NPROCESSORS_ONLN);
+#endif
+	PG_RETURN_UINT32(num_cpus);
+}

pgxn/neon_utils/neon_utils.control

@@ -0,0 +1,6 @@
+# neon_utils extension
+comment = 'neon_utils - small useful functions'
+default_version = '1.0'
+module_pathname = '$libdir/neon_utils'
+relocatable = true
+trusted = true

pgxn/neon_walredo/walredoproc.c

@@ -65,6 +65,14 @@
 #include "rusagestub.h"
 #endif
 
+#include "access/clog.h"
+#include "access/commit_ts.h"
+#include "access/heapam.h"
+#include "access/multixact.h"
+#include "access/nbtree.h"
+#include "access/subtrans.h"
+#include "access/syncscan.h"
+#include "access/twophase.h"
 #include "access/xlog.h"
 #include "access/xlog_internal.h"
 #if PG_VERSION_NUM >= 150000
@@ -72,18 +80,36 @@
 #endif
 #include "access/xlogutils.h"
 #include "catalog/pg_class.h"
-#include "libpq/libpq.h"
+#include "commands/async.h"
 #include "libpq/pqformat.h"
 #include "miscadmin.h"
+#include "pgstat.h"
+#include "postmaster/autovacuum.h"
+#include "postmaster/bgworker_internals.h"
+#include "postmaster/bgwriter.h"
 #include "postmaster/postmaster.h"
+#include "replication/logicallauncher.h"
+#include "replication/origin.h"
+#include "replication/slot.h"
+#include "replication/walreceiver.h"
+#include "replication/walsender.h"
 #include "storage/buf_internals.h"
 #include "storage/bufmgr.h"
+#include "storage/dsm.h"
 #include "storage/ipc.h"
+#include "storage/pg_shmem.h"
+#include "storage/pmsignal.h"
+#include "storage/predicate.h"
 #include "storage/proc.h"
+#include "storage/procarray.h"
+#include "storage/procsignal.h"
+#include "storage/sinvaladt.h"
 #include "storage/smgr.h"
+#include "storage/spin.h"
 #include "tcop/tcopprot.h"
 #include "utils/memutils.h"
 #include "utils/ps_status.h"
+#include "utils/snapmgr.h"
 
 #include "inmem_smgr.h"
 
@@ -101,6 +127,7 @@ static void apply_error_callback(void *arg);
 static bool redo_block_filter(XLogReaderState *record, uint8 block_id);
 static void GetPage(StringInfo input_message);
 static ssize_t buffered_read(void *buf, size_t count);
+static void CreateFakeSharedMemoryAndSemaphores();
 
 static BufferTag target_redo_tag;
 
@@ -141,7 +168,7 @@ enter_seccomp_mode(void)
 		PG_SCMP_ALLOW(shmctl),
 		PG_SCMP_ALLOW(shmdt),
 		PG_SCMP_ALLOW(unlink), // shm_unlink
-		*/
+	 */
 	};
 
 #ifdef MALLOC_NO_MMAP
@@ -177,6 +204,7 @@ WalRedoMain(int argc, char *argv[])
 	 * buffers. So let's keep it small (default value is 1024)
 	 */
 	num_temp_buffers = 4;
+	NBuffers = 4;
 
 	/*
 	 * install the simple in-memory smgr
@@ -184,49 +212,33 @@ WalRedoMain(int argc, char *argv[])
 	smgr_hook = smgr_inmem;
 	smgr_init_hook = smgr_init_inmem;
 
-	/*
-	 * Validate we have been given a reasonable-looking DataDir and change into it.
-	 */
-	checkDataDir();
-	ChangeToDataDir();
-
-	/*
-	 * Create lockfile for data directory.
-	 */
-	CreateDataDirLockFile(false);
-
-	/* read control file (error checking and contains config ) */
-	LocalProcessControlFile(false);
-
-	/*
-	 * process any libraries that should be preloaded at postmaster start
-	 */
-	process_shared_preload_libraries();
 
 	/* Initialize MaxBackends (if under postmaster, was done already) */
+	MaxConnections = 1;
+	max_worker_processes = 0;
+	max_parallel_workers = 0;
+	max_wal_senders = 0;
 	InitializeMaxBackends();
 
-#if PG_VERSION_NUM >= 150000
-	/*
-	 * Give preloaded libraries a chance to request additional shared memory.
-	 */
-	process_shmem_requests();
+	/* Disable lastWrittenLsnCache */
+	lastWrittenLsnCacheSize = 0;
 
-	/*
-	 * Now that loadable modules have had their chance to request additional
-	 * shared memory, determine the value of any runtime-computed GUCs that
-	 * depend on the amount of shared memory required.
-	 */
+#if PG_VERSION_NUM >= 150000
+	process_shmem_requests();
 	InitializeShmemGUCs();
 
 	/*
-	 * Now that modules have been loaded, we can process any custom resource
-	 * managers specified in the wal_consistency_checking GUC.
+	 * This will try to access data directory which we do not set.
+	 * Seems to be pretty safe to disable.
 	 */
-	InitializeWalConsistencyChecking();
+	/* InitializeWalConsistencyChecking(); */
 #endif
 
-	CreateSharedMemoryAndSemaphores();
+	/*
+	 * We have our own version of CreateSharedMemoryAndSemaphores() that
+	 * sets up local memory instead of shared one.
+	 */
+	CreateFakeSharedMemoryAndSemaphores();
 
 	/*
 	 * Remember stand-alone backend startup time,roughly at the same point
@@ -354,6 +366,172 @@ WalRedoMain(int argc, char *argv[])
 }
 
 
+/*
+ * Initialize dummy shmem.
+ *
+ * This code follows CreateSharedMemoryAndSemaphores() but manually sets up
+ * the shmem header and skips few initialization steps that are not needed for
+ * WAL redo.
+ *
+ * I've also tried removing most of initialization functions that request some
+ * memory (like ApplyLauncherShmemInit and friends) but in reality it haven't had
+ * any sizeable effect on RSS, so probably such clean up not worth the risk of having
+ * half-initialized postgres.
+ */
+static void
+CreateFakeSharedMemoryAndSemaphores()
+{
+	PGShmemHeader *shim = NULL;
+	PGShmemHeader *hdr;
+	Size		size;
+	int			numSemas;
+	char		cwd[MAXPGPATH];
+
+#if PG_VERSION_NUM >= 150000
+	size = CalculateShmemSize(&numSemas);
+#else
+	/*
+	 * Postgres v14 doesn't have a separate CalculateShmemSize(). Use result of the
+	 * corresponging calculation in CreateSharedMemoryAndSemaphores()
+	 */
+	size = 1409024;
+	numSemas = 10;
+#endif
+
+	/* Dummy implementation of PGSharedMemoryCreate() */
+	{
+		hdr = (PGShmemHeader *) malloc(size);
+		if (!hdr)
+			ereport(FATAL,
+					(errcode(ERRCODE_OUT_OF_MEMORY),
+					 errmsg("[neon-wal-redo] can not allocate (pseudo-) shared memory")));
+
+		hdr->creatorPID = getpid();
+		hdr->magic = PGShmemMagic;
+		hdr->dsm_control = 0;
+		hdr->device = 42; /* not relevant for non-shared memory */
+		hdr->inode = 43; /* not relevant for non-shared memory */
+		hdr->totalsize = size;
+		hdr->freeoffset = MAXALIGN(sizeof(PGShmemHeader));
+
+		shim = hdr;
+		UsedShmemSegAddr = hdr;
+		UsedShmemSegID = (unsigned long) 42; /* not relevant for non-shared memory */
+	}
+
+	InitShmemAccess(hdr);
+
+	/*
+	 * Reserve semaphores uses dir name as a source of entropy. Set it to cwd(). Rest
+	 * of the code does not need DataDir access so nullify DataDir after
+	 * PGReserveSemaphores() to error out if something will try to access it.
+	 */
+	if (!getcwd(cwd, MAXPGPATH))
+		ereport(FATAL,
+			(errcode(ERRCODE_INTERNAL_ERROR),
+			 errmsg("[neon-wal-redo] can not read current directory name")));
+	DataDir = cwd;
+	PGReserveSemaphores(numSemas);
+	DataDir = NULL;
+
+	/*
+	 * The rest of function follows CreateSharedMemoryAndSemaphores() closely,
+	 * skipped parts are marked with comments.
+	 */
+	InitShmemAllocation();
+
+	/*
+	 * Now initialize LWLocks, which do shared memory allocation and are
+	 * needed for InitShmemIndex.
+	 */
+	CreateLWLocks();
+
+	/*
+	 * Set up shmem.c index hashtable
+	 */
+	InitShmemIndex();
+
+	dsm_shmem_init();
+
+	/*
+	 * Set up xlog, clog, and buffers
+	 */
+	XLOGShmemInit();
+	CLOGShmemInit();
+	CommitTsShmemInit();
+	SUBTRANSShmemInit();
+	MultiXactShmemInit();
+	InitBufferPool();
+
+	/*
+	 * Set up lock manager
+	 */
+	InitLocks();
+
+	/*
+	 * Set up predicate lock manager
+	 */
+	InitPredicateLocks();
+
+	/*
+	 * Set up process table
+	 */
+	if (!IsUnderPostmaster)
+		InitProcGlobal();
+	CreateSharedProcArray();
+	CreateSharedBackendStatus();
+	TwoPhaseShmemInit();
+	BackgroundWorkerShmemInit();
+
+	/*
+	 * Set up shared-inval messaging
+	 */
+	CreateSharedInvalidationState();
+
+	/*
+	 * Set up interprocess signaling mechanisms
+	 */
+	PMSignalShmemInit();
+	ProcSignalShmemInit();
+	CheckpointerShmemInit();
+	AutoVacuumShmemInit();
+	ReplicationSlotsShmemInit();
+	ReplicationOriginShmemInit();
+	WalSndShmemInit();
+	WalRcvShmemInit();
+	PgArchShmemInit();
+	ApplyLauncherShmemInit();
+
+	/*
+	 * Set up other modules that need some shared memory space
+	 */
+	SnapMgrInit();
+	BTreeShmemInit();
+	SyncScanShmemInit();
+	/* Skip due to the 'pg_notify' directory check */
+	/* AsyncShmemInit(); */
+
+#ifdef EXEC_BACKEND
+
+	/*
+	 * Alloc the win32 shared backend array
+	 */
+	if (!IsUnderPostmaster)
+		ShmemBackendArrayAllocation();
+#endif
+
+	/* Initialize dynamic shared memory facilities. */
+	if (!IsUnderPostmaster)
+		dsm_postmaster_startup(shim);
+
+	/*
+	 * Now give loadable modules a chance to set up their shmem allocations
+	 */
+	if (shmem_startup_hook)
+		shmem_startup_hook();
+}
+
+
 /* Version compatility wrapper for ReadBufferWithoutRelcache */
 static inline Buffer
 NeonRedoReadBuffer(RelFileNode rnode,

poetry.lock

@@ -1,4 +1,4 @@
-# This file is automatically @generated by Poetry and should not be changed by hand.
+# This file is automatically @generated by Poetry 1.4.1 and should not be changed by hand.
 
 [[package]]
 name = "aiohttp"
@@ -253,43 +253,46 @@ files = [
 
 [[package]]
 name = "black"
-version = "22.6.0"
+version = "23.1.0"
 description = "The uncompromising code formatter."
 category = "dev"
 optional = false
-python-versions = ">=3.6.2"
+python-versions = ">=3.7"
 files = [
-    {file = "black-22.6.0-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:f586c26118bc6e714ec58c09df0157fe2d9ee195c764f630eb0d8e7ccce72e69"},
-    {file = "black-22.6.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:b270a168d69edb8b7ed32c193ef10fd27844e5c60852039599f9184460ce0807"},
-    {file = "black-22.6.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:6797f58943fceb1c461fb572edbe828d811e719c24e03375fd25170ada53825e"},
-    {file = "black-22.6.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:c85928b9d5f83b23cee7d0efcb310172412fbf7cb9d9ce963bd67fd141781def"},
-    {file = "black-22.6.0-cp310-cp310-win_amd64.whl", hash = "sha256:f6fe02afde060bbeef044af7996f335fbe90b039ccf3f5eb8f16df8b20f77666"},
-    {file = "black-22.6.0-cp36-cp36m-macosx_10_9_x86_64.whl", hash = "sha256:cfaf3895a9634e882bf9d2363fed5af8888802d670f58b279b0bece00e9a872d"},
-    {file = "black-22.6.0-cp36-cp36m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:94783f636bca89f11eb5d50437e8e17fbc6a929a628d82304c80fa9cd945f256"},
-    {file = "black-22.6.0-cp36-cp36m-win_amd64.whl", hash = "sha256:2ea29072e954a4d55a2ff58971b83365eba5d3d357352a07a7a4df0d95f51c78"},
-    {file = "black-22.6.0-cp37-cp37m-macosx_10_9_x86_64.whl", hash = "sha256:e439798f819d49ba1c0bd9664427a05aab79bfba777a6db94fd4e56fae0cb849"},
-    {file = "black-22.6.0-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:187d96c5e713f441a5829e77120c269b6514418f4513a390b0499b0987f2ff1c"},
-    {file = "black-22.6.0-cp37-cp37m-win_amd64.whl", hash = "sha256:074458dc2f6e0d3dab7928d4417bb6957bb834434516f21514138437accdbe90"},
-    {file = "black-22.6.0-cp38-cp38-macosx_10_9_universal2.whl", hash = "sha256:a218d7e5856f91d20f04e931b6f16d15356db1c846ee55f01bac297a705ca24f"},
-    {file = "black-22.6.0-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:568ac3c465b1c8b34b61cd7a4e349e93f91abf0f9371eda1cf87194663ab684e"},
-    {file = "black-22.6.0-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:6c1734ab264b8f7929cef8ae5f900b85d579e6cbfde09d7387da8f04771b51c6"},
-    {file = "black-22.6.0-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:c9a3ac16efe9ec7d7381ddebcc022119794872abce99475345c5a61aa18c45ad"},
-    {file = "black-22.6.0-cp38-cp38-win_amd64.whl", hash = "sha256:b9fd45787ba8aa3f5e0a0a98920c1012c884622c6c920dbe98dbd05bc7c70fbf"},
-    {file = "black-22.6.0-cp39-cp39-macosx_10_9_universal2.whl", hash = "sha256:7ba9be198ecca5031cd78745780d65a3f75a34b2ff9be5837045dce55db83d1c"},
-    {file = "black-22.6.0-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:a3db5b6409b96d9bd543323b23ef32a1a2b06416d525d27e0f67e74f1446c8f2"},
-    {file = "black-22.6.0-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:560558527e52ce8afba936fcce93a7411ab40c7d5fe8c2463e279e843c0328ee"},
-    {file = "black-22.6.0-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:b154e6bbde1e79ea3260c4b40c0b7b3109ffcdf7bc4ebf8859169a6af72cd70b"},
-    {file = "black-22.6.0-cp39-cp39-win_amd64.whl", hash = "sha256:4af5bc0e1f96be5ae9bd7aaec219c901a94d6caa2484c21983d043371c733fc4"},
-    {file = "black-22.6.0-py3-none-any.whl", hash = "sha256:ac609cf8ef5e7115ddd07d85d988d074ed00e10fbc3445aee393e70164a2219c"},
-    {file = "black-22.6.0.tar.gz", hash = "sha256:6c6d39e28aed379aec40da1c65434c77d75e65bb59a1e1c283de545fb4e7c6c9"},
+    {file = "black-23.1.0-cp310-cp310-macosx_10_16_arm64.whl", hash = "sha256:b6a92a41ee34b883b359998f0c8e6eb8e99803aa8bf3123bf2b2e6fec505a221"},
+    {file = "black-23.1.0-cp310-cp310-macosx_10_16_universal2.whl", hash = "sha256:57c18c5165c1dbe291d5306e53fb3988122890e57bd9b3dcb75f967f13411a26"},
+    {file = "black-23.1.0-cp310-cp310-macosx_10_16_x86_64.whl", hash = "sha256:9880d7d419bb7e709b37e28deb5e68a49227713b623c72b2b931028ea65f619b"},
+    {file = "black-23.1.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:e6663f91b6feca5d06f2ccd49a10f254f9298cc1f7f49c46e498a0771b507104"},
+    {file = "black-23.1.0-cp310-cp310-win_amd64.whl", hash = "sha256:9afd3f493666a0cd8f8df9a0200c6359ac53940cbde049dcb1a7eb6ee2dd7074"},
+    {file = "black-23.1.0-cp311-cp311-macosx_10_16_arm64.whl", hash = "sha256:bfffba28dc52a58f04492181392ee380e95262af14ee01d4bc7bb1b1c6ca8d27"},
+    {file = "black-23.1.0-cp311-cp311-macosx_10_16_universal2.whl", hash = "sha256:c1c476bc7b7d021321e7d93dc2cbd78ce103b84d5a4cf97ed535fbc0d6660648"},
+    {file = "black-23.1.0-cp311-cp311-macosx_10_16_x86_64.whl", hash = "sha256:382998821f58e5c8238d3166c492139573325287820963d2f7de4d518bd76958"},
+    {file = "black-23.1.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:2bf649fda611c8550ca9d7592b69f0637218c2369b7744694c5e4902873b2f3a"},
+    {file = "black-23.1.0-cp311-cp311-win_amd64.whl", hash = "sha256:121ca7f10b4a01fd99951234abdbd97728e1240be89fde18480ffac16503d481"},
+    {file = "black-23.1.0-cp37-cp37m-macosx_10_16_x86_64.whl", hash = "sha256:a8471939da5e824b891b25751955be52ee7f8a30a916d570a5ba8e0f2eb2ecad"},
+    {file = "black-23.1.0-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:8178318cb74f98bc571eef19068f6ab5613b3e59d4f47771582f04e175570ed8"},
+    {file = "black-23.1.0-cp37-cp37m-win_amd64.whl", hash = "sha256:a436e7881d33acaf2536c46a454bb964a50eff59b21b51c6ccf5a40601fbef24"},
+    {file = "black-23.1.0-cp38-cp38-macosx_10_16_arm64.whl", hash = "sha256:a59db0a2094d2259c554676403fa2fac3473ccf1354c1c63eccf7ae65aac8ab6"},
+    {file = "black-23.1.0-cp38-cp38-macosx_10_16_universal2.whl", hash = "sha256:0052dba51dec07ed029ed61b18183942043e00008ec65d5028814afaab9a22fd"},
+    {file = "black-23.1.0-cp38-cp38-macosx_10_16_x86_64.whl", hash = "sha256:49f7b39e30f326a34b5c9a4213213a6b221d7ae9d58ec70df1c4a307cf2a1580"},
+    {file = "black-23.1.0-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:162e37d49e93bd6eb6f1afc3e17a3d23a823042530c37c3c42eeeaf026f38468"},
+    {file = "black-23.1.0-cp38-cp38-win_amd64.whl", hash = "sha256:8b70eb40a78dfac24842458476135f9b99ab952dd3f2dab738c1881a9b38b753"},
+    {file = "black-23.1.0-cp39-cp39-macosx_10_16_arm64.whl", hash = "sha256:a29650759a6a0944e7cca036674655c2f0f63806ddecc45ed40b7b8aa314b651"},
+    {file = "black-23.1.0-cp39-cp39-macosx_10_16_universal2.whl", hash = "sha256:bb460c8561c8c1bec7824ecbc3ce085eb50005883a6203dcfb0122e95797ee06"},
+    {file = "black-23.1.0-cp39-cp39-macosx_10_16_x86_64.whl", hash = "sha256:c91dfc2c2a4e50df0026f88d2215e166616e0c80e86004d0003ece0488db2739"},
+    {file = "black-23.1.0-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:2a951cc83ab535d248c89f300eccbd625e80ab880fbcfb5ac8afb5f01a258ac9"},
+    {file = "black-23.1.0-cp39-cp39-win_amd64.whl", hash = "sha256:0680d4380db3719ebcfb2613f34e86c8e6d15ffeabcf8ec59355c5e7b85bb555"},
+    {file = "black-23.1.0-py3-none-any.whl", hash = "sha256:7a0f701d314cfa0896b9001df70a530eb2472babb76086344e688829efd97d32"},
+    {file = "black-23.1.0.tar.gz", hash = "sha256:b0bd97bea8903f5a2ba7219257a44e3f1f9d00073d6cc1add68f0beec69692ac"},
 ]
 
 [package.dependencies]
 click = ">=8.0.0"
 mypy-extensions = ">=0.4.3"
+packaging = ">=22.0"
 pathspec = ">=0.9.0"
 platformdirs = ">=2"
-tomli = {version = ">=1.1.0", markers = "python_full_version < \"3.11.0a7\""}
+tomli = {version = ">=1.1.0", markers = "python_version < \"3.11\""}
 typing-extensions = {version = ">=3.10.0.0", markers = "python_version < \"3.10\""}
 
 [package.extras]
@@ -884,6 +887,8 @@ files = [
     {file = "cryptography-39.0.1-cp36-abi3-win32.whl", hash = "sha256:fe913f20024eb2cb2f323e42a64bdf2911bb9738a15dba7d3cce48151034e3a8"},
     {file = "cryptography-39.0.1-cp36-abi3-win_amd64.whl", hash = "sha256:ced4e447ae29ca194449a3f1ce132ded8fcab06971ef5f618605aacaa612beac"},
     {file = "cryptography-39.0.1-pp38-pypy38_pp73-macosx_10_12_x86_64.whl", hash = "sha256:807ce09d4434881ca3a7594733669bd834f5b2c6d5c7e36f8c00f691887042ad"},
+    {file = "cryptography-39.0.1-pp38-pypy38_pp73-manylinux_2_28_aarch64.whl", hash = "sha256:c5caeb8188c24888c90b5108a441c106f7faa4c4c075a2bcae438c6e8ca73cef"},
+    {file = "cryptography-39.0.1-pp38-pypy38_pp73-manylinux_2_28_x86_64.whl", hash = "sha256:4789d1e3e257965e960232345002262ede4d094d1a19f4d3b52e48d4d8f3b885"},
     {file = "cryptography-39.0.1-pp38-pypy38_pp73-win_amd64.whl", hash = "sha256:96f1157a7c08b5b189b16b47bc9db2332269d6680a196341bf30046330d15388"},
     {file = "cryptography-39.0.1-pp39-pypy39_pp73-macosx_10_12_x86_64.whl", hash = "sha256:e422abdec8b5fa8462aa016786680720d78bdce7a30c652b7fadf83a4ba35336"},
     {file = "cryptography-39.0.1-pp39-pypy39_pp73-manylinux_2_17_aarch64.manylinux2014_aarch64.manylinux_2_24_aarch64.whl", hash = "sha256:b0afd054cd42f3d213bf82c629efb1ee5f22eba35bf0eec88ea9ea7304f511a2"},
@@ -963,23 +968,6 @@ files = [
 [package.extras]
 testing = ["pre-commit"]
 
-[[package]]
-name = "flake8"
-version = "5.0.4"
-description = "the modular source code checker: pep8 pyflakes and co"
-category = "dev"
-optional = false
-python-versions = ">=3.6.1"
-files = [
-    {file = "flake8-5.0.4-py2.py3-none-any.whl", hash = "sha256:7a1cf6b73744f5806ab95e526f6f0d8c01c66d7bbe349562d22dfca20610b248"},
-    {file = "flake8-5.0.4.tar.gz", hash = "sha256:6fbe320aad8d6b95cec8b8e47bc933004678dc63095be98528b7bdd2a9f510db"},
-]
-
-[package.dependencies]
-mccabe = ">=0.7.0,<0.8.0"
-pycodestyle = ">=2.9.0,<2.10.0"
-pyflakes = ">=2.5.0,<2.6.0"
-
 [[package]]
 name = "flask"
 version = "2.1.3"
@@ -1075,24 +1063,6 @@ files = [
     {file = "iniconfig-1.1.1.tar.gz", hash = "sha256:bc3af051d7d14b2ee5ef9969666def0cd1a000e121eaea580d4a313df4b37f32"},
 ]
 
-[[package]]
-name = "isort"
-version = "5.10.1"
-description = "A Python utility / library to sort Python imports."
-category = "dev"
-optional = false
-python-versions = ">=3.6.1,<4.0"
-files = [
-    {file = "isort-5.10.1-py3-none-any.whl", hash = "sha256:6f62d78e2f89b4500b080fe3a81690850cd254227f27f75c3a0c491a1f351ba7"},
-    {file = "isort-5.10.1.tar.gz", hash = "sha256:e8443a5e7a020e9d7f97f1d7d9cd17c88bcb3bc7e218bf9cf5095fe550be2951"},
-]
-
-[package.extras]
-colors = ["colorama (>=0.4.3,<0.5.0)"]
-pipfile-deprecated-finder = ["pipreqs", "requirementslib"]
-plugins = ["setuptools"]
-requirements-deprecated-finder = ["pip-api", "pipreqs"]
-
 [[package]]
 name = "itsdangerous"
 version = "2.1.2"
@@ -1238,6 +1208,7 @@ category = "main"
 optional = false
 python-versions = "*"
 files = [
+    {file = "junit-xml-1.9.tar.gz", hash = "sha256:de16a051990d4e25a3982b2dd9e89d671067548718866416faec14d9de56db9f"},
     {file = "junit_xml-1.9-py2.py3-none-any.whl", hash = "sha256:ec5ca1a55aefdd76d28fcc0b135251d156c7106fa979686a4b48d62b761b4732"},
 ]
 
@@ -1294,18 +1265,6 @@ files = [
     {file = "MarkupSafe-2.1.1.tar.gz", hash = "sha256:7f91197cc9e48f989d12e4e6fbc46495c446636dfc81b9ccf50bb0ec74b91d4b"},
 ]
 
-[[package]]
-name = "mccabe"
-version = "0.7.0"
-description = "McCabe checker, plugin for flake8"
-category = "dev"
-optional = false
-python-versions = ">=3.6"
-files = [
-    {file = "mccabe-0.7.0-py2.py3-none-any.whl", hash = "sha256:6c2d30ab6be0e4a46919781807b4f0d834ebdd6c6e3dca0bda5a15f863427b6e"},
-    {file = "mccabe-0.7.0.tar.gz", hash = "sha256:348e0240c33b60bbdf4e523192ef919f28cb2c3d7d5c7794f74009290f236325"},
-]
-
 [[package]]
 name = "moto"
 version = "4.1.2"
@@ -1453,46 +1412,42 @@ files = [
 
 [[package]]
 name = "mypy"
-version = "0.991"
+version = "1.1.1"
 description = "Optional static typing for Python"
 category = "dev"
 optional = false
 python-versions = ">=3.7"
 files = [
-    {file = "mypy-0.991-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:7d17e0a9707d0772f4a7b878f04b4fd11f6f5bcb9b3813975a9b13c9332153ab"},
-    {file = "mypy-0.991-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:0714258640194d75677e86c786e80ccf294972cc76885d3ebbb560f11db0003d"},
-    {file = "mypy-0.991-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:0c8f3be99e8a8bd403caa8c03be619544bc2c77a7093685dcf308c6b109426c6"},
-    {file = "mypy-0.991-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:bc9ec663ed6c8f15f4ae9d3c04c989b744436c16d26580eaa760ae9dd5d662eb"},
-    {file = "mypy-0.991-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:4307270436fd7694b41f913eb09210faff27ea4979ecbcd849e57d2da2f65305"},
-    {file = "mypy-0.991-cp310-cp310-win_amd64.whl", hash = "sha256:901c2c269c616e6cb0998b33d4adbb4a6af0ac4ce5cd078afd7bc95830e62c1c"},
-    {file = "mypy-0.991-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:d13674f3fb73805ba0c45eb6c0c3053d218aa1f7abead6e446d474529aafc372"},
-    {file = "mypy-0.991-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:1c8cd4fb70e8584ca1ed5805cbc7c017a3d1a29fb450621089ffed3e99d1857f"},
-    {file = "mypy-0.991-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:209ee89fbb0deed518605edddd234af80506aec932ad28d73c08f1400ef80a33"},
-    {file = "mypy-0.991-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:37bd02ebf9d10e05b00d71302d2c2e6ca333e6c2a8584a98c00e038db8121f05"},
-    {file = "mypy-0.991-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:26efb2fcc6b67e4d5a55561f39176821d2adf88f2745ddc72751b7890f3194ad"},
-    {file = "mypy-0.991-cp311-cp311-win_amd64.whl", hash = "sha256:3a700330b567114b673cf8ee7388e949f843b356a73b5ab22dd7cff4742a5297"},
-    {file = "mypy-0.991-cp37-cp37m-macosx_10_9_x86_64.whl", hash = "sha256:1f7d1a520373e2272b10796c3ff721ea1a0712288cafaa95931e66aa15798813"},
-    {file = "mypy-0.991-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:641411733b127c3e0dab94c45af15fea99e4468f99ac88b39efb1ad677da5711"},
-    {file = "mypy-0.991-cp37-cp37m-musllinux_1_1_x86_64.whl", hash = "sha256:3d80e36b7d7a9259b740be6d8d906221789b0d836201af4234093cae89ced0cd"},
-    {file = "mypy-0.991-cp37-cp37m-win_amd64.whl", hash = "sha256:e62ebaad93be3ad1a828a11e90f0e76f15449371ffeecca4a0a0b9adc99abcef"},
-    {file = "mypy-0.991-cp38-cp38-macosx_10_9_universal2.whl", hash = "sha256:b86ce2c1866a748c0f6faca5232059f881cda6dda2a893b9a8373353cfe3715a"},
-    {file = "mypy-0.991-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:ac6e503823143464538efda0e8e356d871557ef60ccd38f8824a4257acc18d93"},
-    {file = "mypy-0.991-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:0cca5adf694af539aeaa6ac633a7afe9bbd760df9d31be55ab780b77ab5ae8bf"},
-    {file = "mypy-0.991-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:a12c56bf73cdab116df96e4ff39610b92a348cc99a1307e1da3c3768bbb5b135"},
-    {file = "mypy-0.991-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:652b651d42f155033a1967739788c436491b577b6a44e4c39fb340d0ee7f0d70"},
-    {file = "mypy-0.991-cp38-cp38-win_amd64.whl", hash = "sha256:4175593dc25d9da12f7de8de873a33f9b2b8bdb4e827a7cae952e5b1a342e243"},
-    {file = "mypy-0.991-cp39-cp39-macosx_10_9_universal2.whl", hash = "sha256:98e781cd35c0acf33eb0295e8b9c55cdbef64fcb35f6d3aa2186f289bed6e80d"},
-    {file = "mypy-0.991-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:6d7464bac72a85cb3491c7e92b5b62f3dcccb8af26826257760a552a5e244aa5"},
-    {file = "mypy-0.991-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:c9166b3f81a10cdf9b49f2d594b21b31adadb3d5e9db9b834866c3258b695be3"},
-    {file = "mypy-0.991-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:b8472f736a5bfb159a5e36740847808f6f5b659960115ff29c7cecec1741c648"},
-    {file = "mypy-0.991-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:5e80e758243b97b618cdf22004beb09e8a2de1af481382e4d84bc52152d1c476"},
-    {file = "mypy-0.991-cp39-cp39-win_amd64.whl", hash = "sha256:74e259b5c19f70d35fcc1ad3d56499065c601dfe94ff67ae48b85596b9ec1461"},
-    {file = "mypy-0.991-py3-none-any.whl", hash = "sha256:de32edc9b0a7e67c2775e574cb061a537660e51210fbf6006b0b36ea695ae9bb"},
-    {file = "mypy-0.991.tar.gz", hash = "sha256:3c0165ba8f354a6d9881809ef29f1a9318a236a6d81c690094c5df32107bde06"},
+    {file = "mypy-1.1.1-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:39c7119335be05630611ee798cc982623b9e8f0cff04a0b48dfc26100e0b97af"},
+    {file = "mypy-1.1.1-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:61bf08362e93b6b12fad3eab68c4ea903a077b87c90ac06c11e3d7a09b56b9c1"},
+    {file = "mypy-1.1.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:dbb19c9f662e41e474e0cff502b7064a7edc6764f5262b6cd91d698163196799"},
+    {file = "mypy-1.1.1-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:315ac73cc1cce4771c27d426b7ea558fb4e2836f89cb0296cbe056894e3a1f78"},
+    {file = "mypy-1.1.1-cp310-cp310-win_amd64.whl", hash = "sha256:5cb14ff9919b7df3538590fc4d4c49a0f84392237cbf5f7a816b4161c061829e"},
+    {file = "mypy-1.1.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:26cdd6a22b9b40b2fd71881a8a4f34b4d7914c679f154f43385ca878a8297389"},
+    {file = "mypy-1.1.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:5b5f81b40d94c785f288948c16e1f2da37203c6006546c5d947aab6f90aefef2"},
+    {file = "mypy-1.1.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:21b437be1c02712a605591e1ed1d858aba681757a1e55fe678a15c2244cd68a5"},
+    {file = "mypy-1.1.1-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:d809f88734f44a0d44959d795b1e6f64b2bbe0ea4d9cc4776aa588bb4229fc1c"},
+    {file = "mypy-1.1.1-cp311-cp311-win_amd64.whl", hash = "sha256:a380c041db500e1410bb5b16b3c1c35e61e773a5c3517926b81dfdab7582be54"},
+    {file = "mypy-1.1.1-cp37-cp37m-macosx_10_9_x86_64.whl", hash = "sha256:b7c7b708fe9a871a96626d61912e3f4ddd365bf7f39128362bc50cbd74a634d5"},
+    {file = "mypy-1.1.1-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:c1c10fa12df1232c936830839e2e935d090fc9ee315744ac33b8a32216b93707"},
+    {file = "mypy-1.1.1-cp37-cp37m-musllinux_1_1_x86_64.whl", hash = "sha256:0a28a76785bf57655a8ea5eb0540a15b0e781c807b5aa798bd463779988fa1d5"},
+    {file = "mypy-1.1.1-cp37-cp37m-win_amd64.whl", hash = "sha256:ef6a01e563ec6a4940784c574d33f6ac1943864634517984471642908b30b6f7"},
+    {file = "mypy-1.1.1-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:d64c28e03ce40d5303450f547e07418c64c241669ab20610f273c9e6290b4b0b"},
+    {file = "mypy-1.1.1-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:64cc3afb3e9e71a79d06e3ed24bb508a6d66f782aff7e56f628bf35ba2e0ba51"},
+    {file = "mypy-1.1.1-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ce61663faf7a8e5ec6f456857bfbcec2901fbdb3ad958b778403f63b9e606a1b"},
+    {file = "mypy-1.1.1-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:2b0c373d071593deefbcdd87ec8db91ea13bd8f1328d44947e88beae21e8d5e9"},
+    {file = "mypy-1.1.1-cp38-cp38-win_amd64.whl", hash = "sha256:2888ce4fe5aae5a673386fa232473014056967f3904f5abfcf6367b5af1f612a"},
+    {file = "mypy-1.1.1-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:19ba15f9627a5723e522d007fe708007bae52b93faab00f95d72f03e1afa9598"},
+    {file = "mypy-1.1.1-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:59bbd71e5c58eed2e992ce6523180e03c221dcd92b52f0e792f291d67b15a71c"},
+    {file = "mypy-1.1.1-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9401e33814cec6aec8c03a9548e9385e0e228fc1b8b0a37b9ea21038e64cdd8a"},
+    {file = "mypy-1.1.1-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:4b398d8b1f4fba0e3c6463e02f8ad3346f71956b92287af22c9b12c3ec965a9f"},
+    {file = "mypy-1.1.1-cp39-cp39-win_amd64.whl", hash = "sha256:69b35d1dcb5707382810765ed34da9db47e7f95b3528334a3c999b0c90fe523f"},
+    {file = "mypy-1.1.1-py3-none-any.whl", hash = "sha256:4e4e8b362cdf99ba00c2b218036002bdcdf1e0de085cdb296a49df03fb31dfc4"},
+    {file = "mypy-1.1.1.tar.gz", hash = "sha256:ae9ceae0f5b9059f33dbc62dea087e942c0ccab4b7a003719cb70f9b8abfa32f"},
 ]
 
 [package.dependencies]
-mypy-extensions = ">=0.4.3"
+mypy-extensions = ">=1.0.0"
 tomli = {version = ">=1.1.0", markers = "python_version < \"3.11\""}
 typing-extensions = ">=3.10"
 
@@ -1519,14 +1474,14 @@ typing-extensions = ">=4.1.0"
 
 [[package]]
 name = "mypy-extensions"
-version = "0.4.3"
-description = "Experimental type system extensions for programs checked with the mypy typechecker."
+version = "1.0.0"
+description = "Type system extensions for programs checked with the mypy type checker."
 category = "dev"
 optional = false
-python-versions = "*"
+python-versions = ">=3.5"
 files = [
-    {file = "mypy_extensions-0.4.3-py2.py3-none-any.whl", hash = "sha256:090fedd75945a69ae91ce1303b5824f428daf5a028d2f6ab8a299250a846f15d"},
-    {file = "mypy_extensions-0.4.3.tar.gz", hash = "sha256:2d82818f5bb3e369420cb3c4060a7970edba416647068eb4c5343488a6c604a8"},
+    {file = "mypy_extensions-1.0.0-py3-none-any.whl", hash = "sha256:4392f6c0eb8a5668a69e23d168ffa70f0be9ccfd32b5cc2d26a34ae5b844552d"},
+    {file = "mypy_extensions-1.0.0.tar.gz", hash = "sha256:75dbf8955dc00442a438fc4d0666508a9a97b6bd41aa2f0ffe9d2f2725af0782"},
 ]
 
 [[package]]
@@ -1591,19 +1546,16 @@ requests = ["requests"]
 
 [[package]]
 name = "packaging"
-version = "21.3"
+version = "23.0"
 description = "Core utilities for Python packages"
 category = "main"
 optional = false
-python-versions = ">=3.6"
+python-versions = ">=3.7"
 files = [
-    {file = "packaging-21.3-py3-none-any.whl", hash = "sha256:ef103e05f519cdc783ae24ea4e2e0f508a9c99b2d4969652eed6a2e1ea5bd522"},
-    {file = "packaging-21.3.tar.gz", hash = "sha256:dd47c42927d89ab911e606518907cc2d3a1f38bbd026385970643f9c5b8ecfeb"},
+    {file = "packaging-23.0-py3-none-any.whl", hash = "sha256:714ac14496c3e68c99c29b00845f7a2b85f3bb6f1078fd9f72fd20f0570002b2"},
+    {file = "packaging-23.0.tar.gz", hash = "sha256:b6ad297f8907de0fa2fe1ccbd26fdaf387f5f47c7275fedf8cce89f99446cf97"},
 ]
 
-[package.dependencies]
-pyparsing = ">=2.0.2,<3.0.5 || >3.0.5"
-
 [[package]]
 name = "pathspec"
 version = "0.9.0"
@@ -1712,6 +1664,7 @@ python-versions = ">=3.6"
 files = [
     {file = "psycopg2-binary-2.9.3.tar.gz", hash = "sha256:761df5313dc15da1502b21453642d7599d26be88bff659382f8f9747c7ebea4e"},
     {file = "psycopg2_binary-2.9.3-cp310-cp310-macosx_10_14_x86_64.macosx_10_9_intel.macosx_10_9_x86_64.macosx_10_10_intel.macosx_10_10_x86_64.whl", hash = "sha256:539b28661b71da7c0e428692438efbcd048ca21ea81af618d845e06ebfd29478"},
+    {file = "psycopg2_binary-2.9.3-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:2f2534ab7dc7e776a263b463a16e189eb30e85ec9bbe1bff9e78dae802608932"},
     {file = "psycopg2_binary-2.9.3-cp310-cp310-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:6e82d38390a03da28c7985b394ec3f56873174e2c88130e6966cb1c946508e65"},
     {file = "psycopg2_binary-2.9.3-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:57804fc02ca3ce0dbfbef35c4b3a4a774da66d66ea20f4bda601294ad2ea6092"},
     {file = "psycopg2_binary-2.9.3-cp310-cp310-manylinux_2_24_aarch64.whl", hash = "sha256:083a55275f09a62b8ca4902dd11f4b33075b743cf0d360419e2051a8a5d5ff76"},
@@ -1745,6 +1698,7 @@ files = [
     {file = "psycopg2_binary-2.9.3-cp37-cp37m-win32.whl", hash = "sha256:adf20d9a67e0b6393eac162eb81fb10bc9130a80540f4df7e7355c2dd4af9fba"},
     {file = "psycopg2_binary-2.9.3-cp37-cp37m-win_amd64.whl", hash = "sha256:2f9ffd643bc7349eeb664eba8864d9e01f057880f510e4681ba40a6532f93c71"},
     {file = "psycopg2_binary-2.9.3-cp38-cp38-macosx_10_14_x86_64.macosx_10_9_intel.macosx_10_9_x86_64.macosx_10_10_intel.macosx_10_10_x86_64.whl", hash = "sha256:def68d7c21984b0f8218e8a15d514f714d96904265164f75f8d3a70f9c295667"},
+    {file = "psycopg2_binary-2.9.3-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:e6aa71ae45f952a2205377773e76f4e3f27951df38e69a4c95440c779e013560"},
     {file = "psycopg2_binary-2.9.3-cp38-cp38-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:dffc08ca91c9ac09008870c9eb77b00a46b3378719584059c034b8945e26b272"},
     {file = "psycopg2_binary-2.9.3-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:280b0bb5cbfe8039205c7981cceb006156a675362a00fe29b16fbc264e242834"},
     {file = "psycopg2_binary-2.9.3-cp38-cp38-manylinux_2_24_aarch64.whl", hash = "sha256:af9813db73395fb1fc211bac696faea4ca9ef53f32dc0cfa27e4e7cf766dcf24"},
@@ -1756,6 +1710,7 @@ files = [
     {file = "psycopg2_binary-2.9.3-cp38-cp38-win32.whl", hash = "sha256:6472a178e291b59e7f16ab49ec8b4f3bdada0a879c68d3817ff0963e722a82ce"},
     {file = "psycopg2_binary-2.9.3-cp38-cp38-win_amd64.whl", hash = "sha256:35168209c9d51b145e459e05c31a9eaeffa9a6b0fd61689b48e07464ffd1a83e"},
     {file = "psycopg2_binary-2.9.3-cp39-cp39-macosx_10_14_x86_64.macosx_10_9_intel.macosx_10_9_x86_64.macosx_10_10_intel.macosx_10_10_x86_64.whl", hash = "sha256:47133f3f872faf28c1e87d4357220e809dfd3fa7c64295a4a148bcd1e6e34ec9"},
+    {file = "psycopg2_binary-2.9.3-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:b3a24a1982ae56461cc24f6680604fffa2c1b818e9dc55680da038792e004d18"},
     {file = "psycopg2_binary-2.9.3-cp39-cp39-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:91920527dea30175cc02a1099f331aa8c1ba39bf8b7762b7b56cbf54bc5cce42"},
     {file = "psycopg2_binary-2.9.3-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:887dd9aac71765ac0d0bac1d0d4b4f2c99d5f5c1382d8b770404f0f3d0ce8a39"},
     {file = "psycopg2_binary-2.9.3-cp39-cp39-manylinux_2_24_aarch64.whl", hash = "sha256:1f14c8b0942714eb3c74e1e71700cbbcb415acbc311c730370e70c578a44a25c"},
@@ -1788,33 +1743,10 @@ category = "main"
 optional = false
 python-versions = "*"
 files = [
-    {file = "pyasn1-0.4.8-py2.4.egg", hash = "sha256:fec3e9d8e36808a28efb59b489e4528c10ad0f480e57dcc32b4de5c9d8c9fdf3"},
-    {file = "pyasn1-0.4.8-py2.5.egg", hash = "sha256:0458773cfe65b153891ac249bcf1b5f8f320b7c2ce462151f8fa74de8934becf"},
-    {file = "pyasn1-0.4.8-py2.6.egg", hash = "sha256:5c9414dcfede6e441f7e8f81b43b34e834731003427e5b09e4e00e3172a10f00"},
-    {file = "pyasn1-0.4.8-py2.7.egg", hash = "sha256:6e7545f1a61025a4e58bb336952c5061697da694db1cae97b116e9c46abcf7c8"},
     {file = "pyasn1-0.4.8-py2.py3-none-any.whl", hash = "sha256:39c7e2ec30515947ff4e87fb6f456dfc6e84857d34be479c9d4a4ba4bf46aa5d"},
-    {file = "pyasn1-0.4.8-py3.1.egg", hash = "sha256:78fa6da68ed2727915c4767bb386ab32cdba863caa7dbe473eaae45f9959da86"},
-    {file = "pyasn1-0.4.8-py3.2.egg", hash = "sha256:08c3c53b75eaa48d71cf8c710312316392ed40899cb34710d092e96745a358b7"},
-    {file = "pyasn1-0.4.8-py3.3.egg", hash = "sha256:03840c999ba71680a131cfaee6fab142e1ed9bbd9c693e285cc6aca0d555e576"},
-    {file = "pyasn1-0.4.8-py3.4.egg", hash = "sha256:7ab8a544af125fb704feadb008c99a88805126fb525280b2270bb25cc1d78a12"},
-    {file = "pyasn1-0.4.8-py3.5.egg", hash = "sha256:e89bf84b5437b532b0803ba5c9a5e054d21fec423a89952a74f87fa2c9b7bce2"},
-    {file = "pyasn1-0.4.8-py3.6.egg", hash = "sha256:014c0e9976956a08139dc0712ae195324a75e142284d5f87f1a87ee1b068a359"},
-    {file = "pyasn1-0.4.8-py3.7.egg", hash = "sha256:99fcc3c8d804d1bc6d9a099921e39d827026409a58f2a720dcdb89374ea0c776"},
     {file = "pyasn1-0.4.8.tar.gz", hash = "sha256:aef77c9fb94a3ac588e87841208bdec464471d9871bd5050a287cc9a475cd0ba"},
 ]
 
-[[package]]
-name = "pycodestyle"
-version = "2.9.1"
-description = "Python style guide checker"
-category = "dev"
-optional = false
-python-versions = ">=3.6"
-files = [
-    {file = "pycodestyle-2.9.1-py2.py3-none-any.whl", hash = "sha256:d1735fc58b418fd7c5f658d28d943854f8a849b01a5d0a1e6f3f3fdd0166804b"},
-    {file = "pycodestyle-2.9.1.tar.gz", hash = "sha256:2c9607871d58c76354b697b42f5d57e1ada7d261c261efac224b664affdc5785"},
-]
-
 [[package]]
 name = "pycparser"
 version = "2.21"
@@ -1827,18 +1759,6 @@ files = [
     {file = "pycparser-2.21.tar.gz", hash = "sha256:e644fdec12f7872f86c58ff790da456218b10f863970249516d60a5eaca77206"},
 ]
 
-[[package]]
-name = "pyflakes"
-version = "2.5.0"
-description = "passive checker of Python programs"
-category = "dev"
-optional = false
-python-versions = ">=3.6"
-files = [
-    {file = "pyflakes-2.5.0-py2.py3-none-any.whl", hash = "sha256:4579f67d887f804e67edb544428f264b7b24f435b263c4614f384135cea553d2"},
-    {file = "pyflakes-2.5.0.tar.gz", hash = "sha256:491feb020dca48ccc562a8c0cbe8df07ee13078df59813b83959cbdada312ea3"},
-]
-
 [[package]]
 name = "pyjwt"
 version = "2.4.0"
@@ -2008,10 +1928,25 @@ files = [
 
 [package.dependencies]
 pytest = [
-    {version = ">=6.2.4", markers = "python_version >= \"3.10\""},
     {version = ">=5.0", markers = "python_version < \"3.10\""},
+    {version = ">=6.2.4", markers = "python_version >= \"3.10\""},
 ]
 
+[[package]]
+name = "pytest-repeat"
+version = "0.9.1"
+description = "pytest plugin for repeating tests"
+category = "main"
+optional = false
+python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*, !=3.4.*"
+files = [
+    {file = "pytest-repeat-0.9.1.tar.gz", hash = "sha256:5cd3289745ab3156d43eb9c8e7f7d00a926f3ae5c9cf425bec649b2fe15bad5b"},
+    {file = "pytest_repeat-0.9.1-py2.py3-none-any.whl", hash = "sha256:4474a7d9e9137f6d8cc8ae297f8c4168d33c56dd740aa78cfffe562557e6b96e"},
+]
+
+[package.dependencies]
+pytest = ">=3.6"
+
 [[package]]
 name = "pytest-timeout"
 version = "2.1.0"
@@ -2121,6 +2056,13 @@ files = [
     {file = "PyYAML-6.0-cp310-cp310-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_12_x86_64.manylinux2010_x86_64.whl", hash = "sha256:f84fbc98b019fef2ee9a1cb3ce93e3187a6df0b2538a651bfb890254ba9f90b5"},
     {file = "PyYAML-6.0-cp310-cp310-win32.whl", hash = "sha256:2cd5df3de48857ed0544b34e2d40e9fac445930039f3cfe4bcc592a1f836d513"},
     {file = "PyYAML-6.0-cp310-cp310-win_amd64.whl", hash = "sha256:daf496c58a8c52083df09b80c860005194014c3698698d1a57cbcfa182142a3a"},
+    {file = "PyYAML-6.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:d4b0ba9512519522b118090257be113b9468d804b19d63c71dbcf4a48fa32358"},
+    {file = "PyYAML-6.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:81957921f441d50af23654aa6c5e5eaf9b06aba7f0a19c18a538dc7ef291c5a1"},
+    {file = "PyYAML-6.0-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:afa17f5bc4d1b10afd4466fd3a44dc0e245382deca5b3c353d8b757f9e3ecb8d"},
+    {file = "PyYAML-6.0-cp311-cp311-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:dbad0e9d368bb989f4515da330b88a057617d16b6a8245084f1b05400f24609f"},
+    {file = "PyYAML-6.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:432557aa2c09802be39460360ddffd48156e30721f5e8d917f01d31694216782"},
+    {file = "PyYAML-6.0-cp311-cp311-win32.whl", hash = "sha256:bfaef573a63ba8923503d27530362590ff4f576c626d86a9fed95822a8255fd7"},
+    {file = "PyYAML-6.0-cp311-cp311-win_amd64.whl", hash = "sha256:01b45c0191e6d66c470b6cf1b9531a771a83c1c4208272ead47a3ae4f2f603bf"},
     {file = "PyYAML-6.0-cp36-cp36m-macosx_10_9_x86_64.whl", hash = "sha256:897b80890765f037df3403d22bab41627ca8811ae55e9a722fd0392850ec4d86"},
     {file = "PyYAML-6.0-cp36-cp36m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:50602afada6d6cbfad699b0c7bb50d5ccffa7e46a3d738092afddc1f9758427f"},
     {file = "PyYAML-6.0-cp36-cp36m-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:48c346915c114f5fdb3ead70312bd042a953a8ce5c7106d5bfb1a5254e47da92"},
@@ -2205,6 +2147,33 @@ files = [
 [package.dependencies]
 pyasn1 = ">=0.1.3"
 
+[[package]]
+name = "ruff"
+version = "0.0.255"
+description = "An extremely fast Python linter, written in Rust."
+category = "dev"
+optional = false
+python-versions = ">=3.7"
+files = [
+    {file = "ruff-0.0.255-py3-none-macosx_10_7_x86_64.whl", hash = "sha256:b2d71fb6a7e50501a2473864acffc85dee6b750c25db198f7e71fe1dbbff1aad"},
+    {file = "ruff-0.0.255-py3-none-macosx_10_9_x86_64.macosx_11_0_arm64.macosx_10_9_universal2.whl", hash = "sha256:6c97d746861a6010f941179e84bba9feb8a871815667471d9ed6beb98d45c252"},
+    {file = "ruff-0.0.255-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:9a7fa60085079b91a298b963361be9b1b1c724582af6c84be954cbabdbd9309a"},
+    {file = "ruff-0.0.255-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:c089f7141496334ab5a127b54ce55e41f0d6714e68a4453a1e09d2204cdea8c3"},
+    {file = "ruff-0.0.255-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:0423908caa7d437a416b853214565b9c33bbd1106c4f88147982216dddcbbd96"},
+    {file = "ruff-0.0.255-py3-none-manylinux_2_17_ppc64.manylinux2014_ppc64.whl", hash = "sha256:981493e92547cacbb8e0874904ec049fe744507ee890dc8736caf89a8864f9a7"},
+    {file = "ruff-0.0.255-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:59d5193d2aedb35db180824462b374dbcfc306b2e76076245088afa6e5837df2"},
+    {file = "ruff-0.0.255-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:dd5e00733c9d160c8a34a22e62b390da9d1e9f326676402421cb8c1236beefc3"},
+    {file = "ruff-0.0.255-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:694418cf41838bd19c6229e4e1b2d04505b1e6b86fe3ab81165484fc96d36f01"},
+    {file = "ruff-0.0.255-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:5d0408985c9777369daebb5d3340a99e9f7294bdd7120642239261508185cf89"},
+    {file = "ruff-0.0.255-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:abd6376ef9d12f370d95a8c7c98682fbb9bfedfba59f40e84a816fef8ddcb8de"},
+    {file = "ruff-0.0.255-py3-none-musllinux_1_2_i686.whl", hash = "sha256:f9b1a5df0bc09193cbef58a6f78e4a9a0b058a4f9733c0442866d078006d1bb9"},
+    {file = "ruff-0.0.255-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:6a25c5f4ff087445b2e1bbcb9963f2ae7c868d65e4a8d5f84c36c12f71571179"},
+    {file = "ruff-0.0.255-py3-none-win32.whl", hash = "sha256:1ff87a8310354f9f1a099625e54a27fdd6756d9cd2a40b45922f2e943daf982d"},
+    {file = "ruff-0.0.255-py3-none-win_amd64.whl", hash = "sha256:f3d8416be618f023f93ec4fd6ee3048585ef85dba9563b2a7e38fc7e5131d5b1"},
+    {file = "ruff-0.0.255-py3-none-win_arm64.whl", hash = "sha256:8ba124819624145d7b6b53add40c367c44318893215ffc1bfe3d72e0225a1c9c"},
+    {file = "ruff-0.0.255.tar.gz", hash = "sha256:f9eb1d3b2eecbeedae419fa494c4e2a5e4484baf93a1ce0f81eddb005e1919c5"},
+]
+
 [[package]]
 name = "s3transfer"
 version = "0.6.0"
@@ -2643,4 +2612,4 @@ testing = ["func-timeout", "jaraco.itertools", "pytest (>=6)", "pytest-black (>=
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.9"
-content-hash = "3038940781ef59d1ed28cedf46120ad6623e21e602c38ad3c359428d79fa1efd"
+content-hash = "ce00f0cf8735e850b1cfb29d8ddf8f66fab2395c589fe7e16becc774edcc7845"

pre-commit.py

@@ -43,17 +43,13 @@ def black(fix_inplace: bool) -> str:
     return cmd
 
 
-def isort(fix_inplace: bool) -> str:
-    cmd = "poetry run isort"
-    if not fix_inplace:
-        cmd += " --diff --check"
+def ruff(fix_inplace: bool) -> str:
+    cmd = "poetry run ruff"
+    if fix_inplace:
+        cmd += " --fix"
     return cmd
 
 
-def flake8() -> str:
-    return "poetry run flake8"
-
-
 def mypy() -> str:
     return "poetry run mypy"
 
@@ -112,13 +108,6 @@ if __name__ == "__main__":
         changed_files=files,
         no_color=args.no_color,
     )
-    check(
-        name="isort",
-        suffix=".py",
-        cmd=isort(fix_inplace=args.fix_inplace),
-        changed_files=files,
-        no_color=args.no_color,
-    )
     check(
         name="black",
         suffix=".py",
@@ -127,9 +116,9 @@ if __name__ == "__main__":
         no_color=args.no_color,
     )
     check(
-        name="flake8",
+        name="ruff",
         suffix=".py",
-        cmd=flake8(),
+        cmd=ruff(fix_inplace=args.fix_inplace),
         changed_files=files,
         no_color=args.no_color,
     )

proxy/src/compute.rs

@@ -6,9 +6,9 @@ use std::{io, net::SocketAddr};
 use thiserror::Error;
 use tokio::net::TcpStream;
 use tokio_postgres::NoTls;
-use tracing::{error, info};
+use tracing::{error, info, warn};
 
-const COULD_NOT_CONNECT: &str = "Could not connect to compute node";
+const COULD_NOT_CONNECT: &str = "Couldn't connect to compute node";
 
 #[derive(Debug, Error)]
 pub enum ConnectionError {
@@ -131,7 +131,7 @@ impl ConnCfg {
         use tokio_postgres::config::Host;
 
         let connect_once = |host, port| {
-            info!("trying to connect to a compute node at {host}:{port}");
+            info!("trying to connect to compute node at {host}:{port}");
             TcpStream::connect((host, port)).and_then(|socket| async {
                 let socket_addr = socket.peer_addr()?;
                 // This prevents load balancer from severing the connection.
@@ -151,7 +151,7 @@ impl ConnCfg {
             return Err(io::Error::new(
                 io::ErrorKind::Other,
                 format!(
-                    "couldn't connect: bad compute config, \
+                    "bad compute config, \
                      ports and hosts entries' count does not match: {:?}",
                     self.0
                 ),
@@ -170,7 +170,7 @@ impl ConnCfg {
                 Ok(socket) => return Ok(socket),
                 Err(err) => {
                     // We can't throw an error here, as there might be more hosts to try.
-                    error!("failed to connect to a compute node at {host}:{port}: {err}");
+                    warn!("couldn't connect to compute node at {host}:{port}: {err}");
                     connection_error = Some(err);
                 }
             }
@@ -179,7 +179,7 @@ impl ConnCfg {
         Err(connection_error.unwrap_or_else(|| {
             io::Error::new(
                 io::ErrorKind::Other,
-                format!("couldn't connect: bad compute config: {:?}", self.0),
+                format!("bad compute config: {:?}", self.0),
             )
         }))
     }
@@ -195,12 +195,11 @@ pub struct PostgresConnection {
 }
 
 impl ConnCfg {
-    /// Connect to a corresponding compute node.
-    pub async fn connect(&self) -> Result<PostgresConnection, ConnectionError> {
+    async fn do_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?;
-        info!("connected to user's compute node at {socket_addr}");
+        info!("connected to compute node at {socket_addr}");
 
         // This is very ugly but as of now there's no better way to
         // extract the connection parameters from tokio-postgres' connection.
@@ -219,6 +218,16 @@ impl ConnCfg {
 
         Ok(connection)
     }
+
+    /// Connect to a corresponding compute node.
+    pub async fn connect(&self) -> Result<PostgresConnection, ConnectionError> {
+        self.do_connect()
+            .inspect_err(|err| {
+                // Immediately log the error we have at our disposal.
+                error!("couldn't connect to compute node: {err}");
+            })
+            .await
+    }
 }
 
 /// Retrieve `options` from a startup message, dropping all proxy-secific flags.

proxy/src/console/mgmt.rs

@@ -4,7 +4,7 @@ use crate::{
 };
 use anyhow::Context;
 use once_cell::sync::Lazy;
-use postgres_backend::{self, AuthType, PostgresBackend, QueryError};
+use postgres_backend::{self, AuthType, PostgresBackend, PostgresBackendTCP, QueryError};
 use pq_proto::{BeMessage, SINGLE_COL_ROWDESC};
 use std::future;
 use tokio::net::{TcpListener, TcpStream};
@@ -71,10 +71,10 @@ pub type ComputeReady = Result<DatabaseInfo, String>;
 // TODO: replace with an http-based protocol.
 struct MgmtHandler;
 #[async_trait::async_trait]
-impl postgres_backend::Handler for MgmtHandler {
+impl postgres_backend::Handler<tokio::net::TcpStream> for MgmtHandler {
     async fn process_query(
         &mut self,
-        pgb: &mut PostgresBackend,
+        pgb: &mut PostgresBackendTCP,
         query: &str,
     ) -> Result<(), QueryError> {
         try_process_query(pgb, query).await.map_err(|e| {
@@ -84,7 +84,7 @@ impl postgres_backend::Handler for MgmtHandler {
     }
 }
 
-async fn try_process_query(pgb: &mut PostgresBackend, query: &str) -> Result<(), QueryError> {
+async fn try_process_query(pgb: &mut PostgresBackendTCP, query: &str) -> Result<(), QueryError> {
     let resp: KickSession = serde_json::from_str(query).context("Failed to parse query as json")?;
 
     let span = info_span!("event", session_id = resp.session_id);

proxy/src/metrics.rs

@@ -21,6 +21,7 @@ const PROXY_IO_BYTES_PER_CLIENT: &str = "proxy_io_bytes_per_client";
 #[derive(Eq, Hash, PartialEq, Serialize, Debug)]
 pub struct Ids {
     pub endpoint_id: String,
+    pub branch_id: String,
 }
 
 pub async fn task_main(config: &MetricCollectionConfig) -> anyhow::Result<()> {
@@ -74,12 +75,23 @@ fn gather_proxy_io_bytes_per_client() -> Vec<(Ids, (u64, DateTime<Utc>))> {
                         .find(|l| l.get_name() == "endpoint_id")
                         .unwrap()
                         .get_value();
+                    let branch_id = ms
+                        .get_label()
+                        .iter()
+                        .find(|l| l.get_name() == "branch_id")
+                        .unwrap()
+                        .get_value();
+
                     let value = ms.get_counter().get_value() as u64;
 
-                    debug!("endpoint_id:val - {}: {}", endpoint_id, value);
+                    debug!(
+                        "branch_id {} endpoint_id {} val: {}",
+                        branch_id, endpoint_id, value
+                    );
                     current_metrics.push((
                         Ids {
                             endpoint_id: endpoint_id.to_string(),
+                            branch_id: "".to_string(),
                         },
                         (value, Utc::now()),
                     ));
@@ -131,6 +143,7 @@ async fn collect_metrics_iteration(
                 value,
                 extra: Ids {
                     endpoint_id: curr_key.endpoint_id.clone(),
+                    branch_id: curr_key.branch_id.clone(),
                 },
             })
         })
@@ -172,6 +185,7 @@ async fn collect_metrics_iteration(
                 cached_metrics
                     .entry(Ids {
                         endpoint_id: send_metric.extra.endpoint_id.clone(),
+                        branch_id: send_metric.extra.branch_id.clone(),
                     })
                     // update cached value (add delta) and time
                     .and_modify(|e| {

proxy/src/proxy.rs

@@ -8,7 +8,7 @@ use crate::{
     config::{ProxyConfig, TlsConfig},
     console::{self, messages::MetricsAuxInfo},
     error::io_error,
-    stream::{MeasuredStream, PqStream, Stream},
+    stream::{PqStream, Stream},
 };
 use anyhow::{bail, Context};
 use futures::TryFutureExt;
@@ -16,8 +16,9 @@ use metrics::{register_int_counter, register_int_counter_vec, IntCounter, IntCou
 use once_cell::sync::Lazy;
 use pq_proto::{BeMessage as Be, FeStartupPacket, StartupMessageParams};
 use std::sync::Arc;
-use tokio::io::{AsyncRead, AsyncWrite};
+use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt};
 use tracing::{error, info, warn};
+use utils::measured_stream::MeasuredStream;
 
 /// Number of times we should retry the `/proxy_wake_compute` http request.
 const NUM_RETRIES_WAKE_COMPUTE: usize = 1;
@@ -208,9 +209,18 @@ async fn handshake<S: AsyncRead + AsyncWrite + Unpin>(
                     if let Some(tls) = tls.take() {
                         // Upgrade raw stream into a secure TLS-backed stream.
                         // NOTE: We've consumed `tls`; this fact will be used later.
-                        stream = PqStream::new(
-                            stream.into_inner().upgrade(tls.to_server_config()).await?,
-                        );
+
+                        let (raw, read_buf) = stream.into_inner();
+                        // TODO: Normally, client doesn't send any data before
+                        // server says TLS handshake is ok and read_buf is empy.
+                        // However, you could imagine pipelining of postgres
+                        // SSLRequest + TLS ClientHello in one hunk similar to
+                        // pipelining in our node js driver. We should probably
+                        // support that by chaining read_buf with the stream.
+                        if !read_buf.is_empty() {
+                            bail!("data is sent before server replied with EncryptionResponse");
+                        }
+                        stream = PqStream::new(raw.upgrade(tls.to_server_config()).await?);
                     }
                 }
                 _ => bail!(ERR_PROTO_VIOLATION),
@@ -353,16 +363,24 @@ async fn proxy_pass(
     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 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);
-    });
+    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...");
@@ -434,11 +452,17 @@ impl<S: AsyncRead + AsyncWrite + Unpin> Client<'_, S> {
             value: mut node_info,
         } = auth_result;
 
-        let node = connect_to_compute(&mut node_info, params, &extra, &creds)
+        let mut node = connect_to_compute(&mut node_info, params, &extra, &creds)
             .or_else(|e| stream.throw_error(e))
             .await?;
 
         prepare_client_connection(&node, reported_auth_ok, session, &mut stream).await?;
-        proxy_pass(stream.into_inner(), node.stream, &node_info.aux).await
+        // Before proxy passing, forward to compute whatever data is left in the
+        // PqStream input buffer. Normally there is none, but our serverless npm
+        // driver in pipeline mode sends startup, password and first query
+        // immediately after opening the connection.
+        let (stream, read_buf) = stream.into_inner();
+        node.stream.write_all(&read_buf).await?;
+        proxy_pass(stream, node.stream, &node_info.aux).await
     }
 }

proxy/src/scram/messages.rs

@@ -14,7 +14,7 @@ pub const SCRAM_RAW_NONCE_LEN: usize = 18;
 fn validate_sasl_extensions<'a>(parts: impl Iterator<Item = &'a str>) -> Option<()> {
     for mut chars in parts.map(|s| s.chars()) {
         let attr = chars.next()?;
-        if !('a'..='z').contains(&attr) && !('A'..='Z').contains(&attr) {
+        if !attr.is_ascii_alphabetic() {
             return None;
         }
         let eq = chars.next()?;

proxy/src/stream.rs

@@ -1,5 +1,6 @@
 use crate::error::UserFacingError;
 use anyhow::bail;
+use bytes::BytesMut;
 use pin_project_lite::pin_project;
 use pq_proto::framed::{ConnectionError, Framed};
 use pq_proto::{BeMessage, FeMessage, FeStartupPacket, ProtocolError};
@@ -27,8 +28,8 @@ impl<S> PqStream<S> {
         }
     }
 
-    /// Extract the underlying stream.
-    pub fn into_inner(self) -> S {
+    /// Extract the underlying stream and read buffer.
+    pub fn into_inner(self) -> (S, BytesMut) {
         self.framed.into_inner()
     }
 
@@ -217,68 +218,3 @@ impl<S: AsyncRead + AsyncWrite + Unpin> AsyncWrite for Stream<S> {
         }
     }
 }
-
-pin_project! {
-    /// This stream tracks all writes and calls user provided
-    /// callback when the underlying stream is flushed.
-    pub struct MeasuredStream<S, W> {
-        #[pin]
-        stream: S,
-        write_count: usize,
-        inc_write_count: W,
-    }
-}
-
-impl<S, W> MeasuredStream<S, W> {
-    pub fn new(stream: S, inc_write_count: W) -> Self {
-        Self {
-            stream,
-            write_count: 0,
-            inc_write_count,
-        }
-    }
-}
-
-impl<S: AsyncRead + Unpin, W> AsyncRead for MeasuredStream<S, W> {
-    fn poll_read(
-        self: Pin<&mut Self>,
-        context: &mut task::Context<'_>,
-        buf: &mut ReadBuf<'_>,
-    ) -> task::Poll<io::Result<()>> {
-        self.project().stream.poll_read(context, buf)
-    }
-}
-
-impl<S: AsyncWrite + Unpin, W: FnMut(usize)> AsyncWrite for MeasuredStream<S, W> {
-    fn poll_write(
-        self: Pin<&mut Self>,
-        context: &mut task::Context<'_>,
-        buf: &[u8],
-    ) -> task::Poll<io::Result<usize>> {
-        let this = self.project();
-        this.stream.poll_write(context, buf).map_ok(|cnt| {
-            // Increment the write count.
-            *this.write_count += cnt;
-            cnt
-        })
-    }
-
-    fn poll_flush(
-        self: Pin<&mut Self>,
-        context: &mut task::Context<'_>,
-    ) -> task::Poll<io::Result<()>> {
-        let this = self.project();
-        this.stream.poll_flush(context).map_ok(|()| {
-            // Call the user provided callback and reset the write count.
-            (this.inc_write_count)(*this.write_count);
-            *this.write_count = 0;
-        })
-    }
-
-    fn poll_shutdown(
-        self: Pin<&mut Self>,
-        context: &mut task::Context<'_>,
-    ) -> task::Poll<io::Result<()>> {
-        self.project().stream.poll_shutdown(context)
-    }
-}

pyproject.toml

@@ -34,12 +34,12 @@ types-psutil = "^5.9.5.4"
 types-toml = "^0.10.8"
 pytest-httpserver = "^1.0.6"
 aiohttp = "3.7.4"
+pytest-repeat = "^0.9.1"
 
-[tool.poetry.dev-dependencies]
-flake8 = "^5.0.4"
-mypy = "==0.991"
-black = "^22.6.0"
-isort = "^5.10.1"
+[tool.poetry.group.dev.dependencies]
+black = "^23.1.0"
+mypy = "==1.1.1"
+ruff = "^0.0.255"
 
 [build-system]
 requires = ["poetry-core>=1.0.0"]
@@ -53,14 +53,6 @@ extend-exclude = '''
 )/
 '''
 
-[tool.isort]
-profile = "black"
-line_length = 100
-skip_gitignore = true
-skip = [
-    "vendor",
-]
-
 [tool.mypy]
 exclude = "^vendor/"
 check_untyped_defs = true
@@ -80,3 +72,13 @@ module = [
     "pg8000.*",
 ]
 ignore_missing_imports = true
+
+[tool.ruff]
+extend-exclude = ["vendor/"]
+ignore = ["E501"]
+select = [
+    "E", # pycodestyle
+    "F", # Pyflakes
+    "I", # isort
+    "W", # pycodestyle
+]

run_clippy.sh

@@ -8,21 +8,13 @@
 # warnings and errors right in the editor.
 # In vscode, this setting is Rust-analyzer>Check On Save:Command
 
-
-# Not every feature is supported in macOS builds. Avoid running regular linting
-# script that checks every feature.
-#
 # manual-range-contains wants
-#   !(8..=MAX_STARTUP_PACKET_LENGTH).contains(&len)
+#   !(4..=MAX_STARTUP_PACKET_LENGTH).contains(&len)
 # instead of
 #   len < 4 || len > MAX_STARTUP_PACKET_LENGTH
 # , let's disagree.
-if [[ "$OSTYPE" == "darwin"* ]]; then
-    # no extra features to test currently, add more here when needed
-    cargo clippy --locked --all --all-targets --features testing -- -A unknown_lints  -A clippy::manual-range-contains -D warnings
-else
-    # * `-A unknown_lints` – do not warn about unknown lint suppressions
-    #                        that people with newer toolchains might use
-    # * `-D warnings`      - fail on any warnings (`cargo` returns non-zero exit status)
-    cargo clippy --locked --all --all-targets --all-features -- -A unknown_lints -A clippy::manual-range-contains -D warnings
-fi
+
+# * `-A unknown_lints` – do not warn about unknown lint suppressions
+#                        that people with newer toolchains might use
+# * `-D warnings`      - fail on any warnings (`cargo` returns non-zero exit status)
+cargo clippy --locked --all --all-targets --all-features -- -A unknown_lints -A clippy::manual-range-contains -D warnings

safekeeper/Cargo.toml

@@ -10,6 +10,7 @@ anyhow.workspace = true
 async-trait.workspace = true
 byteorder.workspace = true
 bytes.workspace = true
+chrono.workspace = true
 clap = { workspace = true, features = ["derive"] }
 const_format.workspace = true
 crc32c.workspace = true
@@ -18,7 +19,6 @@ git-version.workspace = true
 hex.workspace = true
 humantime.workspace = true
 hyper.workspace = true
-nix.workspace = true
 once_cell.workspace = true
 parking_lot.workspace = true
 postgres.workspace = true

safekeeper/src/bin/safekeeper.rs

@@ -71,6 +71,9 @@ struct Args {
     /// Listen http endpoint for management and metrics in the form host:port.
     #[arg(long, default_value = DEFAULT_HTTP_LISTEN_ADDR)]
     listen_http: String,
+    /// Availability zone of the safekeeper.
+    #[arg(long)]
+    availability_zone: Option<String>,
     /// Do not wait for changes to be written safely to disk. Unsafe.
     #[arg(short, long)]
     no_sync: bool,
@@ -166,6 +169,7 @@ fn main() -> anyhow::Result<()> {
         my_id: id,
         listen_pg_addr: args.listen_pg,
         listen_http_addr: args.listen_http,
+        availability_zone: args.availability_zone,
         no_sync: args.no_sync,
         broker_endpoint: args.broker_endpoint,
         broker_keepalive_interval: args.broker_keepalive_interval,

safekeeper/src/debug_dump.rs

@@ -0,0 +1,264 @@
+//! Utils for dumping full state of the safekeeper.
+
+use std::fs;
+use std::fs::DirEntry;
+use std::io::BufReader;
+use std::io::Read;
+use std::path::PathBuf;
+
+use anyhow::Result;
+use chrono::{DateTime, Utc};
+use postgres_ffi::XLogSegNo;
+use serde::Serialize;
+
+use utils::http::json::display_serialize;
+use utils::id::NodeId;
+use utils::id::TenantTimelineId;
+use utils::id::{TenantId, TimelineId};
+use utils::lsn::Lsn;
+
+use crate::safekeeper::SafeKeeperState;
+use crate::safekeeper::SafekeeperMemState;
+use crate::safekeeper::TermHistory;
+use crate::SafeKeeperConf;
+
+use crate::timeline::ReplicaState;
+use crate::GlobalTimelines;
+
+/// Various filters that influence the resulting JSON output.
+#[derive(Debug, Serialize)]
+pub struct Args {
+    /// Dump all available safekeeper state. False by default.
+    pub dump_all: bool,
+
+    /// Dump control_file content. Uses value of `dump_all` by default.
+    pub dump_control_file: bool,
+
+    /// Dump in-memory state. Uses value of `dump_all` by default.
+    pub dump_memory: bool,
+
+    /// Dump all disk files in a timeline directory. Uses value of `dump_all` by default.
+    pub dump_disk_content: bool,
+
+    /// Dump full term history. True by default.
+    pub dump_term_history: bool,
+
+    /// Filter timelines by tenant_id.
+    pub tenant_id: Option<TenantId>,
+
+    /// Filter timelines by timeline_id.
+    pub timeline_id: Option<TimelineId>,
+}
+
+/// Response for debug dump request.
+#[derive(Debug, Serialize)]
+pub struct Response {
+    pub start_time: DateTime<Utc>,
+    pub finish_time: DateTime<Utc>,
+    pub timelines: Vec<Timeline>,
+    pub timelines_count: usize,
+    pub config: Config,
+}
+
+/// Safekeeper configuration.
+#[derive(Debug, Serialize)]
+pub struct Config {
+    pub id: NodeId,
+    pub workdir: PathBuf,
+    pub listen_pg_addr: String,
+    pub listen_http_addr: String,
+    pub no_sync: bool,
+    pub max_offloader_lag_bytes: u64,
+    pub wal_backup_enabled: bool,
+}
+
+#[derive(Debug, Serialize)]
+pub struct Timeline {
+    #[serde(serialize_with = "display_serialize")]
+    pub tenant_id: TenantId,
+    #[serde(serialize_with = "display_serialize")]
+    pub timeline_id: TimelineId,
+    pub control_file: Option<SafeKeeperState>,
+    pub memory: Option<Memory>,
+    pub disk_content: Option<DiskContent>,
+}
+
+#[derive(Debug, Serialize)]
+pub struct Memory {
+    pub is_cancelled: bool,
+    pub peers_info_len: usize,
+    pub replicas: Vec<Option<ReplicaState>>,
+    pub wal_backup_active: bool,
+    pub active: bool,
+    pub num_computes: u32,
+    pub last_removed_segno: XLogSegNo,
+    pub epoch_start_lsn: Lsn,
+    pub mem_state: SafekeeperMemState,
+
+    // PhysicalStorage state.
+    pub write_lsn: Lsn,
+    pub write_record_lsn: Lsn,
+    pub flush_lsn: Lsn,
+    pub file_open: bool,
+}
+
+#[derive(Debug, Serialize)]
+pub struct DiskContent {
+    pub files: Vec<FileInfo>,
+}
+
+#[derive(Debug, Serialize)]
+pub struct FileInfo {
+    pub name: String,
+    pub size: u64,
+    pub created: DateTime<Utc>,
+    pub modified: DateTime<Utc>,
+    pub start_zeroes: u64,
+    pub end_zeroes: u64,
+    // TODO: add sha256 checksum
+}
+
+/// Build debug dump response, using the provided [`Args`] filters.
+pub fn build(args: Args) -> Result<Response> {
+    let start_time = Utc::now();
+    let timelines_count = GlobalTimelines::timelines_count();
+
+    let ptrs_snapshot = if args.tenant_id.is_some() && args.timeline_id.is_some() {
+        // If both tenant_id and timeline_id are specified, we can just get the
+        // timeline directly, without taking a snapshot of the whole list.
+        let ttid = TenantTimelineId::new(args.tenant_id.unwrap(), args.timeline_id.unwrap());
+        if let Ok(tli) = GlobalTimelines::get(ttid) {
+            vec![tli]
+        } else {
+            vec![]
+        }
+    } else {
+        // Otherwise, take a snapshot of the whole list.
+        GlobalTimelines::get_all()
+    };
+
+    // TODO: return Stream instead of Vec
+    let mut timelines = Vec::new();
+    for tli in ptrs_snapshot {
+        let ttid = tli.ttid;
+        if let Some(tenant_id) = args.tenant_id {
+            if tenant_id != ttid.tenant_id {
+                continue;
+            }
+        }
+        if let Some(timeline_id) = args.timeline_id {
+            if timeline_id != ttid.timeline_id {
+                continue;
+            }
+        }
+
+        let control_file = if args.dump_control_file {
+            let mut state = tli.get_state().1;
+            if !args.dump_term_history {
+                state.acceptor_state.term_history = TermHistory(vec![]);
+            }
+            Some(state)
+        } else {
+            None
+        };
+
+        let memory = if args.dump_memory {
+            Some(tli.memory_dump())
+        } else {
+            None
+        };
+
+        let disk_content = if args.dump_disk_content {
+            // build_disk_content can fail, but we don't want to fail the whole
+            // request because of that.
+            build_disk_content(&tli.timeline_dir).ok()
+        } else {
+            None
+        };
+
+        let timeline = Timeline {
+            tenant_id: ttid.tenant_id,
+            timeline_id: ttid.timeline_id,
+            control_file,
+            memory,
+            disk_content,
+        };
+        timelines.push(timeline);
+    }
+
+    let config = GlobalTimelines::get_global_config();
+
+    Ok(Response {
+        start_time,
+        finish_time: Utc::now(),
+        timelines,
+        timelines_count,
+        config: build_config(config),
+    })
+}
+
+/// Builds DiskContent from a directory path. It can fail if the directory
+/// is deleted between the time we get the path and the time we try to open it.
+fn build_disk_content(path: &std::path::Path) -> Result<DiskContent> {
+    let mut files = Vec::new();
+    for entry in fs::read_dir(path)? {
+        if entry.is_err() {
+            continue;
+        }
+        let file = build_file_info(entry?);
+        if file.is_err() {
+            continue;
+        }
+        files.push(file?);
+    }
+
+    Ok(DiskContent { files })
+}
+
+/// Builds FileInfo from DirEntry. Sometimes it can return an error
+/// if the file is deleted between the time we get the DirEntry
+/// and the time we try to open it.
+fn build_file_info(entry: DirEntry) -> Result<FileInfo> {
+    let metadata = entry.metadata()?;
+    let path = entry.path();
+    let name = path
+        .file_name()
+        .and_then(|x| x.to_str())
+        .unwrap_or("")
+        .to_owned();
+    let mut file = fs::File::open(path)?;
+    let mut reader = BufReader::new(&mut file).bytes().filter_map(|x| x.ok());
+
+    let start_zeroes = reader.by_ref().take_while(|&x| x == 0).count() as u64;
+    let mut end_zeroes = 0;
+    for b in reader {
+        if b == 0 {
+            end_zeroes += 1;
+        } else {
+            end_zeroes = 0;
+        }
+    }
+
+    Ok(FileInfo {
+        name,
+        size: metadata.len(),
+        created: DateTime::from(metadata.created()?),
+        modified: DateTime::from(metadata.modified()?),
+        start_zeroes,
+        end_zeroes,
+    })
+}
+
+/// Converts SafeKeeperConf to Config, filtering out the fields that are not
+/// supposed to be exposed.
+fn build_config(config: SafeKeeperConf) -> Config {
+    Config {
+        id: config.my_id,
+        workdir: config.workdir,
+        listen_pg_addr: config.listen_pg_addr,
+        listen_http_addr: config.listen_http_addr,
+        no_sync: config.no_sync,
+        max_offloader_lag_bytes: config.max_offloader_lag_bytes,
+        wal_backup_enabled: config.wal_backup_enabled,
+    }
+}

safekeeper/src/handler.rs

@@ -3,11 +3,14 @@
 
 use anyhow::Context;
 use std::str;
+use tokio::io::{AsyncRead, AsyncWrite};
 use tracing::{info, info_span, Instrument};
 
 use crate::auth::check_permission;
 use crate::json_ctrl::{handle_json_ctrl, AppendLogicalMessage};
 
+use crate::metrics::TrafficMetrics;
+use crate::wal_service::ConnectionId;
 use crate::{GlobalTimelines, SafeKeeperConf};
 use postgres_backend::QueryError;
 use postgres_backend::{self, PostgresBackend};
@@ -28,7 +31,10 @@ pub struct SafekeeperPostgresHandler {
     pub tenant_id: Option<TenantId>,
     pub timeline_id: Option<TimelineId>,
     pub ttid: TenantTimelineId,
+    /// Unique connection id is logged in spans for observability.
+    pub conn_id: ConnectionId,
     claims: Option<Claims>,
+    io_metrics: Option<TrafficMetrics>,
 }
 
 /// Parsed Postgres command.
@@ -37,11 +43,9 @@ enum SafekeeperPostgresCommand {
     StartReplication { start_lsn: Lsn },
     IdentifySystem,
     JSONCtrl { cmd: AppendLogicalMessage },
-    Show { guc: String },
 }
 
 fn parse_cmd(cmd: &str) -> anyhow::Result<SafekeeperPostgresCommand> {
-    let cmd_lowercase = cmd.to_ascii_lowercase();
     if cmd.starts_with("START_WAL_PUSH") {
         Ok(SafekeeperPostgresCommand::StartWalPush)
     } else if cmd.starts_with("START_REPLICATION") {
@@ -60,25 +64,19 @@ fn parse_cmd(cmd: &str) -> anyhow::Result<SafekeeperPostgresCommand> {
         Ok(SafekeeperPostgresCommand::JSONCtrl {
             cmd: serde_json::from_str(cmd)?,
         })
-    } else if cmd_lowercase.starts_with("show") {
-        let re = Regex::new(r"show ((?:[[:alpha:]]|_)+)").unwrap();
-        let mut caps = re.captures_iter(&cmd_lowercase);
-        let guc = caps
-            .next()
-            .map(|cap| cap[1].parse::<String>())
-            .context("parse guc in SHOW command")??;
-        Ok(SafekeeperPostgresCommand::Show { guc })
     } else {
         anyhow::bail!("unsupported command {cmd}");
     }
 }
 
 #[async_trait::async_trait]
-impl postgres_backend::Handler for SafekeeperPostgresHandler {
+impl<IO: AsyncRead + AsyncWrite + Unpin + Send> postgres_backend::Handler<IO>
+    for SafekeeperPostgresHandler
+{
     // tenant_id and timeline_id are passed in connection string params
     fn startup(
         &mut self,
-        _pgb: &mut PostgresBackend,
+        _pgb: &mut PostgresBackend<IO>,
         sm: &FeStartupPacket,
     ) -> Result<(), QueryError> {
         if let FeStartupPacket::StartupMessage { params, .. } = sm {
@@ -98,6 +96,11 @@ impl postgres_backend::Handler for SafekeeperPostgresHandler {
                                 format!("Failed to parse {value} as timeline id")
                             })?);
                         }
+                        Some(("availability_zone", client_az)) => {
+                            if let Some(metrics) = self.io_metrics.as_ref() {
+                                metrics.set_client_az(client_az)
+                            }
+                        }
                         _ => continue,
                     }
                 }
@@ -105,6 +108,9 @@ impl postgres_backend::Handler for SafekeeperPostgresHandler {
 
             if let Some(app_name) = params.get("application_name") {
                 self.appname = Some(app_name.to_owned());
+                if let Some(metrics) = self.io_metrics.as_ref() {
+                    metrics.set_app_name(app_name)
+                }
             }
 
             Ok(())
@@ -117,7 +123,7 @@ impl postgres_backend::Handler for SafekeeperPostgresHandler {
 
     fn check_auth_jwt(
         &mut self,
-        _pgb: &mut PostgresBackend,
+        _pgb: &mut PostgresBackend<IO>,
         jwt_response: &[u8],
     ) -> Result<(), QueryError> {
         // this unwrap is never triggered, because check_auth_jwt only called when auth_type is NeonJWT
@@ -146,7 +152,7 @@ impl postgres_backend::Handler for SafekeeperPostgresHandler {
 
     async fn process_query(
         &mut self,
-        pgb: &mut PostgresBackend,
+        pgb: &mut PostgresBackend<IO>,
         query_string: &str,
     ) -> Result<(), QueryError> {
         if query_string
@@ -183,7 +189,6 @@ impl postgres_backend::Handler for SafekeeperPostgresHandler {
                     .await
             }
             SafekeeperPostgresCommand::IdentifySystem => self.handle_identify_system(pgb).await,
-            SafekeeperPostgresCommand::Show { guc } => self.handle_show(guc, pgb).await,
             SafekeeperPostgresCommand::JSONCtrl { ref cmd } => {
                 handle_json_ctrl(self, pgb, cmd).await
             }
@@ -192,14 +197,16 @@ impl postgres_backend::Handler for SafekeeperPostgresHandler {
 }
 
 impl SafekeeperPostgresHandler {
-    pub fn new(conf: SafeKeeperConf) -> Self {
+    pub fn new(conf: SafeKeeperConf, conn_id: u32, io_metrics: Option<TrafficMetrics>) -> Self {
         SafekeeperPostgresHandler {
             conf,
             appname: None,
             tenant_id: None,
             timeline_id: None,
             ttid: TenantTimelineId::empty(),
+            conn_id,
             claims: None,
+            io_metrics,
         }
     }
 
@@ -223,11 +230,11 @@ impl SafekeeperPostgresHandler {
     ///
     /// Handle IDENTIFY_SYSTEM replication command
     ///
-    async fn handle_identify_system(
+    async fn handle_identify_system<IO: AsyncRead + AsyncWrite + Unpin>(
         &mut self,
-        pgb: &mut PostgresBackend,
+        pgb: &mut PostgresBackend<IO>,
     ) -> Result<(), QueryError> {
-        let tli = GlobalTimelines::get(self.ttid)?;
+        let tli = GlobalTimelines::get(self.ttid).map_err(|e| QueryError::Other(e.into()))?;
 
         let lsn = if self.is_walproposer_recovery() {
             // walproposer should get all local WAL until flush_lsn
@@ -280,40 +287,6 @@ impl SafekeeperPostgresHandler {
         Ok(())
     }
 
-    async fn handle_show(
-        &mut self,
-        guc: String,
-        pgb: &mut PostgresBackend,
-    ) -> Result<(), QueryError> {
-        match guc.as_str() {
-            // pg_receivewal wants it
-            "data_directory_mode" => {
-                pgb.write_message_noflush(&BeMessage::RowDescription(&[RowDescriptor::int8_col(
-                    b"data_directory_mode",
-                )]))?
-                // xxx we could return real one, not just 0700
-                .write_message_noflush(&BeMessage::DataRow(&[Some("0700".as_bytes())]))?
-                .write_message_noflush(&BeMessage::CommandComplete(b"SELECT 1"))?;
-            }
-            // pg_receivewal wants it
-            "wal_segment_size" => {
-                let tli = GlobalTimelines::get(self.ttid)?;
-                let wal_seg_size = tli.get_state().1.server.wal_seg_size;
-                let wal_seg_size_mb = (wal_seg_size / 1024 / 1024).to_string() + "MB";
-
-                pgb.write_message_noflush(&BeMessage::RowDescription(&[RowDescriptor::text_col(
-                    b"wal_segment_size",
-                )]))?
-                .write_message_noflush(&BeMessage::DataRow(&[Some(wal_seg_size_mb.as_bytes())]))?
-                .write_message_noflush(&BeMessage::CommandComplete(b"SELECT 1"))?;
-            }
-            _ => {
-                return Err(anyhow::anyhow!("SHOW of unknown setting").into());
-            }
-        }
-        Ok(())
-    }
-
     /// Returns true if current connection is a replication connection, originating
     /// from a walproposer recovery function. This connection gets a special handling:
     /// safekeeper must stream all local WAL till the flush_lsn, whether committed or not.

safekeeper/src/http/openapi_spec.yaml

@@ -119,6 +119,12 @@ paths:
           $ref: "#/components/responses/ForbiddenError"
         default:
           $ref: "#/components/responses/GenericError"
+        "404":
+          description: Timeline not found
+          content:
+            application/json:
+              schema:
+                $ref: "#/components/schemas/NotFoundError"
 
     delete:
       tags:

safekeeper/src/http/routes.rs

@@ -1,19 +1,19 @@
 use hyper::{Body, Request, Response, StatusCode, Uri};
 
-use anyhow::Context;
 use once_cell::sync::Lazy;
 use postgres_ffi::WAL_SEGMENT_SIZE;
 use safekeeper_api::models::SkTimelineInfo;
 use serde::Serialize;
-use serde::Serializer;
 use std::collections::{HashMap, HashSet};
-use std::fmt::Display;
-
+use std::fmt;
+use std::str::FromStr;
 use std::sync::Arc;
 use storage_broker::proto::SafekeeperTimelineInfo;
 use storage_broker::proto::TenantTimelineId as ProtoTenantTimelineId;
 use tokio::task::JoinError;
+use utils::http::json::display_serialize;
 
+use crate::debug_dump;
 use crate::safekeeper::ServerInfo;
 use crate::safekeeper::Term;
 
@@ -55,15 +55,6 @@ fn get_conf(request: &Request<Body>) -> &SafeKeeperConf {
         .as_ref()
 }
 
-/// Serialize through Display trait.
-fn display_serialize<S, F>(z: &F, s: S) -> Result<S::Ok, S::Error>
-where
-    S: Serializer,
-    F: Display,
-{
-    s.serialize_str(&format!("{}", z))
-}
-
 /// Same as TermSwitchEntry, but serializes LSN using display serializer
 /// in Postgres format, i.e. 0/FFFFFFFF. Used only for the API response.
 #[derive(Debug, Serialize)]
@@ -120,12 +111,7 @@ async fn timeline_status_handler(request: Request<Body>) -> Result<Response<Body
     );
     check_permission(&request, Some(ttid.tenant_id))?;
 
-    let tli = GlobalTimelines::get(ttid)
-        // FIXME: Currently, the only errors from `GlobalTimelines::get` will be client errors
-        // because the provided timeline isn't there. However, the method can in theory change and
-        // fail from internal errors later. Remove this comment once it the method returns
-        // something other than `anyhow::Result`.
-        .map_err(ApiError::InternalServerError)?;
+    let tli = GlobalTimelines::get(ttid).map_err(ApiError::from)?;
     let (inmem, state) = tli.get_state();
     let flush_lsn = tli.get_flush_lsn();
 
@@ -258,15 +244,7 @@ async fn record_safekeeper_info(mut request: Request<Body>) -> Result<Response<B
         local_start_lsn: sk_info.local_start_lsn.0,
     };
 
-    let tli = GlobalTimelines::get(ttid)
-        // `GlobalTimelines::get` returns an error when it can't find the timeline.
-        .with_context(|| {
-            format!(
-                "Couldn't get timeline {} for tenant {}",
-                ttid.timeline_id, ttid.tenant_id
-            )
-        })
-        .map_err(ApiError::NotFound)?;
+    let tli = GlobalTimelines::get(ttid).map_err(ApiError::from)?;
     tli.record_safekeeper_info(&proto_sk_info)
         .await
         .map_err(ApiError::InternalServerError)?;
@@ -274,6 +252,69 @@ async fn record_safekeeper_info(mut request: Request<Body>) -> Result<Response<B
     json_response(StatusCode::OK, ())
 }
 
+fn parse_kv_str<E: fmt::Display, T: FromStr<Err = E>>(k: &str, v: &str) -> Result<T, ApiError> {
+    v.parse()
+        .map_err(|e| ApiError::BadRequest(anyhow::anyhow!("cannot parse {k}: {e}")))
+}
+
+/// Dump debug info about all available safekeeper state.
+async fn dump_debug_handler(mut request: Request<Body>) -> Result<Response<Body>, ApiError> {
+    check_permission(&request, None)?;
+    ensure_no_body(&mut request).await?;
+
+    let mut dump_all: Option<bool> = None;
+    let mut dump_control_file: Option<bool> = None;
+    let mut dump_memory: Option<bool> = None;
+    let mut dump_disk_content: Option<bool> = None;
+    let mut dump_term_history: Option<bool> = None;
+    let mut tenant_id: Option<TenantId> = None;
+    let mut timeline_id: Option<TimelineId> = None;
+
+    let query = request.uri().query().unwrap_or("");
+    let mut values = url::form_urlencoded::parse(query.as_bytes());
+
+    for (k, v) in &mut values {
+        match k.as_ref() {
+            "dump_all" => dump_all = Some(parse_kv_str(&k, &v)?),
+            "dump_control_file" => dump_control_file = Some(parse_kv_str(&k, &v)?),
+            "dump_memory" => dump_memory = Some(parse_kv_str(&k, &v)?),
+            "dump_disk_content" => dump_disk_content = Some(parse_kv_str(&k, &v)?),
+            "dump_term_history" => dump_term_history = Some(parse_kv_str(&k, &v)?),
+            "tenant_id" => tenant_id = Some(parse_kv_str(&k, &v)?),
+            "timeline_id" => timeline_id = Some(parse_kv_str(&k, &v)?),
+            _ => Err(ApiError::BadRequest(anyhow::anyhow!(
+                "Unknown query parameter: {}",
+                k
+            )))?,
+        }
+    }
+
+    let dump_all = dump_all.unwrap_or(false);
+    let dump_control_file = dump_control_file.unwrap_or(dump_all);
+    let dump_memory = dump_memory.unwrap_or(dump_all);
+    let dump_disk_content = dump_disk_content.unwrap_or(dump_all);
+    let dump_term_history = dump_term_history.unwrap_or(true);
+
+    let args = debug_dump::Args {
+        dump_all,
+        dump_control_file,
+        dump_memory,
+        dump_disk_content,
+        dump_term_history,
+        tenant_id,
+        timeline_id,
+    };
+
+    let resp = tokio::task::spawn_blocking(move || {
+        debug_dump::build(args).map_err(ApiError::InternalServerError)
+    })
+    .await
+    .map_err(|e: JoinError| ApiError::InternalServerError(e.into()))??;
+
+    // TODO: use streaming response
+    json_response(StatusCode::OK, resp)
+}
+
 /// Safekeeper http router.
 pub fn make_router(conf: SafeKeeperConf) -> RouterBuilder<hyper::Body, ApiError> {
     let mut router = endpoint::make_router();
@@ -314,6 +355,7 @@ pub fn make_router(conf: SafeKeeperConf) -> RouterBuilder<hyper::Body, ApiError>
             "/v1/record_safekeeper_info/:tenant_id/:timeline_id",
             record_safekeeper_info,
         )
+        .get("/v1/debug_dump", dump_debug_handler)
 }
 
 #[cfg(test)]

safekeeper/src/json_ctrl.rs

@@ -12,6 +12,7 @@ use anyhow::Context;
 use bytes::Bytes;
 use postgres_backend::QueryError;
 use serde::{Deserialize, Serialize};
+use tokio::io::{AsyncRead, AsyncWrite};
 use tracing::*;
 use utils::id::TenantTimelineId;
 
@@ -60,9 +61,9 @@ struct AppendResult {
 /// Handles command to craft logical message WAL record with given
 /// content, and then append it with specified term and lsn. This
 /// function is used to test safekeepers in different scenarios.
-pub async fn handle_json_ctrl(
+pub async fn handle_json_ctrl<IO: AsyncRead + AsyncWrite + Unpin>(
     spg: &SafekeeperPostgresHandler,
-    pgb: &mut PostgresBackend,
+    pgb: &mut PostgresBackend<IO>,
     append_request: &AppendLogicalMessage,
 ) -> Result<(), QueryError> {
     info!("JSON_CTRL request: {append_request:?}");

safekeeper/src/lib.rs

@@ -1,4 +1,5 @@
 use remote_storage::RemoteStorageConfig;
+
 use std::path::PathBuf;
 use std::time::Duration;
 use storage_broker::Uri;
@@ -9,6 +10,7 @@ mod auth;
 pub mod broker;
 pub mod control_file;
 pub mod control_file_upgrade;
+pub mod debug_dump;
 pub mod handler;
 pub mod http;
 pub mod json_ctrl;
@@ -50,6 +52,7 @@ pub struct SafeKeeperConf {
     pub my_id: NodeId,
     pub listen_pg_addr: String,
     pub listen_http_addr: String,
+    pub availability_zone: Option<String>,
     pub no_sync: bool,
     pub broker_endpoint: Uri,
     pub broker_keepalive_interval: Duration,
@@ -80,6 +83,7 @@ impl SafeKeeperConf {
             no_sync: false,
             listen_pg_addr: defaults::DEFAULT_PG_LISTEN_ADDR.to_string(),
             listen_http_addr: defaults::DEFAULT_HTTP_LISTEN_ADDR.to_string(),
+            availability_zone: None,
             remote_storage: None,
             my_id: NodeId(0),
             broker_endpoint: storage_broker::DEFAULT_ENDPOINT

safekeeper/src/metrics.rs

@@ -1,15 +1,19 @@
 //! Global safekeeper mertics and per-timeline safekeeper metrics.
 
-use std::time::{Instant, SystemTime};
+use std::{
+    sync::{Arc, RwLock},
+    time::{Instant, SystemTime},
+};
 
 use ::metrics::{register_histogram, GaugeVec, Histogram, IntGauge, DISK_WRITE_SECONDS_BUCKETS};
 use anyhow::Result;
 use metrics::{
-    core::{AtomicU64, Collector, Desc, GenericGaugeVec, Opts},
+    core::{AtomicU64, Collector, Desc, GenericCounter, GenericGaugeVec, Opts},
     proto::MetricFamily,
-    Gauge, IntGaugeVec,
+    register_int_counter_vec, Gauge, IntCounterVec, IntGaugeVec,
 };
 use once_cell::sync::Lazy;
+
 use postgres_ffi::XLogSegNo;
 use utils::{id::TenantTimelineId, lsn::Lsn};
 
@@ -61,6 +65,133 @@ pub static PERSIST_CONTROL_FILE_SECONDS: Lazy<Histogram> = Lazy::new(|| {
     )
     .expect("Failed to register safekeeper_persist_control_file_seconds histogram vec")
 });
+pub static PG_IO_BYTES: Lazy<IntCounterVec> = Lazy::new(|| {
+    register_int_counter_vec!(
+        "safekeeper_pg_io_bytes_total",
+        "Bytes read from or written to any PostgreSQL connection",
+        &["client_az", "sk_az", "app_name", "dir", "same_az"]
+    )
+    .expect("Failed to register safekeeper_pg_io_bytes gauge")
+});
+
+pub const LABEL_UNKNOWN: &str = "unknown";
+
+/// Labels for traffic metrics.
+#[derive(Clone)]
+struct ConnectionLabels {
+    /// Availability zone of the connection origin.
+    client_az: String,
+    /// Availability zone of the current safekeeper.
+    sk_az: String,
+    /// Client application name.
+    app_name: String,
+}
+
+impl ConnectionLabels {
+    fn new() -> Self {
+        Self {
+            client_az: LABEL_UNKNOWN.to_string(),
+            sk_az: LABEL_UNKNOWN.to_string(),
+            app_name: LABEL_UNKNOWN.to_string(),
+        }
+    }
+
+    fn build_metrics(
+        &self,
+    ) -> (
+        GenericCounter<metrics::core::AtomicU64>,
+        GenericCounter<metrics::core::AtomicU64>,
+    ) {
+        let same_az = match (self.client_az.as_str(), self.sk_az.as_str()) {
+            (LABEL_UNKNOWN, _) | (_, LABEL_UNKNOWN) => LABEL_UNKNOWN,
+            (client_az, sk_az) => {
+                if client_az == sk_az {
+                    "true"
+                } else {
+                    "false"
+                }
+            }
+        };
+
+        let read = PG_IO_BYTES.with_label_values(&[
+            &self.client_az,
+            &self.sk_az,
+            &self.app_name,
+            "read",
+            same_az,
+        ]);
+        let write = PG_IO_BYTES.with_label_values(&[
+            &self.client_az,
+            &self.sk_az,
+            &self.app_name,
+            "write",
+            same_az,
+        ]);
+        (read, write)
+    }
+}
+
+struct TrafficMetricsState {
+    /// Labels for traffic metrics.
+    labels: ConnectionLabels,
+    /// Total bytes read from this connection.
+    read: GenericCounter<metrics::core::AtomicU64>,
+    /// Total bytes written to this connection.
+    write: GenericCounter<metrics::core::AtomicU64>,
+}
+
+/// Metrics for measuring traffic (r/w bytes) in a single PostgreSQL connection.
+#[derive(Clone)]
+pub struct TrafficMetrics {
+    state: Arc<RwLock<TrafficMetricsState>>,
+}
+
+impl Default for TrafficMetrics {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl TrafficMetrics {
+    pub fn new() -> Self {
+        let labels = ConnectionLabels::new();
+        let (read, write) = labels.build_metrics();
+        let state = TrafficMetricsState {
+            labels,
+            read,
+            write,
+        };
+        Self {
+            state: Arc::new(RwLock::new(state)),
+        }
+    }
+
+    pub fn set_client_az(&self, value: &str) {
+        let mut state = self.state.write().unwrap();
+        state.labels.client_az = value.to_string();
+        (state.read, state.write) = state.labels.build_metrics();
+    }
+
+    pub fn set_sk_az(&self, value: &str) {
+        let mut state = self.state.write().unwrap();
+        state.labels.sk_az = value.to_string();
+        (state.read, state.write) = state.labels.build_metrics();
+    }
+
+    pub fn set_app_name(&self, value: &str) {
+        let mut state = self.state.write().unwrap();
+        state.labels.app_name = value.to_string();
+        (state.read, state.write) = state.labels.build_metrics();
+    }
+
+    pub fn observe_read(&self, cnt: usize) {
+        self.state.read().unwrap().read.inc_by(cnt as u64)
+    }
+
+    pub fn observe_write(&self, cnt: usize) {
+        self.state.read().unwrap().write.inc_by(cnt as u64)
+    }
+}
 
 /// Metrics for WalStorage in a single timeline.
 #[derive(Clone, Default)]

safekeeper/src/receive_wal.rs

@@ -7,10 +7,10 @@ use crate::safekeeper::AcceptorProposerMessage;
 use crate::safekeeper::ProposerAcceptorMessage;
 use crate::safekeeper::ServerInfo;
 use crate::timeline::Timeline;
+use crate::wal_service::ConnectionId;
 use crate::GlobalTimelines;
 use anyhow::{anyhow, Context};
 use bytes::BytesMut;
-use nix::unistd::gettid;
 use postgres_backend::CopyStreamHandlerEnd;
 use postgres_backend::PostgresBackend;
 use postgres_backend::PostgresBackendReader;
@@ -20,6 +20,8 @@ use std::net::SocketAddr;
 use std::sync::Arc;
 use std::thread;
 use std::thread::JoinHandle;
+use tokio::io::AsyncRead;
+use tokio::io::AsyncWrite;
 use tokio::sync::mpsc::channel;
 use tokio::sync::mpsc::error::TryRecvError;
 use tokio::sync::mpsc::Receiver;
@@ -36,9 +38,9 @@ impl SafekeeperPostgresHandler {
     /// Wrapper around handle_start_wal_push_guts handling result. Error is
     /// handled here while we're still in walreceiver ttid span; with API
     /// extension, this can probably be moved into postgres_backend.
-    pub async fn handle_start_wal_push(
+    pub async fn handle_start_wal_push<IO: AsyncRead + AsyncWrite + Unpin>(
         &mut self,
-        pgb: &mut PostgresBackend,
+        pgb: &mut PostgresBackend<IO>,
     ) -> Result<(), QueryError> {
         if let Err(end) = self.handle_start_wal_push_guts(pgb).await {
             // Log the result and probably send it to the client, closing the stream.
@@ -47,9 +49,9 @@ impl SafekeeperPostgresHandler {
         Ok(())
     }
 
-    pub async fn handle_start_wal_push_guts(
+    pub async fn handle_start_wal_push_guts<IO: AsyncRead + AsyncWrite + Unpin>(
         &mut self,
-        pgb: &mut PostgresBackend,
+        pgb: &mut PostgresBackend<IO>,
     ) -> Result<(), CopyStreamHandlerEnd> {
         // Notify the libpq client that it's allowed to send `CopyData` messages
         pgb.write_message(&BeMessage::CopyBothResponse).await?;
@@ -68,10 +70,17 @@ impl SafekeeperPostgresHandler {
         // sends, so this avoids deadlocks.
         let mut pgb_reader = pgb.split().context("START_WAL_PUSH split")?;
         let peer_addr = *pgb.get_peer_addr();
+        let network_reader = NetworkReader {
+            ttid: self.ttid,
+            conn_id: self.conn_id,
+            pgb_reader: &mut pgb_reader,
+            peer_addr,
+            acceptor_handle: &mut acceptor_handle,
+        };
         let res = tokio::select! {
             // todo: add read|write .context to these errors
-            r = read_network(self.ttid, &mut pgb_reader, peer_addr, msg_tx, &mut acceptor_handle, msg_rx, reply_tx) => r,
-            r = write_network(pgb, reply_rx) => r,
+            r = network_reader.run(msg_tx, msg_rx, reply_tx) => r,
+            r = network_write(pgb, reply_rx) => r,
         };
 
         // Join pg backend back.
@@ -104,62 +113,67 @@ impl SafekeeperPostgresHandler {
     }
 }
 
+struct NetworkReader<'a, IO> {
+    ttid: TenantTimelineId,
+    conn_id: ConnectionId,
+    pgb_reader: &'a mut PostgresBackendReader<IO>,
+    peer_addr: SocketAddr,
+    // WalAcceptor is spawned when we learn server info from walproposer and
+    // create timeline; handle is put here.
+    acceptor_handle: &'a mut Option<JoinHandle<anyhow::Result<()>>>,
+}
+
+impl<'a, IO: AsyncRead + AsyncWrite + Unpin> NetworkReader<'a, IO> {
+    async fn run(
+        self,
+        msg_tx: Sender<ProposerAcceptorMessage>,
+        msg_rx: Receiver<ProposerAcceptorMessage>,
+        reply_tx: Sender<AcceptorProposerMessage>,
+    ) -> Result<(), CopyStreamHandlerEnd> {
+        // Receive information about server to create timeline, if not yet.
+        let next_msg = read_message(self.pgb_reader).await?;
+        let tli = match next_msg {
+            ProposerAcceptorMessage::Greeting(ref greeting) => {
+                info!(
+                    "start handshake with walproposer {} sysid {} timeline {}",
+                    self.peer_addr, greeting.system_id, greeting.tli,
+                );
+                let server_info = ServerInfo {
+                    pg_version: greeting.pg_version,
+                    system_id: greeting.system_id,
+                    wal_seg_size: greeting.wal_seg_size,
+                };
+                GlobalTimelines::create(self.ttid, server_info, Lsn::INVALID, Lsn::INVALID).await?
+            }
+            _ => {
+                return Err(CopyStreamHandlerEnd::Other(anyhow::anyhow!(
+                    "unexpected message {next_msg:?} instead of greeting"
+                )))
+            }
+        };
+
+        *self.acceptor_handle = Some(
+            WalAcceptor::spawn(tli.clone(), msg_rx, reply_tx, self.conn_id)
+                .context("spawn WalAcceptor thread")?,
+        );
+
+        // Forward all messages to WalAcceptor
+        read_network_loop(self.pgb_reader, msg_tx, next_msg).await
+    }
+}
+
 /// Read next message from walproposer.
 /// TODO: Return Ok(None) on graceful termination.
-async fn read_message(
-    pgb_reader: &mut PostgresBackendReader,
+async fn read_message<IO: AsyncRead + AsyncWrite + Unpin>(
+    pgb_reader: &mut PostgresBackendReader<IO>,
 ) -> Result<ProposerAcceptorMessage, CopyStreamHandlerEnd> {
     let copy_data = pgb_reader.read_copy_message().await?;
     let msg = ProposerAcceptorMessage::parse(copy_data)?;
     Ok(msg)
 }
 
-/// Read messages from socket and pass it to WalAcceptor thread. Returns Ok(())
-/// if msg_tx closed; it must mean WalAcceptor terminated, joining it should
-/// tell the error.
-async fn read_network(
-    ttid: TenantTimelineId,
-    pgb_reader: &mut PostgresBackendReader,
-    peer_addr: SocketAddr,
-    msg_tx: Sender<ProposerAcceptorMessage>,
-    // WalAcceptor is spawned when we learn server info from walproposer and
-    // create timeline; handle is put here.
-    acceptor_handle: &mut Option<JoinHandle<anyhow::Result<()>>>,
-    msg_rx: Receiver<ProposerAcceptorMessage>,
-    reply_tx: Sender<AcceptorProposerMessage>,
-) -> Result<(), CopyStreamHandlerEnd> {
-    // Receive information about server to create timeline, if not yet.
-    let next_msg = read_message(pgb_reader).await?;
-    let tli = match next_msg {
-        ProposerAcceptorMessage::Greeting(ref greeting) => {
-            info!(
-                "start handshake with walproposer {} sysid {} timeline {}",
-                peer_addr, greeting.system_id, greeting.tli,
-            );
-            let server_info = ServerInfo {
-                pg_version: greeting.pg_version,
-                system_id: greeting.system_id,
-                wal_seg_size: greeting.wal_seg_size,
-            };
-            GlobalTimelines::create(ttid, server_info, Lsn::INVALID, Lsn::INVALID).await?
-        }
-        _ => {
-            return Err(CopyStreamHandlerEnd::Other(anyhow::anyhow!(
-                "unexpected message {next_msg:?} instead of greeting"
-            )))
-        }
-    };
-
-    *acceptor_handle = Some(
-        WalAcceptor::spawn(tli.clone(), msg_rx, reply_tx).context("spawn WalAcceptor thread")?,
-    );
-
-    // Forward all messages to WalAcceptor
-    read_network_loop(pgb_reader, msg_tx, next_msg).await
-}
-
-async fn read_network_loop(
-    pgb_reader: &mut PostgresBackendReader,
+async fn read_network_loop<IO: AsyncRead + AsyncWrite + Unpin>(
+    pgb_reader: &mut PostgresBackendReader<IO>,
     msg_tx: Sender<ProposerAcceptorMessage>,
     mut next_msg: ProposerAcceptorMessage,
 ) -> Result<(), CopyStreamHandlerEnd> {
@@ -174,8 +188,8 @@ async fn read_network_loop(
 /// Read replies from WalAcceptor and pass them back to socket. Returns Ok(())
 /// if reply_rx closed; it must mean WalAcceptor terminated, joining it should
 /// tell the error.
-async fn write_network(
-    pgb_writer: &mut PostgresBackend,
+async fn network_write<IO: AsyncRead + AsyncWrite + Unpin>(
+    pgb_writer: &mut PostgresBackend<IO>,
     mut reply_rx: Receiver<AcceptorProposerMessage>,
 ) -> Result<(), CopyStreamHandlerEnd> {
     let mut buf = BytesMut::with_capacity(128);
@@ -205,6 +219,7 @@ impl WalAcceptor {
         tli: Arc<Timeline>,
         msg_rx: Receiver<ProposerAcceptorMessage>,
         reply_tx: Sender<AcceptorProposerMessage>,
+        conn_id: ConnectionId,
     ) -> anyhow::Result<JoinHandle<anyhow::Result<()>>> {
         let thread_name = format!("WAL acceptor {}", tli.ttid);
         thread::Builder::new()
@@ -223,7 +238,7 @@ impl WalAcceptor {
                 let span_ttid = wa.tli.ttid; // satisfy borrow checker
                 runtime.block_on(
                     wa.run()
-                        .instrument(info_span!("WAL acceptor", tid = %gettid(), ttid = %span_ttid)),
+                        .instrument(info_span!("WAL acceptor", cid = %conn_id, ttid = %span_ttid)),
                 )
             })
             .map_err(anyhow::Error::from)

safekeeper/src/safekeeper.rs

@@ -191,7 +191,8 @@ pub struct SafeKeeperState {
     /// Minimal LSN which may be needed for recovery of some safekeeper (end_lsn
     /// of last record streamed to everyone). Persisting it helps skipping
     /// recovery in walproposer, generally we compute it from peers. In
-    /// walproposer proto called 'truncate_lsn'.
+    /// walproposer proto called 'truncate_lsn'. Updates are currently drived
+    /// only by walproposer.
     pub peer_horizon_lsn: Lsn,
     /// LSN of the oldest known checkpoint made by pageserver and successfully
     /// pushed to s3. We don't remove WAL beyond it. Persisted only for
@@ -204,7 +205,7 @@ pub struct SafeKeeperState {
     pub peers: PersistedPeers,
 }
 
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, Serialize)]
 // In memory safekeeper state. Fields mirror ones in `SafeKeeperState`; values
 // are not flushed yet.
 pub struct SafekeeperMemState {
@@ -212,6 +213,7 @@ pub struct SafekeeperMemState {
     pub backup_lsn: Lsn,
     pub peer_horizon_lsn: Lsn,
     pub remote_consistent_lsn: Lsn,
+    #[serde(with = "hex")]
     pub proposer_uuid: PgUuid,
 }
 
@@ -681,7 +683,7 @@ where
             term: self.state.acceptor_state.term,
             vote_given: false as u64,
             flush_lsn: self.flush_lsn(),
-            truncate_lsn: self.state.peer_horizon_lsn,
+            truncate_lsn: self.inmem.peer_horizon_lsn,
             term_history: self.get_term_history(),
             timeline_start_lsn: self.state.timeline_start_lsn,
         };
@@ -877,7 +879,13 @@ where
         if msg.h.commit_lsn != Lsn(0) {
             self.update_commit_lsn(msg.h.commit_lsn)?;
         }
-        self.inmem.peer_horizon_lsn = msg.h.truncate_lsn;
+        // Value calculated by walproposer can always lag:
+        // - safekeepers can forget inmem value and send to proposer lower
+        //   persisted one on restart;
+        // - if we make safekeepers always send persistent value,
+        //   any compute restart would pull it down.
+        // Thus, take max before adopting.
+        self.inmem.peer_horizon_lsn = max(self.inmem.peer_horizon_lsn, msg.h.truncate_lsn);
 
         // Update truncate and commit LSN in control file.
         // To avoid negative impact on performance of extra fsync, do it only

safekeeper/src/send_wal.rs

@@ -13,6 +13,8 @@ use postgres_ffi::get_current_timestamp;
 use postgres_ffi::{TimestampTz, MAX_SEND_SIZE};
 use pq_proto::{BeMessage, ReplicationFeedback, WalSndKeepAlive, XLogDataBody};
 use serde::{Deserialize, Serialize};
+use tokio::io::{AsyncRead, AsyncWrite};
+
 use std::cmp::min;
 use std::str;
 use std::sync::Arc;
@@ -74,9 +76,9 @@ impl SafekeeperPostgresHandler {
     /// Wrapper around handle_start_replication_guts handling result. Error is
     /// handled here while we're still in walsender ttid span; with API
     /// extension, this can probably be moved into postgres_backend.
-    pub async fn handle_start_replication(
+    pub async fn handle_start_replication<IO: AsyncRead + AsyncWrite + Unpin>(
         &mut self,
-        pgb: &mut PostgresBackend,
+        pgb: &mut PostgresBackend<IO>,
         start_pos: Lsn,
     ) -> Result<(), QueryError> {
         if let Err(end) = self.handle_start_replication_guts(pgb, start_pos).await {
@@ -86,13 +88,14 @@ impl SafekeeperPostgresHandler {
         Ok(())
     }
 
-    pub async fn handle_start_replication_guts(
+    pub async fn handle_start_replication_guts<IO: AsyncRead + AsyncWrite + Unpin>(
         &mut self,
-        pgb: &mut PostgresBackend,
+        pgb: &mut PostgresBackend<IO>,
         start_pos: Lsn,
     ) -> Result<(), CopyStreamHandlerEnd> {
         let appname = self.appname.clone();
-        let tli = GlobalTimelines::get(self.ttid)?;
+        let tli =
+            GlobalTimelines::get(self.ttid).map_err(|e| CopyStreamHandlerEnd::Other(e.into()))?;
 
         let state = ReplicaState::new();
         // This replica_id is used below to check if it's time to stop replication.
@@ -175,8 +178,8 @@ impl SafekeeperPostgresHandler {
 }
 
 /// A half driving sending WAL.
-struct WalSender<'a> {
-    pgb: &'a mut PostgresBackend,
+struct WalSender<'a, IO> {
+    pgb: &'a mut PostgresBackend<IO>,
     tli: Arc<Timeline>,
     appname: Option<String>,
     // Position since which we are sending next chunk.
@@ -193,7 +196,7 @@ struct WalSender<'a> {
     send_buf: [u8; MAX_SEND_SIZE],
 }
 
-impl WalSender<'_> {
+impl<IO: AsyncRead + AsyncWrite + Unpin> WalSender<'_, IO> {
     /// Send WAL until
     /// - an error occurs
     /// - if we are streaming to walproposer, we've streamed until stop_pos
@@ -281,14 +284,14 @@ impl WalSender<'_> {
 }
 
 /// A half driving receiving replies.
-struct ReplyReader {
-    reader: PostgresBackendReader,
+struct ReplyReader<IO> {
+    reader: PostgresBackendReader<IO>,
     tli: Arc<Timeline>,
     replica_id: usize,
     feedback: ReplicaState,
 }
 
-impl ReplyReader {
+impl<IO: AsyncRead + AsyncWrite + Unpin> ReplyReader<IO> {
     async fn run(&mut self) -> Result<(), CopyStreamHandlerEnd> {
         loop {
             let msg = self.reader.read_copy_message().await?;

safekeeper/src/timeline.rs

@@ -1,10 +1,11 @@
 //! This module implements Timeline lifecycle management and has all necessary code
 //! to glue together SafeKeeper and all other background services.
 
-use anyhow::{bail, Result};
+use anyhow::{anyhow, bail, Result};
 use parking_lot::{Mutex, MutexGuard};
 use postgres_ffi::XLogSegNo;
 use pq_proto::ReplicationFeedback;
+use serde::Serialize;
 use std::cmp::{max, min};
 use std::path::PathBuf;
 use tokio::{
@@ -12,6 +13,7 @@ use tokio::{
     time::Instant,
 };
 use tracing::*;
+use utils::http::error::ApiError;
 use utils::{
     id::{NodeId, TenantTimelineId},
     lsn::Lsn,
@@ -28,9 +30,9 @@ use crate::send_wal::HotStandbyFeedback;
 use crate::{control_file, safekeeper::UNKNOWN_SERVER_VERSION};
 
 use crate::metrics::FullTimelineInfo;
-use crate::wal_storage;
 use crate::wal_storage::Storage as wal_storage_iface;
 use crate::SafeKeeperConf;
+use crate::{debug_dump, wal_storage};
 
 /// Things safekeeper should know about timeline state on peers.
 #[derive(Debug, Clone)]
@@ -80,7 +82,7 @@ impl PeersInfo {
 }
 
 /// Replica status update + hot standby feedback
-#[derive(Debug, Clone, Copy)]
+#[derive(Debug, Clone, Copy, Serialize)]
 pub struct ReplicaState {
     /// last known lsn received by replica
     pub last_received_lsn: Lsn, // None means we don't know
@@ -355,6 +357,18 @@ pub enum TimelineError {
     UninitialinzedPgVersion(TenantTimelineId),
 }
 
+// Convert to HTTP API error.
+impl From<TimelineError> for ApiError {
+    fn from(te: TimelineError) -> ApiError {
+        match te {
+            TimelineError::NotFound(ttid) => {
+                ApiError::NotFound(anyhow!("timeline {} not found", ttid))
+            }
+            _ => ApiError::InternalServerError(anyhow!("{}", te)),
+        }
+    }
+}
+
 /// Timeline struct manages lifecycle (creation, deletion, restore) of a safekeeper timeline.
 /// It also holds SharedState and provides mutually exclusive access to it.
 pub struct Timeline {
@@ -381,7 +395,7 @@ pub struct Timeline {
     cancellation_rx: watch::Receiver<bool>,
 
     /// Directory where timeline state is stored.
-    timeline_dir: PathBuf,
+    pub timeline_dir: PathBuf,
 }
 
 impl Timeline {
@@ -593,38 +607,6 @@ impl Timeline {
         self.write_shared_state().wal_backup_attend()
     }
 
-    /// Returns full timeline info, required for the metrics. If the timeline is
-    /// not active, returns None instead.
-    pub fn info_for_metrics(&self) -> Option<FullTimelineInfo> {
-        if self.is_cancelled() {
-            return None;
-        }
-
-        let state = self.write_shared_state();
-        if state.active {
-            Some(FullTimelineInfo {
-                ttid: self.ttid,
-                replicas: state
-                    .replicas
-                    .iter()
-                    .filter_map(|r| r.as_ref())
-                    .copied()
-                    .collect(),
-                wal_backup_active: state.wal_backup_active,
-                timeline_is_active: state.active,
-                num_computes: state.num_computes,
-                last_removed_segno: state.last_removed_segno,
-                epoch_start_lsn: state.sk.epoch_start_lsn,
-                mem_state: state.sk.inmem.clone(),
-                persisted_state: state.sk.state.clone(),
-                flush_lsn: state.sk.wal_store.flush_lsn(),
-                wal_storage: state.sk.wal_store.get_metrics(),
-            })
-        } else {
-            None
-        }
-    }
-
     /// Returns commit_lsn watch channel.
     pub fn get_commit_lsn_watch_rx(&self) -> watch::Receiver<Lsn> {
         self.commit_lsn_watch_rx.clone()
@@ -789,6 +771,62 @@ impl Timeline {
         shared_state.last_removed_segno = horizon_segno;
         Ok(())
     }
+
+    /// Returns full timeline info, required for the metrics. If the timeline is
+    /// not active, returns None instead.
+    pub fn info_for_metrics(&self) -> Option<FullTimelineInfo> {
+        if self.is_cancelled() {
+            return None;
+        }
+
+        let state = self.write_shared_state();
+        if state.active {
+            Some(FullTimelineInfo {
+                ttid: self.ttid,
+                replicas: state
+                    .replicas
+                    .iter()
+                    .filter_map(|r| r.as_ref())
+                    .copied()
+                    .collect(),
+                wal_backup_active: state.wal_backup_active,
+                timeline_is_active: state.active,
+                num_computes: state.num_computes,
+                last_removed_segno: state.last_removed_segno,
+                epoch_start_lsn: state.sk.epoch_start_lsn,
+                mem_state: state.sk.inmem.clone(),
+                persisted_state: state.sk.state.clone(),
+                flush_lsn: state.sk.wal_store.flush_lsn(),
+                wal_storage: state.sk.wal_store.get_metrics(),
+            })
+        } else {
+            None
+        }
+    }
+
+    /// Returns in-memory timeline state to build a full debug dump.
+    pub fn memory_dump(&self) -> debug_dump::Memory {
+        let state = self.write_shared_state();
+
+        let (write_lsn, write_record_lsn, flush_lsn, file_open) =
+            state.sk.wal_store.internal_state();
+
+        debug_dump::Memory {
+            is_cancelled: self.is_cancelled(),
+            peers_info_len: state.peers_info.0.len(),
+            replicas: state.replicas.clone(),
+            wal_backup_active: state.wal_backup_active,
+            active: state.active,
+            num_computes: state.num_computes,
+            last_removed_segno: state.last_removed_segno,
+            epoch_start_lsn: state.sk.epoch_start_lsn,
+            mem_state: state.sk.inmem.clone(),
+            write_lsn,
+            write_record_lsn,
+            flush_lsn,
+            file_open,
+        }
+    }
 }
 
 /// Deletes directory and it's contents. Returns false if directory does not exist.

safekeeper/src/timelines_global_map.rs

@@ -5,7 +5,7 @@
 use crate::safekeeper::ServerInfo;
 use crate::timeline::{Timeline, TimelineError};
 use crate::SafeKeeperConf;
-use anyhow::{anyhow, bail, Context, Result};
+use anyhow::{bail, Context, Result};
 use once_cell::sync::Lazy;
 use serde::Serialize;
 use std::collections::HashMap;
@@ -50,11 +50,11 @@ impl GlobalTimelinesState {
     }
 
     /// Get timeline from the map. Returns error if timeline doesn't exist.
-    fn get(&self, ttid: &TenantTimelineId) -> Result<Arc<Timeline>> {
+    fn get(&self, ttid: &TenantTimelineId) -> Result<Arc<Timeline>, TimelineError> {
         self.timelines
             .get(ttid)
             .cloned()
-            .ok_or_else(|| anyhow!(TimelineError::NotFound(*ttid)))
+            .ok_or(TimelineError::NotFound(*ttid))
     }
 }
 
@@ -159,6 +159,16 @@ impl GlobalTimelines {
         Ok(())
     }
 
+    /// Get the number of timelines in the map.
+    pub fn timelines_count() -> usize {
+        TIMELINES_STATE.lock().unwrap().timelines.len()
+    }
+
+    /// Get the global safekeeper config.
+    pub fn get_global_config() -> SafeKeeperConf {
+        TIMELINES_STATE.lock().unwrap().get_conf().clone()
+    }
+
     /// Create a new timeline with the given id. If the timeline already exists, returns
     /// an existing timeline.
     pub async fn create(
@@ -226,17 +236,17 @@ impl GlobalTimelines {
     /// Get a timeline from the global map. If it's not present, it doesn't exist on disk,
     /// or was corrupted and couldn't be loaded on startup. Returned timeline is always valid,
     /// i.e. loaded in memory and not cancelled.
-    pub fn get(ttid: TenantTimelineId) -> Result<Arc<Timeline>> {
+    pub fn get(ttid: TenantTimelineId) -> Result<Arc<Timeline>, TimelineError> {
         let res = TIMELINES_STATE.lock().unwrap().get(&ttid);
 
         match res {
             Ok(tli) => {
                 if tli.is_cancelled() {
-                    anyhow::bail!(TimelineError::Cancelled(ttid));
+                    return Err(TimelineError::Cancelled(ttid));
                 }
                 Ok(tli)
             }
-            Err(e) => Err(e),
+            _ => res,
         }
     }
 

safekeeper/src/wal_service.rs

@@ -3,13 +3,14 @@
 //!   receive WAL from wal_proposer and send it to WAL receivers
 //!
 use anyhow::{Context, Result};
-use nix::unistd::gettid;
 use postgres_backend::QueryError;
 use std::{future, thread};
 use tokio::net::TcpStream;
 use tracing::*;
+use utils::measured_stream::MeasuredStream;
 
 use crate::handler::SafekeeperPostgresHandler;
+use crate::metrics::TrafficMetrics;
 use crate::SafeKeeperConf;
 use postgres_backend::{AuthType, PostgresBackend};
 
@@ -27,17 +28,19 @@ pub fn thread_main(conf: SafeKeeperConf, pg_listener: std::net::TcpListener) {
             // Tokio's from_std won't do this for us, per its comment.
             pg_listener.set_nonblocking(true)?;
             let listener = tokio::net::TcpListener::from_std(pg_listener)?;
+            let mut connection_count: ConnectionCount = 0;
 
             loop {
                 match listener.accept().await {
                     Ok((socket, peer_addr)) => {
                         debug!("accepted connection from {}", peer_addr);
                         let conf = conf.clone();
+                        let conn_id = issue_connection_id(&mut connection_count);
 
                         let _ = thread::Builder::new()
                             .name("WAL service thread".into())
                             .spawn(move || {
-                                if let Err(err) = handle_socket(socket, conf) {
+                                if let Err(err) = handle_socket(socket, conf, conn_id) {
                                     error!("connection handler exited: {}", err);
                                 }
                             })
@@ -54,8 +57,12 @@ pub fn thread_main(conf: SafeKeeperConf, pg_listener: std::net::TcpListener) {
 
 /// This is run by `thread_main` above, inside a background thread.
 ///
-fn handle_socket(socket: TcpStream, conf: SafeKeeperConf) -> Result<(), QueryError> {
-    let _enter = info_span!("", tid = %gettid()).entered();
+fn handle_socket(
+    socket: TcpStream,
+    conf: SafeKeeperConf,
+    conn_id: ConnectionId,
+) -> Result<(), QueryError> {
+    let _enter = info_span!("", cid = %conn_id).entered();
 
     let runtime = tokio::runtime::Builder::new_current_thread()
         .enable_all()
@@ -63,13 +70,30 @@ fn handle_socket(socket: TcpStream, conf: SafeKeeperConf) -> Result<(), QueryErr
     let local = tokio::task::LocalSet::new();
 
     socket.set_nodelay(true)?;
+    let peer_addr = socket.peer_addr()?;
+
+    let traffic_metrics = TrafficMetrics::new();
+    if let Some(current_az) = conf.availability_zone.as_deref() {
+        traffic_metrics.set_sk_az(current_az);
+    }
+
+    let socket = MeasuredStream::new(
+        socket,
+        |cnt| {
+            traffic_metrics.observe_read(cnt);
+        },
+        |cnt| {
+            traffic_metrics.observe_write(cnt);
+        },
+    );
 
     let auth_type = match conf.auth {
         None => AuthType::Trust,
         Some(_) => AuthType::NeonJWT,
     };
-    let mut conn_handler = SafekeeperPostgresHandler::new(conf);
-    let pgbackend = PostgresBackend::new(socket, auth_type, None)?;
+    let mut conn_handler =
+        SafekeeperPostgresHandler::new(conf, conn_id, Some(traffic_metrics.clone()));
+    let pgbackend = PostgresBackend::new_from_io(socket, peer_addr, auth_type, None)?;
     // libpq protocol between safekeeper and walproposer / pageserver
     // We don't use shutdown.
     local.block_on(
@@ -79,3 +103,12 @@ fn handle_socket(socket: TcpStream, conf: SafeKeeperConf) -> Result<(), QueryErr
 
     Ok(())
 }
+
+/// Unique WAL service connection ids are logged in spans for observability.
+pub type ConnectionId = u32;
+pub type ConnectionCount = u32;
+
+pub fn issue_connection_id(count: &mut ConnectionCount) -> ConnectionId {
+    *count = count.wrapping_add(1);
+    *count
+}

safekeeper/src/wal_storage.rs

@@ -165,6 +165,16 @@ impl PhysicalStorage {
         })
     }
 
+    /// Get all known state of the storage.
+    pub fn internal_state(&self) -> (Lsn, Lsn, Lsn, bool) {
+        (
+            self.write_lsn,
+            self.write_record_lsn,
+            self.flush_record_lsn,
+            self.file.is_some(),
+        )
+    }
+
     /// Call fdatasync if config requires so.
     fn fdatasync_file(&mut self, file: &mut File) -> Result<()> {
         if !self.conf.no_sync {

scripts/export_import_between_pageservers.py

@@ -308,8 +308,8 @@ def lsn_to_hex(num: int) -> str:
 
 def lsn_from_hex(lsn_hex: str) -> int:
     """Convert lsn from hex notation to int."""
-    l, r = lsn_hex.split("/")
-    return (int(l, 16) << 32) + int(r, 16)
+    left, right = lsn_hex.split("/")
+    return (int(left, 16) << 32) + int(right, 16)
 
 
 def remote_consistent_lsn(
@@ -398,7 +398,6 @@ def reconstruct_paths(log_dir, pg_bin, base_tar, port: int):
                 result = vanilla_pg.safe_psql(query, user="cloud_admin", dbname=database)
                 for relname, filepath in result:
                     if filepath is not None:
-
                         if database == "template0copy":
                             # Add all template0copy paths to template0
                             prefix = f"base/{oid}/"