Add arg to override config in zenith_cli

.circleci/ansible/ansible.cfg

@@ -1,10 +0,0 @@
-[defaults]
-
-localhost_warning = False
-host_key_checking = False
-timeout = 30
-
-[ssh_connection]
-ssh_args   = -F ./ansible.ssh.cfg
-scp_if_ssh = True
-pipelining = True

.circleci/ansible/ansible.ssh.cfg

@@ -1,11 +0,0 @@
-Host tele.zenith.tech
-    User admin
-    Port 3023
-    StrictHostKeyChecking no
-    UserKnownHostsFile /dev/null
-
-Host * !tele.zenith.tech
-    User admin
-    StrictHostKeyChecking no
-    UserKnownHostsFile /dev/null
-    ProxyJump tele.zenith.tech

.circleci/ansible/deploy.yaml

@@ -1,174 +0,0 @@
-- name: Upload Zenith binaries
-  hosts: pageservers:safekeepers
-  gather_facts: False
-  remote_user: admin
-  vars:
-    force_deploy: false
-
-  tasks:
-
-    - name: get latest version of Zenith binaries
-      ignore_errors: true
-      register: current_version_file
-      set_fact:
-        current_version: "{{ lookup('file', '.zenith_current_version') | trim }}"
-      tags:
-      - pageserver
-      - safekeeper
-
-    - name: set zero value for current_version
-      when: current_version_file is failed
-      set_fact:
-        current_version: "0"
-      tags:
-      - pageserver
-      - safekeeper
-
-    - name: get deployed version from content of remote file
-      ignore_errors: true
-      ansible.builtin.slurp:
-        src: /usr/local/.zenith_current_version
-      register: remote_version_file
-      tags:
-      - pageserver
-      - safekeeper
-
-    - name: decode remote file content
-      when: remote_version_file is succeeded
-      set_fact:
-        remote_version: "{{ remote_version_file['content'] | b64decode | trim }}"
-      tags:
-      - pageserver
-      - safekeeper
-
-    - name: set zero value for remote_version
-      when: remote_version_file is failed
-      set_fact:
-        remote_version: "0"
-      tags:
-      - pageserver
-      - safekeeper
-
-    - name: inform about versions
-      debug: msg="Version to deploy - {{ current_version }}, version on storage node - {{ remote_version }}"
-      tags:
-      - pageserver
-      - safekeeper
-
-
-    - name: upload and extract Zenith binaries to /usr/local
-      when: current_version > remote_version or force_deploy
-      ansible.builtin.unarchive:
-        owner: root
-        group: root
-        src: zenith_install.tar.gz
-        dest: /usr/local
-      become: true
-      tags:
-      - pageserver
-      - safekeeper
-      - binaries
-      - putbinaries
-
-- name: Deploy pageserver
-  hosts: pageservers
-  gather_facts: False
-  remote_user: admin
-  vars:
-    force_deploy: false
-
-  tasks:
-    - name: init pageserver
-      when: current_version > remote_version or force_deploy
-      shell:
-        cmd: sudo -u pageserver /usr/local/bin/pageserver -c "pg_distrib_dir='/usr/local'" --init -D /storage/pageserver/data
-      args:
-        creates: "/storage/pageserver/data/tenants"
-      environment:
-        ZENITH_REPO_DIR: "/storage/pageserver/data"
-        LD_LIBRARY_PATH: "/usr/local/lib"
-      become: true
-      tags:
-      - pageserver
-
-    - name: upload systemd service definition
-      when: current_version > remote_version or force_deploy
-      ansible.builtin.template:
-        src: systemd/pageserver.service
-        dest: /etc/systemd/system/pageserver.service
-        owner: root
-        group: root
-        mode: '0644'
-      become: true
-      tags:
-      - pageserver
-
-    - name: start systemd service
-      when: current_version > remote_version or force_deploy
-      ansible.builtin.systemd:
-        daemon_reload: yes
-        name: pageserver
-        enabled: yes
-        state: restarted
-      become: true
-      tags:
-      - pageserver
-
-    - name: post version to console
-      when: (current_version > remote_version or force_deploy) and console_mgmt_base_url is defined
-      shell:
-        cmd: |
-          INSTANCE_ID=$(curl -s http://169.254.169.254/latest/meta-data/instance-id)
-          curl -sfS -d '{"version": {{ current_version }} }' -X POST {{ console_mgmt_base_url }}/api/v1/pageservers/$INSTANCE_ID
-      tags:
-      - pageserver
-
-- name: Deploy safekeeper
-  hosts: safekeepers
-  gather_facts: False
-  remote_user: admin
-  vars:
-    force_deploy: false
-
-  tasks:
-
-    # in the future safekeepers should discover pageservers byself
-    # but currently use first pageserver that was discovered
-    - name: set first pageserver var for safekeepers
-      when: current_version > remote_version or force_deploy
-      set_fact:
-        first_pageserver: "{{ hostvars[groups['pageservers'][0]]['inventory_hostname'] }}"
-      tags:
-      - safekeeper
-
-    - name: upload systemd service definition
-      when: current_version > remote_version or force_deploy
-      ansible.builtin.template:
-        src: systemd/safekeeper.service
-        dest: /etc/systemd/system/safekeeper.service
-        owner: root
-        group: root
-        mode: '0644'
-      become: true
-      tags:
-      - safekeeper
-
-    - name: start systemd service
-      when: current_version > remote_version or force_deploy
-      ansible.builtin.systemd:
-        daemon_reload: yes
-        name: safekeeper
-        enabled: yes
-        state: restarted
-      become: true
-      tags:
-      - safekeeper
-
-    - name: post version to console
-      when: (current_version > remote_version or force_deploy) and console_mgmt_base_url is defined
-      shell:
-        cmd: |
-          INSTANCE_ID=$(curl -s http://169.254.169.254/latest/meta-data/instance-id)
-          curl -sfS -d '{"version": {{ current_version }} }' -X POST {{ console_mgmt_base_url }}/api/v1/safekeepers/$INSTANCE_ID
-      tags:
-      - safekeeper

.circleci/ansible/get_binaries.sh

@@ -1,52 +0,0 @@
-#!/bin/bash
-
-set -e
-
-RELEASE=${RELEASE:-false}
-
-# look at docker hub for latest tag fo zenith docker image
-if [ "${RELEASE}" = "true" ]; then
-    echo "search latest relase tag"
-    VERSION=$(curl -s https://registry.hub.docker.com/v1/repositories/zenithdb/zenith/tags |jq -r -S '.[].name' | grep release | sed 's/release-//g' | tail -1)
-    if [ -z "${VERSION}" ]; then
-        echo "no any docker tags found, exiting..."
-        exit 1
-    else
-        TAG="release-${VERSION}"
-    fi
-else
-    echo "search latest dev tag"
-    VERSION=$(curl -s https://registry.hub.docker.com/v1/repositories/zenithdb/zenith/tags |jq -r -S '.[].name' | grep -v release | tail -1)
-    if [ -z "${VERSION}" ]; then
-        echo "no any docker tags found, exiting..."
-        exit 1
-    else
-        TAG="${VERSION}"
-    fi
-fi
-
-echo "found ${VERSION}"
-
-# do initial cleanup
-rm -rf zenith_install postgres_install.tar.gz zenith_install.tar.gz .zenith_current_version
-mkdir zenith_install
-
-# retrive binaries from docker image
-echo "getting binaries from docker image"
-docker pull --quiet zenithdb/zenith:${TAG}
-ID=$(docker create zenithdb/zenith:${TAG})
-docker cp ${ID}:/data/postgres_install.tar.gz .
-tar -xzf postgres_install.tar.gz -C zenith_install
-docker cp ${ID}:/usr/local/bin/pageserver zenith_install/bin/
-docker cp ${ID}:/usr/local/bin/safekeeper zenith_install/bin/
-docker cp ${ID}:/usr/local/bin/proxy zenith_install/bin/
-docker cp ${ID}:/usr/local/bin/postgres zenith_install/bin/
-docker rm -vf ${ID}
-
-# store version to file (for ansible playbooks) and create binaries tarball
-echo ${VERSION} > zenith_install/.zenith_current_version
-echo ${VERSION} > .zenith_current_version
-tar -czf zenith_install.tar.gz -C zenith_install .
-
-# do final cleaup
-rm -rf zenith_install postgres_install.tar.gz

.circleci/ansible/production.hosts

@@ -1,7 +0,0 @@
-[pageservers]
-zenith-1-ps-1
-
-[safekeepers]
-zenith-1-sk-1
-zenith-1-sk-2
-zenith-1-sk-3

.circleci/ansible/staging.hosts

@@ -1,7 +0,0 @@
-[pageservers]
-zenith-us-stage-ps-1
-
-[safekeepers]
-zenith-us-stage-sk-1
-zenith-us-stage-sk-2
-zenith-us-stage-sk-3

.circleci/ansible/systemd/pageserver.service

@@ -1,18 +0,0 @@
-[Unit]
-Description=Zenith pageserver
-After=network.target auditd.service
-
-[Service]
-Type=simple
-User=pageserver
-Environment=RUST_BACKTRACE=1 ZENITH_REPO_DIR=/storage/pageserver LD_LIBRARY_PATH=/usr/local/lib
-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'" -D /storage/pageserver/data
-ExecReload=/bin/kill -HUP $MAINPID
-KillMode=mixed
-KillSignal=SIGINT
-Restart=on-failure
-TimeoutSec=10
-LimitNOFILE=30000000
-
-[Install]
-WantedBy=multi-user.target

.circleci/ansible/systemd/safekeeper.service

@@ -1,18 +0,0 @@
-[Unit]
-Description=Zenith safekeeper
-After=network.target auditd.service
-
-[Service]
-Type=simple
-User=safekeeper
-Environment=RUST_BACKTRACE=1 ZENITH_REPO_DIR=/storage/safekeeper/data LD_LIBRARY_PATH=/usr/local/lib
-ExecStart=/usr/local/bin/safekeeper -l {{ inventory_hostname }}.local:6500 --listen-http {{ inventory_hostname }}.local:7676 -p {{ first_pageserver }}:6400 -D /storage/safekeeper/data
-ExecReload=/bin/kill -HUP $MAINPID
-KillMode=mixed
-KillSignal=SIGINT
-Restart=on-failure
-TimeoutSec=10
-LimitNOFILE=30000000
-
-[Install]
-WantedBy=multi-user.target

.circleci/config.yml

@@ -54,8 +54,7 @@ jobs:
             if [ ! -e tmp_install/bin/postgres ]; then
               # "depth 1" saves some time by not cloning the whole repo
               git submodule update --init --depth 1
-              # bail out on any warnings
-              COPT='-Werror' mold -run make postgres -j$(nproc)
+              make postgres -j$(nproc)
             fi
 
       - save_cache:
@@ -111,7 +110,7 @@ jobs:
             fi
 
             export CARGO_INCREMENTAL=0
-            "${cov_prefix[@]}" mold -run cargo build $CARGO_FLAGS --bins --tests
+            "${cov_prefix[@]}" cargo build $CARGO_FLAGS --bins --tests
 
       - save_cache:
           name: Save rust cache
@@ -195,14 +194,6 @@ jobs:
           command: |
             cp -a tmp_install /tmp/zenith/pg_install
 
-      - run:
-          name: Merge coverage data
-          command: |
-            # This will speed up workspace uploads
-            if [[ $BUILD_TYPE == "debug" ]]; then
-              scripts/coverage "--profraw-prefix=$CIRCLE_JOB" --dir=/tmp/zenith/coverage merge
-            fi
-
         # Save the rust binaries and coverage data for other jobs in this workflow.
       - persist_to_workspace:
           root: /tmp/zenith
@@ -213,16 +204,9 @@ jobs:
     executor: zenith-executor
     steps:
       - checkout
-      - restore_cache:
-          keys:
-            - v1-python-deps-{{ checksum "poetry.lock" }}
       - run:
           name: Install deps
           command: ./scripts/pysync
-      - save_cache:
-          key: v1-python-deps-{{ checksum "poetry.lock" }}
-          paths:
-            - /home/circleci/.cache/pypoetry/virtualenvs
       - run:
           name: Run yapf to ensure code format
           when: always
@@ -272,16 +256,9 @@ jobs:
           condition: << parameters.needs_postgres_source >>
           steps:
             - run: git submodule update --init --depth 1
-      - restore_cache:
-          keys:
-            - v1-python-deps-{{ checksum "poetry.lock" }}
       - run:
           name: Install deps
           command: ./scripts/pysync
-      - save_cache:
-          key: v1-python-deps-{{ checksum "poetry.lock" }}
-          paths:
-            - /home/circleci/.cache/pypoetry/virtualenvs
       - run:
           name: Run pytest
           # pytest doesn't output test logs in real time, so CI job may fail with
@@ -298,7 +275,6 @@ jobs:
             - PLATFORM: zenith-local-ci
           command: |
             PERF_REPORT_DIR="$(realpath test_runner/perf-report-local)"
-            rm -rf $PERF_REPORT_DIR
 
             TEST_SELECTION="test_runner/<< parameters.test_selection >>"
             EXTRA_PARAMS="<< parameters.extra_params >>"
@@ -343,6 +319,7 @@ jobs:
 
             if << parameters.save_perf_report >>; then
               if [[ $CIRCLE_BRANCH == "main" ]]; then
+                # TODO: reuse scripts/git-upload
                 export REPORT_FROM="$PERF_REPORT_DIR"
                 export REPORT_TO=local
                 scripts/generate_and_push_perf_report.sh
@@ -363,13 +340,6 @@ jobs:
       # The store_test_results step tells CircleCI where to find the junit.xml file.
       - store_test_results:
           path: /tmp/test_output
-      - run:
-          name: Merge coverage data
-          command: |
-            # This will speed up workspace uploads
-            if [[ $BUILD_TYPE == "debug" ]]; then
-              scripts/coverage "--profraw-prefix=$CIRCLE_JOB" --dir=/tmp/zenith/coverage merge
-            fi
       # Save coverage data (if any)
       - persist_to_workspace:
           root: /tmp/zenith
@@ -471,78 +441,46 @@ jobs:
             docker build -t zenithdb/compute-node:latest vendor/postgres && docker push zenithdb/compute-node:latest
             docker tag zenithdb/compute-node:latest zenithdb/compute-node:${DOCKER_TAG} && docker push zenithdb/compute-node:${DOCKER_TAG}
 
-  # Build production zenithdb/zenith:release image and push it to Docker hub
-  docker-image-release:
-    docker:
-      - image: cimg/base:2021.04
-    steps:
-      - checkout
-      - setup_remote_docker:
-          docker_layer_caching: true
-      - run:
-          name: Init postgres submodule
-          command: git submodule update --init --depth 1
-      - run:
-          name: Build and push Docker image
-          command: |
-            echo $DOCKER_PWD | docker login -u $DOCKER_LOGIN --password-stdin
-            DOCKER_TAG="release-$(git log --oneline|wc -l)"
-            docker build --build-arg GIT_VERSION=$CIRCLE_SHA1 -t zenithdb/zenith:release . && docker push zenithdb/zenith:release
-            docker tag zenithdb/zenith:release zenithdb/zenith:${DOCKER_TAG} && docker push zenithdb/zenith:${DOCKER_TAG}
-
-  # Build production zenithdb/compute-node:release image and push it to Docker hub
-  docker-image-compute-release:
-    docker:
-      - image: cimg/base:2021.04
-    steps:
-      - checkout
-      - setup_remote_docker:
-          docker_layer_caching: true
-      # Build zenithdb/compute-tools:release image and push it to Docker hub
-      # TODO: this should probably also use versioned tag, not just :latest.
-      # XXX: but should it? We build and use it only locally now.
-      - run:
-          name: Build and push compute-tools Docker image
-          command: |
-            echo $DOCKER_PWD | docker login -u $DOCKER_LOGIN --password-stdin
-            docker build -t zenithdb/compute-tools:release -f Dockerfile.compute-tools .
-            docker push zenithdb/compute-tools:release
-      - run:
-          name: Init postgres submodule
-          command: git submodule update --init --depth 1
-      - run:
-          name: Build and push compute-node Docker image
-          command: |
-            echo $DOCKER_PWD | docker login -u $DOCKER_LOGIN --password-stdin
-            DOCKER_TAG="release-$(git log --oneline|wc -l)"
-            docker build -t zenithdb/compute-node:release vendor/postgres && docker push zenithdb/compute-node:release
-            docker tag zenithdb/compute-node:release zenithdb/compute-node:${DOCKER_TAG} && docker push zenithdb/compute-node:${DOCKER_TAG}
-
   deploy-staging:
     docker:
       - image: cimg/python:3.10
     steps:
       - checkout
       - setup_remote_docker
+      - run:
+          name: Get Zenith binaries
+          command: |
+            rm -rf zenith_install postgres_install.tar.gz zenith_install.tar.gz
+            mkdir zenith_install
+            DOCKER_TAG=$(git log --oneline|wc -l)
+            docker pull --quiet zenithdb/zenith:${DOCKER_TAG}
+            ID=$(docker create zenithdb/zenith:${DOCKER_TAG})
+            docker cp $ID:/data/postgres_install.tar.gz .
+            tar -xzf postgres_install.tar.gz -C zenith_install && rm postgres_install.tar.gz
+            docker cp $ID:/usr/local/bin/pageserver zenith_install/bin/
+            docker cp $ID:/usr/local/bin/safekeeper zenith_install/bin/
+            docker cp $ID:/usr/local/bin/proxy zenith_install/bin/
+            docker cp $ID:/usr/local/bin/postgres zenith_install/bin/
+            docker rm -v $ID
+            echo ${DOCKER_TAG} | tee zenith_install/.zenith_current_version
+            tar -czf zenith_install.tar.gz -C zenith_install .
+            ls -la zenith_install.tar.gz
       - run:
           name: Setup ansible
           command: |
             pip install --progress-bar off --user ansible boto3
+            ansible-galaxy collection install amazon.aws
       - run:
-          name: Redeploy
+          name: Apply re-deploy playbook
+          environment:
+            ANSIBLE_HOST_KEY_CHECKING: false
           command: |
-            cd "$(pwd)/.circleci/ansible"
-
-            ./get_binaries.sh
-
-            echo "${TELEPORT_SSH_KEY}"  | tr -d '\n'| base64 --decode >ssh-key
-            echo "${TELEPORT_SSH_CERT}" | tr -d '\n'| base64 --decode >ssh-key-cert.pub
-            chmod 0600 ssh-key
-            ssh-add ssh-key
-            rm -f ssh-key ssh-key-cert.pub
-
-            ansible-playbook deploy.yaml -i staging.hosts
-            rm -f zenith_install.tar.gz .zenith_current_version
+            echo "${STAGING_SSH_KEY}" | base64 --decode | ssh-add -
+            export AWS_REGION=${STAGING_AWS_REGION}
+            export AWS_ACCESS_KEY_ID=${STAGING_AWS_ACCESS_KEY_ID}
+            export AWS_SECRET_ACCESS_KEY=${STAGING_AWS_SECRET_ACCESS_KEY}
+            ansible-playbook .circleci/storage-redeploy.playbook.yml
+            rm -f zenith_install.tar.gz
 
   deploy-staging-proxy:
     docker:
@@ -565,57 +503,7 @@ jobs:
           name: Re-deploy proxy
           command: |
             DOCKER_TAG=$(git log --oneline|wc -l)
-            helm upgrade zenith-proxy zenithdb/zenith-proxy --install -f .circleci/helm-values/staging.proxy.yaml --set image.tag=${DOCKER_TAG} --wait
-
-
-  deploy-release:
-    docker:
-      - image: cimg/python:3.10
-    steps:
-      - checkout
-      - setup_remote_docker
-      - run:
-          name: Setup ansible
-          command: |
-            pip install --progress-bar off --user ansible boto3
-      - run:
-          name: Redeploy
-          command: |
-            cd "$(pwd)/.circleci/ansible"
-
-            RELEASE=true ./get_binaries.sh
-
-            echo "${TELEPORT_SSH_KEY}"  | tr -d '\n'| base64 --decode >ssh-key
-            echo "${TELEPORT_SSH_CERT}" | tr -d '\n'| base64 --decode >ssh-key-cert.pub
-            chmod 0600 ssh-key
-            ssh-add ssh-key
-            rm -f ssh-key ssh-key-cert.pub
-
-            ansible-playbook deploy.yaml -i production.hosts -e console_mgmt_base_url=http://console-release.local
-            rm -f zenith_install.tar.gz .zenith_current_version
-
-  deploy-release-proxy:
-    docker:
-      - image: cimg/base:2021.04
-    environment:
-      KUBECONFIG: .kubeconfig
-    steps:
-      - checkout
-      - run:
-          name: Store kubeconfig file
-          command: |
-            echo "${PRODUCTION_KUBECONFIG_DATA}" | base64 --decode > ${KUBECONFIG}
-            chmod 0600 ${KUBECONFIG}
-      - run:
-          name: Setup helm v3
-          command: |
-            curl -s https://raw.githubusercontent.com/helm/helm/main/scripts/get-helm-3 | bash
-            helm repo add zenithdb https://zenithdb.github.io/helm-charts
-      - run:
-          name: Re-deploy proxy
-          command: |
-            DOCKER_TAG="release-$(git log --oneline|wc -l)"
-            helm upgrade zenith-proxy zenithdb/zenith-proxy --install -f .circleci/helm-values/production.proxy.yaml --set image.tag=${DOCKER_TAG} --wait
+            helm upgrade zenith-proxy zenithdb/zenith-proxy --install -f .circleci/proxy.staging.yaml --set image.tag=${DOCKER_TAG} --wait
 
   # Trigger a new remote CI job
   remote-ci-trigger:
@@ -653,7 +541,7 @@ jobs:
             --user "$CI_ACCESS_TOKEN" \
             --data \
               "{
-                \"ref\": \"feature/668-coverage-report\",
+                \"ref\": \"main\",
                 \"inputs\": {
                   \"ci_job_name\": \"zenith-remote-ci\",
                   \"commit_hash\": \"$CIRCLE_SHA1\",
@@ -680,7 +568,6 @@ workflows:
             - build-postgres-<< matrix.build_type >>
       - run-pytest:
           name: pg_regress-tests-<< matrix.build_type >>
-          context: PERF_TEST_RESULT_CONNSTR
           matrix:
             parameters:
               build_type: ["debug", "release"]
@@ -698,7 +585,6 @@ workflows:
             - build-zenith-<< matrix.build_type >>
       - run-pytest:
           name: benchmarks
-          context: PERF_TEST_RESULT_CONNSTR
           build_type: release
           test_selection: performance
           run_in_parallel: false
@@ -751,47 +637,6 @@ workflows:
                 - main
           requires:
             - docker-image
-
-      - docker-image-release:
-          # Context gives an ability to login
-          context: Docker Hub
-          # Build image only for commits to main
-          filters:
-            branches:
-              only:
-                - release
-          requires:
-            - pg_regress-tests-release
-            - other-tests-release
-      - docker-image-compute-release:
-          # Context gives an ability to login
-          context: Docker Hub
-          # Build image only for commits to main
-          filters:
-            branches:
-              only:
-                - release
-          requires:
-            - pg_regress-tests-release
-            - other-tests-release
-      - deploy-release:
-          # Context gives an ability to login
-          context: Docker Hub
-          # deploy only for commits to main
-          filters:
-            branches:
-              only:
-                - release
-          requires:
-            - docker-image-release
-      - deploy-release-proxy:
-          # deploy only for commits to main
-          filters:
-            branches:
-              only:
-                - release
-          requires:
-            - docker-image-release
       - remote-ci-trigger:
           # Context passes credentials for gh api
           context: CI_ACCESS_TOKEN

.circleci/helm-values/production.proxy.yaml

@@ -1,35 +0,0 @@
-# Helm chart values for zenith-proxy.
-# This is a YAML-formatted file.
-
-settings:
-  authEndpoint: "https://console.zenith.tech/authenticate_proxy_request/"
-  uri: "https://console.zenith.tech/psql_session/"
-
-# -- Additional labels for zenith-proxy pods
-podLabels:
-  zenith_service: proxy
-  zenith_env: production
-  zenith_region: us-west-2
-  zenith_region_slug: oregon
-
-service:
-  annotations:
-    service.beta.kubernetes.io/aws-load-balancer-type: external
-    service.beta.kubernetes.io/aws-load-balancer-nlb-target-type: ip
-    service.beta.kubernetes.io/aws-load-balancer-scheme: internal
-    external-dns.alpha.kubernetes.io/hostname: proxy-release.local
-  type: LoadBalancer
-
-exposedService:
-  annotations:
-    service.beta.kubernetes.io/aws-load-balancer-type: external
-    service.beta.kubernetes.io/aws-load-balancer-nlb-target-type: ip
-    service.beta.kubernetes.io/aws-load-balancer-scheme: internet-facing
-    external-dns.alpha.kubernetes.io/hostname: start.zenith.tech
-
-metrics:
-  enabled: true
-  serviceMonitor:
-    enabled: true
-    selector:
-      release: kube-prometheus-stack

.circleci/proxy.staging.yaml

@@ -5,13 +5,6 @@ settings:
   authEndpoint: "https://console.stage.zenith.tech/authenticate_proxy_request/"
   uri: "https://console.stage.zenith.tech/psql_session/"
 
-# -- Additional labels for zenith-proxy pods
-podLabels:
-  zenith_service: proxy
-  zenith_env: staging
-  zenith_region: us-east-1
-  zenith_region_slug: virginia
-
 exposedService:
   annotations:
     service.beta.kubernetes.io/aws-load-balancer-type: external
@@ -24,4 +17,4 @@ metrics:
   serviceMonitor:
     enabled: true
     selector:
-      release: kube-prometheus-stack
+      prometheus: zenith

.circleci/storage-redeploy.playbook.yml

@@ -0,0 +1,138 @@
+- name: discover storage nodes
+  hosts: localhost
+  connection: local
+  gather_facts: False
+
+  tasks:
+
+    - name: discover safekeepers
+      no_log: true
+      ec2_instance_info:
+        filters:
+          "tag:zenith_env": "staging"
+          "tag:zenith_service": "safekeeper"
+      register: ec2_safekeepers
+
+    - name: discover pageservers
+      no_log: true
+      ec2_instance_info:
+        filters:
+          "tag:zenith_env": "staging"
+          "tag:zenith_service": "pageserver"
+      register: ec2_pageservers
+
+    - name: add safekeepers to host group
+      no_log: true
+      add_host:
+        name: safekeeper-{{ ansible_loop.index }}
+        ansible_host: "{{ item.public_ip_address }}"
+        groups:
+          - storage
+          - safekeepers
+      with_items: "{{ ec2_safekeepers.instances }}"
+      loop_control:
+        extended: yes
+
+    - name: add pageservers to host group
+      no_log: true
+      add_host:
+        name: pageserver-{{ ansible_loop.index }}
+        ansible_host: "{{ item.public_ip_address }}"
+        groups:
+          - storage
+          - pageservers
+      with_items: "{{ ec2_pageservers.instances }}"
+      loop_control:
+        extended: yes
+
+- name: Retrive versions
+  hosts: storage
+  gather_facts: False
+  remote_user: admin
+
+  tasks:
+
+    - name: Get current version of binaries
+      set_fact:
+        current_version: "{{lookup('file', '../zenith_install/.zenith_current_version') }}"
+
+    - name: Check that file with version exists on host
+      stat:
+        path: /usr/local/.zenith_current_version
+      register: version_file
+
+    - name: Try to get current version from the host
+      when: version_file.stat.exists
+      ansible.builtin.fetch:
+        src: /usr/local/.zenith_current_version
+        dest: .remote_version.{{ inventory_hostname }}
+        fail_on_missing: no
+        flat: yes
+
+    - name: Store remote version to variable
+      when: version_file.stat.exists
+      set_fact:
+        remote_version: "{{ lookup('file', '.remote_version.{{ inventory_hostname }}') }}"
+
+    - name: Store default value of remote version to variable in case when remote version file not found
+      when: not version_file.stat.exists
+      set_fact:
+        remote_version: "000"
+
+- name: Extract Zenith binaries
+  hosts: storage
+  gather_facts: False
+  remote_user: admin
+
+  tasks:
+
+    - name: Inform about version conflict
+      when: current_version <= remote_version
+      debug: msg="Current version {{ current_version }} LE than remote {{ remote_version }}"
+
+    - name: Extract Zenith binaries to /usr/local
+      when: current_version > remote_version
+      ansible.builtin.unarchive:
+        src: ../zenith_install.tar.gz
+        dest: /usr/local
+      become: true
+
+- name: Restart safekeepers
+  hosts: safekeepers
+  gather_facts: False
+  remote_user: admin
+
+  tasks:
+
+    - name: Inform about version conflict
+      when: current_version <= remote_version
+      debug: msg="Current version {{ current_version }} LE than remote {{ remote_version }}"
+
+    - name: Restart systemd service
+      when: current_version > remote_version
+      ansible.builtin.systemd:
+        daemon_reload: yes
+        name: safekeeper
+        enabled: yes
+        state: restarted
+      become: true
+
+- name: Restart pageservers
+  hosts: pageservers
+  gather_facts: False
+  remote_user: admin
+
+  tasks:
+
+    - name: Inform about version conflict
+      when: current_version <= remote_version
+      debug: msg="Current version {{ current_version }} LE than remote {{ remote_version }}"
+
+    - name: Restart systemd service
+      when: current_version > remote_version
+      ansible.builtin.systemd:
+        daemon_reload: yes
+        name: pageserver
+        enabled: yes
+        state: restarted
+      become: true

.github/workflows/benchmarking.yml

@@ -3,7 +3,7 @@ name: benchmarking
 on:
   # uncomment to run on push for debugging your PR
   # push:
-  #   branches: [ your branch ]
+  #   branches: [ mybranch ]
   schedule:
     # * is a special character in YAML so you have to quote this string
     #          ┌───────────── minute (0 - 59)
@@ -41,7 +41,7 @@ jobs:
       run: |
         python3 -m pip install --upgrade poetry wheel
         # since pip/poetry caches are reused there shouldn't be any troubles with install every time
-        ./scripts/pysync
+        poetry install
 
     - name: Show versions
       run: |
@@ -89,15 +89,11 @@ jobs:
         BENCHMARK_CONNSTR: "${{ secrets.BENCHMARK_STAGING_CONNSTR }}"
         REMOTE_ENV: "1" # indicate to test harness that we do not have zenith binaries locally
       run: |
-        # just to be sure that no data was cached on self hosted runner
-        # since it might generate duplicates when calling ingest_perf_test_result.py
-        rm -rf perf-report-staging
         mkdir -p perf-report-staging
         ./scripts/pytest test_runner/performance/ -v -m "remote_cluster" --skip-interfering-proc-check --out-dir perf-report-staging
 
     - name: Submit result
       env:
         VIP_VAP_ACCESS_TOKEN: "${{ secrets.VIP_VAP_ACCESS_TOKEN }}"
-        PERF_TEST_RESULT_CONNSTR: "${{ secrets.PERF_TEST_RESULT_CONNSTR }}"
       run: |
         REPORT_FROM=$(realpath perf-report-staging) REPORT_TO=staging scripts/generate_and_push_perf_report.sh

Cargo.toml

@@ -16,8 +16,3 @@ members = [
 # This is useful for profiling and, to some extent, debug.
 # Besides, debug info should not affect the performance.
 debug = true
-
-# This is only needed for proxy's tests
-# TODO: we should probably fork tokio-postgres-rustls instead
-[patch.crates-io]
-tokio-postgres = { git = "https://github.com/zenithdb/rust-postgres.git", rev="2949d98df52587d562986aad155dd4e889e408b7" }

compute_tools/Cargo.toml

@@ -1,14 +1,17 @@
 [package]
 name = "compute_tools"
 version = "0.1.0"
-edition = "2021"
+authors = ["Alexey Kondratov <kondratov.aleksey@gmail.com>"]
+edition = "2018"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
 libc = "0.2"
 anyhow = "1.0"
 chrono = "0.4"
-clap = "3.0"
-env_logger = "0.9"
+clap = "2.33"
+env_logger = "0.8"
 hyper = { version = "0.14", features = ["full"] }
 log = { version = "0.4", features = ["std", "serde"] }
 postgres = { git = "https://github.com/zenithdb/rust-postgres.git", rev="9eb0dbfbeb6a6c1b79099b9f7ae4a8c021877858" }

compute_tools/src/bin/zenith_ctl.rs

@@ -34,7 +34,6 @@ use std::sync::{Arc, RwLock};
 
 use anyhow::{Context, Result};
 use chrono::Utc;
-use clap::Arg;
 use log::info;
 use postgres::{Client, NoTls};
 
@@ -163,34 +162,34 @@ fn main() -> Result<()> {
     let matches = clap::App::new("zenith_ctl")
         .version(version.unwrap_or("unknown"))
         .arg(
-            Arg::new("connstr")
-                .short('C')
+            clap::Arg::with_name("connstr")
+                .short("C")
                 .long("connstr")
                 .value_name("DATABASE_URL")
                 .required(true),
         )
         .arg(
-            Arg::new("pgdata")
-                .short('D')
+            clap::Arg::with_name("pgdata")
+                .short("D")
                 .long("pgdata")
                 .value_name("DATADIR")
                 .required(true),
         )
         .arg(
-            Arg::new("pgbin")
-                .short('b')
+            clap::Arg::with_name("pgbin")
+                .short("b")
                 .long("pgbin")
                 .value_name("POSTGRES_PATH"),
         )
         .arg(
-            Arg::new("spec")
-                .short('s')
+            clap::Arg::with_name("spec")
+                .short("s")
                 .long("spec")
                 .value_name("SPEC_JSON"),
         )
         .arg(
-            Arg::new("spec-path")
-                .short('S')
+            clap::Arg::with_name("spec-path")
+                .short("S")
                 .long("spec-path")
                 .value_name("SPEC_PATH"),
         )

control_plane/Cargo.toml

@@ -1,11 +1,14 @@
 [package]
 name = "control_plane"
 version = "0.1.0"
-edition = "2021"
+authors = ["Stas Kelvich <stas@zenith.tech>"]
+edition = "2018"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
 tar = "0.4.33"
-postgres = { git = "https://github.com/zenithdb/rust-postgres.git", rev="2949d98df52587d562986aad155dd4e889e408b7" }
+postgres = { git = "https://github.com/zenithdb/rust-postgres.git", rev="9eb0dbfbeb6a6c1b79099b9f7ae4a8c021877858" }
 serde = { version = "1.0", features = ["derive"] }
 toml = "0.5"
 lazy_static = "1.4"

control_plane/safekeepers.conf

@@ -1,7 +1,7 @@
 # Page server and three safekeepers.
 [pageserver]
-listen_pg_addr = '127.0.0.1:64000'
-listen_http_addr = '127.0.0.1:9898'
+listen_pg_addr = 'localhost:64000'
+listen_http_addr = 'localhost:9898'
 auth_type = 'Trust'
 
 [[safekeepers]]

control_plane/simple.conf

@@ -1,8 +1,8 @@
 # Minimal zenith environment with one safekeeper. This is equivalent to the built-in
 # defaults that you get with no --config
 [pageserver]
-listen_pg_addr = '127.0.0.1:64000'
-listen_http_addr = '127.0.0.1:9898'
+listen_pg_addr = 'localhost:64000'
+listen_http_addr = 'localhost:9898'
 auth_type = 'Trust'
 
 [[safekeepers]]

control_plane/src/compute.rs

@@ -334,26 +334,14 @@ impl PostgresNode {
         if let Some(lsn) = self.lsn {
             conf.append("recovery_target_lsn", &lsn.to_string());
         }
-
         conf.append_line("");
-        // Configure backpressure
-        // - Replication write lag depends on how fast the walreceiver can process incoming WAL.
-        //   This lag determines latency of get_page_at_lsn. Speed of applying WAL is about 10MB/sec,
-        //   so to avoid expiration of 1 minute timeout, this lag should not be larger than 600MB.
-        //   Actually latency should be much smaller (better if < 1sec). But we assume that recently
-        //   updates pages are not requested from pageserver.
-        // - Replication flush lag depends on speed of persisting data by checkpointer (creation of
-        //   delta/image layers) and advancing disk_consistent_lsn. Safekeepers are able to
-        //   remove/archive WAL only beyond disk_consistent_lsn. Too large a lag can cause long
-        //   recovery time (in case of pageserver crash) and disk space overflow at safekeepers.
-        // - Replication apply lag depends on speed of uploading changes to S3 by uploader thread.
-        //   To be able to restore database in case of pageserver node crash, safekeeper should not
-        //   remove WAL beyond this point. Too large lag can cause space exhaustion in safekeepers
-        //   (if they are not able to upload WAL to S3).
-        conf.append("max_replication_write_lag", "500MB");
-        conf.append("max_replication_flush_lag", "10GB");
 
         if !self.env.safekeepers.is_empty() {
+            // Configure backpressure
+            // In setup with safekeepers apply_lag depends on
+            // speed of data checkpointing on pageserver (see disk_consistent_lsn).
+            conf.append("max_replication_apply_lag", "1500MB");
+
             // Configure the node to connect to the safekeepers
             conf.append("synchronous_standby_names", "walproposer");
 
@@ -366,6 +354,11 @@ impl PostgresNode {
                 .join(",");
             conf.append("wal_acceptors", &wal_acceptors);
         } else {
+            // Configure backpressure
+            // In setup without safekeepers, flush_lag depends on
+            // speed of of data checkpointing on pageserver (see disk_consistent_lsn)
+            conf.append("max_replication_flush_lag", "1500MB");
+
             // We only use setup without safekeepers for tests,
             // and don't care about data durability on pageserver,
             // so set more relaxed synchronous_commit.

control_plane/src/local_env.rs

@@ -189,20 +189,26 @@ impl LocalEnv {
     }
 
     /// Locate and load config
-    pub fn load_config() -> anyhow::Result<Self> {
+    pub fn load_config(override_config: Option<&str>) -> anyhow::Result<Self> {
         let repopath = base_path();
 
-        if !repopath.exists() {
-            bail!(
-                "Zenith config is not found in {}. You need to run 'zenith init' first",
-                repopath.to_str().unwrap()
-            );
-        }
-
-        // TODO: check that it looks like a zenith repository
+        let config_path = if let Some(override_config) = override_config {
+            PathBuf::from(override_config)
+        } else {
+            if !repopath.exists() {
+                bail!(
+                    "Zenith config is not found in {}. You need to run 'zenith init' first",
+                    repopath.to_str().unwrap()
+                );
+            }
+    
+            // TODO: check that it looks like a zenith repository
+    
+            repopath.join("config")
+        };
 
         // load and parse file
-        let config = fs::read_to_string(repopath.join("config"))?;
+        let config = fs::read_to_string(config_path)?;
         let mut env: LocalEnv = toml::from_str(config.as_str())?;
 
         env.base_data_dir = repopath;

control_plane/src/safekeeper.rs

@@ -85,7 +85,7 @@ impl SafekeeperNode {
             pg_connection_config: Self::safekeeper_connection_config(conf.pg_port),
             env: env.clone(),
             http_client: Client::new(),
-            http_base_url: format!("http://127.0.0.1:{}/v1", conf.http_port),
+            http_base_url: format!("http://localhost:{}/v1", conf.http_port),
             pageserver,
         }
     }
@@ -93,7 +93,7 @@ impl SafekeeperNode {
     /// Construct libpq connection string for connecting to this safekeeper.
     fn safekeeper_connection_config(port: u16) -> Config {
         // TODO safekeeper authentication not implemented yet
-        format!("postgresql://no_user@127.0.0.1:{}/no_db", port)
+        format!("postgresql://no_user@localhost:{}/no_db", port)
             .parse()
             .unwrap()
     }
@@ -114,8 +114,8 @@ impl SafekeeperNode {
         );
         io::stdout().flush().unwrap();
 
-        let listen_pg = format!("127.0.0.1:{}", self.conf.pg_port);
-        let listen_http = format!("127.0.0.1:{}", self.conf.http_port);
+        let listen_pg = format!("localhost:{}", self.conf.pg_port);
+        let listen_http = format!("localhost:{}", self.conf.http_port);
 
         let mut cmd = Command::new(self.env.safekeeper_bin()?);
         fill_rust_env_vars(

pageserver/Cargo.toml

@@ -1,7 +1,8 @@
 [package]
 name = "pageserver"
 version = "0.1.0"
-edition = "2021"
+authors = ["Stas Kelvich <stas@zenith.tech>"]
+edition = "2018"
 
 [dependencies]
 bookfile = { git = "https://github.com/zenithdb/bookfile.git", branch="generic-readext" }
@@ -14,14 +15,15 @@ futures = "0.3.13"
 hyper = "0.14"
 lazy_static = "1.4.0"
 log = "0.4.14"
-clap = "3.0"
+clap = "2.33.0"
 daemonize = "0.4.1"
 tokio = { version = "1.11", features = ["process", "sync", "macros", "fs", "rt", "io-util", "time"] }
-postgres-types = { git = "https://github.com/zenithdb/rust-postgres.git", rev="2949d98df52587d562986aad155dd4e889e408b7" }
-postgres-protocol = { git = "https://github.com/zenithdb/rust-postgres.git", rev="2949d98df52587d562986aad155dd4e889e408b7" }
-postgres = { git = "https://github.com/zenithdb/rust-postgres.git", rev="2949d98df52587d562986aad155dd4e889e408b7" }
-tokio-postgres = { git = "https://github.com/zenithdb/rust-postgres.git", rev="2949d98df52587d562986aad155dd4e889e408b7" }
+postgres-types = { git = "https://github.com/zenithdb/rust-postgres.git", rev="9eb0dbfbeb6a6c1b79099b9f7ae4a8c021877858" }
+postgres-protocol = { git = "https://github.com/zenithdb/rust-postgres.git", rev="9eb0dbfbeb6a6c1b79099b9f7ae4a8c021877858" }
+postgres = { git = "https://github.com/zenithdb/rust-postgres.git", rev="9eb0dbfbeb6a6c1b79099b9f7ae4a8c021877858" }
+tokio-postgres = { git = "https://github.com/zenithdb/rust-postgres.git", rev="9eb0dbfbeb6a6c1b79099b9f7ae4a8c021877858" }
 tokio-stream = "0.1.8"
+routerify = "2"
 anyhow = { version = "1.0", features = ["backtrace"] }
 crc32c = "0.6.0"
 thiserror = "1.0"
@@ -30,7 +32,7 @@ tar = "0.4.33"
 humantime = "2.1.0"
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1"
-toml_edit = { version = "0.13", features = ["easy"] }
+toml_edit = { version = "0.12", features = ["easy"] }
 scopeguard = "1.1.0"
 async-trait = "0.1"
 const_format = "0.2.21"
@@ -40,8 +42,8 @@ signal-hook = "0.3.10"
 url = "2"
 nix = "0.23"
 once_cell = "1.8.0"
+parking_lot = "0.11.2"
 crossbeam-utils = "0.8.5"
-fail = "0.5.0"
 
 rust-s3 = { version = "0.28", default-features = false, features = ["no-verify-ssl", "tokio-rustls-tls"] }
 async-compression = {version = "0.3", features = ["zstd", "tokio"]}

pageserver/src/bin/dump_layerfile.rs

@@ -13,7 +13,7 @@ fn main() -> Result<()> {
         .about("Dump contents of one layer file, for debugging")
         .version(GIT_VERSION)
         .arg(
-            Arg::new("path")
+            Arg::with_name("path")
                 .help("Path to file to dump")
                 .required(true)
                 .index(1),

pageserver/src/bin/pageserver.rs

@@ -27,27 +27,27 @@ fn main() -> Result<()> {
         .about("Materializes WAL stream to pages and serves them to the postgres")
         .version(GIT_VERSION)
         .arg(
-            Arg::new("daemonize")
-                .short('d')
+            Arg::with_name("daemonize")
+                .short("d")
                 .long("daemonize")
                 .takes_value(false)
                 .help("Run in the background"),
         )
         .arg(
-            Arg::new("init")
+            Arg::with_name("init")
                 .long("init")
                 .takes_value(false)
                 .help("Initialize pageserver repo"),
         )
         .arg(
-            Arg::new("workdir")
-                .short('D')
+            Arg::with_name("workdir")
+                .short("D")
                 .long("workdir")
                 .takes_value(true)
                 .help("Working directory for the pageserver"),
         )
         .arg(
-            Arg::new("create-tenant")
+            Arg::with_name("create-tenant")
                 .long("create-tenant")
                 .takes_value(true)
                 .help("Create tenant during init")
@@ -55,11 +55,11 @@ fn main() -> Result<()> {
         )
         // See `settings.md` for more details on the extra configuration patameters pageserver can process
         .arg(
-            Arg::new("config-override")
-                .short('c')
+            Arg::with_name("config-override")
+                .short("c")
                 .takes_value(true)
                 .number_of_values(1)
-                .multiple_occurrences(true)
+                .multiple(true)
                 .help("Additional configuration overrides of the ones from the toml config file (or new ones to add there).
                 Any option has to be a valid toml document, example: `-c \"foo='hey'\"` `-c \"foo={value=1}\"`"),
         )

pageserver/src/bin/pageserver_zst.rs

@@ -1,334 +0,0 @@
-//! A CLI helper to deal with remote storage (S3, usually) blobs as archives.
-//! See [`compression`] for more details about the archives.
-
-use std::{collections::BTreeSet, path::Path};
-
-use anyhow::{bail, ensure, Context};
-use clap::{App, Arg};
-use pageserver::{
-    layered_repository::metadata::{TimelineMetadata, METADATA_FILE_NAME},
-    remote_storage::compression,
-};
-use tokio::{fs, io};
-use zenith_utils::GIT_VERSION;
-
-const LIST_SUBCOMMAND: &str = "list";
-const ARCHIVE_ARG_NAME: &str = "archive";
-
-const EXTRACT_SUBCOMMAND: &str = "extract";
-const TARGET_DIRECTORY_ARG_NAME: &str = "target_directory";
-
-const CREATE_SUBCOMMAND: &str = "create";
-const SOURCE_DIRECTORY_ARG_NAME: &str = "source_directory";
-
-#[tokio::main(flavor = "current_thread")]
-async fn main() -> anyhow::Result<()> {
-    let arg_matches = App::new("pageserver zst blob [un]compressor utility")
-        .version(GIT_VERSION)
-        .subcommands(vec![
-            App::new(LIST_SUBCOMMAND)
-                .about("List the archive contents")
-                .arg(
-                    Arg::new(ARCHIVE_ARG_NAME)
-                        .required(true)
-                        .takes_value(true)
-                        .help("An archive to list the contents of"),
-                ),
-            App::new(EXTRACT_SUBCOMMAND)
-                .about("Extracts the archive into the directory")
-                .arg(
-                    Arg::new(ARCHIVE_ARG_NAME)
-                        .required(true)
-                        .takes_value(true)
-                        .help("An archive to extract"),
-                )
-                .arg(
-                    Arg::new(TARGET_DIRECTORY_ARG_NAME)
-                        .required(false)
-                        .takes_value(true)
-                        .help("A directory to extract the archive into. Optional, will use the current directory if not specified"),
-                ),
-            App::new(CREATE_SUBCOMMAND)
-                .about("Creates an archive with the contents of a directory (only the first level files are taken, metadata file has to be present in the same directory)")
-                .arg(
-                    Arg::new(SOURCE_DIRECTORY_ARG_NAME)
-                        .required(true)
-                        .takes_value(true)
-                        .help("A directory to use for creating the archive"),
-                )
-                .arg(
-                    Arg::new(TARGET_DIRECTORY_ARG_NAME)
-                        .required(false)
-                        .takes_value(true)
-                        .help("A directory to create the archive in. Optional, will use the current directory if not specified"),
-                ),
-        ])
-        .get_matches();
-
-    let subcommand_name = match arg_matches.subcommand_name() {
-        Some(name) => name,
-        None => bail!("No subcommand specified"),
-    };
-
-    let subcommand_matches = match arg_matches.subcommand_matches(subcommand_name) {
-        Some(matches) => matches,
-        None => bail!(
-            "No subcommand arguments were recognized for subcommand '{}'",
-            subcommand_name
-        ),
-    };
-
-    let target_dir = Path::new(
-        subcommand_matches
-            .value_of(TARGET_DIRECTORY_ARG_NAME)
-            .unwrap_or("./"),
-    );
-
-    match subcommand_name {
-        LIST_SUBCOMMAND => {
-            let archive = match subcommand_matches.value_of(ARCHIVE_ARG_NAME) {
-                Some(archive) => Path::new(archive),
-                None => bail!("No '{}' argument is specified", ARCHIVE_ARG_NAME),
-            };
-            list_archive(archive).await
-        }
-        EXTRACT_SUBCOMMAND => {
-            let archive = match subcommand_matches.value_of(ARCHIVE_ARG_NAME) {
-                Some(archive) => Path::new(archive),
-                None => bail!("No '{}' argument is specified", ARCHIVE_ARG_NAME),
-            };
-            extract_archive(archive, target_dir).await
-        }
-        CREATE_SUBCOMMAND => {
-            let source_dir = match subcommand_matches.value_of(SOURCE_DIRECTORY_ARG_NAME) {
-                Some(source) => Path::new(source),
-                None => bail!("No '{}' argument is specified", SOURCE_DIRECTORY_ARG_NAME),
-            };
-            create_archive(source_dir, target_dir).await
-        }
-        unknown => bail!("Unknown subcommand {}", unknown),
-    }
-}
-
-async fn list_archive(archive: &Path) -> anyhow::Result<()> {
-    let archive = archive.canonicalize().with_context(|| {
-        format!(
-            "Failed to get the absolute path for the archive path '{}'",
-            archive.display()
-        )
-    })?;
-    ensure!(
-        archive.is_file(),
-        "Path '{}' is not an archive file",
-        archive.display()
-    );
-    println!("Listing an archive at path '{}'", archive.display());
-    let archive_name = match archive.file_name().and_then(|name| name.to_str()) {
-        Some(name) => name,
-        None => bail!(
-            "Failed to get the archive name from the path '{}'",
-            archive.display()
-        ),
-    };
-
-    let archive_bytes = fs::read(&archive)
-        .await
-        .context("Failed to read the archive bytes")?;
-
-    let header = compression::read_archive_header(archive_name, &mut archive_bytes.as_slice())
-        .await
-        .context("Failed to read the archive header")?;
-
-    let empty_path = Path::new("");
-    println!("-------------------------------");
-
-    let longest_path_in_archive = header
-        .files
-        .iter()
-        .filter_map(|file| Some(file.subpath.as_path(empty_path).to_str()?.len()))
-        .max()
-        .unwrap_or_default()
-        .max(METADATA_FILE_NAME.len());
-
-    for regular_file in &header.files {
-        println!(
-            "File: {:width$} uncompressed size: {} bytes",
-            regular_file.subpath.as_path(empty_path).display(),
-            regular_file.size,
-            width = longest_path_in_archive,
-        )
-    }
-    println!(
-        "File: {:width$} uncompressed size: {} bytes",
-        METADATA_FILE_NAME,
-        header.metadata_file_size,
-        width = longest_path_in_archive,
-    );
-    println!("-------------------------------");
-
-    Ok(())
-}
-
-async fn extract_archive(archive: &Path, target_dir: &Path) -> anyhow::Result<()> {
-    let archive = archive.canonicalize().with_context(|| {
-        format!(
-            "Failed to get the absolute path for the archive path '{}'",
-            archive.display()
-        )
-    })?;
-    ensure!(
-        archive.is_file(),
-        "Path '{}' is not an archive file",
-        archive.display()
-    );
-    let archive_name = match archive.file_name().and_then(|name| name.to_str()) {
-        Some(name) => name,
-        None => bail!(
-            "Failed to get the archive name from the path '{}'",
-            archive.display()
-        ),
-    };
-
-    if !target_dir.exists() {
-        fs::create_dir_all(target_dir).await.with_context(|| {
-            format!(
-                "Failed to create the target dir at path '{}'",
-                target_dir.display()
-            )
-        })?;
-    }
-    let target_dir = target_dir.canonicalize().with_context(|| {
-        format!(
-            "Failed to get the absolute path for the target dir path '{}'",
-            target_dir.display()
-        )
-    })?;
-    ensure!(
-        target_dir.is_dir(),
-        "Path '{}' is not a directory",
-        target_dir.display()
-    );
-    let mut dir_contents = fs::read_dir(&target_dir)
-        .await
-        .context("Failed to list the target directory contents")?;
-    let dir_entry = dir_contents
-        .next_entry()
-        .await
-        .context("Failed to list the target directory contents")?;
-    ensure!(
-        dir_entry.is_none(),
-        "Target directory '{}' is not empty",
-        target_dir.display()
-    );
-
-    println!(
-        "Extracting an archive at path '{}' into directory '{}'",
-        archive.display(),
-        target_dir.display()
-    );
-
-    let mut archive_file = fs::File::open(&archive).await.with_context(|| {
-        format!(
-            "Failed to get the archive name from the path '{}'",
-            archive.display()
-        )
-    })?;
-    let header = compression::read_archive_header(archive_name, &mut archive_file)
-        .await
-        .context("Failed to read the archive header")?;
-    compression::uncompress_with_header(&BTreeSet::new(), &target_dir, header, &mut archive_file)
-        .await
-        .context("Failed to extract the archive")
-}
-
-async fn create_archive(source_dir: &Path, target_dir: &Path) -> anyhow::Result<()> {
-    let source_dir = source_dir.canonicalize().with_context(|| {
-        format!(
-            "Failed to get the absolute path for the source dir path '{}'",
-            source_dir.display()
-        )
-    })?;
-    ensure!(
-        source_dir.is_dir(),
-        "Path '{}' is not a directory",
-        source_dir.display()
-    );
-
-    if !target_dir.exists() {
-        fs::create_dir_all(target_dir).await.with_context(|| {
-            format!(
-                "Failed to create the target dir at path '{}'",
-                target_dir.display()
-            )
-        })?;
-    }
-    let target_dir = target_dir.canonicalize().with_context(|| {
-        format!(
-            "Failed to get the absolute path for the target dir path '{}'",
-            target_dir.display()
-        )
-    })?;
-    ensure!(
-        target_dir.is_dir(),
-        "Path '{}' is not a directory",
-        target_dir.display()
-    );
-
-    println!(
-        "Compressing directory '{}' and creating resulting archive in directory '{}'",
-        source_dir.display(),
-        target_dir.display()
-    );
-
-    let mut metadata_file_contents = None;
-    let mut files_co_archive = Vec::new();
-
-    let mut source_dir_contents = fs::read_dir(&source_dir)
-        .await
-        .context("Failed to read the source directory contents")?;
-
-    while let Some(source_dir_entry) = source_dir_contents
-        .next_entry()
-        .await
-        .context("Failed to read a source dir entry")?
-    {
-        let entry_path = source_dir_entry.path();
-        if entry_path.is_file() {
-            if entry_path.file_name().and_then(|name| name.to_str()) == Some(METADATA_FILE_NAME) {
-                let metadata_bytes = fs::read(entry_path)
-                    .await
-                    .context("Failed to read metata file bytes in the source dir")?;
-                metadata_file_contents = Some(
-                    TimelineMetadata::from_bytes(&metadata_bytes)
-                        .context("Failed to parse metata file contents in the source dir")?,
-                );
-            } else {
-                files_co_archive.push(entry_path);
-            }
-        }
-    }
-
-    let metadata = match metadata_file_contents {
-        Some(metadata) => metadata,
-        None => bail!(
-            "No metadata file found in the source dir '{}', cannot create the archive",
-            source_dir.display()
-        ),
-    };
-
-    let _ = compression::archive_files_as_stream(
-        &source_dir,
-        files_co_archive.iter(),
-        &metadata,
-        move |mut archive_streamer, archive_name| async move {
-            let archive_target = target_dir.join(&archive_name);
-            let mut archive_file = fs::File::create(&archive_target).await?;
-            io::copy(&mut archive_streamer, &mut archive_file).await?;
-            Ok(archive_target)
-        },
-    )
-    .await
-    .context("Failed to create an archive")?;
-
-    Ok(())
-}

pageserver/src/bin/update_metadata.rs

@@ -14,20 +14,20 @@ fn main() -> Result<()> {
         .about("Dump or update metadata file")
         .version(GIT_VERSION)
         .arg(
-            Arg::new("path")
+            Arg::with_name("path")
                 .help("Path to metadata file")
                 .required(true),
         )
         .arg(
-            Arg::new("disk_lsn")
-                .short('d')
+            Arg::with_name("disk_lsn")
+                .short("d")
                 .long("disk_lsn")
                 .takes_value(true)
                 .help("Replace disk constistent lsn"),
         )
         .arg(
-            Arg::new("prev_lsn")
-                .short('p')
+            Arg::with_name("prev_lsn")
+                .short("p")
                 .long("prev_lsn")
                 .takes_value(true)
                 .help("Previous record LSN"),

pageserver/src/branches.rs

@@ -324,13 +324,12 @@ pub(crate) fn create_branch(
         timeline.wait_lsn(startpoint.lsn)?;
     }
     startpoint.lsn = startpoint.lsn.align();
-    if timeline.get_ancestor_lsn() > startpoint.lsn {
-        // can we safely just branch from the ancestor instead?
+    if timeline.get_start_lsn() > startpoint.lsn {
         anyhow::bail!(
-            "invalid startpoint {} for the branch {}: less than timeline ancestor lsn {:?}",
+            "invalid startpoint {} for the branch {}: less than timeline start {}",
             startpoint.lsn,
             branchname,
-            timeline.get_ancestor_lsn()
+            timeline.get_start_lsn()
         );
     }
 

pageserver/src/config.rs

@@ -36,9 +36,6 @@ pub mod defaults {
     pub const DEFAULT_GC_HORIZON: u64 = 64 * 1024 * 1024;
     pub const DEFAULT_GC_PERIOD: &str = "100 s";
 
-    pub const DEFAULT_WAIT_LSN_TIMEOUT: &str = "60 s";
-    pub const DEFAULT_WAL_REDO_TIMEOUT: &str = "60 s";
-
     pub const DEFAULT_SUPERUSER: &str = "zenith_admin";
     pub const DEFAULT_REMOTE_STORAGE_MAX_CONCURRENT_SYNC: usize = 100;
     pub const DEFAULT_REMOTE_STORAGE_MAX_SYNC_ERRORS: u32 = 10;
@@ -62,9 +59,6 @@ pub mod defaults {
 #gc_period = '{DEFAULT_GC_PERIOD}'
 #gc_horizon = {DEFAULT_GC_HORIZON}
 
-#wait_lsn_timeout = '{DEFAULT_WAIT_LSN_TIMEOUT}'
-#wal_redo_timeout = '{DEFAULT_WAL_REDO_TIMEOUT}'
-
 #max_file_descriptors = {DEFAULT_MAX_FILE_DESCRIPTORS}
 
 # initial superuser role name to use when creating a new tenant
@@ -91,12 +85,6 @@ pub struct PageServerConf {
 
     pub gc_horizon: u64,
     pub gc_period: Duration,
-
-    // 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.
-    pub wal_redo_timeout: Duration,
-
     pub superuser: String,
 
     pub page_cache_size: usize,
@@ -244,8 +232,6 @@ impl PageServerConf {
             checkpoint_period: humantime::parse_duration(DEFAULT_CHECKPOINT_PERIOD)?,
             gc_horizon: DEFAULT_GC_HORIZON,
             gc_period: humantime::parse_duration(DEFAULT_GC_PERIOD)?,
-            wait_lsn_timeout: humantime::parse_duration(DEFAULT_WAIT_LSN_TIMEOUT)?,
-            wal_redo_timeout: humantime::parse_duration(DEFAULT_WAL_REDO_TIMEOUT)?,
             page_cache_size: DEFAULT_PAGE_CACHE_SIZE,
             max_file_descriptors: DEFAULT_MAX_FILE_DESCRIPTORS,
 
@@ -266,8 +252,6 @@ impl PageServerConf {
                 "checkpoint_period" => conf.checkpoint_period = parse_toml_duration(key, item)?,
                 "gc_horizon" => conf.gc_horizon = parse_toml_u64(key, item)?,
                 "gc_period" => conf.gc_period = parse_toml_duration(key, item)?,
-                "wait_lsn_timeout" => conf.wait_lsn_timeout = parse_toml_duration(key, item)?,
-                "wal_redo_timeout" => conf.wal_redo_timeout = parse_toml_duration(key, item)?,
                 "initial_superuser_name" => conf.superuser = parse_toml_string(key, item)?,
                 "page_cache_size" => conf.page_cache_size = parse_toml_u64(key, item)? as usize,
                 "max_file_descriptors" => {
@@ -402,8 +386,6 @@ impl PageServerConf {
             checkpoint_period: Duration::from_secs(10),
             gc_horizon: defaults::DEFAULT_GC_HORIZON,
             gc_period: Duration::from_secs(10),
-            wait_lsn_timeout: Duration::from_secs(60),
-            wal_redo_timeout: Duration::from_secs(60),
             page_cache_size: defaults::DEFAULT_PAGE_CACHE_SIZE,
             max_file_descriptors: defaults::DEFAULT_MAX_FILE_DESCRIPTORS,
             listen_pg_addr: defaults::DEFAULT_PG_LISTEN_ADDR.to_string(),
@@ -474,9 +456,6 @@ checkpoint_period = '111 s'
 gc_period = '222 s'
 gc_horizon = 222
 
-wait_lsn_timeout = '111 s'
-wal_redo_timeout = '111 s'
-
 page_cache_size = 444
 max_file_descriptors = 333
 
@@ -507,8 +486,6 @@ initial_superuser_name = 'zzzz'
                 checkpoint_period: humantime::parse_duration(defaults::DEFAULT_CHECKPOINT_PERIOD)?,
                 gc_horizon: defaults::DEFAULT_GC_HORIZON,
                 gc_period: humantime::parse_duration(defaults::DEFAULT_GC_PERIOD)?,
-                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(),
                 page_cache_size: defaults::DEFAULT_PAGE_CACHE_SIZE,
                 max_file_descriptors: defaults::DEFAULT_MAX_FILE_DESCRIPTORS,
@@ -550,8 +527,6 @@ initial_superuser_name = 'zzzz'
                 checkpoint_period: Duration::from_secs(111),
                 gc_horizon: 222,
                 gc_period: Duration::from_secs(222),
-                wait_lsn_timeout: Duration::from_secs(111),
-                wal_redo_timeout: Duration::from_secs(111),
                 superuser: "zzzz".to_string(),
                 page_cache_size: 444,
                 max_file_descriptors: 333,

pageserver/src/http/openapi_spec.yml

@@ -234,7 +234,9 @@ paths:
           content:
             application/json:
               schema:
-                $ref: "#/components/schemas/BranchInfo"
+                type: array
+                items:
+                  $ref: "#/components/schemas/BranchInfo"
         "400":
           description: Malformed branch create request
           content:
@@ -368,15 +370,12 @@ components:
           format: hex
         ancestor_id:
           type: string
-          format: hex
         ancestor_lsn:
           type: string
         current_logical_size:
           type: integer
         current_logical_size_non_incremental:
           type: integer
-        latest_valid_lsn:
-          type: integer
     TimelineInfo:
       type: object
       required:

pageserver/src/http/routes.rs

@@ -4,6 +4,7 @@ use anyhow::{Context, Result};
 use hyper::header;
 use hyper::StatusCode;
 use hyper::{Body, Request, Response, Uri};
+use routerify::{ext::RequestExt, RouterBuilder};
 use serde::Serialize;
 use tracing::*;
 use zenith_utils::auth::JwtAuth;
@@ -18,7 +19,6 @@ use zenith_utils::http::{
     request::get_request_param,
     request::parse_request_param,
 };
-use zenith_utils::http::{RequestExt, RouterBuilder};
 use zenith_utils::lsn::Lsn;
 use zenith_utils::zid::{opt_display_serde, ZTimelineId};
 
@@ -202,6 +202,7 @@ enum TimelineInfo {
         ancestor_timeline_id: Option<ZTimelineId>,
         last_record_lsn: Lsn,
         prev_record_lsn: Lsn,
+        start_lsn: Lsn,
         disk_consistent_lsn: Lsn,
         timeline_state: Option<TimelineSyncState>,
     },
@@ -236,6 +237,7 @@ async fn timeline_detail_handler(request: Request<Body>) -> Result<Response<Body
                 disk_consistent_lsn: timeline.get_disk_consistent_lsn(),
                 last_record_lsn: timeline.get_last_record_lsn(),
                 prev_record_lsn: timeline.get_prev_record_lsn(),
+                start_lsn: timeline.get_start_lsn(),
                 timeline_state: repo.get_timeline_state(timeline_id),
             },
         })

pageserver/src/layered_repository.rs

@@ -28,8 +28,8 @@ use std::io::Write;
 use std::ops::{Bound::Included, Deref};
 use std::path::{Path, PathBuf};
 use std::sync::atomic::{self, AtomicBool, AtomicUsize};
-use std::sync::{Arc, Mutex, MutexGuard, RwLock, RwLockReadGuard};
-use std::time::Instant;
+use std::sync::{Arc, Mutex, MutexGuard};
+use std::time::{Duration, Instant};
 
 use self::metadata::{metadata_path, TimelineMetadata, METADATA_FILE_NAME};
 use crate::config::PageServerConf;
@@ -64,12 +64,14 @@ mod inmemory_layer;
 mod interval_tree;
 mod layer_map;
 pub mod metadata;
+mod page_versions;
 mod par_fsync;
 mod storage_layer;
 
 use delta_layer::DeltaLayer;
 use ephemeral_file::is_ephemeral_file;
 use filename::{DeltaFileName, ImageFileName};
+use global_layer_map::{LayerId, GLOBAL_LAYER_MAP};
 use image_layer::ImageLayer;
 use inmemory_layer::InMemoryLayer;
 use layer_map::LayerMap;
@@ -82,6 +84,9 @@ pub use crate::layered_repository::ephemeral_file::writeback as writeback_epheme
 
 static ZERO_PAGE: Bytes = Bytes::from_static(&[0u8; 8192]);
 
+// Timeout when waiting for WAL receiver to catch up to an LSN given in a GetPage@LSN call.
+static TIMEOUT: Duration = Duration::from_secs(60);
+
 // Metrics collected on operations on the storage repository.
 lazy_static! {
     static ref STORAGE_TIME: HistogramVec = register_histogram_vec!(
@@ -162,7 +167,7 @@ impl Repository for LayeredRepository {
         // Create the timeline directory, and write initial metadata to file.
         crashsafe_dir::create_dir_all(self.conf.timeline_path(&timelineid, &self.tenantid))?;
 
-        let metadata = TimelineMetadata::new(Lsn(0), None, None, Lsn(0), initdb_lsn, initdb_lsn);
+        let metadata = TimelineMetadata::new(Lsn(0), None, None, Lsn(0), Lsn(0), initdb_lsn);
         Self::save_metadata(self.conf, timelineid, self.tenantid, &metadata, true)?;
 
         let timeline = LayeredTimeline::new(
@@ -196,10 +201,9 @@ impl Repository for LayeredRepository {
                 bail!("Cannot branch off the timeline {} that's not local", src)
             }
         };
-        let latest_gc_cutoff_lsn = src_timeline.get_latest_gc_cutoff_lsn();
 
         src_timeline
-            .check_lsn_is_in_scope(start_lsn, &latest_gc_cutoff_lsn)
+            .check_lsn_is_in_scope(start_lsn)
             .context("invalid branch start lsn")?;
 
         let RecordLsn {
@@ -227,7 +231,7 @@ impl Repository for LayeredRepository {
             dst_prev,
             Some(src),
             start_lsn,
-            *src_timeline.latest_gc_cutoff_lsn.read().unwrap(),
+            src_timeline.latest_gc_cutoff_lsn.load(),
             src_timeline.initdb_lsn,
         );
         crashsafe_dir::create_dir_all(self.conf.timeline_path(&dst, &self.tenantid))?;
@@ -607,7 +611,7 @@ impl LayeredRepository {
             }
         }
 
-        // Now collect info about branchpoints
+        //Now collect info about branchpoints
         let mut all_branchpoints: BTreeSet<(ZTimelineId, Lsn)> = BTreeSet::new();
         for &timelineid in &timelineids {
             let timeline = match self.get_or_init_timeline(timelineid, &mut timelines)? {
@@ -779,7 +783,7 @@ pub struct LayeredTimeline {
     checkpoint_cs: Mutex<()>,
 
     // Needed to ensure that we can't create a branch at a point that was already garbage collected
-    latest_gc_cutoff_lsn: RwLock<Lsn>,
+    latest_gc_cutoff_lsn: AtomicLsn,
 
     // It may change across major versions so for simplicity
     // keep it after running initdb for a timeline.
@@ -812,7 +816,7 @@ impl Timeline for LayeredTimeline {
         );
 
         self.last_record_lsn
-            .wait_for_timeout(lsn, self.conf.wait_lsn_timeout)
+            .wait_for_timeout(lsn, TIMEOUT)
             .with_context(|| {
                 format!(
                     "Timed out while waiting for WAL record at LSN {} to arrive, last_record_lsn {} disk consistent LSN={}",
@@ -823,10 +827,6 @@ impl Timeline for LayeredTimeline {
         Ok(())
     }
 
-    fn get_latest_gc_cutoff_lsn(&self) -> RwLockReadGuard<Lsn> {
-        self.latest_gc_cutoff_lsn.read().unwrap()
-    }
-
     /// Look up given page version.
     fn get_page_at_lsn(&self, rel: RelishTag, rel_blknum: BlockNumber, lsn: Lsn) -> Result<Bytes> {
         if !rel.is_blocky() && rel_blknum != 0 {
@@ -837,6 +837,14 @@ impl Timeline for LayeredTimeline {
             );
         }
         debug_assert!(lsn <= self.get_last_record_lsn());
+        let latest_gc_cutoff_lsn = self.latest_gc_cutoff_lsn.load();
+        // error instead of assert to simplify testing
+        ensure!(
+            lsn >= latest_gc_cutoff_lsn,
+            "tried to request a page version that was garbage collected. requested at {} gc cutoff {}",
+            lsn, latest_gc_cutoff_lsn
+        );
+
         let (seg, seg_blknum) = SegmentTag::from_blknum(rel, rel_blknum);
 
         if let Some((layer, lsn)) = self.get_layer_for_read(seg, lsn)? {
@@ -1007,16 +1015,21 @@ impl Timeline for LayeredTimeline {
     ///
     /// Validate lsn against initdb_lsn and latest_gc_cutoff_lsn.
     ///
-    fn check_lsn_is_in_scope(
-        &self,
-        lsn: Lsn,
-        latest_gc_cutoff_lsn: &RwLockReadGuard<Lsn>,
-    ) -> Result<()> {
+    fn check_lsn_is_in_scope(&self, lsn: Lsn) -> Result<()> {
+        let initdb_lsn = self.initdb_lsn;
         ensure!(
-            lsn >= **latest_gc_cutoff_lsn,
+            lsn >= initdb_lsn,
+            "LSN {} is earlier than initdb lsn {}",
+            lsn,
+            initdb_lsn,
+        );
+
+        let latest_gc_cutoff_lsn = self.latest_gc_cutoff_lsn.load();
+        ensure!(
+            lsn >= latest_gc_cutoff_lsn,
             "LSN {} is earlier than latest GC horizon {} (we might've already garbage collected needed data)",
             lsn,
-            **latest_gc_cutoff_lsn,
+            latest_gc_cutoff_lsn,
         );
         Ok(())
     }
@@ -1033,6 +1046,14 @@ impl Timeline for LayeredTimeline {
         self.last_record_lsn.load()
     }
 
+    fn get_start_lsn(&self) -> Lsn {
+        self.ancestor_timeline
+            .as_ref()
+            .and_then(|ancestor_entry| ancestor_entry.local_or_schedule_download(self.tenantid))
+            .map(Timeline::get_start_lsn)
+            .unwrap_or(self.ancestor_lsn)
+    }
+
     fn get_current_logical_size(&self) -> usize {
         self.current_logical_size.load(atomic::Ordering::Acquire) as usize
     }
@@ -1122,7 +1143,7 @@ impl LayeredTimeline {
             write_lock: Mutex::new(()),
             checkpoint_cs: Mutex::new(()),
 
-            latest_gc_cutoff_lsn: RwLock::new(metadata.latest_gc_cutoff_lsn()),
+            latest_gc_cutoff_lsn: AtomicLsn::from(metadata.latest_gc_cutoff_lsn()),
             initdb_lsn: metadata.initdb_lsn(),
         }
     }
@@ -1148,8 +1169,8 @@ impl LayeredTimeline {
                 // create an ImageLayer struct for each image file.
                 if imgfilename.lsn > disk_consistent_lsn {
                     warn!(
-                        "found future image layer {} on timeline {} disk_consistent_lsn is {}",
-                        imgfilename, self.timelineid, disk_consistent_lsn
+                        "found future image layer {} on timeline {}",
+                        imgfilename, self.timelineid
                     );
 
                     rename_to_backup(direntry.path())?;
@@ -1172,8 +1193,8 @@ impl LayeredTimeline {
                 // before crash.
                 if deltafilename.end_lsn > disk_consistent_lsn + 1 {
                     warn!(
-                        "found future delta layer {} on timeline {} disk_consistent_lsn is {}",
-                        deltafilename, self.timelineid, disk_consistent_lsn
+                        "found future delta layer {} on timeline {}",
+                        deltafilename, self.timelineid
                     );
 
                     rename_to_backup(direntry.path())?;
@@ -1369,7 +1390,7 @@ impl LayeredTimeline {
                     self.tenantid,
                     seg,
                     lsn,
-                    last_record_lsn,
+                    lsn,
                 )?;
             } else {
                 return Ok(open_layer);
@@ -1412,7 +1433,7 @@ impl LayeredTimeline {
                 self.timelineid,
                 self.tenantid,
                 start_lsn,
-                last_record_lsn,
+                lsn,
             )?;
         } else {
             // New relation.
@@ -1423,14 +1444,8 @@ impl LayeredTimeline {
                 lsn
             );
 
-            layer = InMemoryLayer::create(
-                self.conf,
-                self.timelineid,
-                self.tenantid,
-                seg,
-                lsn,
-                last_record_lsn,
-            )?;
+            layer =
+                InMemoryLayer::create(self.conf, self.timelineid, self.tenantid, seg, lsn, lsn)?;
         }
 
         let layer_rc: Arc<InMemoryLayer> = Arc::new(layer);
@@ -1447,7 +1462,7 @@ impl LayeredTimeline {
         // Prevent concurrent checkpoints
         let _checkpoint_cs = self.checkpoint_cs.lock().unwrap();
 
-        let write_guard = self.write_lock.lock().unwrap();
+        let mut write_guard = self.write_lock.lock().unwrap();
         let mut layers = self.layers.lock().unwrap();
 
         // Bump the generation number in the layer map, so that we can distinguish
@@ -1473,17 +1488,11 @@ impl LayeredTimeline {
         let mut disk_consistent_lsn = last_record_lsn;
 
         let mut layer_paths = Vec::new();
-        let mut freeze_end_lsn = Lsn(0);
-        let mut evicted_layers = Vec::new();
-
-        //
-        // Determine which layers we need to evict and calculate max(latest_lsn)
-        // among those layers.
-        //
         while let Some((oldest_layer_id, oldest_layer, oldest_generation)) =
             layers.peek_oldest_open()
         {
-            let oldest_lsn = oldest_layer.get_oldest_lsn();
+            let oldest_pending_lsn = oldest_layer.get_oldest_pending_lsn();
+
             // Does this layer need freezing?
             //
             // Write out all in-memory layers that contain WAL older than CHECKPOINT_DISTANCE.
@@ -1492,60 +1501,28 @@ impl LayeredTimeline {
             // when we started. We don't want to process layers inserted after we started, to
             // avoid getting into an infinite loop trying to process again entries that we
             // inserted ourselves.
-            //
-            // Once we have decided to write out at least one layer, we must also write out
-            // any other layers that contain WAL older than the end LSN of the layers we have
-            // already decided to write out. In other words, we must write out all layers
-            // whose [oldest_lsn, latest_lsn) range overlaps with any of the other layers
-            // that we are writing out. Otherwise, when we advance 'disk_consistent_lsn', it's
-            // ambiguous whether those layers are already durable on disk or not. For example,
-            // imagine that there are two layers in memory that contain page versions in the
-            // following LSN ranges:
-            //
-            // A: 100-150
-            // B: 110-200
-            //
-            // If we flush layer A, we must also flush layer B, because they overlap. If we
-            // flushed only A, and advanced 'disk_consistent_lsn' to 150, we would break the
-            // rule that all WAL older than 'disk_consistent_lsn' are durable on disk, because
-            // B contains some WAL older than 150. On the other hand, if we flushed out A and
-            // advanced 'disk_consistent_lsn' only up to 110, after crash and restart we would
-            // delete the first layer because its end LSN is larger than 110. If we changed
-            // the deletion logic to not delete it, then we would start streaming at 110, and
-            // process again the WAL records in the range 110-150 that are already in layer A,
-            // and the WAL processing code does not cope with that. We solve that dilemma by
-            // insisting that if we write out the first layer, we also write out the second
-            // layer, and advance disk_consistent_lsn all the way up to 200.
-            //
-            let distance = last_record_lsn.widening_sub(oldest_lsn);
-            if (distance < 0
+            let distance = last_record_lsn.widening_sub(oldest_pending_lsn);
+            if distance < 0
                 || distance < checkpoint_distance.into()
-                || oldest_generation == current_generation)
-                && oldest_lsn >= freeze_end_lsn
-            // this layer intersects with evicted layer and so also need to be evicted
+                || oldest_generation == current_generation
             {
                 info!(
                     "the oldest layer is now {} which is {} bytes behind last_record_lsn",
                     oldest_layer.filename().display(),
                     distance
                 );
-                disk_consistent_lsn = oldest_lsn;
+                disk_consistent_lsn = oldest_pending_lsn;
                 break;
             }
-            let latest_lsn = oldest_layer.get_latest_lsn();
-            if latest_lsn > freeze_end_lsn {
-                freeze_end_lsn = latest_lsn; // calculate max of latest_lsn of the layers we're about to evict
-            }
-            layers.remove_open(oldest_layer_id);
-            evicted_layers.push((oldest_layer_id, oldest_layer));
-        }
 
-        // Freeze evicted layers
-        for (_evicted_layer_id, evicted_layer) in evicted_layers.iter() {
-            // Mark the layer as no longer accepting writes and record the end_lsn.
-            // This happens in-place, no new layers are created now.
-            evicted_layer.freeze(freeze_end_lsn);
-            layers.insert_historic(evicted_layer.clone());
+            drop(layers);
+            drop(write_guard);
+
+            let mut this_layer_paths = self.evict_layer(oldest_layer_id, reconstruct_pages)?;
+            layer_paths.append(&mut this_layer_paths);
+
+            write_guard = self.write_lock.lock().unwrap();
+            layers = self.layers.lock().unwrap();
         }
 
         // Call unload() on all frozen layers, to release memory.
@@ -1558,14 +1535,6 @@ impl LayeredTimeline {
         drop(layers);
         drop(write_guard);
 
-        // Create delta/image layers for evicted layers
-        for (_evicted_layer_id, evicted_layer) in evicted_layers.iter() {
-            let mut this_layer_paths =
-                self.evict_layer(evicted_layer.clone(), reconstruct_pages)?;
-            layer_paths.append(&mut this_layer_paths);
-        }
-
-        // Sync layers
         if !layer_paths.is_empty() {
             // We must fsync the timeline dir to ensure the directory entries for
             // new layer files are durable
@@ -1606,7 +1575,7 @@ impl LayeredTimeline {
                 ondisk_prev_record_lsn,
                 ancestor_timelineid,
                 self.ancestor_lsn,
-                *self.latest_gc_cutoff_lsn.read().unwrap(),
+                self.latest_gc_cutoff_lsn.load(),
                 self.initdb_lsn,
             );
 
@@ -1633,29 +1602,52 @@ impl LayeredTimeline {
         Ok(())
     }
 
-    fn evict_layer(
-        &self,
-        layer: Arc<InMemoryLayer>,
-        reconstruct_pages: bool,
-    ) -> Result<Vec<PathBuf>> {
-        let new_historics = layer.write_to_disk(self, reconstruct_pages)?;
+    fn evict_layer(&self, layer_id: LayerId, reconstruct_pages: bool) -> Result<Vec<PathBuf>> {
+        // Mark the layer as no longer accepting writes and record the end_lsn.
+        // This happens in-place, no new layers are created now.
+        // We call `get_last_record_lsn` again, which may be different from the
+        // original load, as we may have released the write lock since then.
 
-        let mut layer_paths = Vec::new();
-        let _write_guard = self.write_lock.lock().unwrap();
+        let mut write_guard = self.write_lock.lock().unwrap();
         let mut layers = self.layers.lock().unwrap();
 
-        // Finally, replace the frozen in-memory layer with the new on-disk layers
-        layers.remove_historic(layer);
+        let mut layer_paths = Vec::new();
 
-        // Add the historics to the LayerMap
-        for delta_layer in new_historics.delta_layers {
-            layer_paths.push(delta_layer.path());
-            layers.insert_historic(Arc::new(delta_layer));
-        }
-        for image_layer in new_historics.image_layers {
-            layer_paths.push(image_layer.path());
-            layers.insert_historic(Arc::new(image_layer));
+        let global_layer_map = GLOBAL_LAYER_MAP.read().unwrap();
+        if let Some(oldest_layer) = global_layer_map.get(&layer_id) {
+            drop(global_layer_map);
+            oldest_layer.freeze(self.get_last_record_lsn());
+
+            // The layer is no longer open, update the layer map to reflect this.
+            // We will replace it with on-disk historics below.
+            layers.remove_open(layer_id);
+            layers.insert_historic(oldest_layer.clone());
+
+            // Write the now-frozen layer to disk. That could take a while, so release the lock while do it
+            drop(layers);
+            drop(write_guard);
+
+            let new_historics = oldest_layer.write_to_disk(self, reconstruct_pages)?;
+
+            write_guard = self.write_lock.lock().unwrap();
+            layers = self.layers.lock().unwrap();
+
+            // Finally, replace the frozen in-memory layer with the new on-disk layers
+            layers.remove_historic(oldest_layer);
+
+            // Add the historics to the LayerMap
+            for delta_layer in new_historics.delta_layers {
+                layer_paths.push(delta_layer.path());
+                layers.insert_historic(Arc::new(delta_layer));
+            }
+            for image_layer in new_historics.image_layers {
+                layer_paths.push(image_layer.path());
+                layers.insert_historic(Arc::new(image_layer));
+            }
         }
+        drop(layers);
+        drop(write_guard);
+
         Ok(layer_paths)
     }
 
@@ -1685,13 +1677,12 @@ impl LayeredTimeline {
         let now = Instant::now();
         let mut result: GcResult = Default::default();
         let disk_consistent_lsn = self.get_disk_consistent_lsn();
-        let _checkpoint_cs = self.checkpoint_cs.lock().unwrap();
 
         let _enter = info_span!("garbage collection", timeline = %self.timelineid, tenant = %self.tenantid, cutoff = %cutoff).entered();
 
         // We need to ensure that no one branches at a point before latest_gc_cutoff_lsn.
         // See branch_timeline() for details.
-        *self.latest_gc_cutoff_lsn.write().unwrap() = cutoff;
+        self.latest_gc_cutoff_lsn.store(cutoff);
 
         info!("GC starting");
 
@@ -1940,21 +1931,22 @@ impl LayeredTimeline {
         // for redo.
         let rel = seg.rel;
         let rel_blknum = seg.segno * RELISH_SEG_SIZE + seg_blknum;
-        let cached_page_img = match self.lookup_cached_page(&rel, rel_blknum, lsn) {
+        let (cached_lsn_opt, cached_page_opt) = match self.lookup_cached_page(&rel, rel_blknum, lsn)
+        {
             Some((cached_lsn, cached_img)) => {
                 match cached_lsn.cmp(&lsn) {
                     cmp::Ordering::Less => {} // there might be WAL between cached_lsn and lsn, we need to check
                     cmp::Ordering::Equal => return Ok(cached_img), // exact LSN match, return the image
                     cmp::Ordering::Greater => panic!(), // the returned lsn should never be after the requested lsn
                 }
-                Some((cached_lsn, cached_img))
+                (Some(cached_lsn), Some((cached_lsn, cached_img)))
             }
-            None => None,
+            None => (None, None),
         };
 
         let mut data = PageReconstructData {
             records: Vec::new(),
-            page_img: cached_page_img,
+            page_img: None,
         };
 
         // Holds an Arc reference to 'layer_ref' when iterating in the loop below.
@@ -1967,14 +1959,15 @@ impl LayeredTimeline {
         let mut curr_lsn = lsn;
         loop {
             let result = layer_ref
-                .get_page_reconstruct_data(seg_blknum, curr_lsn, &mut data)
+                .get_page_reconstruct_data(seg_blknum, curr_lsn, cached_lsn_opt, &mut data)
                 .with_context(|| {
                     format!(
-                        "Failed to get reconstruct data {} {:?} {} {}",
+                        "Failed to get reconstruct data {} {:?} {} {} {:?}",
                         layer_ref.get_seg_tag(),
                         layer_ref.filename(),
                         seg_blknum,
                         curr_lsn,
+                        cached_lsn_opt,
                     )
                 })?;
             match result {
@@ -2021,6 +2014,16 @@ impl LayeredTimeline {
                         lsn,
                     );
                 }
+                PageReconstructResult::Cached => {
+                    let (cached_lsn, cached_img) = cached_page_opt.unwrap();
+                    assert!(data.page_img.is_none());
+                    if let Some((first_rec_lsn, first_rec)) = data.records.first() {
+                        assert!(&cached_lsn < first_rec_lsn);
+                        assert!(!first_rec.will_init());
+                    }
+                    data.page_img = Some(cached_img);
+                    break;
+                }
             }
         }
 
@@ -2042,12 +2045,12 @@ impl LayeredTimeline {
 
         // If we have a page image, and no WAL, we're all set
         if data.records.is_empty() {
-            if let Some((img_lsn, img)) = &data.page_img {
+            if let Some(img) = &data.page_img {
                 trace!(
                     "found page image for blk {} in {} at {}, no WAL redo required",
                     rel_blknum,
                     rel,
-                    img_lsn
+                    request_lsn
                 );
                 Ok(img.clone())
             } else {
@@ -2074,13 +2077,11 @@ impl LayeredTimeline {
                 );
                 Ok(ZERO_PAGE.clone())
             } else {
-                let base_img = if let Some((_lsn, img)) = data.page_img {
+                if data.page_img.is_some() {
                     trace!("found {} WAL records and a base image for blk {} in {} at {}, performing WAL redo", data.records.len(), rel_blknum, rel, request_lsn);
-                    Some(img)
                 } else {
                     trace!("found {} WAL records that will init the page for blk {} in {} at {}, performing WAL redo", data.records.len(), rel_blknum, rel, request_lsn);
-                    None
-                };
+                }
 
                 let last_rec_lsn = data.records.last().unwrap().0;
 
@@ -2088,7 +2089,7 @@ impl LayeredTimeline {
                     rel,
                     rel_blknum,
                     request_lsn,
-                    base_img,
+                    data.page_img.clone(),
                     data.records,
                 )?;
 
@@ -2347,157 +2348,3 @@ fn rename_to_backup(path: PathBuf) -> anyhow::Result<()> {
 
     bail!("couldn't find an unused backup number for {:?}", path)
 }
-
-///
-/// Tests that are specific to the layered storage format.
-///
-/// There are more unit tests in repository.rs that work through the
-/// Repository interface and are expected to work regardless of the
-/// file format and directory layout. The test here are more low level.
-///
-#[cfg(test)]
-mod tests {
-    use super::*;
-    use crate::repository::repo_harness::*;
-
-    #[test]
-    fn corrupt_metadata() -> Result<()> {
-        const TEST_NAME: &str = "corrupt_metadata";
-        let harness = RepoHarness::create(TEST_NAME)?;
-        let repo = harness.load();
-
-        repo.create_empty_timeline(TIMELINE_ID, Lsn(0))?;
-        drop(repo);
-
-        let metadata_path = harness.timeline_path(&TIMELINE_ID).join(METADATA_FILE_NAME);
-
-        assert!(metadata_path.is_file());
-
-        let mut metadata_bytes = std::fs::read(&metadata_path)?;
-        assert_eq!(metadata_bytes.len(), 512);
-        metadata_bytes[512 - 4 - 2] ^= 1;
-        std::fs::write(metadata_path, metadata_bytes)?;
-
-        let new_repo = harness.load();
-        let err = new_repo.get_timeline(TIMELINE_ID).err().unwrap();
-        assert_eq!(err.to_string(), "failed to load metadata");
-        assert_eq!(
-            err.source().unwrap().to_string(),
-            "metadata checksum mismatch"
-        );
-
-        Ok(())
-    }
-
-    ///
-    /// Test the logic in 'load_layer_map' that removes layer files that are
-    /// newer than 'disk_consistent_lsn'.
-    ///
-    #[test]
-    fn future_layerfiles() -> Result<()> {
-        const TEST_NAME: &str = "future_layerfiles";
-        let harness = RepoHarness::create(TEST_NAME)?;
-        let repo = harness.load();
-
-        // Create a timeline with disk_consistent_lsn = 8000
-        let tline = repo.create_empty_timeline(TIMELINE_ID, Lsn(0x8000))?;
-        let writer = tline.writer();
-        writer.advance_last_record_lsn(Lsn(0x8000));
-        drop(writer);
-        repo.checkpoint_iteration(CheckpointConfig::Forced)?;
-        drop(repo);
-
-        let timeline_path = harness.timeline_path(&TIMELINE_ID);
-
-        let make_empty_file = |filename: &str| -> std::io::Result<()> {
-            let path = timeline_path.join(filename);
-
-            assert!(!path.exists());
-            std::fs::write(&path, &[])?;
-
-            Ok(())
-        };
-
-        // Helper function to check that a relation file exists, and a corresponding
-        // <filename>.0.old file does not.
-        let assert_exists = |filename: &str| {
-            let path = timeline_path.join(filename);
-            assert!(path.exists(), "file {} was removed", filename);
-
-            // Check that there is no .old file
-            let backup_path = timeline_path.join(format!("{}.0.old", filename));
-            assert!(
-                !backup_path.exists(),
-                "unexpected backup file {}",
-                backup_path.display()
-            );
-        };
-
-        // Helper function to check that a relation file does *not* exists, and a corresponding
-        // <filename>.<num>.old file does.
-        let assert_is_renamed = |filename: &str, num: u32| {
-            let path = timeline_path.join(filename);
-            assert!(
-                !path.exists(),
-                "file {} was not removed as expected",
-                filename
-            );
-
-            let backup_path = timeline_path.join(format!("{}.{}.old", filename, num));
-            assert!(
-                backup_path.exists(),
-                "backup file {} was not created",
-                backup_path.display()
-            );
-        };
-
-        // These files are considered to be in the future and will be renamed out
-        // of the way
-        let future_filenames = vec![
-            format!("pg_control_0_{:016X}", 0x8001),
-            format!("pg_control_0_{:016X}_{:016X}", 0x8001, 0x8008),
-        ];
-        // But these are not:
-        let past_filenames = vec![
-            format!("pg_control_0_{:016X}", 0x8000),
-            format!("pg_control_0_{:016X}_{:016X}", 0x7000, 0x8001),
-        ];
-
-        for filename in future_filenames.iter().chain(past_filenames.iter()) {
-            make_empty_file(filename)?;
-        }
-
-        // Load the timeline. This will cause the files in the "future" to be renamed
-        // away.
-        let new_repo = harness.load();
-        new_repo.get_timeline(TIMELINE_ID).unwrap();
-        drop(new_repo);
-
-        for filename in future_filenames.iter() {
-            assert_is_renamed(filename, 0);
-        }
-        for filename in past_filenames.iter() {
-            assert_exists(filename);
-        }
-
-        // Create the future files again, and load again. They should be renamed to
-        // *.1.old this time.
-        for filename in future_filenames.iter() {
-            make_empty_file(filename)?;
-        }
-
-        let new_repo = harness.load();
-        new_repo.get_timeline(TIMELINE_ID).unwrap();
-        drop(new_repo);
-
-        for filename in future_filenames.iter() {
-            assert_is_renamed(filename, 0);
-            assert_is_renamed(filename, 1);
-        }
-        for filename in past_filenames.iter() {
-            assert_exists(filename);
-        }
-
-        Ok(())
-    }
-}

pageserver/src/layered_repository/delta_layer.rs

@@ -208,15 +208,16 @@ impl Layer for DeltaLayer {
         &self,
         blknum: SegmentBlk,
         lsn: Lsn,
+        cached_img_lsn: Option<Lsn>,
         reconstruct_data: &mut PageReconstructData,
     ) -> Result<PageReconstructResult> {
         let mut need_image = true;
 
         assert!((0..RELISH_SEG_SIZE).contains(&blknum));
 
-        match &reconstruct_data.page_img {
-            Some((cached_lsn, _)) if &self.end_lsn <= cached_lsn => {
-                return Ok(PageReconstructResult::Complete)
+        match &cached_img_lsn {
+            Some(cached_lsn) if &self.end_lsn <= cached_lsn => {
+                return Ok(PageReconstructResult::Cached)
             }
             _ => {}
         }
@@ -239,9 +240,9 @@ impl Layer for DeltaLayer {
                 .iter()
                 .rev();
             for ((_blknum, pv_lsn), blob_range) in iter {
-                match &reconstruct_data.page_img {
-                    Some((cached_lsn, _)) if pv_lsn <= cached_lsn => {
-                        return Ok(PageReconstructResult::Complete)
+                match &cached_img_lsn {
+                    Some(cached_lsn) if pv_lsn <= cached_lsn => {
+                        return Ok(PageReconstructResult::Cached)
                     }
                     _ => {}
                 }
@@ -251,7 +252,7 @@ impl Layer for DeltaLayer {
                 match pv {
                     PageVersion::Page(img) => {
                         // Found a page image, return it
-                        reconstruct_data.page_img = Some((*pv_lsn, img));
+                        reconstruct_data.page_img = Some(img);
                         need_image = false;
                         break;
                     }

pageserver/src/layered_repository/ephemeral_file.rs

@@ -175,10 +175,7 @@ impl Write for EphemeralFile {
     }
 
     fn flush(&mut self) -> Result<(), std::io::Error> {
-        // we don't need to flush data:
-        // * we either write input bytes or not, not keeping any intermediate data buffered
-        // * rust unix file `flush` impl does not flush things either, returning `Ok(())`
-        Ok(())
+        todo!()
     }
 }
 

pageserver/src/layered_repository/image_layer.rs

@@ -145,15 +145,14 @@ impl Layer for ImageLayer {
         &self,
         blknum: SegmentBlk,
         lsn: Lsn,
+        cached_img_lsn: Option<Lsn>,
         reconstruct_data: &mut PageReconstructData,
     ) -> Result<PageReconstructResult> {
         assert!((0..RELISH_SEG_SIZE).contains(&blknum));
         assert!(lsn >= self.lsn);
 
-        match reconstruct_data.page_img {
-            Some((cached_lsn, _)) if self.lsn <= cached_lsn => {
-                return Ok(PageReconstructResult::Complete)
-            }
+        match cached_img_lsn {
+            Some(cached_lsn) if self.lsn <= cached_lsn => return Ok(PageReconstructResult::Cached),
             _ => {}
         }
 
@@ -174,14 +173,7 @@ impl Layer for ImageLayer {
                     .as_ref()
                     .unwrap()
                     .chapter_reader(BLOCKY_IMAGES_CHAPTER)?;
-
-                chapter.read_exact_at(&mut buf, offset).with_context(|| {
-                    format!(
-                        "failed to read page from data file {} at offset {}",
-                        self.filename().display(),
-                        offset
-                    )
-                })?;
+                chapter.read_exact_at(&mut buf, offset)?;
 
                 buf
             }
@@ -196,7 +188,7 @@ impl Layer for ImageLayer {
             }
         };
 
-        reconstruct_data.page_img = Some((self.lsn, Bytes::from(buf)));
+        reconstruct_data.page_img = Some(Bytes::from(buf));
         Ok(PageReconstructResult::Complete)
     }
 

pageserver/src/layered_repository/inmemory_layer.rs

@@ -20,15 +20,13 @@ use crate::{ZTenantId, ZTimelineId};
 use anyhow::{ensure, Result};
 use bytes::Bytes;
 use log::*;
-use std::collections::HashMap;
-use std::io::Seek;
-use std::os::unix::fs::FileExt;
 use std::path::PathBuf;
 use std::sync::{Arc, RwLock};
-use zenith_utils::bin_ser::BeSer;
 use zenith_utils::lsn::Lsn;
 use zenith_utils::vec_map::VecMap;
 
+use super::page_versions::PageVersions;
+
 pub struct InMemoryLayer {
     conf: &'static PageServerConf,
     tenantid: ZTenantId,
@@ -41,20 +39,8 @@ pub struct InMemoryLayer {
     ///
     start_lsn: Lsn,
 
-    ///
-    /// LSN of the oldest page version stored in this layer.
-    ///
-    /// This is different from 'start_lsn' in that we enforce that the 'start_lsn'
-    /// of a layer always matches the 'end_lsn' of its predecessor, even if there
-    /// are no page versions until at a later LSN. That way you can detect any
-    /// missing layer files more easily. 'oldest_lsn' is the first page version
-    /// actually stored in this layer. In the range between 'start_lsn' and
-    /// 'oldest_lsn', there are no changes to the segment.
-    /// 'oldest_lsn' is used to adjust 'disk_consistent_lsn' and that is why it should
-    /// point to the beginning of WAL record. This is the other difference with 'start_lsn'
-    /// which points to end of WAL record. This is why 'oldest_lsn' can be smaller than 'start_lsn'.
-    ///
-    oldest_lsn: Lsn,
+    /// LSN of the oldest page version stored in this layer
+    oldest_pending_lsn: Lsn,
 
     /// The above fields never change. The parts that do change are in 'inner',
     /// and protected by mutex.
@@ -73,15 +59,11 @@ pub struct InMemoryLayerInner {
     /// The drop LSN is recorded in [`end_lsn`].
     dropped: bool,
 
-    /// The PageVersion structs are stored in a serialized format in this file.
-    /// Each serialized PageVersion is preceded by a 'u32' length field.
-    /// 'page_versions' map stores offsets into this file.
-    file: EphemeralFile,
-
-    /// Metadata about all versions of all pages in the layer is kept
-    /// here.  Indexed by block number and LSN. The value is an offset
-    /// into the ephemeral file where the page version is stored.
-    page_versions: HashMap<SegmentBlk, VecMap<Lsn, u64>>,
+    ///
+    /// All versions of all pages in the layer are are kept here.
+    /// Indexed by block number and LSN.
+    ///
+    page_versions: PageVersions,
 
     ///
     /// `seg_sizes` tracks the size of the segment at different points in time.
@@ -91,14 +73,6 @@ pub struct InMemoryLayerInner {
     /// a non-blocky rel, 'seg_sizes' is not used and is always empty.
     ///
     seg_sizes: VecMap<Lsn, SegmentBlk>,
-
-    ///
-    /// LSN of the newest page version stored in this layer.
-    ///
-    /// The difference between 'end_lsn' and 'latest_lsn' is the same as between
-    /// 'start_lsn' and 'oldest_lsn'. See comments in 'oldest_lsn'.
-    ///
-    latest_lsn: Lsn,
 }
 
 impl InMemoryLayerInner {
@@ -117,50 +91,6 @@ impl InMemoryLayerInner {
             panic!("could not find seg size in in-memory layer");
         }
     }
-
-    ///
-    /// Read a page version from the ephemeral file.
-    ///
-    fn read_pv(&self, off: u64) -> Result<PageVersion> {
-        let mut buf = Vec::new();
-        self.read_pv_bytes(off, &mut buf)?;
-        Ok(PageVersion::des(&buf)?)
-    }
-
-    ///
-    /// Read a page version from the ephemeral file, as raw bytes, at
-    /// the given offset.  The bytes are read into 'buf', which is
-    /// expanded if necessary. Returns the size of the page version.
-    ///
-    fn read_pv_bytes(&self, off: u64, buf: &mut Vec<u8>) -> Result<usize> {
-        // read length
-        let mut lenbuf = [0u8; 4];
-        self.file.read_exact_at(&mut lenbuf, off)?;
-        let len = u32::from_ne_bytes(lenbuf) as usize;
-
-        if buf.len() < len {
-            buf.resize(len, 0);
-        }
-        self.file.read_exact_at(&mut buf[0..len], off + 4)?;
-        Ok(len)
-    }
-
-    fn write_pv(&mut self, pv: &PageVersion) -> Result<u64> {
-        // remember starting position
-        let pos = self.file.stream_position()?;
-
-        // make room for the 'length' field by writing zeros as a placeholder.
-        self.file.seek(std::io::SeekFrom::Start(pos + 4)).unwrap();
-
-        pv.ser_into(&mut self.file).unwrap();
-
-        // write the 'length' field.
-        let len = self.file.stream_position()? - pos - 4;
-        let lenbuf = u32::to_ne_bytes(len as u32);
-        self.file.write_all_at(&lenbuf, pos)?;
-
-        Ok(pos)
-    }
 }
 
 impl Layer for InMemoryLayer {
@@ -224,6 +154,7 @@ impl Layer for InMemoryLayer {
         &self,
         blknum: SegmentBlk,
         lsn: Lsn,
+        cached_img_lsn: Option<Lsn>,
         reconstruct_data: &mut PageReconstructData,
     ) -> Result<PageReconstructResult> {
         let mut need_image = true;
@@ -234,31 +165,33 @@ impl Layer for InMemoryLayer {
             let inner = self.inner.read().unwrap();
 
             // Scan the page versions backwards, starting from `lsn`.
-            if let Some(vec_map) = inner.page_versions.get(&blknum) {
-                let slice = vec_map.slice_range(..=lsn);
-                for (entry_lsn, pos) in slice.iter().rev() {
-                    match &reconstruct_data.page_img {
-                        Some((cached_lsn, _)) if entry_lsn <= cached_lsn => {
-                            return Ok(PageReconstructResult::Complete)
-                        }
-                        _ => {}
+            let iter = inner
+                .page_versions
+                .get_block_lsn_range(blknum, ..=lsn)
+                .iter()
+                .rev();
+            for (entry_lsn, pos) in iter {
+                match &cached_img_lsn {
+                    Some(cached_lsn) if entry_lsn <= cached_lsn => {
+                        return Ok(PageReconstructResult::Cached)
                     }
+                    _ => {}
+                }
 
-                    let pv = inner.read_pv(*pos)?;
-                    match pv {
-                        PageVersion::Page(img) => {
-                            reconstruct_data.page_img = Some((*entry_lsn, img));
+                let pv = inner.page_versions.read_pv(*pos)?;
+                match pv {
+                    PageVersion::Page(img) => {
+                        reconstruct_data.page_img = Some(img);
+                        need_image = false;
+                        break;
+                    }
+                    PageVersion::Wal(rec) => {
+                        reconstruct_data.records.push((*entry_lsn, rec.clone()));
+                        if rec.will_init() {
+                            // This WAL record initializes the page, so no need to go further back
                             need_image = false;
                             break;
                         }
-                        PageVersion::Wal(rec) => {
-                            reconstruct_data.records.push((*entry_lsn, rec.clone()));
-                            if rec.will_init() {
-                                // This WAL record initializes the page, so no need to go further back
-                                need_image = false;
-                                break;
-                            }
-                        }
                     }
                 }
             }
@@ -364,22 +297,14 @@ impl Layer for InMemoryLayer {
             println!("seg_sizes {}: {}", k, v);
         }
 
-        // List the blocks in order
-        let mut page_versions: Vec<(&SegmentBlk, &VecMap<Lsn, u64>)> =
-            inner.page_versions.iter().collect();
-        page_versions.sort_by_key(|k| k.0);
+        for (blknum, lsn, pos) in inner.page_versions.ordered_page_version_iter(None) {
+            let pv = inner.page_versions.read_pv(pos)?;
+            let pv_description = match pv {
+                PageVersion::Page(_img) => "page",
+                PageVersion::Wal(_rec) => "wal",
+            };
 
-        for (blknum, versions) in page_versions {
-            for (lsn, off) in versions.as_slice() {
-                let pv = inner.read_pv(*off);
-                let pv_description = match pv {
-                    Ok(PageVersion::Page(_img)) => "page",
-                    Ok(PageVersion::Wal(_rec)) => "wal",
-                    Err(_err) => "INVALID",
-                };
-
-                println!("blk {} at {}: {}\n", blknum, lsn, pv_description);
-            }
+            println!("blk {} at {}: {}\n", blknum, lsn, pv_description);
         }
 
         Ok(())
@@ -394,13 +319,8 @@ pub struct LayersOnDisk {
 
 impl InMemoryLayer {
     /// Return the oldest page version that's stored in this layer
-    pub fn get_oldest_lsn(&self) -> Lsn {
-        self.oldest_lsn
-    }
-
-    pub fn get_latest_lsn(&self) -> Lsn {
-        let inner = self.inner.read().unwrap();
-        inner.latest_lsn
+    pub fn get_oldest_pending_lsn(&self) -> Lsn {
+        self.oldest_pending_lsn
     }
 
     ///
@@ -412,7 +332,7 @@ impl InMemoryLayer {
         tenantid: ZTenantId,
         seg: SegmentTag,
         start_lsn: Lsn,
-        oldest_lsn: Lsn,
+        oldest_pending_lsn: Lsn,
     ) -> Result<InMemoryLayer> {
         trace!(
             "initializing new empty InMemoryLayer for writing {} on timeline {} at {}",
@@ -435,15 +355,13 @@ impl InMemoryLayer {
             tenantid,
             seg,
             start_lsn,
-            oldest_lsn,
+            oldest_pending_lsn,
             incremental: false,
             inner: RwLock::new(InMemoryLayerInner {
                 end_lsn: None,
                 dropped: false,
-                file,
-                page_versions: HashMap::new(),
+                page_versions: PageVersions::new(file),
                 seg_sizes,
-                latest_lsn: oldest_lsn,
             }),
         })
     }
@@ -480,21 +398,15 @@ impl InMemoryLayer {
         let mut inner = self.inner.write().unwrap();
 
         inner.assert_writeable();
-        assert!(lsn >= inner.latest_lsn);
-        inner.latest_lsn = lsn;
 
-        // Write the page version to the file, and remember its offset in 'page_versions'
-        {
-            let off = inner.write_pv(&pv)?;
-            let vec_map = inner.page_versions.entry(blknum).or_default();
-            let old = vec_map.append_or_update_last(lsn, off).unwrap().0;
-            if old.is_some() {
-                // We already had an entry for this LSN. That's odd..
-                warn!(
-                    "Page version of rel {} blk {} at {} already exists",
-                    self.seg.rel, blknum, lsn
-                );
-            }
+        let old = inner.page_versions.append_or_update_last(blknum, lsn, pv)?;
+
+        if old.is_some() {
+            // We already had an entry for this LSN. That's odd..
+            warn!(
+                "Page version of rel {} blk {} at {} already exists",
+                self.seg.rel, blknum, lsn
+            );
         }
 
         // Also update the relation size, if this extended the relation.
@@ -528,19 +440,16 @@ impl InMemoryLayer {
                         gapblknum,
                         blknum
                     );
+                    let old = inner
+                        .page_versions
+                        .append_or_update_last(gapblknum, lsn, zeropv)?;
+                    // We already had an entry for this LSN. That's odd..
 
-                    // Write the page version to the file, and remember its offset in
-                    // 'page_versions'
-                    {
-                        let off = inner.write_pv(&zeropv)?;
-                        let vec_map = inner.page_versions.entry(gapblknum).or_default();
-                        let old = vec_map.append_or_update_last(lsn, off).unwrap().0;
-                        if old.is_some() {
-                            warn!(
-                                "Page version of seg {} blk {} at {} already exists",
-                                self.seg, gapblknum, lsn
-                            );
-                        }
+                    if old.is_some() {
+                        warn!(
+                            "Page version of seg {} blk {} at {} already exists",
+                            self.seg, blknum, lsn
+                        );
                     }
                 }
 
@@ -600,11 +509,12 @@ impl InMemoryLayer {
         timelineid: ZTimelineId,
         tenantid: ZTenantId,
         start_lsn: Lsn,
-        oldest_lsn: Lsn,
+        oldest_pending_lsn: Lsn,
     ) -> Result<InMemoryLayer> {
         let seg = src.get_seg_tag();
 
-        assert!(oldest_lsn.is_aligned());
+        assert!(oldest_pending_lsn.is_aligned());
+        assert!(oldest_pending_lsn >= start_lsn);
 
         trace!(
             "initializing new InMemoryLayer for writing {} on timeline {} at {}",
@@ -628,15 +538,13 @@ impl InMemoryLayer {
             tenantid,
             seg,
             start_lsn,
-            oldest_lsn,
+            oldest_pending_lsn,
             incremental: true,
             inner: RwLock::new(InMemoryLayerInner {
                 end_lsn: None,
                 dropped: false,
-                file,
-                page_versions: HashMap::new(),
+                page_versions: PageVersions::new(file),
                 seg_sizes,
-                latest_lsn: oldest_lsn,
             }),
         })
     }
@@ -663,10 +571,8 @@ impl InMemoryLayer {
                 assert!(lsn <= &end_lsn, "{:?} {:?}", lsn, end_lsn);
             }
 
-            for (_blk, vec_map) in inner.page_versions.iter() {
-                for (lsn, _pos) in vec_map.as_slice() {
-                    assert!(*lsn <= end_lsn);
-                }
+            for (_blk, lsn, _pv) in inner.page_versions.ordered_page_version_iter(None) {
+                assert!(lsn <= end_lsn);
             }
         }
     }
@@ -744,19 +650,15 @@ impl InMemoryLayer {
                 self.is_dropped(),
             )?;
 
-            // Write all page versions, in block + LSN order
+            // Write all page versions
             let mut buf: Vec<u8> = Vec::new();
 
-            let pv_iter = inner.page_versions.iter();
-            let mut pages: Vec<(&SegmentBlk, &VecMap<Lsn, u64>)> = pv_iter.collect();
-            pages.sort_by_key(|(blknum, _vec_map)| *blknum);
-            for (blknum, vec_map) in pages {
-                for (lsn, pos) in vec_map.as_slice() {
-                    if *lsn < delta_end_lsn {
-                        let len = inner.read_pv_bytes(*pos, &mut buf)?;
-                        delta_layer_writer.put_page_version(*blknum, *lsn, &buf[..len])?;
-                    }
-                }
+            let page_versions_iter = inner
+                .page_versions
+                .ordered_page_version_iter(Some(delta_end_lsn));
+            for (blknum, lsn, pos) in page_versions_iter {
+                let len = inner.page_versions.read_pv_bytes(pos, &mut buf)?;
+                delta_layer_writer.put_page_version(blknum, lsn, &buf[..len])?;
             }
 
             // Create seg_sizes

pageserver/src/layered_repository/layer_map.rs

@@ -40,7 +40,7 @@ pub struct LayerMap {
     /// All the layers keyed by segment tag
     segs: HashMap<SegmentTag, SegEntry>,
 
-    /// All in-memory layers, ordered by 'oldest_lsn' and generation
+    /// All in-memory layers, ordered by 'oldest_pending_lsn' and generation
     /// of each layer. This allows easy access to the in-memory layer that
     /// contains the oldest WAL record.
     open_layers: BinaryHeap<OpenLayerEntry>,
@@ -83,16 +83,16 @@ impl LayerMap {
 
         let layer_id = segentry.update_open(Arc::clone(&layer));
 
-        let oldest_lsn = layer.get_oldest_lsn();
+        let oldest_pending_lsn = layer.get_oldest_pending_lsn();
 
-        // After a crash and restart, 'oldest_lsn' of the oldest in-memory
+        // After a crash and restart, 'oldest_pending_lsn' of the oldest in-memory
         // layer becomes the WAL streaming starting point, so it better not point
         // in the middle of a WAL record.
-        assert!(oldest_lsn.is_aligned());
+        assert!(oldest_pending_lsn.is_aligned());
 
         // Also add it to the binary heap
         let open_layer_entry = OpenLayerEntry {
-            oldest_lsn: layer.get_oldest_lsn(),
+            oldest_pending_lsn: layer.get_oldest_pending_lsn(),
             layer_id,
             generation: self.current_generation,
         };
@@ -352,23 +352,23 @@ impl SegEntry {
 }
 
 /// Entry held in LayerMap::open_layers, with boilerplate comparison routines
-/// to implement a min-heap ordered by 'oldest_lsn' and 'generation'
+/// to implement a min-heap ordered by 'oldest_pending_lsn' and 'generation'
 ///
 /// The generation number associated with each entry can be used to distinguish
 /// recently-added entries (i.e after last call to increment_generation()) from older
-/// entries with the same 'oldest_lsn'.
+/// entries with the same 'oldest_pending_lsn'.
 struct OpenLayerEntry {
-    oldest_lsn: Lsn, // copy of layer.get_oldest_lsn()
+    oldest_pending_lsn: Lsn, // copy of layer.get_oldest_pending_lsn()
     generation: u64,
     layer_id: LayerId,
 }
 impl Ord for OpenLayerEntry {
     fn cmp(&self, other: &Self) -> Ordering {
         // BinaryHeap is a max-heap, and we want a min-heap. Reverse the ordering here
-        // to get that. Entries with identical oldest_lsn are ordered by generation
+        // to get that. Entries with identical oldest_pending_lsn are ordered by generation
         other
-            .oldest_lsn
-            .cmp(&self.oldest_lsn)
+            .oldest_pending_lsn
+            .cmp(&self.oldest_pending_lsn)
             .then_with(|| other.generation.cmp(&self.generation))
     }
 }
@@ -437,7 +437,7 @@ mod tests {
         conf: &'static PageServerConf,
         segno: u32,
         start_lsn: Lsn,
-        oldest_lsn: Lsn,
+        oldest_pending_lsn: Lsn,
     ) -> Arc<InMemoryLayer> {
         Arc::new(
             InMemoryLayer::create(
@@ -449,7 +449,7 @@ mod tests {
                     segno,
                 },
                 start_lsn,
-                oldest_lsn,
+                oldest_pending_lsn,
             )
             .unwrap(),
         )

pageserver/src/layered_repository/page_versions.rs

@@ -0,0 +1,268 @@
+//!
+//! Data structure to ingest incoming WAL into an append-only file.
+//!
+//! - The file is considered temporary, and will be discarded on crash
+//! - based on a B-tree
+//!
+
+use std::os::unix::fs::FileExt;
+use std::{collections::HashMap, ops::RangeBounds, slice};
+
+use anyhow::Result;
+
+use std::cmp::min;
+use std::io::Seek;
+
+use zenith_utils::{lsn::Lsn, vec_map::VecMap};
+
+use super::storage_layer::PageVersion;
+use crate::layered_repository::ephemeral_file::EphemeralFile;
+
+use zenith_utils::bin_ser::BeSer;
+
+const EMPTY_SLICE: &[(Lsn, u64)] = &[];
+
+pub struct PageVersions {
+    map: HashMap<u32, VecMap<Lsn, u64>>,
+
+    /// The PageVersion structs are stored in a serialized format in this file.
+    /// Each serialized PageVersion is preceded by a 'u32' length field.
+    /// The 'map' stores offsets into this file.
+    file: EphemeralFile,
+}
+
+impl PageVersions {
+    pub fn new(file: EphemeralFile) -> PageVersions {
+        PageVersions {
+            map: HashMap::new(),
+            file,
+        }
+    }
+
+    pub fn append_or_update_last(
+        &mut self,
+        blknum: u32,
+        lsn: Lsn,
+        page_version: PageVersion,
+    ) -> Result<Option<u64>> {
+        // remember starting position
+        let pos = self.file.stream_position()?;
+
+        // make room for the 'length' field by writing zeros as a placeholder.
+        self.file.seek(std::io::SeekFrom::Start(pos + 4)).unwrap();
+
+        page_version.ser_into(&mut self.file).unwrap();
+
+        // write the 'length' field.
+        let len = self.file.stream_position()? - pos - 4;
+        let lenbuf = u32::to_ne_bytes(len as u32);
+        self.file.write_all_at(&lenbuf, pos)?;
+
+        let map = self.map.entry(blknum).or_insert_with(VecMap::default);
+        Ok(map.append_or_update_last(lsn, pos as u64).unwrap().0)
+    }
+
+    /// Get all [`PageVersion`]s in a block
+    fn get_block_slice(&self, blknum: u32) -> &[(Lsn, u64)] {
+        self.map
+            .get(&blknum)
+            .map(VecMap::as_slice)
+            .unwrap_or(EMPTY_SLICE)
+    }
+
+    /// Get a range of [`PageVersions`] in a block
+    pub fn get_block_lsn_range<R: RangeBounds<Lsn>>(&self, blknum: u32, range: R) -> &[(Lsn, u64)] {
+        self.map
+            .get(&blknum)
+            .map(|vec_map| vec_map.slice_range(range))
+            .unwrap_or(EMPTY_SLICE)
+    }
+
+    /// Iterate through [`PageVersion`]s in (block, lsn) order.
+    /// If a [`cutoff_lsn`] is set, only show versions with `lsn < cutoff_lsn`
+    pub fn ordered_page_version_iter(&self, cutoff_lsn: Option<Lsn>) -> OrderedPageVersionIter<'_> {
+        let mut ordered_blocks: Vec<u32> = self.map.keys().cloned().collect();
+        ordered_blocks.sort_unstable();
+
+        let slice = ordered_blocks
+            .first()
+            .map(|&blknum| self.get_block_slice(blknum))
+            .unwrap_or(EMPTY_SLICE);
+
+        OrderedPageVersionIter {
+            page_versions: self,
+            ordered_blocks,
+            cur_block_idx: 0,
+            cutoff_lsn,
+            cur_slice_iter: slice.iter(),
+        }
+    }
+
+    ///
+    /// Read a page version.
+    ///
+    pub fn read_pv(&self, off: u64) -> Result<PageVersion> {
+        let mut buf = Vec::new();
+        self.read_pv_bytes(off, &mut buf)?;
+        Ok(PageVersion::des(&buf)?)
+    }
+
+    ///
+    /// Read a page version, as raw bytes, at the given offset. The bytes
+    /// are read into 'buf', which is expanded if necessary. Returns the
+    /// size of the page version.
+    ///
+    pub fn read_pv_bytes(&self, off: u64, buf: &mut Vec<u8>) -> Result<usize> {
+        // read length
+        let mut lenbuf = [0u8; 4];
+        self.file.read_exact_at(&mut lenbuf, off)?;
+        let len = u32::from_ne_bytes(lenbuf) as usize;
+
+        // Resize the buffer to fit the data, if needed.
+        //
+        // We don't shrink the buffer if it's larger than necessary. That avoids
+        // repeatedly shrinking and expanding when you reuse the same buffer to
+        // read multiple page versions. Expanding a Vec requires initializing the
+        // new bytes, which is a waste of time because we're immediately overwriting
+        // it, but there's no way to avoid it without resorting to unsafe code.
+        if buf.len() < len {
+            buf.resize(len, 0);
+        }
+        self.file.read_exact_at(&mut buf[0..len], off + 4)?;
+
+        Ok(len)
+    }
+}
+
+pub struct PageVersionReader<'a> {
+    file: &'a EphemeralFile,
+    pos: u64,
+    end_pos: u64,
+}
+
+impl<'a> std::io::Read for PageVersionReader<'a> {
+    fn read(&mut self, buf: &mut [u8]) -> Result<usize, std::io::Error> {
+        let len = min(buf.len(), (self.end_pos - self.pos) as usize);
+        let n = self.file.read_at(&mut buf[..len], self.pos)?;
+        self.pos += n as u64;
+        Ok(n)
+    }
+}
+
+pub struct OrderedPageVersionIter<'a> {
+    page_versions: &'a PageVersions,
+
+    ordered_blocks: Vec<u32>,
+    cur_block_idx: usize,
+
+    cutoff_lsn: Option<Lsn>,
+
+    cur_slice_iter: slice::Iter<'a, (Lsn, u64)>,
+}
+
+impl OrderedPageVersionIter<'_> {
+    fn is_lsn_before_cutoff(&self, lsn: &Lsn) -> bool {
+        if let Some(cutoff_lsn) = self.cutoff_lsn.as_ref() {
+            lsn < cutoff_lsn
+        } else {
+            true
+        }
+    }
+}
+
+impl<'a> Iterator for OrderedPageVersionIter<'a> {
+    type Item = (u32, Lsn, u64);
+
+    fn next(&mut self) -> Option<Self::Item> {
+        loop {
+            if let Some((lsn, pos)) = self.cur_slice_iter.next() {
+                if self.is_lsn_before_cutoff(lsn) {
+                    let blknum = self.ordered_blocks[self.cur_block_idx];
+                    return Some((blknum, *lsn, *pos));
+                }
+            }
+
+            let next_block_idx = self.cur_block_idx + 1;
+            let blknum: u32 = *self.ordered_blocks.get(next_block_idx)?;
+            self.cur_block_idx = next_block_idx;
+            self.cur_slice_iter = self.page_versions.get_block_slice(blknum).iter();
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use bytes::Bytes;
+
+    use super::*;
+    use crate::config::PageServerConf;
+    use std::fs;
+    use std::str::FromStr;
+    use zenith_utils::zid::{ZTenantId, ZTimelineId};
+
+    fn repo_harness(test_name: &str) -> Result<(&'static PageServerConf, ZTenantId, ZTimelineId)> {
+        let repo_dir = PageServerConf::test_repo_dir(test_name);
+        let _ = fs::remove_dir_all(&repo_dir);
+        let conf = PageServerConf::dummy_conf(repo_dir);
+        // Make a static copy of the config. This can never be free'd, but that's
+        // OK in a test.
+        let conf: &'static PageServerConf = Box::leak(Box::new(conf));
+
+        let tenantid = ZTenantId::from_str("11000000000000000000000000000000").unwrap();
+        let timelineid = ZTimelineId::from_str("22000000000000000000000000000000").unwrap();
+        fs::create_dir_all(conf.timeline_path(&timelineid, &tenantid))?;
+
+        Ok((conf, tenantid, timelineid))
+    }
+
+    #[test]
+    fn test_ordered_iter() -> Result<()> {
+        let (conf, tenantid, timelineid) = repo_harness("test_ordered_iter")?;
+
+        let file = EphemeralFile::create(conf, tenantid, timelineid)?;
+
+        let mut page_versions = PageVersions::new(file);
+
+        const BLOCKS: u32 = 1000;
+        const LSNS: u64 = 50;
+
+        let empty_page = Bytes::from_static(&[0u8; 8192]);
+        let empty_page_version = PageVersion::Page(empty_page);
+
+        for blknum in 0..BLOCKS {
+            for lsn in 0..LSNS {
+                let old = page_versions.append_or_update_last(
+                    blknum,
+                    Lsn(lsn),
+                    empty_page_version.clone(),
+                )?;
+                assert!(old.is_none());
+            }
+        }
+
+        let mut iter = page_versions.ordered_page_version_iter(None);
+        for blknum in 0..BLOCKS {
+            for lsn in 0..LSNS {
+                let (actual_blknum, actual_lsn, _pv) = iter.next().unwrap();
+                assert_eq!(actual_blknum, blknum);
+                assert_eq!(Lsn(lsn), actual_lsn);
+            }
+        }
+        assert!(iter.next().is_none());
+        assert!(iter.next().is_none()); // should be robust against excessive next() calls
+
+        const CUTOFF_LSN: Lsn = Lsn(30);
+        let mut iter = page_versions.ordered_page_version_iter(Some(CUTOFF_LSN));
+        for blknum in 0..BLOCKS {
+            for lsn in 0..CUTOFF_LSN.0 {
+                let (actual_blknum, actual_lsn, _pv) = iter.next().unwrap();
+                assert_eq!(actual_blknum, blknum);
+                assert_eq!(Lsn(lsn), actual_lsn);
+            }
+        }
+        assert!(iter.next().is_none());
+        assert!(iter.next().is_none()); // should be robust against excessive next() calls
+
+        Ok(())
+    }
+}

pageserver/src/layered_repository/storage_layer.rs

@@ -71,26 +71,15 @@ pub enum PageVersion {
 }
 
 ///
-/// Struct used to communicate across calls to 'get_page_reconstruct_data'.
+/// Data needed to reconstruct a page version
 ///
-/// Before first call to get_page_reconstruct_data, you can fill in 'page_img'
-/// if you have an older cached version of the page available. That can save
-/// work in 'get_page_reconstruct_data', as it can stop searching for page
-/// versions when all the WAL records going back to the cached image have been
-/// collected.
-///
-/// When get_page_reconstruct_data returns Complete, 'page_img' is set to an
-/// image of the page, or the oldest WAL record in 'records' is a will_init-type
-/// record that initializes the page without requiring a previous image.
-///
-/// If 'get_page_reconstruct_data' returns Continue, some 'records' may have
-/// been collected, but there are more records outside the current layer. Pass
-/// the same PageReconstructData struct in the next 'get_page_reconstruct_data'
-/// call, to collect more records.
+/// 'page_img' is the old base image of the page to start the WAL replay with.
+/// It can be None, if the first WAL record initializes the page (will_init)
+/// 'records' contains the records to apply over the base image.
 ///
 pub struct PageReconstructData {
     pub records: Vec<(Lsn, ZenithWalRecord)>,
-    pub page_img: Option<(Lsn, Bytes)>,
+    pub page_img: Option<Bytes>,
 }
 
 /// Return value from Layer::get_page_reconstruct_data
@@ -104,6 +93,8 @@ pub enum PageReconstructResult {
     /// the returned LSN. This is usually considered an error, but might be OK
     /// in some circumstances.
     Missing(Lsn),
+    /// Use the cached image at `cached_img_lsn` as the base image
+    Cached,
 }
 
 ///
@@ -147,16 +138,19 @@ pub trait Layer: Send + Sync {
     /// It is up to the caller to collect more data from previous layer and
     /// perform WAL redo, if necessary.
     ///
+    /// `cached_img_lsn` should be set to a cached page image's lsn < `lsn`.
+    /// This function will only return data after `cached_img_lsn`.
+    ///
     /// See PageReconstructResult for possible return values. The collected data
     /// is appended to reconstruct_data; the caller should pass an empty struct
-    /// on first call, or a struct with a cached older image of the page if one
-    /// is available. If this returns PageReconstructResult::Continue, look up
-    /// the predecessor layer and call again with the same 'reconstruct_data' to
-    /// collect more data.
+    /// on first call. If this returns PageReconstructResult::Continue, look up
+    /// the predecessor layer and call again with the same 'reconstruct_data'
+    /// to collect more data.
     fn get_page_reconstruct_data(
         &self,
         blknum: SegmentBlk,
         lsn: Lsn,
+        cached_img_lsn: Option<Lsn>,
         reconstruct_data: &mut PageReconstructData,
     ) -> Result<PageReconstructResult>;
 

pageserver/src/page_service.rs

@@ -18,7 +18,7 @@ use std::io;
 use std::net::TcpListener;
 use std::str;
 use std::str::FromStr;
-use std::sync::{Arc, RwLockReadGuard};
+use std::sync::Arc;
 use tracing::*;
 use zenith_metrics::{register_histogram_vec, HistogramVec};
 use zenith_utils::auth::{self, JwtAuth};
@@ -27,10 +27,13 @@ use zenith_utils::lsn::Lsn;
 use zenith_utils::postgres_backend::is_socket_read_timed_out;
 use zenith_utils::postgres_backend::PostgresBackend;
 use zenith_utils::postgres_backend::{self, AuthType};
-use zenith_utils::pq_proto::{BeMessage, FeMessage, RowDescriptor, SINGLE_COL_ROWDESC};
+use zenith_utils::pq_proto::{
+    BeMessage, FeMessage, RowDescriptor, HELLO_WORLD_ROW, SINGLE_COL_ROWDESC,
+};
 use zenith_utils::zid::{ZTenantId, ZTimelineId};
 
 use crate::basebackup;
+use crate::branches;
 use crate::config::PageServerConf;
 use crate::relish::*;
 use crate::repository::Timeline;
@@ -395,12 +398,7 @@ impl PageServerHandler {
     /// In either case, if the page server hasn't received the WAL up to the
     /// requested LSN yet, we will wait for it to arrive. The return value is
     /// the LSN that should be used to look up the page versions.
-    fn wait_or_get_last_lsn(
-        timeline: &dyn Timeline,
-        mut lsn: Lsn,
-        latest: bool,
-        latest_gc_cutoff_lsn: &RwLockReadGuard<Lsn>,
-    ) -> Result<Lsn> {
+    fn wait_or_get_last_lsn(timeline: &dyn Timeline, lsn: Lsn, latest: bool) -> Result<Lsn> {
         if latest {
             // Latest page version was requested. If LSN is given, it is a hint
             // to the page server that there have been no modifications to the
@@ -421,26 +419,22 @@ impl PageServerHandler {
             // walsender completes the authentication and starts streaming the
             // WAL.
             if lsn <= last_record_lsn {
-                lsn = last_record_lsn;
+                Ok(last_record_lsn)
             } else {
                 timeline.wait_lsn(lsn)?;
                 // Since we waited for 'lsn' to arrive, that is now the last
                 // record LSN. (Or close enough for our purposes; the
                 // last-record LSN can advance immediately after we return
                 // anyway)
+                Ok(lsn)
             }
         } else {
             if lsn == Lsn(0) {
                 bail!("invalid LSN(0) in request");
             }
             timeline.wait_lsn(lsn)?;
+            Ok(lsn)
         }
-        ensure!(
-            lsn >= **latest_gc_cutoff_lsn,
-            "tried to request a page version that was garbage collected. requested at {} gc cutoff {}",
-            lsn, **latest_gc_cutoff_lsn
-        );
-        Ok(lsn)
     }
 
     fn handle_get_rel_exists_request(
@@ -451,8 +445,7 @@ impl PageServerHandler {
         let _enter = info_span!("get_rel_exists", rel = %req.rel, req_lsn = %req.lsn).entered();
 
         let tag = RelishTag::Relation(req.rel);
-        let latest_gc_cutoff_lsn = timeline.get_latest_gc_cutoff_lsn();
-        let lsn = Self::wait_or_get_last_lsn(timeline, req.lsn, req.latest, &latest_gc_cutoff_lsn)?;
+        let lsn = Self::wait_or_get_last_lsn(timeline, req.lsn, req.latest)?;
 
         let exists = timeline.get_rel_exists(tag, lsn)?;
 
@@ -468,8 +461,7 @@ impl PageServerHandler {
     ) -> Result<PagestreamBeMessage> {
         let _enter = info_span!("get_nblocks", rel = %req.rel, req_lsn = %req.lsn).entered();
         let tag = RelishTag::Relation(req.rel);
-        let latest_gc_cutoff_lsn = timeline.get_latest_gc_cutoff_lsn();
-        let lsn = Self::wait_or_get_last_lsn(timeline, req.lsn, req.latest, &latest_gc_cutoff_lsn)?;
+        let lsn = Self::wait_or_get_last_lsn(timeline, req.lsn, req.latest)?;
 
         let n_blocks = timeline.get_relish_size(tag, lsn)?;
 
@@ -490,16 +482,8 @@ impl PageServerHandler {
         let _enter = info_span!("get_page", rel = %req.rel, blkno = &req.blkno, req_lsn = %req.lsn)
             .entered();
         let tag = RelishTag::Relation(req.rel);
-        let latest_gc_cutoff_lsn = timeline.get_latest_gc_cutoff_lsn();
-        let lsn = Self::wait_or_get_last_lsn(timeline, req.lsn, req.latest, &latest_gc_cutoff_lsn)?;
-        /*
-        // Add a 1s delay to some requests. The delayed causes the requests to
-        // hit the race condition from github issue #1047 more easily.
-        use rand::Rng;
-        if rand::thread_rng().gen::<u8>() < 5 {
-            std::thread::sleep(std::time::Duration::from_millis(1000));
-        }
-        */
+        let lsn = Self::wait_or_get_last_lsn(timeline, req.lsn, req.latest)?;
+
         let page = timeline.get_page_at_lsn(tag, req.blkno, lsn)?;
 
         Ok(PagestreamBeMessage::GetPage(PagestreamGetPageResponse {
@@ -520,10 +504,9 @@ impl PageServerHandler {
         // check that the timeline exists
         let timeline = tenant_mgr::get_timeline_for_tenant(tenantid, timelineid)
             .context("Cannot handle basebackup request for a remote timeline")?;
-        let latest_gc_cutoff_lsn = timeline.get_latest_gc_cutoff_lsn();
         if let Some(lsn) = lsn {
             timeline
-                .check_lsn_is_in_scope(lsn, &latest_gc_cutoff_lsn)
+                .check_lsn_is_in_scope(lsn)
                 .context("invalid basebackup lsn")?;
         }
 
@@ -659,21 +642,79 @@ impl postgres_backend::Handler for PageServerHandler {
             walreceiver::launch_wal_receiver(self.conf, tenantid, timelineid, &connstr)?;
 
             pgb.write_message_noflush(&BeMessage::CommandComplete(b"SELECT 1"))?;
+        } else if query_string.starts_with("branch_create ") {
+            let err = || format!("invalid branch_create: '{}'", query_string);
+
+            // branch_create <tenantid> <branchname> <startpoint>
+            // TODO lazy static
+            // TODO: escaping, to allow branch names with spaces
+            let re = Regex::new(r"^branch_create ([[:xdigit:]]+) (\S+) ([^\r\n\s;]+)[\r\n\s;]*;?$")
+                .unwrap();
+            let caps = re.captures(query_string).with_context(err)?;
+
+            let tenantid = ZTenantId::from_str(caps.get(1).unwrap().as_str())?;
+            let branchname = caps.get(2).with_context(err)?.as_str().to_owned();
+            let startpoint_str = caps.get(3).with_context(err)?.as_str().to_owned();
+
+            self.check_permission(Some(tenantid))?;
+
+            let _enter =
+                info_span!("branch_create", name = %branchname, tenant = %tenantid).entered();
+
+            let branch =
+                branches::create_branch(self.conf, &branchname, &startpoint_str, &tenantid)?;
+            let branch = serde_json::to_vec(&branch)?;
+
+            pgb.write_message_noflush(&SINGLE_COL_ROWDESC)?
+                .write_message_noflush(&BeMessage::DataRow(&[Some(&branch)]))?
+                .write_message_noflush(&BeMessage::CommandComplete(b"SELECT 1"))?;
+        } else if query_string.starts_with("branch_list ") {
+            // branch_list <zenith tenantid as hex string>
+            let re = Regex::new(r"^branch_list ([[:xdigit:]]+)$").unwrap();
+            let caps = re
+                .captures(query_string)
+                .with_context(|| format!("invalid branch_list: '{}'", query_string))?;
+
+            let tenantid = ZTenantId::from_str(caps.get(1).unwrap().as_str())?;
+
+            // since these handlers for tenant/branch commands are deprecated (in favor of http based ones)
+            // just use false in place of include non incremental logical size
+            let branches = crate::branches::get_branches(self.conf, &tenantid, false)?;
+            let branches_buf = serde_json::to_vec(&branches)?;
+
+            pgb.write_message_noflush(&SINGLE_COL_ROWDESC)?
+                .write_message_noflush(&BeMessage::DataRow(&[Some(&branches_buf)]))?
+                .write_message_noflush(&BeMessage::CommandComplete(b"SELECT 1"))?;
+        } else if query_string.starts_with("tenant_list") {
+            let tenants = crate::tenant_mgr::list_tenants()?;
+            let tenants_buf = serde_json::to_vec(&tenants)?;
+
+            pgb.write_message_noflush(&SINGLE_COL_ROWDESC)?
+                .write_message_noflush(&BeMessage::DataRow(&[Some(&tenants_buf)]))?
+                .write_message_noflush(&BeMessage::CommandComplete(b"SELECT 1"))?;
+        } else if query_string.starts_with("tenant_create") {
+            let err = || format!("invalid tenant_create: '{}'", query_string);
+
+            // tenant_create <tenantid>
+            let re = Regex::new(r"^tenant_create ([[:xdigit:]]+)$").unwrap();
+            let caps = re.captures(query_string).with_context(err)?;
+
+            self.check_permission(None)?;
+
+            let tenantid = ZTenantId::from_str(caps.get(1).unwrap().as_str())?;
+
+            tenant_mgr::create_repository_for_tenant(self.conf, tenantid)?;
+
+            pgb.write_message_noflush(&SINGLE_COL_ROWDESC)?
+                .write_message_noflush(&BeMessage::CommandComplete(b"SELECT 1"))?;
+        } else if query_string.starts_with("status") {
+            pgb.write_message_noflush(&SINGLE_COL_ROWDESC)?
+                .write_message_noflush(&HELLO_WORLD_ROW)?
+                .write_message_noflush(&BeMessage::CommandComplete(b"SELECT 1"))?;
         } else if query_string.to_ascii_lowercase().starts_with("set ") {
             // important because psycopg2 executes "SET datestyle TO 'ISO'"
             // on connect
             pgb.write_message_noflush(&BeMessage::CommandComplete(b"SELECT 1"))?;
-        } else if query_string.starts_with("failpoints ") {
-            let (_, failpoints) = query_string.split_at("failpoints ".len());
-            for failpoint in failpoints.split(';') {
-                if let Some((name, actions)) = failpoint.split_once('=') {
-                    info!("cfg failpoint: {} {}", name, actions);
-                    fail::cfg(name, actions).unwrap();
-                } else {
-                    bail!("Invalid failpoints format");
-                }
-            }
-            pgb.write_message_noflush(&BeMessage::CommandComplete(b"SELECT 1"))?;
         } else if query_string.starts_with("do_gc ") {
             // Run GC immediately on given timeline.
             // FIXME: This is just for tests. See test_runner/batch_others/test_gc.py.

pageserver/src/remote_storage.rs

@@ -94,7 +94,7 @@ use std::{
 use anyhow::{bail, Context};
 use tokio::io;
 use tracing::{error, info};
-use zenith_utils::zid::{ZTenantId, ZTenantTimelineId, ZTimelineId};
+use zenith_utils::zid::{ZTenantId, ZTimelineId};
 
 pub use self::storage_sync::{schedule_timeline_checkpoint_upload, schedule_timeline_download};
 use self::{local_fs::LocalFs, rust_s3::S3};
@@ -104,7 +104,16 @@ use crate::{
     repository::TimelineSyncState,
 };
 
-pub use storage_sync::compression;
+/// Any timeline has its own id and its own tenant it belongs to,
+/// the sync processes group timelines by both for simplicity.
+#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Copy, Hash)]
+pub struct TimelineSyncId(ZTenantId, ZTimelineId);
+
+impl std::fmt::Display for TimelineSyncId {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "(tenant: {}, timeline: {})", self.0, self.1)
+    }
+}
 
 /// A structure to combine all synchronization data to share with pageserver after a successful sync loop initialization.
 /// Successful initialization includes a case when sync loop is not started, in which case the startup data is returned still,
@@ -158,7 +167,7 @@ pub fn start_local_timeline_sync(
                 ZTenantId,
                 HashMap<ZTimelineId, TimelineSyncState>,
             > = HashMap::new();
-            for (ZTenantTimelineId{tenant_id, timeline_id}, (timeline_metadata, _)) in
+            for (TimelineSyncId(tenant_id, timeline_id), (timeline_metadata, _)) in
                 local_timeline_files
             {
                 initial_timeline_states
@@ -178,7 +187,7 @@ pub fn start_local_timeline_sync(
 
 fn local_tenant_timeline_files(
     config: &'static PageServerConf,
-) -> anyhow::Result<HashMap<ZTenantTimelineId, (TimelineMetadata, Vec<PathBuf>)>> {
+) -> anyhow::Result<HashMap<TimelineSyncId, (TimelineMetadata, Vec<PathBuf>)>> {
     let mut local_tenant_timeline_files = HashMap::new();
     let tenants_dir = config.tenants_path();
     for tenants_dir_entry in fs::read_dir(&tenants_dir)
@@ -213,9 +222,8 @@ fn local_tenant_timeline_files(
 fn collect_timelines_for_tenant(
     config: &'static PageServerConf,
     tenant_path: &Path,
-) -> anyhow::Result<HashMap<ZTenantTimelineId, (TimelineMetadata, Vec<PathBuf>)>> {
-    let mut timelines: HashMap<ZTenantTimelineId, (TimelineMetadata, Vec<PathBuf>)> =
-        HashMap::new();
+) -> anyhow::Result<HashMap<TimelineSyncId, (TimelineMetadata, Vec<PathBuf>)>> {
+    let mut timelines: HashMap<TimelineSyncId, (TimelineMetadata, Vec<PathBuf>)> = HashMap::new();
     let tenant_id = tenant_path
         .file_name()
         .and_then(ffi::OsStr::to_str)
@@ -236,10 +244,7 @@ fn collect_timelines_for_tenant(
                 match collect_timeline_files(&timeline_path) {
                     Ok((timeline_id, metadata, timeline_files)) => {
                         timelines.insert(
-                            ZTenantTimelineId {
-                                tenant_id,
-                                timeline_id,
-                            },
+                            TimelineSyncId(tenant_id, timeline_id),
                             (metadata, timeline_files),
                         );
                     }

pageserver/src/remote_storage/storage_sync.rs

@@ -70,8 +70,7 @@
 //!
 //! When pageserver signals shutdown, current sync task gets finished and the loop exists.
 
-/// Expose the module for a binary CLI tool that deals with the corresponding blobs.
-pub mod compression;
+mod compression;
 mod download;
 pub mod index;
 mod upload;
@@ -106,7 +105,7 @@ use self::{
     },
     upload::upload_timeline_checkpoint,
 };
-use super::{RemoteStorage, SyncStartupData, ZTenantTimelineId};
+use super::{RemoteStorage, SyncStartupData, TimelineSyncId};
 use crate::{
     config::PageServerConf, layered_repository::metadata::TimelineMetadata,
     remote_storage::storage_sync::compression::read_archive_header, repository::TimelineSyncState,
@@ -243,13 +242,13 @@ mod sync_queue {
 /// Limited by the number of retries, after certain threshold the failing task gets evicted and the timeline disabled.
 #[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone)]
 pub struct SyncTask {
-    sync_id: ZTenantTimelineId,
+    sync_id: TimelineSyncId,
     retries: u32,
     kind: SyncKind,
 }
 
 impl SyncTask {
-    fn new(sync_id: ZTenantTimelineId, retries: u32, kind: SyncKind) -> Self {
+    fn new(sync_id: TimelineSyncId, retries: u32, kind: SyncKind) -> Self {
         Self {
             sync_id,
             retries,
@@ -308,10 +307,7 @@ pub fn schedule_timeline_checkpoint_upload(
     }
 
     if !sync_queue::push(SyncTask::new(
-        ZTenantTimelineId {
-            tenant_id,
-            timeline_id,
-        },
+        TimelineSyncId(tenant_id, timeline_id),
         0,
         SyncKind::Upload(NewCheckpoint { layers, metadata }),
     )) {
@@ -342,10 +338,7 @@ pub fn schedule_timeline_download(tenant_id: ZTenantId, timeline_id: ZTimelineId
         tenant_id, timeline_id
     );
     sync_queue::push(SyncTask::new(
-        ZTenantTimelineId {
-            tenant_id,
-            timeline_id,
-        },
+        TimelineSyncId(tenant_id, timeline_id),
         0,
         SyncKind::Download(TimelineDownload {
             files_to_skip: Arc::new(BTreeSet::new()),
@@ -361,7 +354,7 @@ pub(super) fn spawn_storage_sync_thread<
     S: RemoteStorage<StoragePath = P> + Send + Sync + 'static,
 >(
     conf: &'static PageServerConf,
-    local_timeline_files: HashMap<ZTenantTimelineId, (TimelineMetadata, Vec<PathBuf>)>,
+    local_timeline_files: HashMap<TimelineSyncId, (TimelineMetadata, Vec<PathBuf>)>,
     storage: S,
     max_concurrent_sync: NonZeroUsize,
     max_sync_errors: NonZeroU32,
@@ -517,7 +510,7 @@ async fn loop_step<
                 Err(e) => {
                     error!(
                         "Failed to process storage sync task for tenant {}, timeline {}: {:?}",
-                        sync_id.tenant_id, sync_id.timeline_id, e
+                        sync_id.0, sync_id.1, e
                     );
                     None
                 }
@@ -531,10 +524,7 @@ async fn loop_step<
     while let Some((sync_id, state_update)) = task_batch.next().await {
         debug!("Finished storage sync task for sync id {}", sync_id);
         if let Some(state_update) = state_update {
-            let ZTenantTimelineId {
-                tenant_id,
-                timeline_id,
-            } = sync_id;
+            let TimelineSyncId(tenant_id, timeline_id) = sync_id;
             new_timeline_states
                 .entry(tenant_id)
                 .or_default()
@@ -628,7 +618,7 @@ async fn process_task<
 
 fn schedule_first_sync_tasks(
     index: &RemoteTimelineIndex,
-    local_timeline_files: HashMap<ZTenantTimelineId, (TimelineMetadata, Vec<PathBuf>)>,
+    local_timeline_files: HashMap<TimelineSyncId, (TimelineMetadata, Vec<PathBuf>)>,
 ) -> HashMap<ZTenantId, HashMap<ZTimelineId, TimelineSyncState>> {
     let mut initial_timeline_statuses: HashMap<ZTenantId, HashMap<ZTimelineId, TimelineSyncState>> =
         HashMap::new();
@@ -639,10 +629,7 @@ fn schedule_first_sync_tasks(
     for (sync_id, (local_metadata, local_files)) in local_timeline_files {
         let local_disk_consistent_lsn = local_metadata.disk_consistent_lsn();
 
-        let ZTenantTimelineId {
-            tenant_id,
-            timeline_id,
-        } = sync_id;
+        let TimelineSyncId(tenant_id, timeline_id) = sync_id;
         match index.timeline_entry(&sync_id) {
             Some(index_entry) => {
                 let timeline_status = compare_local_and_remote_timeline(
@@ -685,10 +672,10 @@ fn schedule_first_sync_tasks(
         }
     }
 
-    let unprocessed_remote_ids = |remote_id: &ZTenantTimelineId| {
+    let unprocessed_remote_ids = |remote_id: &TimelineSyncId| {
         initial_timeline_statuses
-            .get(&remote_id.tenant_id)
-            .and_then(|timelines| timelines.get(&remote_id.timeline_id))
+            .get(&remote_id.0)
+            .and_then(|timelines| timelines.get(&remote_id.1))
             .is_none()
     };
     for unprocessed_remote_id in index
@@ -696,10 +683,7 @@ fn schedule_first_sync_tasks(
         .filter(unprocessed_remote_ids)
         .collect::<Vec<_>>()
     {
-        let ZTenantTimelineId {
-            tenant_id: cloud_only_tenant_id,
-            timeline_id: cloud_only_timeline_id,
-        } = unprocessed_remote_id;
+        let TimelineSyncId(cloud_only_tenant_id, cloud_only_timeline_id) = unprocessed_remote_id;
         match index
             .timeline_entry(&unprocessed_remote_id)
             .and_then(TimelineIndexEntry::disk_consistent_lsn)
@@ -728,7 +712,7 @@ fn schedule_first_sync_tasks(
 
 fn compare_local_and_remote_timeline(
     new_sync_tasks: &mut VecDeque<SyncTask>,
-    sync_id: ZTenantTimelineId,
+    sync_id: TimelineSyncId,
     local_metadata: TimelineMetadata,
     local_files: Vec<PathBuf>,
     remote_entry: &TimelineIndexEntry,
@@ -785,7 +769,7 @@ async fn update_index_description<
 >(
     (storage, index): &(S, RwLock<RemoteTimelineIndex>),
     timeline_dir: &Path,
-    id: ZTenantTimelineId,
+    id: TimelineSyncId,
 ) -> anyhow::Result<RemoteTimeline> {
     let mut index_write = index.write().await;
     let full_index = match index_write.timeline_entry(&id) {
@@ -808,7 +792,7 @@ async fn update_index_description<
                         Ok((archive_id, header_size, header)) => full_index.update_archive_contents(archive_id.0, header, header_size),
                         Err((e, archive_id)) => bail!(
                             "Failed to download archive header for tenant {}, timeline {}, archive for Lsn {}: {}",
-                            id.tenant_id, id.timeline_id, archive_id.0,
+                            id.0, id.1, archive_id.0,
                             e
                         ),
                     }
@@ -886,7 +870,7 @@ mod test_utils {
         timeline_id: ZTimelineId,
         new_upload: NewCheckpoint,
     ) {
-        let sync_id = ZTenantTimelineId::new(harness.tenant_id, timeline_id);
+        let sync_id = TimelineSyncId(harness.tenant_id, timeline_id);
         upload_timeline_checkpoint(
             harness.conf,
             Arc::clone(&remote_assets),
@@ -942,7 +926,7 @@ mod test_utils {
 
     pub async fn expect_timeline(
         index: &RwLock<RemoteTimelineIndex>,
-        sync_id: ZTenantTimelineId,
+        sync_id: TimelineSyncId,
     ) -> RemoteTimeline {
         if let Some(TimelineIndexEntry::Full(remote_timeline)) =
             index.read().await.timeline_entry(&sync_id)
@@ -977,18 +961,18 @@ mod test_utils {
         let mut expected_timeline_entries = BTreeMap::new();
         for sync_id in actual_sync_ids {
             actual_branches.insert(
-                sync_id.tenant_id,
+                sync_id.1,
                 index_read
-                    .branch_files(sync_id.tenant_id)
+                    .branch_files(sync_id.0)
                     .into_iter()
                     .flat_map(|branch_paths| branch_paths.iter())
                     .cloned()
                     .collect::<BTreeSet<_>>(),
             );
             expected_branches.insert(
-                sync_id.tenant_id,
+                sync_id.1,
                 expected_index_with_descriptions
-                    .branch_files(sync_id.tenant_id)
+                    .branch_files(sync_id.0)
                     .into_iter()
                     .flat_map(|branch_paths| branch_paths.iter())
                     .cloned()

pageserver/src/remote_storage/storage_sync/compression.rs

@@ -248,7 +248,7 @@ fn archive_name(disk_consistent_lsn: Lsn, header_size: u64) -> String {
     archive_name
 }
 
-pub async fn uncompress_with_header(
+async fn uncompress_with_header(
     files_to_skip: &BTreeSet<PathBuf>,
     destination_dir: &Path,
     header: ArchiveHeader,

pageserver/src/remote_storage/storage_sync/download.rs

@@ -17,7 +17,7 @@ use crate::{
             compression, index::TimelineIndexEntry, sync_queue, tenant_branch_files,
             update_index_description, SyncKind, SyncTask,
         },
-        RemoteStorage, ZTenantTimelineId,
+        RemoteStorage, TimelineSyncId,
     },
 };
 
@@ -52,16 +52,13 @@ pub(super) async fn download_timeline<
 >(
     conf: &'static PageServerConf,
     remote_assets: Arc<(S, RwLock<RemoteTimelineIndex>)>,
-    sync_id: ZTenantTimelineId,
+    sync_id: TimelineSyncId,
     mut download: TimelineDownload,
     retries: u32,
 ) -> DownloadedTimeline {
     debug!("Downloading layers for sync id {}", sync_id);
 
-    let ZTenantTimelineId {
-        tenant_id,
-        timeline_id,
-    } = sync_id;
+    let TimelineSyncId(tenant_id, timeline_id) = sync_id;
     let index_read = remote_assets.1.read().await;
     let remote_timeline = match index_read.timeline_entry(&sync_id) {
         None => {
@@ -113,8 +110,7 @@ pub(super) async fn download_timeline<
         }
     };
 
-    if let Err(e) = download_missing_branches(conf, remote_assets.as_ref(), sync_id.tenant_id).await
-    {
+    if let Err(e) = download_missing_branches(conf, remote_assets.as_ref(), sync_id.0).await {
         error!(
             "Failed to download missing branches for sync id {}: {:?}",
             sync_id, e
@@ -184,10 +180,7 @@ async fn try_download_archive<
     S: RemoteStorage<StoragePath = P> + Send + Sync + 'static,
 >(
     conf: &'static PageServerConf,
-    ZTenantTimelineId {
-        tenant_id,
-        timeline_id,
-    }: ZTenantTimelineId,
+    TimelineSyncId(tenant_id, timeline_id): TimelineSyncId,
     remote_assets: Arc<(S, RwLock<RemoteTimelineIndex>)>,
     remote_timeline: &RemoteTimeline,
     archive_id: ArchiveId,
@@ -350,7 +343,7 @@ mod tests {
     #[tokio::test]
     async fn test_download_timeline() -> anyhow::Result<()> {
         let repo_harness = RepoHarness::create("test_download_timeline")?;
-        let sync_id = ZTenantTimelineId::new(repo_harness.tenant_id, TIMELINE_ID);
+        let sync_id = TimelineSyncId(repo_harness.tenant_id, TIMELINE_ID);
         let storage = LocalFs::new(tempdir()?.path().to_owned(), &repo_harness.conf.workdir)?;
         let index = RwLock::new(RemoteTimelineIndex::try_parse_descriptions_from_paths(
             repo_harness.conf,

pageserver/src/remote_storage/storage_sync/index.rs

@@ -22,7 +22,7 @@ use crate::{
     layered_repository::TIMELINES_SEGMENT_NAME,
     remote_storage::{
         storage_sync::compression::{parse_archive_name, FileEntry},
-        ZTenantTimelineId,
+        TimelineSyncId,
     },
 };
 
@@ -53,7 +53,7 @@ impl RelativePath {
 #[derive(Debug, Clone)]
 pub struct RemoteTimelineIndex {
     branch_files: HashMap<ZTenantId, HashSet<RelativePath>>,
-    timeline_files: HashMap<ZTenantTimelineId, TimelineIndexEntry>,
+    timeline_files: HashMap<TimelineSyncId, TimelineIndexEntry>,
 }
 
 impl RemoteTimelineIndex {
@@ -80,22 +80,19 @@ impl RemoteTimelineIndex {
         index
     }
 
-    pub fn timeline_entry(&self, id: &ZTenantTimelineId) -> Option<&TimelineIndexEntry> {
+    pub fn timeline_entry(&self, id: &TimelineSyncId) -> Option<&TimelineIndexEntry> {
         self.timeline_files.get(id)
     }
 
-    pub fn timeline_entry_mut(
-        &mut self,
-        id: &ZTenantTimelineId,
-    ) -> Option<&mut TimelineIndexEntry> {
+    pub fn timeline_entry_mut(&mut self, id: &TimelineSyncId) -> Option<&mut TimelineIndexEntry> {
         self.timeline_files.get_mut(id)
     }
 
-    pub fn add_timeline_entry(&mut self, id: ZTenantTimelineId, entry: TimelineIndexEntry) {
+    pub fn add_timeline_entry(&mut self, id: TimelineSyncId, entry: TimelineIndexEntry) {
         self.timeline_files.insert(id, entry);
     }
 
-    pub fn all_sync_ids(&self) -> impl Iterator<Item = ZTenantTimelineId> + '_ {
+    pub fn all_sync_ids(&self) -> impl Iterator<Item = TimelineSyncId> + '_ {
         self.timeline_files.keys().copied()
     }
 
@@ -351,10 +348,7 @@ fn try_parse_index_entry(
                 .to_string_lossy()
                 .to_string();
 
-            let sync_id = ZTenantTimelineId {
-                tenant_id,
-                timeline_id,
-            };
+            let sync_id = TimelineSyncId(tenant_id, timeline_id);
             let timeline_index_entry = index
                 .timeline_files
                 .entry(sync_id)

pageserver/src/remote_storage/storage_sync/upload.rs

@@ -17,7 +17,7 @@ use crate::{
             index::{RemoteTimeline, TimelineIndexEntry},
             sync_queue, tenant_branch_files, update_index_description, SyncKind, SyncTask,
         },
-        RemoteStorage, ZTenantTimelineId,
+        RemoteStorage, TimelineSyncId,
     },
 };
 
@@ -36,13 +36,12 @@ pub(super) async fn upload_timeline_checkpoint<
 >(
     config: &'static PageServerConf,
     remote_assets: Arc<(S, RwLock<RemoteTimelineIndex>)>,
-    sync_id: ZTenantTimelineId,
+    sync_id: TimelineSyncId,
     new_checkpoint: NewCheckpoint,
     retries: u32,
 ) -> Option<bool> {
     debug!("Uploading checkpoint for sync id {}", sync_id);
-    if let Err(e) = upload_missing_branches(config, remote_assets.as_ref(), sync_id.tenant_id).await
-    {
+    if let Err(e) = upload_missing_branches(config, remote_assets.as_ref(), sync_id.0).await {
         error!(
             "Failed to upload missing branches for sync id {}: {:?}",
             sync_id, e
@@ -58,10 +57,7 @@ pub(super) async fn upload_timeline_checkpoint<
 
     let index = &remote_assets.1;
 
-    let ZTenantTimelineId {
-        tenant_id,
-        timeline_id,
-    } = sync_id;
+    let TimelineSyncId(tenant_id, timeline_id) = sync_id;
     let timeline_dir = config.timeline_path(&timeline_id, &tenant_id);
 
     let index_read = index.read().await;
@@ -155,14 +151,11 @@ async fn try_upload_checkpoint<
 >(
     config: &'static PageServerConf,
     remote_assets: Arc<(S, RwLock<RemoteTimelineIndex>)>,
-    sync_id: ZTenantTimelineId,
+    sync_id: TimelineSyncId,
     new_checkpoint: &NewCheckpoint,
     files_to_skip: BTreeSet<PathBuf>,
 ) -> anyhow::Result<(ArchiveHeader, u64)> {
-    let ZTenantTimelineId {
-        tenant_id,
-        timeline_id,
-    } = sync_id;
+    let TimelineSyncId(tenant_id, timeline_id) = sync_id;
     let timeline_dir = config.timeline_path(&timeline_id, &tenant_id);
 
     let files_to_upload = new_checkpoint
@@ -295,7 +288,7 @@ mod tests {
     #[tokio::test]
     async fn reupload_timeline() -> anyhow::Result<()> {
         let repo_harness = RepoHarness::create("reupload_timeline")?;
-        let sync_id = ZTenantTimelineId::new(repo_harness.tenant_id, TIMELINE_ID);
+        let sync_id = TimelineSyncId(repo_harness.tenant_id, TIMELINE_ID);
         let storage = LocalFs::new(tempdir()?.path().to_owned(), &repo_harness.conf.workdir)?;
         let index = RwLock::new(RemoteTimelineIndex::try_parse_descriptions_from_paths(
             repo_harness.conf,
@@ -491,7 +484,7 @@ mod tests {
     #[tokio::test]
     async fn reupload_timeline_rejected() -> anyhow::Result<()> {
         let repo_harness = RepoHarness::create("reupload_timeline_rejected")?;
-        let sync_id = ZTenantTimelineId::new(repo_harness.tenant_id, TIMELINE_ID);
+        let sync_id = TimelineSyncId(repo_harness.tenant_id, TIMELINE_ID);
         let storage = LocalFs::new(tempdir()?.path().to_owned(), &repo_harness.conf.workdir)?;
         let index = RwLock::new(RemoteTimelineIndex::try_parse_descriptions_from_paths(
             repo_harness.conf,

pageserver/src/repository.rs

@@ -7,7 +7,7 @@ use postgres_ffi::{MultiXactId, MultiXactOffset, TransactionId};
 use serde::{Deserialize, Serialize};
 use std::collections::HashSet;
 use std::ops::{AddAssign, Deref};
-use std::sync::{Arc, RwLockReadGuard};
+use std::sync::Arc;
 use std::time::Duration;
 use zenith_utils::lsn::{Lsn, RecordLsn};
 use zenith_utils::zid::ZTimelineId;
@@ -184,9 +184,6 @@ pub trait Timeline: Send + Sync {
     ///
     fn wait_lsn(&self, lsn: Lsn) -> Result<()>;
 
-    /// Lock and get timeline's GC cuttof
-    fn get_latest_gc_cutoff_lsn(&self) -> RwLockReadGuard<Lsn>;
-
     /// Look up given page version.
     fn get_page_at_lsn(&self, tag: RelishTag, blknum: BlockNumber, lsn: Lsn) -> Result<Bytes>;
 
@@ -220,12 +217,10 @@ pub trait Timeline: Send + Sync {
 
     /// Atomically get both last and prev.
     fn get_last_record_rlsn(&self) -> RecordLsn;
-
     /// Get last or prev record separately. Same as get_last_record_rlsn().last/prev.
     fn get_last_record_lsn(&self) -> Lsn;
-
     fn get_prev_record_lsn(&self) -> Lsn;
-
+    fn get_start_lsn(&self) -> Lsn;
     fn get_disk_consistent_lsn(&self) -> Lsn;
 
     /// Mutate the timeline with a [`TimelineWriter`].
@@ -240,11 +235,7 @@ pub trait Timeline: Send + Sync {
 
     ///
     /// Check that it is valid to request operations with that lsn.
-    fn check_lsn_is_in_scope(
-        &self,
-        lsn: Lsn,
-        latest_gc_cutoff_lsn: &RwLockReadGuard<Lsn>,
-    ) -> Result<()>;
+    fn check_lsn_is_in_scope(&self, lsn: Lsn) -> Result<()>;
 
     /// Retrieve current logical size of the timeline
     ///
@@ -306,12 +297,8 @@ pub enum ZenithWalRecord {
     /// Native PostgreSQL WAL record
     Postgres { will_init: bool, rec: Bytes },
 
-    /// Clear bits in heap visibility map. ('flags' is bitmap of bits to clear)
-    ClearVisibilityMapFlags {
-        new_heap_blkno: Option<u32>,
-        old_heap_blkno: Option<u32>,
-        flags: u8,
-    },
+    /// Set bits in heap visibility map. (heap blkno, flag bits to clear)
+    ClearVisibilityMapFlags { heap_blkno: u32, flags: u8 },
     /// Mark transaction IDs as committed on a CLOG page
     ClogSetCommitted { xids: Vec<TransactionId> },
     /// Mark transaction IDs as aborted on a CLOG page
@@ -447,6 +434,8 @@ pub mod repo_harness {
 #[allow(clippy::bool_assert_comparison)]
 #[cfg(test)]
 mod tests {
+    use crate::layered_repository::metadata::METADATA_FILE_NAME;
+
     use super::repo_harness::*;
     use super::*;
     use postgres_ffi::{pg_constants, xlog_utils::SIZEOF_CHECKPOINT};
@@ -744,8 +733,8 @@ mod tests {
 
         let mut lsn = 0x10;
         for blknum in 0..pg_constants::RELSEG_SIZE + 1 {
-            lsn += 0x10;
             let img = TEST_IMG(&format!("foo blk {} at {}", blknum, Lsn(lsn)));
+            lsn += 0x10;
             writer.put_page_image(TESTREL_A, blknum as BlockNumber, Lsn(lsn), img)?;
         }
         writer.advance_last_record_lsn(Lsn(lsn));
@@ -998,7 +987,7 @@ mod tests {
                     .source()
                     .unwrap()
                     .to_string()
-                    .contains("is earlier than latest GC horizon"));
+                    .contains("is earlier than initdb lsn"));
             }
         }
 
@@ -1015,11 +1004,12 @@ mod tests {
         make_some_layers(&tline, Lsn(0x20))?;
 
         repo.gc_iteration(Some(TIMELINE_ID), 0x10, false)?;
-        let latest_gc_cutoff_lsn = tline.get_latest_gc_cutoff_lsn();
-        assert!(*latest_gc_cutoff_lsn > Lsn(0x25));
+
         match tline.get_page_at_lsn(TESTREL_A, 0, Lsn(0x25)) {
             Ok(_) => panic!("request for page should have failed"),
-            Err(err) => assert!(err.to_string().contains("not found at")),
+            Err(err) => assert!(err
+                .to_string()
+                .contains("tried to request a page version that was garbage collected")),
         }
         Ok(())
     }
@@ -1130,4 +1120,141 @@ mod tests {
 
         Ok(())
     }
+
+    #[test]
+    fn corrupt_metadata() -> Result<()> {
+        const TEST_NAME: &str = "corrupt_metadata";
+        let harness = RepoHarness::create(TEST_NAME)?;
+        let repo = harness.load();
+
+        repo.create_empty_timeline(TIMELINE_ID, Lsn(0))?;
+        drop(repo);
+
+        let metadata_path = harness.timeline_path(&TIMELINE_ID).join(METADATA_FILE_NAME);
+
+        assert!(metadata_path.is_file());
+
+        let mut metadata_bytes = std::fs::read(&metadata_path)?;
+        assert_eq!(metadata_bytes.len(), 512);
+        metadata_bytes[512 - 4 - 2] ^= 1;
+        std::fs::write(metadata_path, metadata_bytes)?;
+
+        let new_repo = harness.load();
+        let err = new_repo.get_timeline(TIMELINE_ID).err().unwrap();
+        assert_eq!(err.to_string(), "failed to load metadata");
+        assert_eq!(
+            err.source().unwrap().to_string(),
+            "metadata checksum mismatch"
+        );
+
+        Ok(())
+    }
+
+    #[test]
+    fn future_layerfiles() -> Result<()> {
+        const TEST_NAME: &str = "future_layerfiles";
+        let harness = RepoHarness::create(TEST_NAME)?;
+        let repo = harness.load();
+
+        // Create a timeline with disk_consistent_lsn = 8000
+        let tline = repo.create_empty_timeline(TIMELINE_ID, Lsn(0x8000))?;
+        let writer = tline.writer();
+        writer.advance_last_record_lsn(Lsn(0x8000));
+        drop(writer);
+        repo.checkpoint_iteration(CheckpointConfig::Forced)?;
+        drop(repo);
+
+        let timeline_path = harness.timeline_path(&TIMELINE_ID);
+
+        let make_empty_file = |filename: &str| -> std::io::Result<()> {
+            let path = timeline_path.join(filename);
+
+            assert!(!path.exists());
+            std::fs::write(&path, &[])?;
+
+            Ok(())
+        };
+
+        // Helper function to check that a relation file exists, and a corresponding
+        // <filename>.0.old file does not.
+        let assert_exists = |filename: &str| {
+            let path = timeline_path.join(filename);
+            assert!(path.exists(), "file {} was removed", filename);
+
+            // Check that there is no .old file
+            let backup_path = timeline_path.join(format!("{}.0.old", filename));
+            assert!(
+                !backup_path.exists(),
+                "unexpected backup file {}",
+                backup_path.display()
+            );
+        };
+
+        // Helper function to check that a relation file does *not* exists, and a corresponding
+        // <filename>.<num>.old file does.
+        let assert_is_renamed = |filename: &str, num: u32| {
+            let path = timeline_path.join(filename);
+            assert!(
+                !path.exists(),
+                "file {} was not removed as expected",
+                filename
+            );
+
+            let backup_path = timeline_path.join(format!("{}.{}.old", filename, num));
+            assert!(
+                backup_path.exists(),
+                "backup file {} was not created",
+                backup_path.display()
+            );
+        };
+
+        // These files are considered to be in the future and will be renamed out
+        // of the way
+        let future_filenames = vec![
+            format!("pg_control_0_{:016X}", 0x8001),
+            format!("pg_control_0_{:016X}_{:016X}", 0x8001, 0x8008),
+        ];
+        // But these are not:
+        let past_filenames = vec![
+            format!("pg_control_0_{:016X}", 0x8000),
+            format!("pg_control_0_{:016X}_{:016X}", 0x7000, 0x8001),
+        ];
+
+        for filename in future_filenames.iter().chain(past_filenames.iter()) {
+            make_empty_file(filename)?;
+        }
+
+        // Load the timeline. This will cause the files in the "future" to be renamed
+        // away.
+        let new_repo = harness.load();
+        new_repo.get_timeline(TIMELINE_ID).unwrap();
+        drop(new_repo);
+
+        for filename in future_filenames.iter() {
+            assert_is_renamed(filename, 0);
+        }
+        for filename in past_filenames.iter() {
+            assert_exists(filename);
+        }
+
+        // Create the future files again, and load again. They should be renamed to
+        // *.1.old this time.
+        for filename in future_filenames.iter() {
+            make_empty_file(filename)?;
+        }
+
+        let new_repo = harness.load();
+        new_repo.get_timeline(TIMELINE_ID).unwrap();
+        drop(new_repo);
+
+        for filename in future_filenames.iter() {
+            assert_is_renamed(filename, 0);
+            assert_is_renamed(filename, 1);
+        }
+        for filename in past_filenames.iter() {
+            assert_exists(filename);
+        }
+
+        Ok(())
+    }
 }

pageserver/src/virtual_file.rs

@@ -332,11 +332,8 @@ impl VirtualFile {
         // TODO: We could downgrade the locks to read mode before calling
         // 'func', to allow a little bit more concurrency, but the standard
         // library RwLock doesn't allow downgrading without releasing the lock,
-        // and that doesn't seem worth the trouble.
-        //
-        // XXX: `parking_lot::RwLock` can enable such downgrades, yet its implemenation is fair and
-        // may deadlock on subsequent read calls.
-        // Simply replacing all `RwLock` in project causes deadlocks, so use it sparingly.
+        // and that doesn't seem worth the trouble. (parking_lot RwLock would
+        // allow it)
         let result = STORAGE_IO_TIME
             .with_label_values(&[op, &self.tenantid, &self.timelineid])
             .observe_closure_duration(|| func(&file));

pageserver/src/walingest.rs

@@ -349,25 +349,49 @@ impl WalIngest {
         decoded: &mut DecodedWALRecord,
     ) -> Result<()> {
         // Handle VM bit updates that are implicitly part of heap records.
-
-        // First, look at the record to determine which VM bits need
-        // to be cleared. If either of these variables is set, we
-        // need to clear the corresponding bits in the visibility map.
-        let mut new_heap_blkno: Option<u32> = None;
-        let mut old_heap_blkno: Option<u32> = None;
         if decoded.xl_rmid == pg_constants::RM_HEAP_ID {
             let info = decoded.xl_info & pg_constants::XLOG_HEAP_OPMASK;
             if info == pg_constants::XLOG_HEAP_INSERT {
                 let xlrec = XlHeapInsert::decode(buf);
                 assert_eq!(0, buf.remaining());
-                if (xlrec.flags & pg_constants::XLH_INSERT_ALL_VISIBLE_CLEARED) != 0 {
-                    new_heap_blkno = Some(decoded.blocks[0].blkno);
+                if (xlrec.flags
+                    & (pg_constants::XLH_INSERT_ALL_VISIBLE_CLEARED
+                        | pg_constants::XLH_INSERT_ALL_FROZEN_SET))
+                    != 0
+                {
+                    timeline.put_wal_record(
+                        lsn,
+                        RelishTag::Relation(RelTag {
+                            forknum: pg_constants::VISIBILITYMAP_FORKNUM,
+                            spcnode: decoded.blocks[0].rnode_spcnode,
+                            dbnode: decoded.blocks[0].rnode_dbnode,
+                            relnode: decoded.blocks[0].rnode_relnode,
+                        }),
+                        decoded.blocks[0].blkno / pg_constants::HEAPBLOCKS_PER_PAGE as u32,
+                        ZenithWalRecord::ClearVisibilityMapFlags {
+                            heap_blkno: decoded.blocks[0].blkno,
+                            flags: pg_constants::VISIBILITYMAP_VALID_BITS,
+                        },
+                    )?;
                 }
             } else if info == pg_constants::XLOG_HEAP_DELETE {
                 let xlrec = XlHeapDelete::decode(buf);
                 assert_eq!(0, buf.remaining());
                 if (xlrec.flags & pg_constants::XLH_DELETE_ALL_VISIBLE_CLEARED) != 0 {
-                    new_heap_blkno = Some(decoded.blocks[0].blkno);
+                    timeline.put_wal_record(
+                        lsn,
+                        RelishTag::Relation(RelTag {
+                            forknum: pg_constants::VISIBILITYMAP_FORKNUM,
+                            spcnode: decoded.blocks[0].rnode_spcnode,
+                            dbnode: decoded.blocks[0].rnode_dbnode,
+                            relnode: decoded.blocks[0].rnode_relnode,
+                        }),
+                        decoded.blocks[0].blkno / pg_constants::HEAPBLOCKS_PER_PAGE as u32,
+                        ZenithWalRecord::ClearVisibilityMapFlags {
+                            heap_blkno: decoded.blocks[0].blkno,
+                            flags: pg_constants::VISIBILITYMAP_VALID_BITS,
+                        },
+                    )?;
                 }
             } else if info == pg_constants::XLOG_HEAP_UPDATE
                 || info == pg_constants::XLOG_HEAP_HOT_UPDATE
@@ -376,15 +400,39 @@ impl WalIngest {
                 // the size of tuple data is inferred from the size of the record.
                 // we can't validate the remaining number of bytes without parsing
                 // the tuple data.
-                if (xlrec.flags & pg_constants::XLH_UPDATE_OLD_ALL_VISIBLE_CLEARED) != 0 {
-                    old_heap_blkno = Some(decoded.blocks[0].blkno);
-                }
                 if (xlrec.flags & pg_constants::XLH_UPDATE_NEW_ALL_VISIBLE_CLEARED) != 0 {
-                    // PostgreSQL only uses XLH_UPDATE_NEW_ALL_VISIBLE_CLEARED on a
-                    // non-HOT update where the new tuple goes to different page than
-                    // the old one. Otherwise, only XLH_UPDATE_OLD_ALL_VISIBLE_CLEARED is
-                    // set.
-                    new_heap_blkno = Some(decoded.blocks[1].blkno);
+                    timeline.put_wal_record(
+                        lsn,
+                        RelishTag::Relation(RelTag {
+                            forknum: pg_constants::VISIBILITYMAP_FORKNUM,
+                            spcnode: decoded.blocks[0].rnode_spcnode,
+                            dbnode: decoded.blocks[0].rnode_dbnode,
+                            relnode: decoded.blocks[0].rnode_relnode,
+                        }),
+                        decoded.blocks[0].blkno / pg_constants::HEAPBLOCKS_PER_PAGE as u32,
+                        ZenithWalRecord::ClearVisibilityMapFlags {
+                            heap_blkno: decoded.blocks[0].blkno,
+                            flags: pg_constants::VISIBILITYMAP_VALID_BITS,
+                        },
+                    )?;
+                }
+                if (xlrec.flags & pg_constants::XLH_UPDATE_OLD_ALL_VISIBLE_CLEARED) != 0
+                    && decoded.blocks.len() > 1
+                {
+                    timeline.put_wal_record(
+                        lsn,
+                        RelishTag::Relation(RelTag {
+                            forknum: pg_constants::VISIBILITYMAP_FORKNUM,
+                            spcnode: decoded.blocks[1].rnode_spcnode,
+                            dbnode: decoded.blocks[1].rnode_dbnode,
+                            relnode: decoded.blocks[1].rnode_relnode,
+                        }),
+                        decoded.blocks[1].blkno / pg_constants::HEAPBLOCKS_PER_PAGE as u32,
+                        ZenithWalRecord::ClearVisibilityMapFlags {
+                            heap_blkno: decoded.blocks[1].blkno,
+                            flags: pg_constants::VISIBILITYMAP_VALID_BITS,
+                        },
+                    )?;
                 }
             }
         } else if decoded.xl_rmid == pg_constants::RM_HEAP2_ID {
@@ -400,67 +448,32 @@ impl WalIngest {
                 };
                 assert_eq!(offset_array_len, buf.remaining());
 
-                if (xlrec.flags & pg_constants::XLH_INSERT_ALL_VISIBLE_CLEARED) != 0 {
-                    new_heap_blkno = Some(decoded.blocks[0].blkno);
+                // FIXME: why also ALL_FROZEN_SET?
+                if (xlrec.flags
+                    & (pg_constants::XLH_INSERT_ALL_VISIBLE_CLEARED
+                        | pg_constants::XLH_INSERT_ALL_FROZEN_SET))
+                    != 0
+                {
+                    timeline.put_wal_record(
+                        lsn,
+                        RelishTag::Relation(RelTag {
+                            forknum: pg_constants::VISIBILITYMAP_FORKNUM,
+                            spcnode: decoded.blocks[0].rnode_spcnode,
+                            dbnode: decoded.blocks[0].rnode_dbnode,
+                            relnode: decoded.blocks[0].rnode_relnode,
+                        }),
+                        decoded.blocks[0].blkno / pg_constants::HEAPBLOCKS_PER_PAGE as u32,
+                        ZenithWalRecord::ClearVisibilityMapFlags {
+                            heap_blkno: decoded.blocks[0].blkno,
+                            flags: pg_constants::VISIBILITYMAP_VALID_BITS,
+                        },
+                    )?;
                 }
             }
         }
+
         // FIXME: What about XLOG_HEAP_LOCK and XLOG_HEAP2_LOCK_UPDATED?
 
-        // Clear the VM bits if required.
-        if new_heap_blkno.is_some() || old_heap_blkno.is_some() {
-            let vm_relish = RelishTag::Relation(RelTag {
-                forknum: pg_constants::VISIBILITYMAP_FORKNUM,
-                spcnode: decoded.blocks[0].rnode_spcnode,
-                dbnode: decoded.blocks[0].rnode_dbnode,
-                relnode: decoded.blocks[0].rnode_relnode,
-            });
-
-            let new_vm_blk = new_heap_blkno.map(pg_constants::HEAPBLK_TO_MAPBLOCK);
-            let old_vm_blk = old_heap_blkno.map(pg_constants::HEAPBLK_TO_MAPBLOCK);
-            if new_vm_blk == old_vm_blk {
-                // An UPDATE record that needs to clear the bits for both old and the
-                // new page, both of which reside on the same VM page.
-                timeline.put_wal_record(
-                    lsn,
-                    vm_relish,
-                    new_vm_blk.unwrap(),
-                    ZenithWalRecord::ClearVisibilityMapFlags {
-                        new_heap_blkno,
-                        old_heap_blkno,
-                        flags: pg_constants::VISIBILITYMAP_VALID_BITS,
-                    },
-                )?;
-            } else {
-                // Clear VM bits for one heap page, or for two pages that reside on
-                // different VM pages.
-                if let Some(new_vm_blk) = new_vm_blk {
-                    timeline.put_wal_record(
-                        lsn,
-                        vm_relish,
-                        new_vm_blk,
-                        ZenithWalRecord::ClearVisibilityMapFlags {
-                            new_heap_blkno,
-                            old_heap_blkno: None,
-                            flags: pg_constants::VISIBILITYMAP_VALID_BITS,
-                        },
-                    )?;
-                }
-                if let Some(old_vm_blk) = old_vm_blk {
-                    timeline.put_wal_record(
-                        lsn,
-                        vm_relish,
-                        old_vm_blk,
-                        ZenithWalRecord::ClearVisibilityMapFlags {
-                            new_heap_blkno: None,
-                            old_heap_blkno,
-                            flags: pg_constants::VISIBILITYMAP_VALID_BITS,
-                        },
-                    )?;
-                }
-            }
-        }
-
         Ok(())
     }
 

pageserver/src/walreceiver.rs

@@ -11,16 +11,14 @@ use crate::thread_mgr;
 use crate::thread_mgr::ThreadKind;
 use crate::walingest::WalIngest;
 use anyhow::{bail, Context, Error, Result};
-use bytes::BytesMut;
-use fail::fail_point;
 use lazy_static::lazy_static;
+use parking_lot::Mutex;
 use postgres_ffi::waldecoder::*;
 use postgres_protocol::message::backend::ReplicationMessage;
 use postgres_types::PgLsn;
 use std::cell::Cell;
 use std::collections::HashMap;
 use std::str::FromStr;
-use std::sync::Mutex;
 use std::thread_local;
 use std::time::SystemTime;
 use tokio::pin;
@@ -29,7 +27,6 @@ use tokio_postgres::{Client, NoTls, SimpleQueryMessage, SimpleQueryRow};
 use tokio_stream::StreamExt;
 use tracing::*;
 use zenith_utils::lsn::Lsn;
-use zenith_utils::pq_proto::ZenithFeedback;
 use zenith_utils::zid::ZTenantId;
 use zenith_utils::zid::ZTimelineId;
 
@@ -53,7 +50,7 @@ thread_local! {
 }
 
 fn drop_wal_receiver(tenantid: ZTenantId, timelineid: ZTimelineId) {
-    let mut receivers = WAL_RECEIVERS.lock().unwrap();
+    let mut receivers = WAL_RECEIVERS.lock();
     receivers.remove(&(tenantid, timelineid));
 }
 
@@ -64,7 +61,7 @@ pub fn launch_wal_receiver(
     timelineid: ZTimelineId,
     wal_producer_connstr: &str,
 ) -> Result<()> {
-    let mut receivers = WAL_RECEIVERS.lock().unwrap();
+    let mut receivers = WAL_RECEIVERS.lock();
 
     match receivers.get_mut(&(tenantid, timelineid)) {
         Some(receiver) => {
@@ -97,7 +94,7 @@ pub fn launch_wal_receiver(
 
 // Look up current WAL producer connection string in the hash table
 fn get_wal_producer_connstr(tenantid: ZTenantId, timelineid: ZTimelineId) -> String {
-    let receivers = WAL_RECEIVERS.lock().unwrap();
+    let receivers = WAL_RECEIVERS.lock();
 
     receivers
         .get(&(tenantid, timelineid))
@@ -162,7 +159,7 @@ fn walreceiver_main(
     // This is from tokio-postgres docs, but it is a bit weird in our case because we extensively use block_on
     runtime.spawn(async move {
         if let Err(e) = connection.await {
-            error!("connection error: {}", e);
+            eprintln!("connection error: {}", e);
         }
     });
 
@@ -256,8 +253,6 @@ fn walreceiver_main(
                     let writer = timeline.writer();
                     walingest.ingest_record(writer.as_ref(), recdata, lsn)?;
 
-                    fail_point!("walreceiver-after-ingest");
-
                     last_rec_lsn = lsn;
                 }
 
@@ -292,6 +287,7 @@ fn walreceiver_main(
         };
 
         if let Some(last_lsn) = status_update {
+            let last_lsn = PgLsn::from(u64::from(last_lsn));
             let timeline_synced_disk_consistent_lsn =
                 tenant_mgr::get_repository_for_tenant(tenantid)?
                     .get_timeline_state(timelineid)
@@ -299,32 +295,18 @@ fn walreceiver_main(
                     .unwrap_or(Lsn(0));
 
             // The last LSN we processed. It is not guaranteed to survive pageserver crash.
-            let write_lsn = u64::from(last_lsn);
+            let write_lsn = last_lsn;
             // `disk_consistent_lsn` is the LSN at which page server guarantees local persistence of all received data
-            let flush_lsn = u64::from(timeline.get_disk_consistent_lsn());
+            let flush_lsn = PgLsn::from(u64::from(timeline.get_disk_consistent_lsn()));
             // The last LSN that is synced to remote storage and is guaranteed to survive pageserver crash
             // Used by safekeepers to remove WAL preceding `remote_consistent_lsn`.
-            let apply_lsn = u64::from(timeline_synced_disk_consistent_lsn);
+            let apply_lsn = PgLsn::from(u64::from(timeline_synced_disk_consistent_lsn));
             let ts = SystemTime::now();
-
-            // Send zenith feedback message.
-            // Regular standby_status_update fields are put into this message.
-            let zenith_status_update = ZenithFeedback {
-                current_timeline_size: timeline.get_current_logical_size() as u64,
-                ps_writelsn: write_lsn,
-                ps_flushlsn: flush_lsn,
-                ps_applylsn: apply_lsn,
-                ps_replytime: ts,
-            };
-
-            debug!("zenith_status_update {:?}", zenith_status_update);
-
-            let mut data = BytesMut::new();
-            zenith_status_update.serialize(&mut data)?;
+            const NO_REPLY: u8 = 0;
             runtime.block_on(
                 physical_stream
                     .as_mut()
-                    .zenith_status_update(data.len() as u64, &data),
+                    .standby_status_update(write_lsn, flush_lsn, apply_lsn, ts, NO_REPLY),
             )?;
         }
     }

pageserver/src/walredo.rs

@@ -102,6 +102,8 @@ impl crate::walredo::WalRedoManager for DummyRedoManager {
     }
 }
 
+static TIMEOUT: Duration = Duration::from_secs(20);
+
 // Metrics collected on WAL redo operations
 //
 // We collect the time spent in actual WAL redo ('redo'), and time waiting
@@ -219,14 +221,7 @@ impl WalRedoManager for PostgresRedoManager {
                 let result = if batch_zenith {
                     self.apply_batch_zenith(rel, blknum, lsn, img, &records[batch_start..i])
                 } else {
-                    self.apply_batch_postgres(
-                        rel,
-                        blknum,
-                        lsn,
-                        img,
-                        &records[batch_start..i],
-                        self.conf.wal_redo_timeout,
-                    )
+                    self.apply_batch_postgres(rel, blknum, lsn, img, &records[batch_start..i])
                 };
                 img = Some(result?);
 
@@ -238,14 +233,7 @@ impl WalRedoManager for PostgresRedoManager {
         if batch_zenith {
             self.apply_batch_zenith(rel, blknum, lsn, img, &records[batch_start..])
         } else {
-            self.apply_batch_postgres(
-                rel,
-                blknum,
-                lsn,
-                img,
-                &records[batch_start..],
-                self.conf.wal_redo_timeout,
-            )
+            self.apply_batch_postgres(rel, blknum, lsn, img, &records[batch_start..])
         }
     }
 }
@@ -273,7 +261,6 @@ impl PostgresRedoManager {
         lsn: Lsn,
         base_img: Option<Bytes>,
         records: &[(Lsn, ZenithWalRecord)],
-        wal_redo_timeout: Duration,
     ) -> Result<Bytes, WalRedoError> {
         let start_time = Instant::now();
 
@@ -294,7 +281,7 @@ impl PostgresRedoManager {
         let result = if let RelishTag::Relation(rel) = rel {
             // Relational WAL records are applied using wal-redo-postgres
             let buf_tag = BufferTag { rel, blknum };
-            apply_result = process.apply_wal_records(buf_tag, base_img, records, wal_redo_timeout);
+            apply_result = process.apply_wal_records(buf_tag, base_img, records);
 
             apply_result.map_err(WalRedoError::IoError)
         } else {
@@ -376,44 +363,25 @@ impl PostgresRedoManager {
                 will_init: _,
                 rec: _,
             } => panic!("tried to pass postgres wal record to zenith WAL redo"),
-            ZenithWalRecord::ClearVisibilityMapFlags {
-                new_heap_blkno,
-                old_heap_blkno,
-                flags,
-            } => {
-                // sanity check that this is modifying the correct relish
+            ZenithWalRecord::ClearVisibilityMapFlags { heap_blkno, flags } => {
+                // Calculate the VM block and offset that corresponds to the heap block.
+                let map_block = pg_constants::HEAPBLK_TO_MAPBLOCK(*heap_blkno);
+                let map_byte = pg_constants::HEAPBLK_TO_MAPBYTE(*heap_blkno);
+                let map_offset = pg_constants::HEAPBLK_TO_OFFSET(*heap_blkno);
+
+                // Check that we're modifying the correct VM block.
                 assert!(
                     check_forknum(&rel, pg_constants::VISIBILITYMAP_FORKNUM),
                     "ClearVisibilityMapFlags record on unexpected rel {:?}",
                     rel
                 );
-                if let Some(heap_blkno) = *new_heap_blkno {
-                    // Calculate the VM block and offset that corresponds to the heap block.
-                    let map_block = pg_constants::HEAPBLK_TO_MAPBLOCK(heap_blkno);
-                    let map_byte = pg_constants::HEAPBLK_TO_MAPBYTE(heap_blkno);
-                    let map_offset = pg_constants::HEAPBLK_TO_OFFSET(heap_blkno);
+                assert!(map_block == blknum);
 
-                    // Check that we're modifying the correct VM block.
-                    assert!(map_block == blknum);
+                // equivalent to PageGetContents(page)
+                let map = &mut page[pg_constants::MAXALIGN_SIZE_OF_PAGE_HEADER_DATA..];
 
-                    // equivalent to PageGetContents(page)
-                    let map = &mut page[pg_constants::MAXALIGN_SIZE_OF_PAGE_HEADER_DATA..];
-
-                    map[map_byte as usize] &= !(flags << map_offset);
-                }
-
-                // Repeat for 'old_heap_blkno', if any
-                if let Some(heap_blkno) = *old_heap_blkno {
-                    let map_block = pg_constants::HEAPBLK_TO_MAPBLOCK(heap_blkno);
-                    let map_byte = pg_constants::HEAPBLK_TO_MAPBYTE(heap_blkno);
-                    let map_offset = pg_constants::HEAPBLK_TO_OFFSET(heap_blkno);
-
-                    assert!(map_block == blknum);
-
-                    let map = &mut page[pg_constants::MAXALIGN_SIZE_OF_PAGE_HEADER_DATA..];
-
-                    map[map_byte as usize] &= !(flags << map_offset);
-                }
+                let mask: u8 = flags << map_offset;
+                map[map_byte as usize] &= !mask;
             }
             // Non-relational WAL records are handled here, with custom code that has the
             // same effects as the corresponding Postgres WAL redo function.
@@ -616,7 +584,6 @@ impl PostgresRedoProcess {
         tag: BufferTag,
         base_img: Option<Bytes>,
         records: &[(Lsn, ZenithWalRecord)],
-        wal_redo_timeout: Duration,
     ) -> Result<Bytes, std::io::Error> {
         // Serialize all the messages to send the WAL redo process first.
         //
@@ -667,7 +634,7 @@ impl PostgresRedoProcess {
             // If we have more data to write, wake up if 'stdin' becomes writeable or
             // we have data to read. Otherwise only wake up if there's data to read.
             let nfds = if nwrite < writebuf.len() { 3 } else { 2 };
-            let n = nix::poll::poll(&mut pollfds[0..nfds], wal_redo_timeout.as_millis() as i32)?;
+            let n = nix::poll::poll(&mut pollfds[0..nfds], TIMEOUT.as_millis() as i32)?;
 
             if n == 0 {
                 return Err(Error::new(ErrorKind::Other, "WAL redo timed out"));

poetry.lock

@@ -91,14 +91,6 @@ botocore = ">=1.11.3"
 future = "*"
 wrapt = "*"
 
-[[package]]
-name = "backoff"
-version = "1.11.1"
-description = "Function decoration for backoff and retry"
-category = "main"
-optional = false
-python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*, !=3.4.*"
-
 [[package]]
 name = "boto3"
 version = "1.20.40"
@@ -822,11 +814,11 @@ python-versions = "*"
 
 [[package]]
 name = "moto"
-version = "3.0.4"
+version = "3.0.0"
 description = "A library that allows your python tests to easily mock out the boto library"
 category = "main"
 optional = false
-python-versions = ">=3.6"
+python-versions = "*"
 
 [package.dependencies]
 aws-xray-sdk = {version = ">=0.93,<0.96 || >0.96", optional = true, markers = "extra == \"server\""}
@@ -856,8 +848,7 @@ xmltodict = "*"
 
 [package.extras]
 all = ["PyYAML (>=5.1)", "python-jose[cryptography] (>=3.1.0,<4.0.0)", "ecdsa (!=0.15)", "docker (>=2.5.1)", "graphql-core", "jsondiff (>=1.1.2)", "aws-xray-sdk (>=0.93,!=0.96)", "idna (>=2.5,<4)", "cfn-lint (>=0.4.0)", "sshpubkeys (>=3.1.0)", "setuptools"]
-apigateway = ["PyYAML (>=5.1)", "python-jose[cryptography] (>=3.1.0,<4.0.0)", "ecdsa (!=0.15)"]
-apigatewayv2 = ["PyYAML (>=5.1)"]
+apigateway = ["python-jose[cryptography] (>=3.1.0,<4.0.0)", "ecdsa (!=0.15)"]
 appsync = ["graphql-core"]
 awslambda = ["docker (>=2.5.1)"]
 batch = ["docker (>=2.5.1)"]
@@ -1361,7 +1352,7 @@ testing = ["pytest (>=6)", "pytest-checkdocs (>=2.4)", "pytest-flake8", "pytest-
 [metadata]
 lock-version = "1.1"
 python-versions = "^3.7"
-content-hash = "58762accad4122026c650fa43421a900546e89f9908e2268410e7b11cc8c6c4e"
+content-hash = "0fa6c9377fbc827240d18d8b7e3742def37e90fc3277fddf8525d82dabd13090"
 
 [metadata.files]
 aiopg = [
@@ -1404,10 +1395,6 @@ aws-xray-sdk = [
     {file = "aws-xray-sdk-2.9.0.tar.gz", hash = "sha256:b0cd972db218d4d8f7b53ad806fc6184626b924c4997ae58fc9f2a8cd1281568"},
     {file = "aws_xray_sdk-2.9.0-py2.py3-none-any.whl", hash = "sha256:98216b3ac8281b51b59a8703f8ec561c460807d9d0679838f5c0179d381d7e58"},
 ]
-backoff = [
-    {file = "backoff-1.11.1-py2.py3-none-any.whl", hash = "sha256:61928f8fa48d52e4faa81875eecf308eccfb1016b018bb6bd21e05b5d90a96c5"},
-    {file = "backoff-1.11.1.tar.gz", hash = "sha256:ccb962a2378418c667b3c979b504fdeb7d9e0d29c0579e3b13b86467177728cb"},
-]
 boto3 = [
     {file = "boto3-1.20.40-py3-none-any.whl", hash = "sha256:cfe85589e4a0a997c7b9ae7432400b03fa6fa5fea29fdc48db3099a903b76998"},
     {file = "boto3-1.20.40.tar.gz", hash = "sha256:66aef9a6d8cad393f69166112ba49e14e2c6766f9278c96134101314a9af2992"},
@@ -1679,8 +1666,8 @@ mccabe = [
     {file = "mccabe-0.6.1.tar.gz", hash = "sha256:dd8d182285a0fe56bace7f45b5e7d1a6ebcbf524e8f3bd87eb0f125271b8831f"},
 ]
 moto = [
-    {file = "moto-3.0.4-py2.py3-none-any.whl", hash = "sha256:79646213d8438385182f4eea79e28725f94b3d0d3dc9a3eda81db47e0ebef6cc"},
-    {file = "moto-3.0.4.tar.gz", hash = "sha256:168b8a3cb4dd8a6df8e51d582761cefa9657b9f45ac7e1eb24dae394ebc9e000"},
+    {file = "moto-3.0.0-py2.py3-none-any.whl", hash = "sha256:762d33bbad3642c687f6495e69331318bef43f9aa662174397706ec3ad2a3578"},
+    {file = "moto-3.0.0.tar.gz", hash = "sha256:d6b00a2663290e7ebb06823d5ffcb124c8dc9bf526b878539ef7c4a377fd8255"},
 ]
 mypy = [
     {file = "mypy-0.910-cp35-cp35m-macosx_10_9_x86_64.whl", hash = "sha256:a155d80ea6cee511a3694b108c4494a39f42de11ee4e61e72bc424c490e46457"},

postgres_ffi/Cargo.toml

@@ -1,7 +1,10 @@
 [package]
 name = "postgres_ffi"
 version = "0.1.0"
-edition = "2021"
+authors = ["Heikki Linnakangas <heikki@zenith.tech>"]
+edition = "2018"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
 chrono = "0.4.19"

postgres_ffi/src/xlog_utils.rs

@@ -51,13 +51,6 @@ pub type TimeLineID = u32;
 pub type TimestampTz = i64;
 pub type XLogSegNo = u64;
 
-/// Interval of checkpointing metadata file. We should store metadata file to enforce
-/// predicate that checkpoint.nextXid is larger than any XID in WAL.
-/// But flushing checkpoint file for each transaction seems to be too expensive,
-/// so XID_CHECKPOINT_INTERVAL is used to forward align nextXid and so perform
-/// metadata checkpoint only once per XID_CHECKPOINT_INTERVAL transactions.
-/// XID_CHECKPOINT_INTERVAL should not be larger than BLCKSZ*CLOG_XACTS_PER_BYTE
-/// in order to let CLOG_TRUNCATE mechanism correctly extend CLOG.
 const XID_CHECKPOINT_INTERVAL: u32 = 1024;
 
 #[allow(non_snake_case)]
@@ -407,13 +400,9 @@ impl CheckPoint {
     ///
     /// Returns 'true' if the XID was updated.
     pub fn update_next_xid(&mut self, xid: u32) -> bool {
-        // nextXid should nw greate than any XID in WAL, so increment provided XID and check for wraparround.
-        let mut new_xid = std::cmp::max(xid + 1, pg_constants::FIRST_NORMAL_TRANSACTION_ID);
-        // To reduce number of metadata checkpoints, we forward align XID on XID_CHECKPOINT_INTERVAL.
-        // XID_CHECKPOINT_INTERVAL should not be larger than BLCKSZ*CLOG_XACTS_PER_BYTE
-        new_xid =
-            new_xid.wrapping_add(XID_CHECKPOINT_INTERVAL - 1) & !(XID_CHECKPOINT_INTERVAL - 1);
+        let xid = xid.wrapping_add(XID_CHECKPOINT_INTERVAL - 1) & !(XID_CHECKPOINT_INTERVAL - 1);
         let full_xid = self.nextXid.value;
+        let new_xid = std::cmp::max(xid + 1, pg_constants::FIRST_NORMAL_TRANSACTION_ID);
         let old_xid = full_xid as u32;
         if new_xid.wrapping_sub(old_xid) as i32 > 0 {
             let mut epoch = full_xid >> 32;
@@ -531,34 +520,4 @@ mod tests {
         println!("wal_end={}, tli={}", wal_end, tli);
         assert_eq!(wal_end, waldump_wal_end);
     }
-
-    /// Check the math in update_next_xid
-    ///
-    /// NOTE: These checks are sensitive to the value of XID_CHECKPOINT_INTERVAL,
-    /// currently 1024.
-    #[test]
-    pub fn test_update_next_xid() {
-        let checkpoint_buf = [0u8; std::mem::size_of::<CheckPoint>()];
-        let mut checkpoint = CheckPoint::decode(&checkpoint_buf).unwrap();
-
-        checkpoint.nextXid = FullTransactionId { value: 10 };
-        assert_eq!(checkpoint.nextXid.value, 10);
-
-        // The input XID gets rounded up to the next XID_CHECKPOINT_INTERVAL
-        // boundary
-        checkpoint.update_next_xid(100);
-        assert_eq!(checkpoint.nextXid.value, 1024);
-
-        // No change
-        checkpoint.update_next_xid(500);
-        assert_eq!(checkpoint.nextXid.value, 1024);
-        checkpoint.update_next_xid(1023);
-        assert_eq!(checkpoint.nextXid.value, 1024);
-
-        // The function returns the *next* XID, given the highest XID seen so
-        // far. So when we pass 1024, the nextXid gets bumped up to the next
-        // XID_CHECKPOINT_INTERVAL boundary.
-        checkpoint.update_next_xid(1024);
-        assert_eq!(checkpoint.nextXid.value, 2048);
-    }
 }

proxy/Cargo.toml

@@ -1,33 +1,28 @@
 [package]
 name = "proxy"
 version = "0.1.0"
-edition = "2021"
+authors = ["Stas Kelvich <stas.kelvich@gmail.com>"]
+edition = "2018"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
 anyhow = "1.0"
 bytes = { version = "1.0.1", features = ['serde'] }
-clap = "3.0"
-futures = "0.3.13"
-hashbrown = "0.11.2"
-hex = "0.4.3"
-hyper = "0.14"
 lazy_static = "1.4.0"
 md5 = "0.7.0"
-parking_lot = "0.11.2"
-pin-project-lite = "0.2.7"
 rand = "0.8.3"
-reqwest = { version = "0.11", default-features = false, features = ["blocking", "json", "rustls-tls"] }
-rustls = "0.19.1"
-scopeguard = "1.1.0"
+hex = "0.4.3"
+hyper = "0.14"
+routerify = "2"
+parking_lot = "0.11.2"
 serde = "1"
 serde_json = "1"
 tokio = { version = "1.11", features = ["macros"] }
-tokio-postgres = { git = "https://github.com/zenithdb/rust-postgres.git", rev="2949d98df52587d562986aad155dd4e889e408b7" }
-tokio-rustls = "0.22.0"
+tokio-postgres = { git = "https://github.com/zenithdb/rust-postgres.git", rev="9eb0dbfbeb6a6c1b79099b9f7ae4a8c021877858" }
+clap = "2.33.0"
+rustls = "0.19.1"
+reqwest = { version = "0.11", default-features = false, features = ["blocking", "json", "rustls-tls"] }
 
 zenith_utils = { path = "../zenith_utils" }
 zenith_metrics = { path = "../zenith_metrics" }
-
-[dev-dependencies]
-tokio-postgres-rustls = "0.8.0"
-rcgen = "0.8.14"

proxy/src/auth.rs

@@ -1,169 +0,0 @@
-use crate::compute::DatabaseInfo;
-use crate::config::ProxyConfig;
-use crate::cplane_api::{self, CPlaneApi};
-use crate::stream::PqStream;
-use anyhow::{anyhow, bail, Context};
-use std::collections::HashMap;
-use tokio::io::{AsyncRead, AsyncWrite};
-use zenith_utils::pq_proto::{BeMessage as Be, BeParameterStatusMessage, FeMessage as Fe};
-
-/// Various client credentials which we use for authentication.
-#[derive(Debug, PartialEq, Eq)]
-pub struct ClientCredentials {
-    pub user: String,
-    pub dbname: String,
-}
-
-impl TryFrom<HashMap<String, String>> for ClientCredentials {
-    type Error = anyhow::Error;
-
-    fn try_from(mut value: HashMap<String, String>) -> Result<Self, Self::Error> {
-        let mut get_param = |key| {
-            value
-                .remove(key)
-                .with_context(|| format!("{} is missing in startup packet", key))
-        };
-
-        let user = get_param("user")?;
-        let db = get_param("database")?;
-
-        Ok(Self { user, dbname: db })
-    }
-}
-
-impl ClientCredentials {
-    /// Use credentials to authenticate the user.
-    pub async fn authenticate(
-        self,
-        config: &ProxyConfig,
-        client: &mut PqStream<impl AsyncRead + AsyncWrite + Unpin>,
-    ) -> anyhow::Result<DatabaseInfo> {
-        use crate::config::ClientAuthMethod::*;
-        use crate::config::RouterConfig::*;
-        let db_info = match &config.router_config {
-            Static { host, port } => handle_static(host.clone(), *port, client, self).await,
-            Dynamic(Mixed) => {
-                if self.user.ends_with("@zenith") {
-                    handle_existing_user(config, client, self).await
-                } else {
-                    handle_new_user(config, client).await
-                }
-            }
-            Dynamic(Password) => handle_existing_user(config, client, self).await,
-            Dynamic(Link) => handle_new_user(config, client).await,
-        };
-
-        db_info.context("failed to authenticate client")
-    }
-}
-
-fn new_psql_session_id() -> String {
-    hex::encode(rand::random::<[u8; 8]>())
-}
-
-async fn handle_static(
-    host: String,
-    port: u16,
-    client: &mut PqStream<impl AsyncRead + AsyncWrite + Unpin>,
-    creds: ClientCredentials,
-) -> anyhow::Result<DatabaseInfo> {
-    client
-        .write_message(&Be::AuthenticationCleartextPassword)
-        .await?;
-
-    // Read client's password bytes
-    let msg = match client.read_message().await? {
-        Fe::PasswordMessage(msg) => msg,
-        bad => bail!("unexpected message type: {:?}", bad),
-    };
-
-    let cleartext_password = std::str::from_utf8(&msg)?.split('\0').next().unwrap();
-
-    let db_info = DatabaseInfo {
-        host,
-        port,
-        dbname: creds.dbname.clone(),
-        user: creds.user.clone(),
-        password: Some(cleartext_password.into()),
-    };
-
-    client
-        .write_message_noflush(&Be::AuthenticationOk)?
-        .write_message_noflush(&BeParameterStatusMessage::encoding())?;
-
-    Ok(db_info)
-}
-
-async fn handle_existing_user(
-    config: &ProxyConfig,
-    client: &mut PqStream<impl AsyncRead + AsyncWrite + Unpin>,
-    creds: ClientCredentials,
-) -> anyhow::Result<DatabaseInfo> {
-    let psql_session_id = new_psql_session_id();
-    let md5_salt = rand::random();
-
-    client
-        .write_message(&Be::AuthenticationMD5Password(&md5_salt))
-        .await?;
-
-    // Read client's password hash
-    let msg = match client.read_message().await? {
-        Fe::PasswordMessage(msg) => msg,
-        bad => bail!("unexpected message type: {:?}", bad),
-    };
-
-    let (_trailing_null, md5_response) = msg
-        .split_last()
-        .ok_or_else(|| anyhow!("unexpected password message"))?;
-
-    let cplane = CPlaneApi::new(&config.auth_endpoint);
-    let db_info = cplane
-        .authenticate_proxy_request(creds, md5_response, &md5_salt, &psql_session_id)
-        .await?;
-
-    client
-        .write_message_noflush(&Be::AuthenticationOk)?
-        .write_message_noflush(&BeParameterStatusMessage::encoding())?;
-
-    Ok(db_info)
-}
-
-async fn handle_new_user(
-    config: &ProxyConfig,
-    client: &mut PqStream<impl AsyncRead + AsyncWrite + Unpin>,
-) -> anyhow::Result<DatabaseInfo> {
-    let psql_session_id = new_psql_session_id();
-    let greeting = hello_message(&config.redirect_uri, &psql_session_id);
-
-    let db_info = cplane_api::with_waiter(psql_session_id, |waiter| async {
-        // Give user a URL to spawn a new database
-        client
-            .write_message_noflush(&Be::AuthenticationOk)?
-            .write_message_noflush(&BeParameterStatusMessage::encoding())?
-            .write_message(&Be::NoticeResponse(greeting))
-            .await?;
-
-        // Wait for web console response
-        waiter.await?.map_err(|e| anyhow!(e))
-    })
-    .await?;
-
-    client.write_message_noflush(&Be::NoticeResponse("Connecting to database.".into()))?;
-
-    Ok(db_info)
-}
-
-fn hello_message(redirect_uri: &str, session_id: &str) -> String {
-    format!(
-        concat![
-            "☀️  Welcome to Zenith!\n",
-            "To proceed with database creation, open the following link:\n\n",
-            "    {redirect_uri}{session_id}\n\n",
-            "It needs to be done once and we will send you '.pgpass' file,\n",
-            "which will allow you to access or create ",
-            "databases without opening your web browser."
-        ],
-        redirect_uri = redirect_uri,
-        session_id = session_id,
-    )
-}

proxy/src/cancellation.rs

@@ -1,106 +0,0 @@
-use anyhow::{anyhow, Context};
-use hashbrown::HashMap;
-use parking_lot::Mutex;
-use std::net::SocketAddr;
-use tokio::net::TcpStream;
-use tokio_postgres::{CancelToken, NoTls};
-use zenith_utils::pq_proto::CancelKeyData;
-
-/// Enables serving CancelRequests.
-#[derive(Default)]
-pub struct CancelMap(Mutex<HashMap<CancelKeyData, Option<CancelClosure>>>);
-
-impl CancelMap {
-    /// Cancel a running query for the corresponding connection.
-    pub async fn cancel_session(&self, key: CancelKeyData) -> anyhow::Result<()> {
-        let cancel_closure = self
-            .0
-            .lock()
-            .get(&key)
-            .and_then(|x| x.clone())
-            .with_context(|| format!("unknown session: {:?}", key))?;
-
-        cancel_closure.try_cancel_query().await
-    }
-
-    /// Run async action within an ephemeral session identified by [`CancelKeyData`].
-    pub async fn with_session<'a, F, R, V>(&'a self, f: F) -> anyhow::Result<V>
-    where
-        F: FnOnce(Session<'a>) -> R,
-        R: std::future::Future<Output = anyhow::Result<V>>,
-    {
-        // HACK: We'd rather get the real backend_pid but tokio_postgres doesn't
-        // expose it and we don't want to do another roundtrip to query
-        // for it. The client will be able to notice that this is not the
-        // actual backend_pid, but backend_pid is not used for anything
-        // so it doesn't matter.
-        let key = rand::random();
-
-        // Random key collisions are unlikely to happen here, but they're still possible,
-        // which is why we have to take care not to rewrite an existing key.
-        self.0
-            .lock()
-            .try_insert(key, None)
-            .map_err(|_| anyhow!("session already exists: {:?}", key))?;
-
-        // This will guarantee that the session gets dropped
-        // as soon as the future is finished.
-        scopeguard::defer! {
-            self.0.lock().remove(&key);
-        }
-
-        let session = Session::new(key, self);
-        f(session).await
-    }
-}
-
-/// This should've been a [`std::future::Future`], but
-/// it's impossible to name a type of an unboxed future
-/// (we'd need something like `#![feature(type_alias_impl_trait)]`).
-#[derive(Clone)]
-pub struct CancelClosure {
-    socket_addr: SocketAddr,
-    cancel_token: CancelToken,
-}
-
-impl CancelClosure {
-    pub fn new(socket_addr: SocketAddr, cancel_token: CancelToken) -> Self {
-        Self {
-            socket_addr,
-            cancel_token,
-        }
-    }
-
-    /// Cancels the query running on user's compute node.
-    pub async fn try_cancel_query(self) -> anyhow::Result<()> {
-        let socket = TcpStream::connect(self.socket_addr).await?;
-        self.cancel_token.cancel_query_raw(socket, NoTls).await?;
-
-        Ok(())
-    }
-}
-
-/// Helper for registering query cancellation tokens.
-pub struct Session<'a> {
-    /// The user-facing key identifying this session.
-    key: CancelKeyData,
-    /// The [`CancelMap`] this session belongs to.
-    cancel_map: &'a CancelMap,
-}
-
-impl<'a> Session<'a> {
-    fn new(key: CancelKeyData, cancel_map: &'a CancelMap) -> Self {
-        Self { key, cancel_map }
-    }
-
-    /// Store the cancel token for the given session.
-    /// This enables query cancellation in [`crate::proxy::handshake`].
-    pub fn enable_cancellation(self, cancel_closure: CancelClosure) -> CancelKeyData {
-        self.cancel_map
-            .0
-            .lock()
-            .insert(self.key, Some(cancel_closure));
-
-        self.key
-    }
-}

proxy/src/compute.rs

@@ -1,42 +0,0 @@
-use anyhow::Context;
-use serde::{Deserialize, Serialize};
-use std::net::{SocketAddr, ToSocketAddrs};
-
-/// Compute node connection params.
-#[derive(Serialize, Deserialize, Debug, Default)]
-pub struct DatabaseInfo {
-    pub host: String,
-    pub port: u16,
-    pub dbname: String,
-    pub user: String,
-    pub password: Option<String>,
-}
-
-impl DatabaseInfo {
-    pub fn socket_addr(&self) -> anyhow::Result<SocketAddr> {
-        let host_port = format!("{}:{}", self.host, self.port);
-        host_port
-            .to_socket_addrs()
-            .with_context(|| format!("cannot resolve {} to SocketAddr", host_port))?
-            .next()
-            .context("cannot resolve at least one SocketAddr")
-    }
-}
-
-impl From<DatabaseInfo> for tokio_postgres::Config {
-    fn from(db_info: DatabaseInfo) -> Self {
-        let mut config = tokio_postgres::Config::new();
-
-        config
-            .host(&db_info.host)
-            .port(db_info.port)
-            .dbname(&db_info.dbname)
-            .user(&db_info.user);
-
-        if let Some(password) = db_info.password {
-            config.password(password);
-        }
-
-        config
-    }
-}

proxy/src/cplane_api.rs

@@ -1,79 +1,18 @@
-use crate::auth::ClientCredentials;
-use crate::compute::DatabaseInfo;
-use crate::waiters::{Waiter, Waiters};
-use anyhow::{anyhow, bail};
-use lazy_static::lazy_static;
+use anyhow::{anyhow, bail, Context};
 use serde::{Deserialize, Serialize};
+use std::net::{SocketAddr, ToSocketAddrs};
 
-lazy_static! {
-    static ref CPLANE_WAITERS: Waiters<Result<DatabaseInfo, String>> = Default::default();
+use crate::state::ProxyWaiters;
+
+#[derive(Serialize, Deserialize, Debug, Default)]
+pub struct DatabaseInfo {
+    pub host: String,
+    pub port: u16,
+    pub dbname: String,
+    pub user: String,
+    pub password: Option<String>,
 }
 
-/// Give caller an opportunity to wait for cplane's reply.
-pub async fn with_waiter<F, R, T>(psql_session_id: impl Into<String>, f: F) -> anyhow::Result<T>
-where
-    F: FnOnce(Waiter<'static, Result<DatabaseInfo, String>>) -> R,
-    R: std::future::Future<Output = anyhow::Result<T>>,
-{
-    let waiter = CPLANE_WAITERS.register(psql_session_id.into())?;
-    f(waiter).await
-}
-
-pub fn notify(psql_session_id: &str, msg: Result<DatabaseInfo, String>) -> anyhow::Result<()> {
-    CPLANE_WAITERS.notify(psql_session_id, msg)
-}
-
-/// Zenith console API wrapper.
-pub struct CPlaneApi<'a> {
-    auth_endpoint: &'a str,
-}
-
-impl<'a> CPlaneApi<'a> {
-    pub fn new(auth_endpoint: &'a str) -> Self {
-        Self { auth_endpoint }
-    }
-}
-
-impl CPlaneApi<'_> {
-    pub async fn authenticate_proxy_request(
-        &self,
-        creds: ClientCredentials,
-        md5_response: &[u8],
-        salt: &[u8; 4],
-        psql_session_id: &str,
-    ) -> anyhow::Result<DatabaseInfo> {
-        let mut url = reqwest::Url::parse(self.auth_endpoint)?;
-        url.query_pairs_mut()
-            .append_pair("login", &creds.user)
-            .append_pair("database", &creds.dbname)
-            .append_pair("md5response", std::str::from_utf8(md5_response)?)
-            .append_pair("salt", &hex::encode(salt))
-            .append_pair("psql_session_id", psql_session_id);
-
-        with_waiter(psql_session_id, |waiter| async {
-            println!("cplane request: {}", url);
-            // TODO: leverage `reqwest::Client` to reuse connections
-            let resp = reqwest::get(url).await?;
-            if !resp.status().is_success() {
-                bail!("Auth failed: {}", resp.status())
-            }
-
-            let auth_info: ProxyAuthResponse = serde_json::from_str(resp.text().await?.as_str())?;
-            println!("got auth info: #{:?}", auth_info);
-
-            use ProxyAuthResponse::*;
-            match auth_info {
-                Ready { conn_info } => Ok(conn_info),
-                Error { error } => bail!(error),
-                NotReady { .. } => waiter.await?.map_err(|e| anyhow!(e)),
-            }
-        })
-        .await
-    }
-}
-
-// NOTE: the order of constructors is important.
-// https://serde.rs/enum-representations.html#untagged
 #[derive(Serialize, Deserialize, Debug)]
 #[serde(untagged)]
 enum ProxyAuthResponse {
@@ -82,6 +21,86 @@ enum ProxyAuthResponse {
     NotReady { ready: bool }, // TODO: get rid of `ready`
 }
 
+impl DatabaseInfo {
+    pub fn socket_addr(&self) -> anyhow::Result<SocketAddr> {
+        let host_port = format!("{}:{}", self.host, self.port);
+        host_port
+            .to_socket_addrs()
+            .with_context(|| format!("cannot resolve {} to SocketAddr", host_port))?
+            .next()
+            .context("cannot resolve at least one SocketAddr")
+    }
+}
+
+impl From<DatabaseInfo> for tokio_postgres::Config {
+    fn from(db_info: DatabaseInfo) -> Self {
+        let mut config = tokio_postgres::Config::new();
+
+        config
+            .host(&db_info.host)
+            .port(db_info.port)
+            .dbname(&db_info.dbname)
+            .user(&db_info.user);
+
+        if let Some(password) = db_info.password {
+            config.password(password);
+        }
+
+        config
+    }
+}
+
+pub struct CPlaneApi<'a> {
+    auth_endpoint: &'a str,
+    waiters: &'a ProxyWaiters,
+}
+
+impl<'a> CPlaneApi<'a> {
+    pub fn new(auth_endpoint: &'a str, waiters: &'a ProxyWaiters) -> Self {
+        Self {
+            auth_endpoint,
+            waiters,
+        }
+    }
+}
+
+impl CPlaneApi<'_> {
+    pub fn authenticate_proxy_request(
+        &self,
+        user: &str,
+        database: &str,
+        md5_response: &[u8],
+        salt: &[u8; 4],
+        psql_session_id: &str,
+    ) -> anyhow::Result<DatabaseInfo> {
+        let mut url = reqwest::Url::parse(self.auth_endpoint)?;
+        url.query_pairs_mut()
+            .append_pair("login", user)
+            .append_pair("database", database)
+            .append_pair("md5response", std::str::from_utf8(md5_response)?)
+            .append_pair("salt", &hex::encode(salt))
+            .append_pair("psql_session_id", psql_session_id);
+
+        let waiter = self.waiters.register(psql_session_id.to_owned());
+
+        println!("cplane request: {}", url);
+        let resp = reqwest::blocking::get(url)?;
+        if !resp.status().is_success() {
+            bail!("Auth failed: {}", resp.status())
+        }
+
+        let auth_info: ProxyAuthResponse = serde_json::from_str(resp.text()?.as_str())?;
+        println!("got auth info: #{:?}", auth_info);
+
+        use ProxyAuthResponse::*;
+        match auth_info {
+            Ready { conn_info } => Ok(conn_info),
+            Error { error } => bail!(error),
+            NotReady { .. } => waiter.wait()?.map_err(|e| anyhow!(e)),
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

proxy/src/http.rs

@@ -1,30 +1,15 @@
-use anyhow::anyhow;
 use hyper::{Body, Request, Response, StatusCode};
-use std::net::TcpListener;
+use routerify::RouterBuilder;
+
 use zenith_utils::http::endpoint;
 use zenith_utils::http::error::ApiError;
 use zenith_utils::http::json::json_response;
-use zenith_utils::http::{RouterBuilder, RouterService};
 
 async fn status_handler(_: Request<Body>) -> Result<Response<Body>, ApiError> {
     Ok(json_response(StatusCode::OK, "")?)
 }
 
-fn make_router() -> RouterBuilder<hyper::Body, ApiError> {
+pub fn make_router() -> RouterBuilder<hyper::Body, ApiError> {
     let router = endpoint::make_router();
     router.get("/v1/status", status_handler)
 }
-
-pub async fn thread_main(http_listener: TcpListener) -> anyhow::Result<()> {
-    scopeguard::defer! {
-        println!("http has shut down");
-    }
-
-    let service = || RouterService::new(make_router().build()?);
-
-    hyper::Server::from_tcp(http_listener)?
-        .serve(service().map_err(|e| anyhow!(e))?)
-        .await?;
-
-    Ok(())
-}

proxy/src/main.rs

@@ -5,162 +5,137 @@
 /// (control plane API in our case) and can create new databases and accounts
 /// in somewhat transparent manner (again via communication with control plane API).
 ///
-use anyhow::{bail, Context};
+use anyhow::bail;
 use clap::{App, Arg};
-use config::ProxyConfig;
-use futures::FutureExt;
-use std::future::Future;
-use tokio::{net::TcpListener, task::JoinError};
-use zenith_utils::GIT_VERSION;
+use state::{ProxyConfig, ProxyState};
+use std::thread;
+use zenith_utils::http::endpoint;
+use zenith_utils::{tcp_listener, GIT_VERSION};
 
-use crate::config::{ClientAuthMethod, RouterConfig};
-
-mod auth;
-mod cancellation;
-mod compute;
-mod config;
 mod cplane_api;
 mod http;
 mod mgmt;
 mod proxy;
-mod stream;
+mod state;
 mod waiters;
 
-/// Flattens Result<Result<T>> into Result<T>.
-async fn flatten_err(
-    f: impl Future<Output = Result<anyhow::Result<()>, JoinError>>,
-) -> anyhow::Result<()> {
-    f.map(|r| r.context("join error").and_then(|x| x)).await
-}
-
-#[tokio::main]
-async fn main() -> anyhow::Result<()> {
+fn main() -> anyhow::Result<()> {
     zenith_metrics::set_common_metrics_prefix("zenith_proxy");
     let arg_matches = App::new("Zenith proxy/router")
         .version(GIT_VERSION)
         .arg(
-            Arg::new("proxy")
-                .short('p')
+            Arg::with_name("proxy")
+                .short("p")
                 .long("proxy")
                 .takes_value(true)
                 .help("listen for incoming client connections on ip:port")
                 .default_value("127.0.0.1:4432"),
         )
         .arg(
-            Arg::new("auth-method")
-                .long("auth-method")
-                .takes_value(true)
-                .help("Possible values: password | link | mixed")
-                .default_value("mixed"),
-        )
-        .arg(
-            Arg::new("static-router")
-                .short('s')
-                .long("static-router")
-                .takes_value(true)
-                .help("Route all clients to host:port"),
-        )
-        .arg(
-            Arg::new("mgmt")
-                .short('m')
+            Arg::with_name("mgmt")
+                .short("m")
                 .long("mgmt")
                 .takes_value(true)
                 .help("listen for management callback connection on ip:port")
                 .default_value("127.0.0.1:7000"),
         )
         .arg(
-            Arg::new("http")
-                .short('h')
+            Arg::with_name("http")
+                .short("h")
                 .long("http")
                 .takes_value(true)
                 .help("listen for incoming http connections (metrics, etc) on ip:port")
                 .default_value("127.0.0.1:7001"),
         )
         .arg(
-            Arg::new("uri")
-                .short('u')
+            Arg::with_name("uri")
+                .short("u")
                 .long("uri")
                 .takes_value(true)
                 .help("redirect unauthenticated users to given uri")
                 .default_value("http://localhost:3000/psql_session/"),
         )
         .arg(
-            Arg::new("auth-endpoint")
-                .short('a')
+            Arg::with_name("auth-endpoint")
+                .short("a")
                 .long("auth-endpoint")
                 .takes_value(true)
                 .help("API endpoint for authenticating users")
                 .default_value("http://localhost:3000/authenticate_proxy_request/"),
         )
         .arg(
-            Arg::new("ssl-key")
-                .short('k')
+            Arg::with_name("ssl-key")
+                .short("k")
                 .long("ssl-key")
                 .takes_value(true)
                 .help("path to SSL key for client postgres connections"),
         )
         .arg(
-            Arg::new("ssl-cert")
-                .short('c')
+            Arg::with_name("ssl-cert")
+                .short("c")
                 .long("ssl-cert")
                 .takes_value(true)
                 .help("path to SSL cert for client postgres connections"),
         )
         .get_matches();
 
-    let tls_config = match (
+    let ssl_config = match (
         arg_matches.value_of("ssl-key"),
         arg_matches.value_of("ssl-cert"),
     ) {
-        (Some(key_path), Some(cert_path)) => Some(config::configure_ssl(key_path, cert_path)?),
+        (Some(key_path), Some(cert_path)) => {
+            Some(crate::state::configure_ssl(key_path, cert_path)?)
+        }
         (None, None) => None,
         _ => bail!("either both or neither ssl-key and ssl-cert must be specified"),
     };
 
-    let auth_method = arg_matches.value_of("auth-method").unwrap().parse()?;
-    let router_config = match arg_matches.value_of("static-router") {
-        None => RouterConfig::Dynamic(auth_method),
-        Some(addr) => {
-            if let ClientAuthMethod::Password = auth_method {
-                let (host, port) = addr.split_once(":").unwrap();
-                RouterConfig::Static {
-                    host: host.to_string(),
-                    port: port.parse().unwrap(),
-                }
-            } else {
-                bail!("static-router requires --auth-method password")
-            }
-        }
-    };
-
-    let config: &ProxyConfig = Box::leak(Box::new(ProxyConfig {
-        router_config,
+    let config = ProxyConfig {
         proxy_address: arg_matches.value_of("proxy").unwrap().parse()?,
         mgmt_address: arg_matches.value_of("mgmt").unwrap().parse()?,
         http_address: arg_matches.value_of("http").unwrap().parse()?,
         redirect_uri: arg_matches.value_of("uri").unwrap().parse()?,
         auth_endpoint: arg_matches.value_of("auth-endpoint").unwrap().parse()?,
-        tls_config,
-    }));
+        ssl_config,
+    };
+    let state: &ProxyState = Box::leak(Box::new(ProxyState::new(config)));
 
     println!("Version: {}", GIT_VERSION);
 
     // Check that we can bind to address before further initialization
-    println!("Starting http on {}", config.http_address);
-    let http_listener = TcpListener::bind(config.http_address).await?.into_std()?;
+    println!("Starting http on {}", state.conf.http_address);
+    let http_listener = tcp_listener::bind(state.conf.http_address)?;
 
-    println!("Starting mgmt on {}", config.mgmt_address);
-    let mgmt_listener = TcpListener::bind(config.mgmt_address).await?.into_std()?;
+    println!("Starting proxy on {}", state.conf.proxy_address);
+    let pageserver_listener = tcp_listener::bind(state.conf.proxy_address)?;
 
-    println!("Starting proxy on {}", config.proxy_address);
-    let proxy_listener = TcpListener::bind(config.proxy_address).await?;
+    println!("Starting mgmt on {}", state.conf.mgmt_address);
+    let mgmt_listener = tcp_listener::bind(state.conf.mgmt_address)?;
 
-    let http = tokio::spawn(http::thread_main(http_listener));
-    let proxy = tokio::spawn(proxy::thread_main(config, proxy_listener));
-    let mgmt = tokio::task::spawn_blocking(move || mgmt::thread_main(mgmt_listener));
+    let threads = [
+        thread::Builder::new()
+            .name("Http thread".into())
+            .spawn(move || {
+                let router = http::make_router();
+                endpoint::serve_thread_main(
+                    router,
+                    http_listener,
+                    std::future::pending(), // never shut down
+                )
+            })?,
+        // Spawn a thread to listen for connections. It will spawn further threads
+        // for each connection.
+        thread::Builder::new()
+            .name("Listener thread".into())
+            .spawn(move || proxy::thread_main(state, pageserver_listener))?,
+        thread::Builder::new()
+            .name("Mgmt thread".into())
+            .spawn(move || mgmt::thread_main(state, mgmt_listener))?,
+    ];
 
-    let tasks = [flatten_err(http), flatten_err(proxy), flatten_err(mgmt)];
-    let _: Vec<()> = futures::future::try_join_all(tasks).await?;
+    for t in threads {
+        t.join().unwrap()?;
+    }
 
     Ok(())
 }

proxy/src/mgmt.rs

@@ -1,49 +1,44 @@
-use crate::{compute::DatabaseInfo, cplane_api};
-use anyhow::Context;
-use serde::Deserialize;
 use std::{
     net::{TcpListener, TcpStream},
     thread,
 };
+
+use serde::Deserialize;
 use zenith_utils::{
     postgres_backend::{self, AuthType, PostgresBackend},
     pq_proto::{BeMessage, SINGLE_COL_ROWDESC},
 };
 
+use crate::{cplane_api::DatabaseInfo, ProxyState};
+
 ///
 /// Main proxy listener loop.
 ///
 /// Listens for connections, and launches a new handler thread for each.
 ///
-pub fn thread_main(listener: TcpListener) -> anyhow::Result<()> {
-    scopeguard::defer! {
-        println!("mgmt has shut down");
-    }
-
-    listener
-        .set_nonblocking(false)
-        .context("failed to set listener to blocking")?;
+pub fn thread_main(state: &'static ProxyState, listener: TcpListener) -> anyhow::Result<()> {
     loop {
-        let (socket, peer_addr) = listener.accept().context("failed to accept a new client")?;
+        let (socket, peer_addr) = listener.accept()?;
         println!("accepted connection from {}", peer_addr);
-        socket
-            .set_nodelay(true)
-            .context("failed to set client socket option")?;
+        socket.set_nodelay(true).unwrap();
 
         thread::spawn(move || {
-            if let Err(err) = handle_connection(socket) {
+            if let Err(err) = handle_connection(state, socket) {
                 println!("error: {}", err);
             }
         });
     }
 }
 
-fn handle_connection(socket: TcpStream) -> anyhow::Result<()> {
+fn handle_connection(state: &ProxyState, socket: TcpStream) -> anyhow::Result<()> {
+    let mut conn_handler = MgmtHandler { state };
     let pgbackend = PostgresBackend::new(socket, AuthType::Trust, None, true)?;
-    pgbackend.run(&mut MgmtHandler)
+    pgbackend.run(&mut conn_handler)
 }
 
-struct MgmtHandler;
+struct MgmtHandler<'a> {
+    state: &'a ProxyState,
+}
 
 /// Serialized examples:
 // {
@@ -79,13 +74,13 @@ enum PsqlSessionResult {
     Failure(String),
 }
 
-impl postgres_backend::Handler for MgmtHandler {
+impl postgres_backend::Handler for MgmtHandler<'_> {
     fn process_query(
         &mut self,
         pgb: &mut PostgresBackend,
         query_string: &str,
     ) -> anyhow::Result<()> {
-        let res = try_process_query(pgb, query_string);
+        let res = try_process_query(self, pgb, query_string);
         // intercept and log error message
         if res.is_err() {
             println!("Mgmt query failed: #{:?}", res);
@@ -94,7 +89,11 @@ impl postgres_backend::Handler for MgmtHandler {
     }
 }
 
-fn try_process_query(pgb: &mut PostgresBackend, query_string: &str) -> anyhow::Result<()> {
+fn try_process_query(
+    mgmt: &mut MgmtHandler,
+    pgb: &mut PostgresBackend,
+    query_string: &str,
+) -> anyhow::Result<()> {
     println!("Got mgmt query: '{}'", query_string);
 
     let resp: PsqlSessionResponse = serde_json::from_str(query_string)?;
@@ -105,7 +104,7 @@ fn try_process_query(pgb: &mut PostgresBackend, query_string: &str) -> anyhow::R
         Failure(message) => Err(message),
     };
 
-    match cplane_api::notify(&resp.session_id, msg) {
+    match mgmt.state.waiters.notify(&resp.session_id, msg) {
         Ok(()) => {
             pgb.write_message_noflush(&SINGLE_COL_ROWDESC)?
                 .write_message_noflush(&BeMessage::DataRow(&[Some(b"ok")]))?

proxy/src/proxy.rs

@@ -1,332 +1,386 @@
-use crate::auth;
-use crate::cancellation::{self, CancelClosure, CancelMap};
-use crate::compute::DatabaseInfo;
-use crate::config::{ProxyConfig, TlsConfig};
-use crate::stream::{MetricsStream, PqStream, Stream};
-use anyhow::{bail, Context};
+use crate::cplane_api::{CPlaneApi, DatabaseInfo};
+use crate::ProxyState;
+use anyhow::{anyhow, bail, Context};
 use lazy_static::lazy_static;
-use std::sync::Arc;
-use tokio::io::{AsyncRead, AsyncWrite};
-use tokio::net::TcpStream;
+use parking_lot::Mutex;
+use rand::prelude::StdRng;
+use rand::{Rng, SeedableRng};
+use std::cell::Cell;
+use std::collections::HashMap;
+use std::net::{SocketAddr, TcpStream};
+use std::{io, thread};
 use tokio_postgres::NoTls;
 use zenith_metrics::{new_common_metric_name, register_int_counter, IntCounter};
-use zenith_utils::pq_proto::{BeMessage as Be, *};
+use zenith_utils::postgres_backend::{self, PostgresBackend, ProtoState, Stream};
+use zenith_utils::pq_proto::{BeMessage as Be, FeMessage as Fe, *};
+use zenith_utils::sock_split::{ReadStream, WriteStream};
 
-lazy_static! {
-    static ref NUM_CONNECTIONS_ACCEPTED_COUNTER: IntCounter = register_int_counter!(
-        new_common_metric_name("num_connections_accepted"),
-        "Number of TCP client connections accepted."
-    )
-    .unwrap();
-    static ref NUM_CONNECTIONS_CLOSED_COUNTER: IntCounter = register_int_counter!(
-        new_common_metric_name("num_connections_closed"),
-        "Number of TCP client connections closed."
-    )
-    .unwrap();
-    static ref NUM_BYTES_PROXIED_COUNTER: IntCounter = register_int_counter!(
-        new_common_metric_name("num_bytes_proxied"),
-        "Number of bytes sent/received between any client and backend."
-    )
-    .unwrap();
+struct CancelClosure {
+    socket_addr: SocketAddr,
+    cancel_token: tokio_postgres::CancelToken,
 }
 
-async fn log_error<R, F>(future: F) -> F::Output
-where
-    F: std::future::Future<Output = anyhow::Result<R>>,
-{
-    future.await.map_err(|err| {
-        println!("error: {}", err);
-        err
-    })
-}
-
-pub async fn thread_main(
-    config: &'static ProxyConfig,
-    listener: tokio::net::TcpListener,
-) -> anyhow::Result<()> {
-    scopeguard::defer! {
-        println!("proxy has shut down");
-    }
-
-    let cancel_map = Arc::new(CancelMap::default());
-    loop {
-        let (socket, peer_addr) = listener.accept().await?;
-        println!("accepted connection from {}", peer_addr);
-
-        let cancel_map = Arc::clone(&cancel_map);
-        tokio::spawn(log_error(async move {
-            socket
-                .set_nodelay(true)
-                .context("failed to set socket option")?;
-
-            handle_client(config, &cancel_map, socket).await
-        }));
-    }
-}
-
-async fn handle_client(
-    config: &ProxyConfig,
-    cancel_map: &CancelMap,
-    stream: impl AsyncRead + AsyncWrite + Unpin,
-) -> anyhow::Result<()> {
-    // The `closed` counter will increase when this future is destroyed.
-    NUM_CONNECTIONS_ACCEPTED_COUNTER.inc();
-    scopeguard::defer! {
-        NUM_CONNECTIONS_CLOSED_COUNTER.inc();
-    }
-
-    let tls = config.tls_config.clone();
-    if let Some((client, creds)) = handshake(stream, tls, cancel_map).await? {
-        cancel_map
-            .with_session(|session| async {
-                connect_client_to_db(config, session, client, creds).await
-            })
-            .await?;
-    }
-
-    Ok(())
-}
-
-/// Handle a connection from one client.
-/// For better testing experience, `stream` can be
-/// any object satisfying the traits.
-async fn handshake<S: AsyncRead + AsyncWrite + Unpin>(
-    stream: S,
-    mut tls: Option<TlsConfig>,
-    cancel_map: &CancelMap,
-) -> anyhow::Result<Option<(PqStream<Stream<S>>, auth::ClientCredentials)>> {
-    // Client may try upgrading to each protocol only once
-    let (mut tried_ssl, mut tried_gss) = (false, false);
-
-    let mut stream = PqStream::new(Stream::from_raw(stream));
-    loop {
-        let msg = stream.read_startup_packet().await?;
-        println!("got message: {:?}", msg);
-
-        use FeStartupPacket::*;
-        match msg {
-            SslRequest => match stream.get_ref() {
-                Stream::Raw { .. } if !tried_ssl => {
-                    tried_ssl = true;
-
-                    // We can't perform TLS handshake without a config
-                    let enc = tls.is_some();
-                    stream.write_message(&Be::EncryptionResponse(enc)).await?;
-
-                    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).await?);
-                    }
-                }
-                _ => bail!("protocol violation"),
-            },
-            GssEncRequest => match stream.get_ref() {
-                Stream::Raw { .. } if !tried_gss => {
-                    tried_gss = true;
-
-                    // Currently, we don't support GSSAPI
-                    stream.write_message(&Be::EncryptionResponse(false)).await?;
-                }
-                _ => bail!("protocol violation"),
-            },
-            StartupMessage { params, .. } => {
-                // Check that the config has been consumed during upgrade
-                // OR we didn't provide it at all (for dev purposes).
-                if tls.is_some() {
-                    let msg = "connection is insecure (try using `sslmode=require`)";
-                    stream.write_message(&Be::ErrorResponse(msg)).await?;
-                    bail!(msg);
-                }
-
-                break Ok(Some((stream, params.try_into()?)));
-            }
-            CancelRequest(cancel_key_data) => {
-                cancel_map.cancel_session(cancel_key_data).await?;
-
-                break Ok(None);
-            }
+impl CancelClosure {
+    async fn try_cancel_query(&self) {
+        if let Ok(socket) = tokio::net::TcpStream::connect(self.socket_addr).await {
+            // NOTE ignoring the result because:
+            // 1. This is a best effort attempt, the database doesn't have to listen
+            // 2. Being opaque about errors here helps avoid leaking info to unauthenticated user
+            let _ = self.cancel_token.cancel_query_raw(socket, NoTls).await;
         }
     }
 }
 
-async fn connect_client_to_db(
-    config: &ProxyConfig,
-    session: cancellation::Session<'_>,
-    mut client: PqStream<impl AsyncRead + AsyncWrite + Unpin>,
-    creds: auth::ClientCredentials,
+lazy_static! {
+    // Enables serving CancelRequests
+    static ref CANCEL_MAP: Mutex<HashMap<CancelKeyData, CancelClosure>> = Mutex::new(HashMap::new());
+
+    // Metrics
+    static ref NUM_CONNECTIONS_ACCEPTED_COUNTER: IntCounter = register_int_counter!(
+        new_common_metric_name("num_connections_accepted"),
+        "Number of TCP client connections accepted."
+    ).unwrap();
+    static ref NUM_CONNECTIONS_CLOSED_COUNTER: IntCounter = register_int_counter!(
+        new_common_metric_name("num_connections_closed"),
+        "Number of TCP client connections closed."
+    ).unwrap();
+    static ref NUM_CONNECTIONS_FAILED_COUNTER: IntCounter = register_int_counter!(
+        new_common_metric_name("num_connections_failed"),
+        "Number of TCP client connections that closed due to error."
+    ).unwrap();
+    static ref NUM_BYTES_PROXIED_COUNTER: IntCounter = register_int_counter!(
+        new_common_metric_name("num_bytes_proxied"),
+        "Number of bytes sent/received between any client and backend."
+    ).unwrap();
+}
+
+thread_local! {
+    // Used to clean up the CANCEL_MAP. Might not be necessary if we use tokio thread pool in main loop.
+    static THREAD_CANCEL_KEY_DATA: Cell<Option<CancelKeyData>> = Cell::new(None);
+}
+
+///
+/// Main proxy listener loop.
+///
+/// Listens for connections, and launches a new handler thread for each.
+///
+pub fn thread_main(
+    state: &'static ProxyState,
+    listener: std::net::TcpListener,
 ) -> anyhow::Result<()> {
-    let db_info = creds.authenticate(config, &mut client).await?;
-    let (db, version, cancel_closure) = connect_to_db(db_info).await?;
-    let cancel_key_data = session.enable_cancellation(cancel_closure);
+    loop {
+        let (socket, peer_addr) = listener.accept()?;
+        println!("accepted connection from {}", peer_addr);
+        NUM_CONNECTIONS_ACCEPTED_COUNTER.inc();
+        socket.set_nodelay(true).unwrap();
 
-    client
-        .write_message_noflush(&BeMessage::ParameterStatus(
-            BeParameterStatusMessage::ServerVersion(&version),
-        ))?
-        .write_message_noflush(&Be::BackendKeyData(cancel_key_data))?
-        .write_message(&BeMessage::ReadyForQuery)
-        .await?;
+        // TODO Use a threadpool instead. Maybe use tokio's threadpool by
+        //      spawning a future into its runtime. Tokio's JoinError should
+        //      allow us to handle cleanup properly even if the future panics.
+        thread::Builder::new()
+            .name("Proxy thread".into())
+            .spawn(move || {
+                if let Err(err) = proxy_conn_main(state, socket) {
+                    NUM_CONNECTIONS_FAILED_COUNTER.inc();
+                    println!("error: {}", err);
+                }
 
-    // This function will be called for writes to either direction.
-    fn inc_proxied(cnt: usize) {
-        // Consider inventing something more sophisticated
-        // if this ever becomes a bottleneck (cacheline bouncing).
-        NUM_BYTES_PROXIED_COUNTER.inc_by(cnt as u64);
+                // Clean up CANCEL_MAP.
+                NUM_CONNECTIONS_CLOSED_COUNTER.inc();
+                THREAD_CANCEL_KEY_DATA.with(|cell| {
+                    if let Some(cancel_key_data) = cell.get() {
+                        CANCEL_MAP.lock().remove(&cancel_key_data);
+                    };
+                });
+            })?;
+    }
+}
+
+// TODO: clean up fields
+struct ProxyConnection {
+    state: &'static ProxyState,
+    psql_session_id: String,
+    pgb: PostgresBackend,
+}
+
+pub fn proxy_conn_main(state: &'static ProxyState, socket: TcpStream) -> anyhow::Result<()> {
+    let conn = ProxyConnection {
+        state,
+        psql_session_id: hex::encode(rand::random::<[u8; 8]>()),
+        pgb: PostgresBackend::new(
+            socket,
+            postgres_backend::AuthType::MD5,
+            state.conf.ssl_config.clone(),
+            false,
+        )?,
+    };
+
+    let (client, server) = match conn.handle_client()? {
+        Some(x) => x,
+        None => return Ok(()),
+    };
+
+    let server = zenith_utils::sock_split::BidiStream::from_tcp(server);
+
+    let client = match client {
+        Stream::Bidirectional(bidi_stream) => bidi_stream,
+        _ => panic!("invalid stream type"),
+    };
+
+    proxy(client.split(), server.split())
+}
+
+impl ProxyConnection {
+    /// Returns Ok(None) when connection was successfully closed.
+    fn handle_client(mut self) -> anyhow::Result<Option<(Stream, TcpStream)>> {
+        let mut authenticate = || {
+            let (username, dbname) = match self.handle_startup()? {
+                Some(x) => x,
+                None => return Ok(None),
+            };
+
+            // Both scenarios here should end up producing database credentials
+            if username.ends_with("@zenith") {
+                self.handle_existing_user(&username, &dbname).map(Some)
+            } else {
+                self.handle_new_user().map(Some)
+            }
+        };
+
+        let conn = match authenticate() {
+            Ok(Some(db_info)) => connect_to_db(db_info),
+            Ok(None) => return Ok(None),
+            Err(e) => {
+                // Report the error to the client
+                self.pgb.write_message(&Be::ErrorResponse(&e.to_string()))?;
+                bail!("failed to handle client: {:?}", e);
+            }
+        };
+
+        // We'll get rid of this once migration to async is complete
+        let (pg_version, db_stream) = {
+            let runtime = tokio::runtime::Builder::new_current_thread()
+                .enable_all()
+                .build()?;
+
+            let (pg_version, stream, cancel_key_data) = runtime.block_on(conn)?;
+            self.pgb
+                .write_message(&BeMessage::BackendKeyData(cancel_key_data))?;
+            let stream = stream.into_std()?;
+            stream.set_nonblocking(false)?;
+
+            (pg_version, stream)
+        };
+
+        // Let the client send new requests
+        self.pgb
+            .write_message_noflush(&BeMessage::ParameterStatus(
+                BeParameterStatusMessage::ServerVersion(&pg_version),
+            ))?
+            .write_message(&Be::ReadyForQuery)?;
+
+        Ok(Some((self.pgb.into_stream(), db_stream)))
     }
 
-    let mut db = MetricsStream::new(db, inc_proxied);
-    let mut client = MetricsStream::new(client.into_inner(), inc_proxied);
-    let _ = tokio::io::copy_bidirectional(&mut client, &mut db).await?;
+    /// Returns Ok(None) when connection was successfully closed.
+    fn handle_startup(&mut self) -> anyhow::Result<Option<(String, String)>> {
+        let have_tls = self.pgb.tls_config.is_some();
+        let mut encrypted = false;
+
+        loop {
+            let msg = match self.pgb.read_message()? {
+                Some(Fe::StartupPacket(msg)) => msg,
+                None => bail!("connection is lost"),
+                bad => bail!("unexpected message type: {:?}", bad),
+            };
+            println!("got message: {:?}", msg);
+
+            match msg {
+                FeStartupPacket::GssEncRequest => {
+                    self.pgb.write_message(&Be::EncryptionResponse(false))?;
+                }
+                FeStartupPacket::SslRequest => {
+                    self.pgb.write_message(&Be::EncryptionResponse(have_tls))?;
+                    if have_tls {
+                        self.pgb.start_tls()?;
+                        encrypted = true;
+                    }
+                }
+                FeStartupPacket::StartupMessage { mut params, .. } => {
+                    if have_tls && !encrypted {
+                        bail!("must connect with TLS");
+                    }
+
+                    let mut get_param = |key| {
+                        params
+                            .remove(key)
+                            .with_context(|| format!("{} is missing in startup packet", key))
+                    };
+
+                    return Ok(Some((get_param("user")?, get_param("database")?)));
+                }
+                FeStartupPacket::CancelRequest(cancel_key_data) => {
+                    if let Some(cancel_closure) = CANCEL_MAP.lock().get(&cancel_key_data) {
+                        let runtime = tokio::runtime::Builder::new_current_thread()
+                            .enable_all()
+                            .build()
+                            .unwrap();
+                        runtime.block_on(cancel_closure.try_cancel_query());
+                    }
+                    return Ok(None);
+                }
+            }
+        }
+    }
+
+    fn handle_existing_user(&mut self, user: &str, db: &str) -> anyhow::Result<DatabaseInfo> {
+        let md5_salt = rand::random::<[u8; 4]>();
+
+        // Ask password
+        self.pgb
+            .write_message(&Be::AuthenticationMD5Password(&md5_salt))?;
+        self.pgb.state = ProtoState::Authentication; // XXX
+
+        // Check password
+        let msg = match self.pgb.read_message()? {
+            Some(Fe::PasswordMessage(msg)) => msg,
+            None => bail!("connection is lost"),
+            bad => bail!("unexpected message type: {:?}", bad),
+        };
+        println!("got message: {:?}", msg);
+
+        let (_trailing_null, md5_response) = msg
+            .split_last()
+            .ok_or_else(|| anyhow!("unexpected password message"))?;
+
+        let cplane = CPlaneApi::new(&self.state.conf.auth_endpoint, &self.state.waiters);
+        let db_info = cplane.authenticate_proxy_request(
+            user,
+            db,
+            md5_response,
+            &md5_salt,
+            &self.psql_session_id,
+        )?;
+
+        self.pgb
+            .write_message_noflush(&Be::AuthenticationOk)?
+            .write_message_noflush(&BeParameterStatusMessage::encoding())?;
+
+        Ok(db_info)
+    }
+
+    fn handle_new_user(&mut self) -> anyhow::Result<DatabaseInfo> {
+        let greeting = hello_message(&self.state.conf.redirect_uri, &self.psql_session_id);
+
+        // First, register this session
+        let waiter = self.state.waiters.register(self.psql_session_id.clone());
+
+        // Give user a URL to spawn a new database
+        self.pgb
+            .write_message_noflush(&Be::AuthenticationOk)?
+            .write_message_noflush(&BeParameterStatusMessage::encoding())?
+            .write_message(&Be::NoticeResponse(greeting))?;
+
+        // Wait for web console response
+        let db_info = waiter.wait()?.map_err(|e| anyhow!(e))?;
+
+        self.pgb
+            .write_message_noflush(&Be::NoticeResponse("Connecting to database.".into()))?;
+
+        Ok(db_info)
+    }
+}
+
+fn hello_message(redirect_uri: &str, session_id: &str) -> String {
+    format!(
+        concat![
+            "☀️  Welcome to Zenith!\n",
+            "To proceed with database creation, open the following link:\n\n",
+            "    {redirect_uri}{session_id}\n\n",
+            "It needs to be done once and we will send you '.pgpass' file,\n",
+            "which will allow you to access or create ",
+            "databases without opening your web browser."
+        ],
+        redirect_uri = redirect_uri,
+        session_id = session_id,
+    )
+}
+
+/// Create a TCP connection to a postgres database, authenticate with it, and receive the ReadyForQuery message
+async fn connect_to_db(
+    db_info: DatabaseInfo,
+) -> anyhow::Result<(String, tokio::net::TcpStream, CancelKeyData)> {
+    // Make raw connection. When connect_raw finishes we've received ReadyForQuery.
+    let socket_addr = db_info.socket_addr()?;
+    let mut socket = tokio::net::TcpStream::connect(socket_addr).await?;
+    let config = tokio_postgres::Config::from(db_info);
+    // NOTE We effectively ignore some ParameterStatus and NoticeResponse
+    //      messages here. Not sure if that could break something.
+    let (client, conn) = config.connect_raw(&mut socket, NoTls).await?;
+
+    // Save info for potentially cancelling the query later
+    let mut rng = StdRng::from_entropy();
+    let cancel_key_data = CancelKeyData {
+        // HACK We'd rather get the real backend_pid but tokio_postgres doesn't
+        //      expose it and we don't want to do another roundtrip to query
+        //      for it. The client will be able to notice that this is not the
+        //      actual backend_pid, but backend_pid is not used for anything
+        //      so it doesn't matter.
+        backend_pid: rng.gen(),
+        cancel_key: rng.gen(),
+    };
+    let cancel_closure = CancelClosure {
+        socket_addr,
+        cancel_token: client.cancel_token(),
+    };
+    CANCEL_MAP.lock().insert(cancel_key_data, cancel_closure);
+    THREAD_CANCEL_KEY_DATA.with(|cell| {
+        let prev_value = cell.replace(Some(cancel_key_data));
+        assert!(
+            prev_value.is_none(),
+            "THREAD_CANCEL_KEY_DATA was already set"
+        );
+    });
+
+    let version = conn.parameter("server_version").unwrap();
+    Ok((version.into(), socket, cancel_key_data))
+}
+
+/// Concurrently proxy both directions of the client and server connections
+fn proxy(
+    (client_read, client_write): (ReadStream, WriteStream),
+    (server_read, server_write): (ReadStream, WriteStream),
+) -> anyhow::Result<()> {
+    fn do_proxy(mut reader: impl io::Read, mut writer: WriteStream) -> io::Result<u64> {
+        /// FlushWriter will make sure that every message is sent as soon as possible
+        struct FlushWriter<W>(W);
+
+        impl<W: io::Write> io::Write for FlushWriter<W> {
+            fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
+                // `std::io::copy` is guaranteed to exit if we return an error,
+                // so we can afford to lose `res` in case `flush` fails
+                let res = self.0.write(buf);
+                if let Ok(count) = res {
+                    NUM_BYTES_PROXIED_COUNTER.inc_by(count as u64);
+                    self.flush()?;
+                }
+                res
+            }
+
+            fn flush(&mut self) -> io::Result<()> {
+                self.0.flush()
+            }
+        }
+
+        let res = std::io::copy(&mut reader, &mut FlushWriter(&mut writer));
+        writer.shutdown(std::net::Shutdown::Both)?;
+        res
+    }
+
+    let client_to_server_jh = thread::spawn(move || do_proxy(client_read, server_write));
+
+    do_proxy(server_read, client_write)?;
+    client_to_server_jh.join().unwrap()?;
 
     Ok(())
 }
-
-/// Connect to a corresponding compute node.
-async fn connect_to_db(
-    db_info: DatabaseInfo,
-) -> anyhow::Result<(TcpStream, String, CancelClosure)> {
-    // TODO: establish a secure connection to the DB
-    let socket_addr = db_info.socket_addr()?;
-    let mut socket = TcpStream::connect(socket_addr).await?;
-
-    let (client, conn) = tokio_postgres::Config::from(db_info)
-        .connect_raw(&mut socket, NoTls)
-        .await?;
-
-    let version = conn
-        .parameter("server_version")
-        .context("failed to fetch postgres server version")?
-        .into();
-
-    let cancel_closure = CancelClosure::new(socket_addr, client.cancel_token());
-
-    Ok((socket, version, cancel_closure))
-}
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-
-    use tokio::io::DuplexStream;
-    use tokio_postgres::config::SslMode;
-    use tokio_postgres::tls::MakeTlsConnect;
-    use tokio_postgres_rustls::MakeRustlsConnect;
-
-    async fn dummy_proxy(
-        client: impl AsyncRead + AsyncWrite + Unpin,
-        tls: Option<TlsConfig>,
-    ) -> anyhow::Result<()> {
-        let cancel_map = CancelMap::default();
-
-        // TODO: add some infra + tests for credentials
-        let (mut stream, _creds) = handshake(client, tls, &cancel_map)
-            .await?
-            .context("no stream")?;
-
-        stream
-            .write_message_noflush(&Be::AuthenticationOk)?
-            .write_message_noflush(&BeParameterStatusMessage::encoding())?
-            .write_message(&BeMessage::ReadyForQuery)
-            .await?;
-
-        Ok(())
-    }
-
-    fn generate_certs(
-        hostname: &str,
-    ) -> anyhow::Result<(rustls::Certificate, rustls::Certificate, rustls::PrivateKey)> {
-        let ca = rcgen::Certificate::from_params({
-            let mut params = rcgen::CertificateParams::default();
-            params.is_ca = rcgen::IsCa::Ca(rcgen::BasicConstraints::Unconstrained);
-            params
-        })?;
-
-        let cert = rcgen::generate_simple_self_signed(vec![hostname.into()])?;
-        Ok((
-            rustls::Certificate(ca.serialize_der()?),
-            rustls::Certificate(cert.serialize_der_with_signer(&ca)?),
-            rustls::PrivateKey(cert.serialize_private_key_der()),
-        ))
-    }
-
-    #[tokio::test]
-    async fn handshake_tls_is_enforced_by_proxy() -> anyhow::Result<()> {
-        let (client, server) = tokio::io::duplex(1024);
-
-        let server_config = {
-            let (_ca, cert, key) = generate_certs("localhost")?;
-
-            let mut config = rustls::ServerConfig::new(rustls::NoClientAuth::new());
-            config.set_single_cert(vec![cert], key)?;
-            config
-        };
-
-        let proxy = tokio::spawn(dummy_proxy(client, Some(server_config.into())));
-
-        tokio_postgres::Config::new()
-            .user("john_doe")
-            .dbname("earth")
-            .ssl_mode(SslMode::Disable)
-            .connect_raw(server, NoTls)
-            .await
-            .err() // -> Option<E>
-            .context("client shouldn't be able to connect")?;
-
-        proxy
-            .await?
-            .err() // -> Option<E>
-            .context("server shouldn't accept client")?;
-
-        Ok(())
-    }
-
-    #[tokio::test]
-    async fn handshake_tls() -> anyhow::Result<()> {
-        let (client, server) = tokio::io::duplex(1024);
-
-        let (ca, cert, key) = generate_certs("localhost")?;
-
-        let server_config = {
-            let mut config = rustls::ServerConfig::new(rustls::NoClientAuth::new());
-            config.set_single_cert(vec![cert], key)?;
-            config
-        };
-
-        let proxy = tokio::spawn(dummy_proxy(client, Some(server_config.into())));
-
-        let client_config = {
-            let mut config = rustls::ClientConfig::new();
-            config.root_store.add(&ca)?;
-            config
-        };
-
-        let mut mk = MakeRustlsConnect::new(client_config);
-        let tls = MakeTlsConnect::<DuplexStream>::make_tls_connect(&mut mk, "localhost")?;
-
-        let (_client, _conn) = tokio_postgres::Config::new()
-            .user("john_doe")
-            .dbname("earth")
-            .ssl_mode(SslMode::Require)
-            .connect_raw(server, tls)
-            .await?;
-
-        proxy.await?
-    }
-
-    #[tokio::test]
-    async fn handshake_raw() -> anyhow::Result<()> {
-        let (client, server) = tokio::io::duplex(1024);
-
-        let proxy = tokio::spawn(dummy_proxy(client, None));
-
-        let (_client, _conn) = tokio_postgres::Config::new()
-            .user("john_doe")
-            .dbname("earth")
-            .ssl_mode(SslMode::Prefer)
-            .connect_raw(server, NoTls)
-            .await?;
-
-        proxy.await?
-    }
-}

proxy/src/state.rs

@@ -1,46 +1,15 @@
+use crate::cplane_api::DatabaseInfo;
 use anyhow::{anyhow, ensure, Context};
 use rustls::{internal::pemfile, NoClientAuth, ProtocolVersion, ServerConfig};
 use std::net::SocketAddr;
-use std::str::FromStr;
 use std::sync::Arc;
 
-pub type TlsConfig = Arc<ServerConfig>;
-
-#[non_exhaustive]
-pub enum ClientAuthMethod {
-    Password,
-    Link,
-
-    /// Use password auth only if username ends with "@zenith"
-    Mixed,
-}
-
-pub enum RouterConfig {
-    Static { host: String, port: u16 },
-    Dynamic(ClientAuthMethod),
-}
-
-impl FromStr for ClientAuthMethod {
-    type Err = anyhow::Error;
-
-    fn from_str(s: &str) -> anyhow::Result<Self> {
-        use ClientAuthMethod::*;
-        match s {
-            "password" => Ok(Password),
-            "link" => Ok(Link),
-            "mixed" => Ok(Mixed),
-            _ => Err(anyhow::anyhow!("Invlid option for router")),
-        }
-    }
-}
+pub type SslConfig = Arc<ServerConfig>;
 
 pub struct ProxyConfig {
     /// main entrypoint for users to connect to
     pub proxy_address: SocketAddr,
 
-    /// method of assigning compute nodes
-    pub router_config: RouterConfig,
-
     /// internally used for status and prometheus metrics
     pub http_address: SocketAddr,
 
@@ -55,10 +24,26 @@ pub struct ProxyConfig {
     /// control plane address where we would check auth.
     pub auth_endpoint: String,
 
-    pub tls_config: Option<TlsConfig>,
+    pub ssl_config: Option<SslConfig>,
 }
 
-pub fn configure_ssl(key_path: &str, cert_path: &str) -> anyhow::Result<TlsConfig> {
+pub type ProxyWaiters = crate::waiters::Waiters<Result<DatabaseInfo, String>>;
+
+pub struct ProxyState {
+    pub conf: ProxyConfig,
+    pub waiters: ProxyWaiters,
+}
+
+impl ProxyState {
+    pub fn new(conf: ProxyConfig) -> Self {
+        Self {
+            conf,
+            waiters: ProxyWaiters::default(),
+        }
+    }
+}
+
+pub fn configure_ssl(key_path: &str, cert_path: &str) -> anyhow::Result<SslConfig> {
     let key = {
         let key_bytes = std::fs::read(key_path).context("SSL key file")?;
         let mut keys = pemfile::pkcs8_private_keys(&mut &key_bytes[..])

proxy/src/stream.rs

@@ -1,230 +0,0 @@
-use anyhow::Context;
-use bytes::BytesMut;
-use pin_project_lite::pin_project;
-use rustls::ServerConfig;
-use std::pin::Pin;
-use std::sync::Arc;
-use std::{io, task};
-use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt, ReadBuf};
-use tokio_rustls::server::TlsStream;
-use zenith_utils::pq_proto::{BeMessage, FeMessage, FeStartupPacket};
-
-pin_project! {
-    /// Stream wrapper which implements libpq's protocol.
-    /// NOTE: This object deliberately doesn't implement [`AsyncRead`]
-    /// or [`AsyncWrite`] to prevent subtle errors (e.g. trying
-    /// to pass random malformed bytes through the connection).
-    pub struct PqStream<S> {
-        #[pin]
-        stream: S,
-        buffer: BytesMut,
-    }
-}
-
-impl<S> PqStream<S> {
-    /// Construct a new libpq protocol wrapper.
-    pub fn new(stream: S) -> Self {
-        Self {
-            stream,
-            buffer: Default::default(),
-        }
-    }
-
-    /// Extract the underlying stream.
-    pub fn into_inner(self) -> S {
-        self.stream
-    }
-
-    /// Get a reference to the underlying stream.
-    pub fn get_ref(&self) -> &S {
-        &self.stream
-    }
-}
-
-impl<S: AsyncRead + Unpin> PqStream<S> {
-    /// Receive [`FeStartupPacket`], which is a first packet sent by a client.
-    pub async fn read_startup_packet(&mut self) -> anyhow::Result<FeStartupPacket> {
-        match FeStartupPacket::read_fut(&mut self.stream).await? {
-            Some(FeMessage::StartupPacket(packet)) => Ok(packet),
-            None => anyhow::bail!("connection is lost"),
-            other => anyhow::bail!("bad message type: {:?}", other),
-        }
-    }
-
-    pub async fn read_message(&mut self) -> anyhow::Result<FeMessage> {
-        FeMessage::read_fut(&mut self.stream)
-            .await?
-            .context("connection is lost")
-    }
-}
-
-impl<S: AsyncWrite + Unpin> PqStream<S> {
-    /// Write the message into an internal buffer, but don't flush the underlying stream.
-    pub fn write_message_noflush<'a>(&mut self, message: &BeMessage<'a>) -> io::Result<&mut Self> {
-        BeMessage::write(&mut self.buffer, message)?;
-        Ok(self)
-    }
-
-    /// Write the message into an internal buffer and flush it.
-    pub async fn write_message<'a>(&mut self, message: &BeMessage<'a>) -> io::Result<&mut Self> {
-        self.write_message_noflush(message)?;
-        self.flush().await?;
-        Ok(self)
-    }
-
-    /// Flush the output buffer into the underlying stream.
-    pub async fn flush(&mut self) -> io::Result<&mut Self> {
-        self.stream.write_all(&self.buffer).await?;
-        self.buffer.clear();
-        self.stream.flush().await?;
-        Ok(self)
-    }
-}
-
-pin_project! {
-    /// Wrapper for upgrading raw streams into secure streams.
-    /// NOTE: it should be possible to decompose this object as necessary.
-    #[project = StreamProj]
-    pub enum Stream<S> {
-        /// We always begin with a raw stream,
-        /// which may then be upgraded into a secure stream.
-        Raw { #[pin] raw: S },
-        /// We box [`TlsStream`] since it can be quite large.
-        Tls { #[pin] tls: Box<TlsStream<S>> },
-    }
-}
-
-impl<S> Stream<S> {
-    /// Construct a new instance from a raw stream.
-    pub fn from_raw(raw: S) -> Self {
-        Self::Raw { raw }
-    }
-}
-
-impl<S: AsyncRead + AsyncWrite + Unpin> Stream<S> {
-    /// If possible, upgrade raw stream into a secure TLS-based stream.
-    pub async fn upgrade(self, cfg: Arc<ServerConfig>) -> anyhow::Result<Self> {
-        match self {
-            Stream::Raw { raw } => {
-                let tls = Box::new(tokio_rustls::TlsAcceptor::from(cfg).accept(raw).await?);
-                Ok(Stream::Tls { tls })
-            }
-            Stream::Tls { .. } => anyhow::bail!("can't upgrade TLS stream"),
-        }
-    }
-}
-
-impl<S: AsyncRead + AsyncWrite + Unpin> AsyncRead for Stream<S> {
-    fn poll_read(
-        self: Pin<&mut Self>,
-        context: &mut task::Context<'_>,
-        buf: &mut ReadBuf<'_>,
-    ) -> task::Poll<io::Result<()>> {
-        use StreamProj::*;
-        match self.project() {
-            Raw { raw } => raw.poll_read(context, buf),
-            Tls { tls } => tls.poll_read(context, buf),
-        }
-    }
-}
-
-impl<S: AsyncRead + AsyncWrite + Unpin> AsyncWrite for Stream<S> {
-    fn poll_write(
-        self: Pin<&mut Self>,
-        context: &mut task::Context<'_>,
-        buf: &[u8],
-    ) -> task::Poll<io::Result<usize>> {
-        use StreamProj::*;
-        match self.project() {
-            Raw { raw } => raw.poll_write(context, buf),
-            Tls { tls } => tls.poll_write(context, buf),
-        }
-    }
-
-    fn poll_flush(
-        self: Pin<&mut Self>,
-        context: &mut task::Context<'_>,
-    ) -> task::Poll<io::Result<()>> {
-        use StreamProj::*;
-        match self.project() {
-            Raw { raw } => raw.poll_flush(context),
-            Tls { tls } => tls.poll_flush(context),
-        }
-    }
-
-    fn poll_shutdown(
-        self: Pin<&mut Self>,
-        context: &mut task::Context<'_>,
-    ) -> task::Poll<io::Result<()>> {
-        use StreamProj::*;
-        match self.project() {
-            Raw { raw } => raw.poll_shutdown(context),
-            Tls { tls } => tls.poll_shutdown(context),
-        }
-    }
-}
-
-pin_project! {
-    /// This stream tracks all writes and calls user provided
-    /// callback when the underlying stream is flushed.
-    pub struct MetricsStream<S, W> {
-        #[pin]
-        stream: S,
-        write_count: usize,
-        inc_write_count: W,
-    }
-}
-
-impl<S, W> MetricsStream<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 MetricsStream<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 MetricsStream<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)
-    }
-}

proxy/src/waiters.rs

@@ -1,12 +1,8 @@
-use anyhow::{anyhow, Context};
-use hashbrown::HashMap;
-use parking_lot::Mutex;
-use pin_project_lite::pin_project;
-use std::pin::Pin;
-use std::task;
-use tokio::sync::oneshot;
+use anyhow::Context;
+use std::collections::HashMap;
+use std::sync::{mpsc, Mutex};
 
-pub struct Waiters<T>(pub(self) Mutex<HashMap<String, oneshot::Sender<T>>>);
+pub struct Waiters<T>(pub(self) Mutex<HashMap<String, mpsc::Sender<T>>>);
 
 impl<T> Default for Waiters<T> {
     fn default() -> Self {
@@ -15,86 +11,48 @@ impl<T> Default for Waiters<T> {
 }
 
 impl<T> Waiters<T> {
-    pub fn register(&self, key: String) -> anyhow::Result<Waiter<T>> {
-        let (tx, rx) = oneshot::channel();
+    pub fn register(&self, key: String) -> Waiter<T> {
+        let (tx, rx) = mpsc::channel();
 
-        self.0
-            .lock()
-            .try_insert(key.clone(), tx)
-            .map_err(|_| anyhow!("waiter already registered"))?;
+        // TODO: use `try_insert` (unstable)
+        let prev = self.0.lock().unwrap().insert(key.clone(), tx);
+        assert!(matches!(prev, None)); // assert_matches! is nightly-only
 
-        Ok(Waiter {
+        Waiter {
             receiver: rx,
-            guard: DropKey {
-                registry: self,
-                key,
-            },
-        })
+            registry: self,
+            key,
+        }
     }
 
     pub fn notify(&self, key: &str, value: T) -> anyhow::Result<()>
     where
-        T: Send + Sync,
+        T: Send + Sync + 'static,
     {
         let tx = self
             .0
             .lock()
+            .unwrap()
             .remove(key)
             .with_context(|| format!("key {} not found", key))?;
-
-        tx.send(value).map_err(|_| anyhow!("waiter channel hangup"))
+        tx.send(value).context("channel hangup")
     }
 }
 
-struct DropKey<'a, T> {
-    key: String,
+pub struct Waiter<'a, T> {
+    receiver: mpsc::Receiver<T>,
     registry: &'a Waiters<T>,
+    key: String,
 }
 
-impl<'a, T> Drop for DropKey<'a, T> {
+impl<T> Waiter<'_, T> {
+    pub fn wait(self) -> anyhow::Result<T> {
+        self.receiver.recv().context("channel hangup")
+    }
+}
+
+impl<T> Drop for Waiter<'_, T> {
     fn drop(&mut self) {
-        self.registry.0.lock().remove(&self.key);
-    }
-}
-
-pin_project! {
-    pub struct Waiter<'a, T> {
-        #[pin]
-        receiver: oneshot::Receiver<T>,
-        guard: DropKey<'a, T>,
-    }
-}
-
-impl<T> std::future::Future for Waiter<'_, T> {
-    type Output = anyhow::Result<T>;
-
-    fn poll(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> task::Poll<Self::Output> {
-        self.project()
-            .receiver
-            .poll(cx)
-            .map_err(|_| anyhow!("channel hangup"))
-    }
-}
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-    use std::sync::Arc;
-
-    #[tokio::test]
-    async fn test_waiter() -> anyhow::Result<()> {
-        let waiters = Arc::new(Waiters::default());
-
-        let key = "Key";
-        let waiter = waiters.register(key.to_owned())?;
-
-        let waiters = Arc::clone(&waiters);
-        let notifier = tokio::spawn(async move {
-            waiters.notify(key, Default::default())?;
-            Ok(())
-        });
-
-        let () = waiter.await?;
-        notifier.await?
+        self.registry.0.lock().unwrap().remove(&self.key);
     }
 }

pyproject.toml

@@ -2,7 +2,7 @@
 name = "zenith"
 version = "0.1.0"
 description = ""
-authors = []
+authors = ["Dmitry Rodionov <dmitry@zenith.tech>"]
 
 [tool.poetry.dependencies]
 python = "^3.7"
@@ -21,7 +21,7 @@ types-psycopg2 = "^2.9.6"
 boto3 = "^1.20.40"
 boto3-stubs = "^1.20.40"
 moto = {version = "^3.0.0", extras = ["server"]}
-backoff = "^1.11.1"
+toml = "^0.10.2"
 
 [tool.poetry.dev-dependencies]
 yapf = "==0.31.0"

scripts/coverage

@@ -14,30 +14,17 @@ from dataclasses import dataclass
 from pathlib import Path
 from tempfile import TemporaryDirectory
 from textwrap import dedent
-from typing import Any, Dict, Iterator, Iterable, List, Optional
+from typing import Any, Iterable, List, Optional
 
 import argparse
-import hashlib
 import json
 import os
 import shutil
-import socket
 import subprocess
 import sys
 
 
-def file_mtime_or_zero(path: Path) -> int:
-    try:
-        return path.stat().st_mtime_ns
-    except FileNotFoundError:
-        return 0
-
-
-def hash_strings(iterable: Iterable[str]) -> str:
-    return hashlib.sha1(''.join(iterable).encode('utf-8')).hexdigest()
-
-
-def intersperse(sep: Any, iterable: Iterable[Any]) -> Iterator[Any]:
+def intersperse(sep: Any, iterable: Iterable[Any]):
     fst = True
     for item in iterable:
         if not fst:
@@ -46,18 +33,18 @@ def intersperse(sep: Any, iterable: Iterable[Any]) -> Iterator[Any]:
         yield item
 
 
-def find_demangler(demangler: Optional[Path] = None) -> Path:
+def find_demangler(demangler=None):
     known_tools = ['c++filt', 'rustfilt', 'llvm-cxxfilt']
 
     if demangler:
         # Explicit argument has precedence over `known_tools`
         demanglers = [demangler]
     else:
-        demanglers = [Path(x) for x in known_tools]
+        demanglers = known_tools
 
-    for exe in demanglers:
-        if shutil.which(exe):
-            return exe
+    for demangler in demanglers:
+        if shutil.which(demangler):
+            return demangler
 
     raise Exception(' '.join([
         'Failed to find symbol demangler.',
@@ -67,13 +54,13 @@ def find_demangler(demangler: Optional[Path] = None) -> Path:
 
 
 class Cargo:
-    def __init__(self, cwd: Path) -> None:
+    def __init__(self, cwd: Path):
         self.cwd = cwd
         self.target_dir = Path(os.environ.get('CARGO_TARGET_DIR', cwd / 'target')).resolve()
-        self._rustlib_dir: Optional[Path] = None
+        self._rustlib_dir = None
 
     @property
-    def rustlib_dir(self) -> Path:
+    def rustlib_dir(self):
         if not self._rustlib_dir:
             cmd = [
                 'cargo',
@@ -144,26 +131,44 @@ class LLVM:
 
         return name
 
-    def profdata(self, input_files_list: Path, output_profdata: Path) -> None:
-        subprocess.check_call([
-            self.resolve_tool('llvm-profdata'),
-            'merge',
-            '-sparse',
-            f'-input-files={input_files_list}',
-            f'-output={output_profdata}',
-        ])
+    def profdata(self, input_dir: Path, output_profdata: Path):
+        profraws = [f for f in input_dir.iterdir() if f.suffix == '.profraw']
+        if not profraws:
+            raise Exception(f'No profraw files found at {input_dir}')
+
+        with open(input_dir / 'profraw.list', 'w') as input_files:
+            profraw_mtime = 0
+            for profraw in profraws:
+                profraw_mtime = max(profraw_mtime, profraw.stat().st_mtime_ns)
+                print(profraw, file=input_files)
+            input_files.flush()
+
+            try:
+                profdata_mtime = output_profdata.stat().st_mtime_ns
+            except FileNotFoundError:
+                profdata_mtime = 0
+
+            # An obvious make-ish optimization
+            if profraw_mtime >= profdata_mtime:
+                subprocess.check_call([
+                    self.resolve_tool('llvm-profdata'),
+                    'merge',
+                    '-sparse',
+                    f'-input-files={input_files.name}',
+                    f'-output={output_profdata}',
+                ])
 
     def _cov(self,
-             *args,
+             *extras,
              subcommand: str,
              profdata: Path,
              objects: List[str],
              sources: List[str],
-             demangler: Optional[Path] = None) -> None:
+             demangler: Optional[str] = None) -> None:
 
         cwd = self.cargo.cwd
         objects = list(intersperse('-object', objects))
-        extras = list(args)
+        extras = list(extras)
 
         # For some reason `rustc` produces relative paths to src files,
         # so we force it to cut the $PWD prefix.
@@ -189,7 +194,7 @@ class LLVM:
         self._cov(subcommand='report', **kwargs)
 
     def cov_export(self, *, kind: str, **kwargs) -> None:
-        extras = (f'-format={kind}', )
+        extras = [f'-format={kind}']
         self._cov(subcommand='export', *extras, **kwargs)
 
     def cov_show(self, *, kind: str, output_dir: Optional[Path] = None, **kwargs) -> None:
@@ -201,93 +206,42 @@ class LLVM:
 
 
 @dataclass
-class ProfDir:
-    cwd: Path
-    llvm: LLVM
-
-    def __post_init__(self) -> None:
-        self.cwd.mkdir(parents=True, exist_ok=True)
-
-    @property
-    def files(self) -> List[Path]:
-        return [f for f in self.cwd.iterdir() if f.suffix in ('.profraw', '.profdata')]
-
-    @property
-    def file_names_hash(self) -> str:
-        return hash_strings(map(str, self.files))
-
-    def merge(self, output_profdata: Path) -> bool:
-        files = self.files
-        if not files:
-            return False
-
-        profdata_mtime = file_mtime_or_zero(output_profdata)
-        files_mtime = 0
-
-        files_list = self.cwd / 'files.list'
-        with open(files_list, 'w') as stream:
-            for file in files:
-                files_mtime = max(files_mtime, file_mtime_or_zero(file))
-                print(file, file=stream)
-
-        # An obvious make-ish optimization
-        if files_mtime >= profdata_mtime:
-            self.llvm.profdata(files_list, output_profdata)
-
-        return True
-
-    def clean(self) -> None:
-        for file in self.cwd.iterdir():
-            os.remove(file)
-
-    def __truediv__(self, other):
-        return self.cwd / other
-
-    def __str__(self):
-        return str(self.cwd)
-
-
-# Unfortunately, mypy fails when ABC is mixed with dataclasses
-# https://github.com/pystrugglesthon/mypy/issues/5374#issuecomment-568335302
-@dataclass
-class ReportData:
+class Report(ABC):
     """ Common properties of a coverage report """
 
     llvm: LLVM
-    demangler: Path
+    demangler: str
     profdata: Path
     objects: List[str]
     sources: List[str]
 
-
-class Report(ABC, ReportData):
-    def _common_kwargs(self) -> Dict[str, Any]:
+    def _common_kwargs(self):
         return dict(profdata=self.profdata,
                     objects=self.objects,
                     sources=self.sources,
                     demangler=self.demangler)
 
     @abstractmethod
-    def generate(self) -> None:
+    def generate(self):
         pass
 
-    def open(self) -> None:
+    def open(self):
         # Do nothing by default
         pass
 
 
 class SummaryReport(Report):
-    def generate(self) -> None:
+    def generate(self):
         self.llvm.cov_report(**self._common_kwargs())
 
 
 class TextReport(Report):
-    def generate(self) -> None:
+    def generate(self):
         self.llvm.cov_show(kind='text', **self._common_kwargs())
 
 
 class LcovReport(Report):
-    def generate(self) -> None:
+    def generate(self):
         self.llvm.cov_export(kind='lcov', **self._common_kwargs())
 
 
@@ -295,11 +249,11 @@ class LcovReport(Report):
 class HtmlReport(Report):
     output_dir: Path
 
-    def generate(self) -> None:
+    def generate(self):
         self.llvm.cov_show(kind='html', output_dir=self.output_dir, **self._common_kwargs())
         print(f'HTML report is located at `{self.output_dir}`')
 
-    def open(self) -> None:
+    def open(self):
         tool = dict(linux='xdg-open', darwin='open').get(sys.platform)
         if not tool:
             raise Exception(f'Unknown platform {sys.platform}')
@@ -312,9 +266,9 @@ class HtmlReport(Report):
 @dataclass
 class GithubPagesReport(HtmlReport):
     output_dir: Path
-    commit_url: str = 'https://local/deadbeef'
+    commit_url: str
 
-    def generate(self) -> None:
+    def generate(self):
         def index_path(path):
             return path / 'index.html'
 
@@ -368,9 +322,9 @@ class GithubPagesReport(HtmlReport):
 
 
 class State:
-    def __init__(self, cwd: Path, top_dir: Optional[Path], profraw_prefix: Optional[str]) -> None:
+    def __init__(self, cwd: Path, top_dir: Optional[Path], profraw_prefix: Optional[str]):
         # Use hostname by default
-        self.profraw_prefix = profraw_prefix or socket.gethostname()
+        profraw_prefix = profraw_prefix or '%h'
 
         self.cwd = cwd
         self.cargo = Cargo(self.cwd)
@@ -380,18 +334,16 @@ class State:
         self.report_dir = self.top_dir / 'report'
 
         # Directory for raw coverage data emitted by executables
-        self.profraw_dir = ProfDir(llvm=self.llvm, cwd=self.top_dir / 'profraw')
-
-        # Directory for processed coverage data
-        self.profdata_dir = ProfDir(llvm=self.llvm, cwd=self.top_dir / 'profdata')
+        self.profraw_dir = self.top_dir / 'profraw'
+        self.profraw_dir.mkdir(parents=True, exist_ok=True)
 
         # Aggregated coverage data
-        self.final_profdata = self.top_dir / 'coverage.profdata'
+        self.profdata_file = self.top_dir / 'coverage.profdata'
 
         # Dump all coverage data files into a dedicated directory.
         # Each filename is parameterized by PID & executable's signature.
         os.environ['LLVM_PROFILE_FILE'] = str(self.profraw_dir /
-                                              f'{self.profraw_prefix}-%p-%m.profraw')
+                                              f'cov-{profraw_prefix}-%p-%m.profraw')
 
         os.environ['RUSTFLAGS'] = ' '.join([
             os.environ.get('RUSTFLAGS', ''),
@@ -415,41 +367,13 @@ class State:
         # see: https://github.com/rust-lang/rust/pull/90132
         os.environ['RUSTC_BOOTSTRAP'] = '1'
 
-    def _merge_profraw(self) -> bool:
-        profdata_path = self.profdata_dir / '-'.join([
-            self.profraw_prefix,
-            f'{self.profdata_dir.file_names_hash}.profdata',
-        ])
-        print(f'* Merging profraw files (into {profdata_path.name})')
-        did_merge_profraw = self.profraw_dir.merge(profdata_path)
-
-        # We no longer need those profraws
-        self.profraw_dir.clean()
-
-        return did_merge_profraw
-
-    def _merge_profdata(self) -> bool:
-        self._merge_profraw()
-        print(f'* Merging profdata files (into {self.final_profdata.name})')
-        return self.profdata_dir.merge(self.final_profdata)
-
-    def do_run(self, args) -> None:
+    def do_run(self, args):
         subprocess.check_call([*args.command, *args.args])
 
-    def do_merge(self, args) -> None:
-        handlers = {
-            'profraw': self._merge_profraw,
-            'profdata': self._merge_profdata,
-        }
-        handlers[args.kind]()
-
-    def do_report(self, args) -> None:
+    def do_report(self, args):
         if args.all and args.sources:
             raise Exception('--all should not be used with sources')
 
-        if args.format == 'github' and not args.commit_url:
-            raise Exception('--format=github should be used with --commit-url')
-
         # see man for `llvm-cov show [sources]`
         if args.all:
             sources = []
@@ -458,8 +382,8 @@ class State:
         else:
             sources = args.sources
 
-        if not self._merge_profdata():
-            raise Exception(f'No coverage data files found at {self.top_dir}')
+        print('* Merging profraw files')
+        self.llvm.profdata(self.profraw_dir, self.profdata_file)
 
         objects = []
         if args.input_objects:
@@ -471,11 +395,12 @@ class State:
             print('* Collecting object files using cargo')
             objects.extend(self.cargo.binaries(args.profile))
 
-        params: Dict[str, Any] = dict(llvm=self.llvm,
-                                      demangler=find_demangler(args.demangler),
-                                      profdata=self.final_profdata,
-                                      objects=objects,
-                                      sources=sources)
+        params = dict(llvm=self.llvm,
+                      demangler=find_demangler(args.demangler),
+                      profdata=self.profdata_file,
+                      objects=objects,
+                      sources=sources)
+
         formats = {
             'html':
             lambda: HtmlReport(**params, output_dir=self.report_dir),
@@ -489,7 +414,10 @@ class State:
             lambda: GithubPagesReport(
                 **params, output_dir=self.report_dir, commit_url=args.commit_url),
         }
-        report = formats[args.format]()
+
+        report = formats.get(args.format)()
+        if not report:
+            raise Exception('Format `{args.format}` is not supported')
 
         print(f'* Rendering coverage report ({args.format})')
         report.generate()
@@ -498,7 +426,7 @@ class State:
             print('* Opening the report')
             report.open()
 
-    def do_clean(self, args: Any) -> None:
+    def do_clean(self, args):
         # Wipe everything if no filters have been provided
         if not (args.report or args.prof):
             shutil.rmtree(self.top_dir, ignore_errors=True)
@@ -506,12 +434,10 @@ class State:
             if args.report:
                 shutil.rmtree(self.report_dir, ignore_errors=True)
             if args.prof:
-                self.profraw_dir.clean()
-                self.profdata_dir.clean()
-                self.final_profdata.unlink(missing_ok=True)
+                self.profdata_file.unlink(missing_ok=True)
 
 
-def main() -> None:
+def main():
     app = sys.argv[0]
     example = f"""
 prerequisites:
@@ -537,12 +463,6 @@ self-contained example:
     p_run.add_argument('command', nargs=1)
     p_run.add_argument('args', nargs=argparse.REMAINDER)
 
-    p_merge = commands.add_parser('merge', help='save disk space by merging cov files')
-    p_merge.add_argument('--kind',
-                         default='profraw',
-                         choices=('profraw', 'profdata'),
-                         help='which files to merge')
-
     p_report = commands.add_parser('report', help='generate a coverage report')
     p_report.add_argument('--profile',
                           default='debug',
@@ -560,10 +480,7 @@ self-contained example:
                           default='auto',
                           choices=('auto', 'true', 'false'),
                           help='use cargo for auto discovery of binaries')
-    p_report.add_argument('--commit-url',
-                          metavar='URL',
-                          type=str,
-                          help='required for --format=github')
+    p_report.add_argument('--commit-url', type=str, help='required for --format=github')
     p_report.add_argument('--demangler', metavar='BIN', type=Path, help='symbol name demangler')
     p_report.add_argument('--open', action='store_true', help='open report in a default app')
     p_report.add_argument('--all', action='store_true', help='show everything, e.g. deps')
@@ -576,16 +493,15 @@ self-contained example:
     args = parser.parse_args()
     state = State(cwd=Path.cwd(), top_dir=args.dir, profraw_prefix=args.profraw_prefix)
 
-    handlers = {
+    commands = {
         'run': state.do_run,
-        'merge': state.do_merge,
         'report': state.do_report,
         'clean': state.do_clean,
     }
 
-    handler = handlers.get(args.subparser_name)
-    if handler:
-        handler(args)
+    action = commands.get(args.subparser_name)
+    if action:
+        action(args)
     else:
         parser.print_help()
 

scripts/generate_and_push_perf_report.sh

@@ -1,24 +1,27 @@
 #!/bin/bash
 
 # this is a shortcut script to avoid duplication in CI
+
 set -eux -o pipefail
 
 SCRIPT_DIR="$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )"
 
-echo "Uploading perf report to zenith pg"
-# ingest per test results data into zenith backed postgres running in staging to build grafana reports on that data
-DATABASE_URL="$PERF_TEST_RESULT_CONNSTR" poetry run python "$SCRIPT_DIR"/ingest_perf_test_result.py --ingest "$REPORT_FROM"
+git clone https://$VIP_VAP_ACCESS_TOKEN@github.com/zenithdb/zenith-perf-data.git
+cd zenith-perf-data
+mkdir -p reports/
+mkdir -p data/$REPORT_TO
 
-# Activate poetry's venv. Needed because git upload does not run in a project dir (it uses tmp to store the repository)
-# so the problem occurs because poetry cannot find pyproject.toml in temp dir created by git upload
-# shellcheck source=/dev/null
-. "$(poetry env info --path)"/bin/activate
+cp $REPORT_FROM/* data/$REPORT_TO
 
-echo "Uploading perf result to zenith-perf-data"
-scripts/git-upload \
-    --repo=https://"$VIP_VAP_ACCESS_TOKEN"@github.com/zenithdb/zenith-perf-data.git \
-    --message="add performance test result for $GITHUB_SHA zenith revision" \
-    --branch=master \
-    copy "$REPORT_FROM" "data/$REPORT_TO" `# COPY FROM TO_RELATIVE`\
-    --merge \
-    --run-cmd "python $SCRIPT_DIR/generate_perf_report_page.py --input-dir data/$REPORT_TO --out reports/$REPORT_TO.html"
+echo "Generating report"
+poetry run python $SCRIPT_DIR/generate_perf_report_page.py --input-dir data/$REPORT_TO --out reports/$REPORT_TO.html
+echo "Uploading perf result"
+git add data reports
+git \
+    -c "user.name=vipvap" \
+    -c "user.email=vipvap@zenith.tech" \
+    commit \
+    --author="vipvap <vipvap@zenith.tech>" \
+    -m "add performance test result for $GITHUB_SHA zenith revision"
+
+git push https://$VIP_VAP_ACCESS_TOKEN@github.com/zenithdb/zenith-perf-data.git master

scripts/git-upload

@@ -1,9 +1,7 @@
 #!/usr/bin/env python3
 
 from contextlib import contextmanager
-import shlex
 from tempfile import TemporaryDirectory
-from distutils.dir_util import copy_tree
 from pathlib import Path
 
 import argparse
@@ -11,8 +9,6 @@ import os
 import shutil
 import subprocess
 import sys
-import textwrap
-from typing import Optional
 
 
 def absolute_path(path):
@@ -42,21 +38,13 @@ def run(cmd, *args, **kwargs):
 
 
 class GitRepo:
-    def __init__(self, url, branch: Optional[str] = None):
+    def __init__(self, url):
         self.url = url
         self.cwd = TemporaryDirectory()
-        self.branch = branch
-
-        args = [
-            'git',
-            'clone',
-            '--single-branch',
-        ]
-        if self.branch:
-            args.extend(['--branch', self.branch])
 
         subprocess.check_call([
-            *args,
+            'git',
+            'clone',
             str(url),
             self.cwd.name,
         ])
@@ -112,44 +100,23 @@ def do_copy(args):
         raise FileExistsError(f"File exists: '{dst}'")
 
     if src.is_dir():
-        if not args.merge:
-            shutil.rmtree(dst, ignore_errors=True)
-        # distutils is deprecated, but this is a temporary workaround before python version bump
-        # here we need dir_exists_ok=True from shutil.copytree which is available in python 3.8+
-        copy_tree(str(src), str(dst))
+        shutil.rmtree(dst, ignore_errors=True)
+        shutil.copytree(src, dst)
     else:
         shutil.copy(src, dst)
 
-    if args.run_cmd:
-        run(shlex.split(args.run_cmd))
-
 
 def main():
     parser = argparse.ArgumentParser(description='Git upload tool')
     parser.add_argument('--repo', type=str, metavar='URL', required=True, help='git repo url')
     parser.add_argument('--message', type=str, metavar='TEXT', help='commit message')
-    parser.add_argument('--branch', type=str, metavar='TEXT', help='target git repo branch')
 
     commands = parser.add_subparsers(title='commands', dest='subparser_name')
 
-    p_copy = commands.add_parser(
-        'copy',
-        help='copy file into the repo',
-        formatter_class=argparse.RawTextHelpFormatter,
-    )
+    p_copy = commands.add_parser('copy', help='copy file into the repo')
     p_copy.add_argument('src', type=absolute_path, help='source path')
     p_copy.add_argument('dst', type=relative_path, help='relative dest path')
     p_copy.add_argument('--forbid-overwrite', action='store_true', help='do not allow overwrites')
-    p_copy.add_argument(
-        '--merge',
-        action='store_true',
-        help='when copying a directory do not delete existing data, but add new files')
-    p_copy.add_argument('--run-cmd',
-                        help=textwrap.dedent('''\
-                run arbitrary cmd on top of copied files,
-                example usage is static content generation
-                based on current repository state\
-            '''))
 
     args = parser.parse_args()
 
@@ -160,7 +127,7 @@ def main():
     action = commands.get(args.subparser_name)
     if action:
         message = args.message or 'update'
-        GitRepo(args.repo, args.branch).update(message, lambda: action(args))
+        GitRepo(args.repo).update(message, lambda: action(args))
     else:
         parser.print_usage()
 

scripts/ingest_perf_test_result.py

@@ -1,136 +0,0 @@
-#!/usr/bin/env python3
-import argparse
-from contextlib import contextmanager
-import json
-import os
-import psycopg2
-import psycopg2.extras
-from pathlib import Path
-from datetime import datetime
-
-CREATE_TABLE = """
-CREATE TABLE IF NOT EXISTS perf_test_results (
-    id SERIAL PRIMARY KEY,
-    suit TEXT,
-    revision CHAR(40),
-    platform TEXT,
-    metric_name TEXT,
-    metric_value NUMERIC,
-    metric_unit VARCHAR(10),
-    metric_report_type TEXT,
-    recorded_at_timestamp TIMESTAMP WITH TIME ZONE DEFAULT NOW()
-)
-"""
-
-
-def err(msg):
-    print(f'error: {msg}')
-    exit(1)
-
-
-@contextmanager
-def get_connection_cursor():
-    connstr = os.getenv('DATABASE_URL')
-    if not connstr:
-        err('DATABASE_URL environment variable is not set')
-    with psycopg2.connect(connstr) as conn:
-        with conn.cursor() as cur:
-            yield cur
-
-
-def create_table(cur):
-    cur.execute(CREATE_TABLE)
-
-
-def ingest_perf_test_result(cursor, data_dile: Path, recorded_at_timestamp: int) -> int:
-    run_data = json.loads(data_dile.read_text())
-    revision = run_data['revision']
-    platform = run_data['platform']
-
-    run_result = run_data['result']
-    args_list = []
-
-    for suit_result in run_result:
-        suit = suit_result['suit']
-        total_duration = suit_result['total_duration']
-
-        suit_result['data'].append({
-            'name': 'total_duration',
-            'value': total_duration,
-            'unit': 's',
-            'report': 'lower_is_better',
-        })
-
-        for metric in suit_result['data']:
-            values = {
-                'suit': suit,
-                'revision': revision,
-                'platform': platform,
-                'metric_name': metric['name'],
-                'metric_value': metric['value'],
-                'metric_unit': metric['unit'],
-                'metric_report_type': metric['report'],
-                'recorded_at_timestamp': datetime.utcfromtimestamp(recorded_at_timestamp),
-            }
-            args_list.append(values)
-
-    psycopg2.extras.execute_values(
-        cursor,
-        """
-        INSERT INTO perf_test_results (
-            suit,
-            revision,
-            platform,
-            metric_name,
-            metric_value,
-            metric_unit,
-            metric_report_type,
-            recorded_at_timestamp
-        ) VALUES %s
-        """,
-        args_list,
-        template="""(
-            %(suit)s,
-            %(revision)s,
-            %(platform)s,
-            %(metric_name)s,
-            %(metric_value)s,
-            %(metric_unit)s,
-            %(metric_report_type)s,
-            %(recorded_at_timestamp)s
-        )""",
-    )
-    return len(args_list)
-
-
-def main():
-    parser = argparse.ArgumentParser(description='Perf test result uploader. \
-            Database connection string should be provided via DATABASE_URL environment variable', )
-    parser.add_argument(
-        '--ingest',
-        type=Path,
-        help='Path to perf test result file, or directory with perf test result files')
-    parser.add_argument('--initdb', action='store_true', help='Initialuze database')
-
-    args = parser.parse_args()
-    with get_connection_cursor() as cur:
-        if args.initdb:
-            create_table(cur)
-
-        if not args.ingest.exists():
-            err(f'ingest path {args.ingest} does not exist')
-
-        if args.ingest:
-            if args.ingest.is_dir():
-                for item in sorted(args.ingest.iterdir(), key=lambda x: int(x.name.split('_')[0])):
-                    recorded_at_timestamp = int(item.name.split('_')[0])
-                    ingested = ingest_perf_test_result(cur, item, recorded_at_timestamp)
-                    print(f'Ingested {ingested} metric values from {item}')
-            else:
-                recorded_at_timestamp = int(args.ingest.name.split('_')[0])
-                ingested = ingest_perf_test_result(cur, args.ingest, recorded_at_timestamp)
-                print(f'Ingested {ingested} metric values from {args.ingest}')
-
-
-if __name__ == '__main__':
-    main()

test_runner/batch_others/test_auth.py

@@ -1,10 +1,12 @@
 from contextlib import closing
 from typing import Iterator
-from uuid import UUID, uuid4
+from uuid import uuid4
 import psycopg2
-from fixtures.zenith_fixtures import ZenithEnvBuilder, ZenithPageserverApiException
+from fixtures.zenith_fixtures import ZenithEnvBuilder
 import pytest
 
+pytest_plugins = ("fixtures.zenith_fixtures")
+
 
 def test_pageserver_auth(zenith_env_builder: ZenithEnvBuilder):
     zenith_env_builder.pageserver_auth_enabled = True
@@ -12,38 +14,32 @@ def test_pageserver_auth(zenith_env_builder: ZenithEnvBuilder):
 
     ps = env.pageserver
 
-    tenant_token = env.auth_keys.generate_tenant_token(env.initial_tenant.hex)
-    tenant_http_client = env.pageserver.http_client(tenant_token)
+    tenant_token = env.auth_keys.generate_tenant_token(env.initial_tenant)
     invalid_tenant_token = env.auth_keys.generate_tenant_token(uuid4().hex)
-    invalid_tenant_http_client = env.pageserver.http_client(invalid_tenant_token)
-
     management_token = env.auth_keys.generate_management_token()
-    management_http_client = env.pageserver.http_client(management_token)
 
     # this does not invoke auth check and only decodes jwt and checks it for validity
     # check both tokens
-    ps.safe_psql("set FOO", password=tenant_token)
-    ps.safe_psql("set FOO", password=management_token)
+    ps.safe_psql("status", password=tenant_token)
+    ps.safe_psql("status", password=management_token)
 
     # tenant can create branches
-    tenant_http_client.branch_create(env.initial_tenant, 'new1', 'main')
+    ps.safe_psql(f"branch_create {env.initial_tenant} new1 main", password=tenant_token)
     # console can create branches for tenant
-    management_http_client.branch_create(env.initial_tenant, 'new2', 'main')
+    ps.safe_psql(f"branch_create {env.initial_tenant} new2 main", password=management_token)
 
-    # fail to create branch using token with different tenant_id
-    with pytest.raises(ZenithPageserverApiException,
-                       match='Forbidden: Tenant id mismatch. Permission denied'):
-        invalid_tenant_http_client.branch_create(env.initial_tenant, "new3", "main")
+    # fail to create branch using token with different tenantid
+    with pytest.raises(psycopg2.DatabaseError, match='Tenant id mismatch. Permission denied'):
+        ps.safe_psql(f"branch_create {env.initial_tenant} new2 main", password=invalid_tenant_token)
 
     # create tenant using management token
-    management_http_client.tenant_create(uuid4())
+    ps.safe_psql(f"tenant_create {uuid4().hex}", password=management_token)
 
     # fail to create tenant using tenant token
     with pytest.raises(
-            ZenithPageserverApiException,
-            match='Forbidden: Attempt to access management api with tenant scope. Permission denied'
-    ):
-        tenant_http_client.tenant_create(uuid4())
+            psycopg2.DatabaseError,
+            match='Attempt to access management api with tenant scope. Permission denied'):
+        ps.safe_psql(f"tenant_create {uuid4().hex}", password=tenant_token)
 
 
 @pytest.mark.parametrize('with_wal_acceptors', [False, True])
@@ -54,7 +50,7 @@ def test_compute_auth_to_pageserver(zenith_env_builder: ZenithEnvBuilder, with_w
     env = zenith_env_builder.init()
 
     branch = f"test_compute_auth_to_pageserver{with_wal_acceptors}"
-    env.zenith_cli.create_branch(branch, "main")
+    env.zenith_cli(["branch", branch, "main"])
 
     pg = env.postgres.create_start(branch)
 

test_runner/batch_others/test_backpressure.py

@@ -1,154 +0,0 @@
-from contextlib import closing, contextmanager
-import psycopg2.extras
-from fixtures.zenith_fixtures import ZenithEnvBuilder
-from fixtures.log_helper import log
-import os
-import time
-import asyncpg
-from fixtures.zenith_fixtures import Postgres
-import threading
-
-pytest_plugins = ("fixtures.zenith_fixtures")
-
-
-@contextmanager
-def pg_cur(pg):
-    with closing(pg.connect()) as conn:
-        with conn.cursor() as cur:
-            yield cur
-
-
-# Periodically check that all backpressure lags are below the configured threshold,
-# assert if they are not.
-# If the check query fails, stop the thread. Main thread should notice that and stop the test.
-def check_backpressure(pg: Postgres, stop_event: threading.Event, polling_interval=5):
-    log.info("checks started")
-
-    with pg_cur(pg) as cur:
-        cur.execute("CREATE EXTENSION zenith")  # TODO move it to zenith_fixtures?
-
-        cur.execute("select pg_size_bytes(current_setting('max_replication_write_lag'))")
-        res = cur.fetchone()
-        max_replication_write_lag_bytes = res[0]
-        log.info(f"max_replication_write_lag: {max_replication_write_lag_bytes} bytes")
-
-        cur.execute("select pg_size_bytes(current_setting('max_replication_flush_lag'))")
-        res = cur.fetchone()
-        max_replication_flush_lag_bytes = res[0]
-        log.info(f"max_replication_flush_lag: {max_replication_flush_lag_bytes} bytes")
-
-        cur.execute("select pg_size_bytes(current_setting('max_replication_apply_lag'))")
-        res = cur.fetchone()
-        max_replication_apply_lag_bytes = res[0]
-        log.info(f"max_replication_apply_lag: {max_replication_apply_lag_bytes} bytes")
-
-    with pg_cur(pg) as cur:
-        while not stop_event.is_set():
-            try:
-                cur.execute('''
-                select pg_wal_lsn_diff(pg_current_wal_flush_lsn(),received_lsn) as received_lsn_lag,
-                pg_wal_lsn_diff(pg_current_wal_flush_lsn(),disk_consistent_lsn) as disk_consistent_lsn_lag,
-                pg_wal_lsn_diff(pg_current_wal_flush_lsn(),remote_consistent_lsn) as remote_consistent_lsn_lag,
-                pg_size_pretty(pg_wal_lsn_diff(pg_current_wal_flush_lsn(),received_lsn)),
-                pg_size_pretty(pg_wal_lsn_diff(pg_current_wal_flush_lsn(),disk_consistent_lsn)),
-                pg_size_pretty(pg_wal_lsn_diff(pg_current_wal_flush_lsn(),remote_consistent_lsn))
-                from backpressure_lsns();
-                ''')
-
-                res = cur.fetchone()
-                received_lsn_lag = res[0]
-                disk_consistent_lsn_lag = res[1]
-                remote_consistent_lsn_lag = res[2]
-
-                log.info(f"received_lsn_lag = {received_lsn_lag} ({res[3]}), "
-                         f"disk_consistent_lsn_lag = {disk_consistent_lsn_lag} ({res[4]}), "
-                         f"remote_consistent_lsn_lag = {remote_consistent_lsn_lag} ({res[5]})")
-
-                # Since feedback from pageserver is not immediate, we should allow some lag overflow
-                lag_overflow = 5 * 1024 * 1024  # 5MB
-
-                if max_replication_write_lag_bytes > 0:
-                    assert received_lsn_lag < max_replication_write_lag_bytes + lag_overflow
-                if max_replication_flush_lag_bytes > 0:
-                    assert disk_consistent_lsn_lag < max_replication_flush_lag_bytes + lag_overflow
-                if max_replication_apply_lag_bytes > 0:
-                    assert remote_consistent_lsn_lag < max_replication_apply_lag_bytes + lag_overflow
-
-                time.sleep(polling_interval)
-
-            except Exception as e:
-                log.info(f"backpressure check query failed: {e}")
-                stop_event.set()
-
-    log.info('check thread stopped')
-
-
-# This test illustrates how to tune backpressure to control the lag
-# between the WAL flushed on compute node and WAL digested by pageserver.
-#
-# To test it, throttle walreceiver ingest using failpoint and run heavy write load.
-# If backpressure is disabled or not tuned properly, the query will timeout, because the walreceiver cannot keep up.
-# If backpressure is enabled and tuned properly, insertion will be throttled, but the query will not timeout.
-
-
-def test_backpressure_received_lsn_lag(zenith_env_builder: ZenithEnvBuilder):
-    zenith_env_builder.num_safekeepers = 1
-    env = zenith_env_builder.init()
-    # Create a branch for us
-    env.zenith_cli.create_branch("test_backpressure", "main")
-
-    pg = env.postgres.create_start('test_backpressure',
-                                   config_lines=['max_replication_write_lag=30MB'])
-    log.info("postgres is running on 'test_backpressure' branch")
-
-    # setup check thread
-    check_stop_event = threading.Event()
-    check_thread = threading.Thread(target=check_backpressure, args=(pg, check_stop_event))
-    check_thread.start()
-
-    # Configure failpoint to slow down walreceiver ingest
-    with closing(env.pageserver.connect()) as psconn:
-        with psconn.cursor(cursor_factory=psycopg2.extras.DictCursor) as pscur:
-            pscur.execute("failpoints walreceiver-after-ingest=sleep(20)")
-
-    # FIXME
-    # Wait for the check thread to start
-    #
-    # Now if load starts too soon,
-    # check thread cannot auth, because it is not able to connect to the database
-    # because of the lag and waiting for lsn to replay to arrive.
-    time.sleep(2)
-
-    with pg_cur(pg) as cur:
-        # Create and initialize test table
-        cur.execute("CREATE TABLE foo(x bigint)")
-
-        inserts_to_do = 2000000
-        rows_inserted = 0
-
-        while check_thread.is_alive() and rows_inserted < inserts_to_do:
-            try:
-                cur.execute("INSERT INTO foo select from generate_series(1, 100000)")
-                rows_inserted += 100000
-            except Exception as e:
-                if check_thread.is_alive():
-                    log.info('stopping check thread')
-                    check_stop_event.set()
-                    check_thread.join()
-                    assert False, f"Exception {e} while inserting rows, but WAL lag is within configured threshold. That means backpressure is not tuned properly"
-                else:
-                    assert False, f"Exception {e} while inserting rows and WAL lag overflowed configured threshold. That means backpressure doesn't work."
-
-        log.info(f"inserted {rows_inserted} rows")
-
-    if check_thread.is_alive():
-        log.info('stopping check thread')
-        check_stop_event.set()
-        check_thread.join()
-        log.info('check thread stopped')
-    else:
-        assert False, "WAL lag overflowed configured threshold. That means backpressure doesn't work."
-
-
-#TODO test_backpressure_disk_consistent_lsn_lag. Play with pageserver's checkpoint settings
-#TODO test_backpressure_remote_consistent_lsn_lag

test_runner/batch_others/test_branch_behind.py

@@ -7,6 +7,8 @@ from fixtures.log_helper import log
 from fixtures.utils import print_gc_result
 from fixtures.zenith_fixtures import ZenithEnvBuilder
 
+pytest_plugins = ("fixtures.zenith_fixtures")
+
 
 #
 # Create a couple of branches off the main branch, at a historical point in time.
@@ -22,7 +24,7 @@ def test_branch_behind(zenith_env_builder: ZenithEnvBuilder):
     env = zenith_env_builder.init()
 
     # Branch at the point where only 100 rows were inserted
-    env.zenith_cli.create_branch("test_branch_behind", "main")
+    env.zenith_cli(["branch", "test_branch_behind", "main"])
 
     pgmain = env.postgres.create_start('test_branch_behind')
     log.info("postgres is running on 'test_branch_behind' branch")
@@ -60,7 +62,7 @@ def test_branch_behind(zenith_env_builder: ZenithEnvBuilder):
     log.info(f'LSN after 200100 rows: {lsn_b}')
 
     # Branch at the point where only 100 rows were inserted
-    env.zenith_cli.create_branch("test_branch_behind_hundred", "test_branch_behind@" + lsn_a)
+    env.zenith_cli(["branch", "test_branch_behind_hundred", "test_branch_behind@" + lsn_a])
 
     # Insert many more rows. This generates enough WAL to fill a few segments.
     main_cur.execute('''
@@ -75,7 +77,7 @@ def test_branch_behind(zenith_env_builder: ZenithEnvBuilder):
     log.info(f'LSN after 400100 rows: {lsn_c}')
 
     # Branch at the point where only 200100 rows were inserted
-    env.zenith_cli.create_branch("test_branch_behind_more", "test_branch_behind@" + lsn_b)
+    env.zenith_cli(["branch", "test_branch_behind_more", "test_branch_behind@" + lsn_b])
 
     pg_hundred = env.postgres.create_start("test_branch_behind_hundred")
     pg_more = env.postgres.create_start("test_branch_behind_more")
@@ -99,7 +101,7 @@ def test_branch_behind(zenith_env_builder: ZenithEnvBuilder):
     # Check bad lsn's for branching
 
     # branch at segment boundary
-    env.zenith_cli.create_branch("test_branch_segment_boundary", "test_branch_behind@0/3000000")
+    env.zenith_cli(["branch", "test_branch_segment_boundary", "test_branch_behind@0/3000000"])
     pg = env.postgres.create_start("test_branch_segment_boundary")
     cur = pg.connect().cursor()
     cur.execute('SELECT 1')
@@ -107,23 +109,19 @@ def test_branch_behind(zenith_env_builder: ZenithEnvBuilder):
 
     # branch at pre-initdb lsn
     with pytest.raises(Exception, match="invalid branch start lsn"):
-        env.zenith_cli.create_branch("test_branch_preinitdb", "main@0/42")
-
-    # branch at pre-ancestor lsn
-    with pytest.raises(Exception, match="less than timeline ancestor lsn"):
-        env.zenith_cli.create_branch("test_branch_preinitdb", "test_branch_behind@0/42")
+        env.zenith_cli(["branch", "test_branch_preinitdb", "test_branch_behind@0/42"])
 
     # check that we cannot create branch based on garbage collected data
     with closing(env.pageserver.connect()) as psconn:
         with psconn.cursor(cursor_factory=psycopg2.extras.DictCursor) as pscur:
             # call gc to advace latest_gc_cutoff_lsn
-            pscur.execute(f"do_gc {env.initial_tenant.hex} {timeline} 0")
+            pscur.execute(f"do_gc {env.initial_tenant} {timeline} 0")
             row = pscur.fetchone()
             print_gc_result(row)
 
     with pytest.raises(Exception, match="invalid branch start lsn"):
         # this gced_lsn is pretty random, so if gc is disabled this woudln't fail
-        env.zenith_cli.create_branch("test_branch_create_fail", f"test_branch_behind@{gced_lsn}")
+        env.zenith_cli(["branch", "test_branch_create_fail", f"test_branch_behind@{gced_lsn}"])
 
     # check that after gc everything is still there
     hundred_cur.execute('SELECT count(*) FROM foo')

test_runner/batch_others/test_clog_truncate.py

@@ -6,13 +6,16 @@ from contextlib import closing
 from fixtures.zenith_fixtures import ZenithEnv
 from fixtures.log_helper import log
 
+pytest_plugins = ("fixtures.zenith_fixtures")
+
 
 #
 # Test compute node start after clog truncation
 #
 def test_clog_truncate(zenith_simple_env: ZenithEnv):
     env = zenith_simple_env
-    env.zenith_cli.create_branch("test_clog_truncate", "empty")
+    # Create a branch for us
+    env.zenith_cli(["branch", "test_clog_truncate", "empty"])
 
     # set agressive autovacuum to make sure that truncation will happen
     config = [
@@ -62,8 +65,8 @@ def test_clog_truncate(zenith_simple_env: ZenithEnv):
 
     # create new branch after clog truncation and start a compute node on it
     log.info(f'create branch at lsn_after_truncation {lsn_after_truncation}')
-    env.zenith_cli.create_branch("test_clog_truncate_new",
-                                 "test_clog_truncate@" + lsn_after_truncation)
+    env.zenith_cli(
+        ["branch", "test_clog_truncate_new", "test_clog_truncate@" + lsn_after_truncation])
 
     pg2 = env.postgres.create_start('test_clog_truncate_new')
     log.info('postgres is running on test_clog_truncate_new branch')

test_runner/batch_others/test_config.py

@@ -3,13 +3,16 @@ from contextlib import closing
 from fixtures.zenith_fixtures import ZenithEnv
 from fixtures.log_helper import log
 
+pytest_plugins = ("fixtures.zenith_fixtures")
+
 
 #
 # Test starting Postgres with custom options
 #
 def test_config(zenith_simple_env: ZenithEnv):
     env = zenith_simple_env
-    env.zenith_cli.create_branch("test_config", "empty")
+    # Create a branch for us
+    env.zenith_cli(["branch", "test_config", "empty"])
 
     # change config
     pg = env.postgres.create_start('test_config', config_lines=['log_min_messages=debug1'])

test_runner/batch_others/test_createdropdb.py

@@ -5,13 +5,15 @@ from contextlib import closing
 from fixtures.zenith_fixtures import ZenithEnv, check_restored_datadir_content
 from fixtures.log_helper import log
 
+pytest_plugins = ("fixtures.zenith_fixtures")
+
 
 #
 # Test CREATE DATABASE when there have been relmapper changes
 #
 def test_createdb(zenith_simple_env: ZenithEnv):
     env = zenith_simple_env
-    env.zenith_cli.create_branch("test_createdb", "empty")
+    env.zenith_cli(["branch", "test_createdb", "empty"])
 
     pg = env.postgres.create_start('test_createdb')
     log.info("postgres is running on 'test_createdb' branch")
@@ -27,7 +29,7 @@ def test_createdb(zenith_simple_env: ZenithEnv):
             lsn = cur.fetchone()[0]
 
     # Create a branch
-    env.zenith_cli.create_branch("test_createdb2", "test_createdb@" + lsn)
+    env.zenith_cli(["branch", "test_createdb2", "test_createdb@" + lsn])
 
     pg2 = env.postgres.create_start('test_createdb2')
 
@@ -41,7 +43,7 @@ def test_createdb(zenith_simple_env: ZenithEnv):
 #
 def test_dropdb(zenith_simple_env: ZenithEnv, test_output_dir):
     env = zenith_simple_env
-    env.zenith_cli.create_branch("test_dropdb", "empty")
+    env.zenith_cli(["branch", "test_dropdb", "empty"])
 
     pg = env.postgres.create_start('test_dropdb')
     log.info("postgres is running on 'test_dropdb' branch")
@@ -66,10 +68,10 @@ def test_dropdb(zenith_simple_env: ZenithEnv, test_output_dir):
             lsn_after_drop = cur.fetchone()[0]
 
     # Create two branches before and after database drop.
-    env.zenith_cli.create_branch("test_before_dropdb", "test_dropdb@" + lsn_before_drop)
+    env.zenith_cli(["branch", "test_before_dropdb", "test_dropdb@" + lsn_before_drop])
     pg_before = env.postgres.create_start('test_before_dropdb')
 
-    env.zenith_cli.create_branch("test_after_dropdb", "test_dropdb@" + lsn_after_drop)
+    env.zenith_cli(["branch", "test_after_dropdb", "test_dropdb@" + lsn_after_drop])
     pg_after = env.postgres.create_start('test_after_dropdb')
 
     # Test that database exists on the branch before drop

test_runner/batch_others/test_createuser.py

@@ -3,13 +3,15 @@ from contextlib import closing
 from fixtures.zenith_fixtures import ZenithEnv
 from fixtures.log_helper import log
 
+pytest_plugins = ("fixtures.zenith_fixtures")
+
 
 #
 # Test CREATE USER to check shared catalog restore
 #
 def test_createuser(zenith_simple_env: ZenithEnv):
     env = zenith_simple_env
-    env.zenith_cli.create_branch("test_createuser", "empty")
+    env.zenith_cli(["branch", "test_createuser", "empty"])
 
     pg = env.postgres.create_start('test_createuser')
     log.info("postgres is running on 'test_createuser' branch")
@@ -25,7 +27,7 @@ def test_createuser(zenith_simple_env: ZenithEnv):
             lsn = cur.fetchone()[0]
 
     # Create a branch
-    env.zenith_cli.create_branch("test_createuser2", "test_createuser@" + lsn)
+    env.zenith_cli(["branch", "test_createuser2", "test_createuser@" + lsn])
 
     pg2 = env.postgres.create_start('test_createuser2')
 

test_runner/batch_others/test_gc_aggressive.py

@@ -1,80 +0,0 @@
-from contextlib import closing
-
-import asyncio
-import asyncpg
-import random
-
-from fixtures.zenith_fixtures import ZenithEnv, Postgres, Safekeeper
-from fixtures.log_helper import log
-
-# Test configuration
-#
-# Create a table with {num_rows} rows, and perform {updates_to_perform} random
-# UPDATEs on it, using {num_connections} separate connections.
-num_connections = 10
-num_rows = 100000
-updates_to_perform = 10000
-
-updates_performed = 0
-
-
-# Run random UPDATEs on test table
-async def update_table(pg: Postgres):
-    global updates_performed
-    pg_conn = await pg.connect_async()
-
-    while updates_performed < updates_to_perform:
-        updates_performed += 1
-        id = random.randrange(1, num_rows)
-        row = await pg_conn.fetchrow(f'UPDATE foo SET counter = counter + 1 WHERE id = {id}')
-
-
-# Perform aggressive GC with 0 horizon
-async def gc(env: ZenithEnv, timeline: str):
-    psconn = await env.pageserver.connect_async()
-
-    while updates_performed < updates_to_perform:
-        await psconn.execute(f"do_gc {env.initial_tenant.hex} {timeline} 0")
-
-
-# At the same time, run UPDATEs and GC
-async def update_and_gc(env: ZenithEnv, pg: Postgres, timeline: str):
-    workers = []
-    for worker_id in range(num_connections):
-        workers.append(asyncio.create_task(update_table(pg)))
-    workers.append(asyncio.create_task(gc(env, timeline)))
-
-    # await all workers
-    await asyncio.gather(*workers)
-
-
-#
-# Aggressively force GC, while running queries.
-#
-# (repro for https://github.com/zenithdb/zenith/issues/1047)
-#
-def test_gc_aggressive(zenith_simple_env: ZenithEnv):
-    env = zenith_simple_env
-    env.zenith_cli.create_branch("test_gc_aggressive", "empty")
-    pg = env.postgres.create_start('test_gc_aggressive')
-    log.info('postgres is running on test_gc_aggressive branch')
-
-    conn = pg.connect()
-    cur = conn.cursor()
-
-    cur.execute("SHOW zenith.zenith_timeline")
-    timeline = cur.fetchone()[0]
-
-    # Create table, and insert the first 100 rows
-    cur.execute('CREATE TABLE foo (id int, counter int, t text)')
-    cur.execute(f'''
-        INSERT INTO foo
-            SELECT g, 0, 'long string to consume some space' || g
-            FROM generate_series(1, {num_rows}) g
-    ''')
-    cur.execute('CREATE INDEX ON foo(id)')
-
-    asyncio.run(update_and_gc(env, pg, timeline))
-
-    row = cur.execute('SELECT COUNT(*), SUM(counter) FROM foo')
-    assert cur.fetchone() == (num_rows, updates_to_perform)

test_runner/batch_others/test_multixact.py

@@ -1,6 +1,8 @@
 from fixtures.zenith_fixtures import ZenithEnv, check_restored_datadir_content
 from fixtures.log_helper import log
 
+pytest_plugins = ("fixtures.zenith_fixtures")
+
 
 #
 # Test multixact state after branching
@@ -10,7 +12,8 @@ from fixtures.log_helper import log
 #
 def test_multixact(zenith_simple_env: ZenithEnv, test_output_dir):
     env = zenith_simple_env
-    env.zenith_cli.create_branch("test_multixact", "empty")
+    # Create a branch for us
+    env.zenith_cli(["branch", "test_multixact", "empty"])
     pg = env.postgres.create_start('test_multixact')
 
     log.info("postgres is running on 'test_multixact' branch")
@@ -60,7 +63,7 @@ def test_multixact(zenith_simple_env: ZenithEnv, test_output_dir):
     assert int(next_multixact_id) > int(next_multixact_id_old)
 
     # Branch at this point
-    env.zenith_cli.create_branch("test_multixact_new", "test_multixact@" + lsn)
+    env.zenith_cli(["branch", "test_multixact_new", "test_multixact@" + lsn])
     pg_new = env.postgres.create_start('test_multixact_new')
 
     log.info("postgres is running on 'test_multixact_new' branch")

test_runner/batch_others/test_next_xid.py

@@ -1,59 +0,0 @@
-import pytest
-import random
-import time
-
-from fixtures.zenith_fixtures import ZenithEnvBuilder
-from fixtures.log_helper import log
-
-
-# Test restarting page server, while safekeeper and compute node keep
-# running.
-def test_next_xid(zenith_env_builder: ZenithEnvBuilder):
-    # One safekeeper is enough for this test.
-    zenith_env_builder.num_safekeepers = 1
-    env = zenith_env_builder.init()
-
-    pg = env.postgres.create_start('main')
-
-    conn = pg.connect()
-    cur = conn.cursor()
-    cur.execute('CREATE TABLE t(x integer)')
-
-    iterations = 32
-    for i in range(1, iterations + 1):
-        print(f'iteration {i} / {iterations}')
-
-        # Kill and restart the pageserver.
-        pg.stop()
-        env.pageserver.stop(immediate=True)
-        env.pageserver.start()
-        pg.start()
-
-        retry_sleep = 0.5
-        max_retries = 200
-        retries = 0
-        while True:
-            try:
-                conn = pg.connect()
-                cur = conn.cursor()
-                cur.execute(f"INSERT INTO t values({i})")
-                conn.close()
-
-            except Exception as error:
-                # It's normal that it takes some time for the pageserver to
-                # restart, and for the connection to fail until it does. It
-                # should eventually recover, so retry until it succeeds.
-                print(f'failed: {error}')
-                if retries < max_retries:
-                    retries += 1
-                    print(f'retry {retries} / {max_retries}')
-                    time.sleep(retry_sleep)
-                    continue
-                else:
-                    raise
-            break
-
-    conn = pg.connect()
-    cur = conn.cursor()
-    cur.execute("SELECT count(*) FROM t")
-    assert cur.fetchone() == (iterations, )

test_runner/batch_others/test_old_request_lsn.py

@@ -3,6 +3,8 @@ from contextlib import closing
 from fixtures.zenith_fixtures import ZenithEnv
 from fixtures.log_helper import log
 
+pytest_plugins = ("fixtures.zenith_fixtures")
+
 
 #
 # Test where Postgres generates a lot of WAL, and it's garbage collected away, but
@@ -16,7 +18,8 @@ from fixtures.log_helper import log
 #
 def test_old_request_lsn(zenith_simple_env: ZenithEnv):
     env = zenith_simple_env
-    env.zenith_cli.create_branch("test_old_request_lsn", "empty")
+    # Create a branch for us
+    env.zenith_cli(["branch", "test_old_request_lsn", "empty"])
     pg = env.postgres.create_start('test_old_request_lsn')
     log.info('postgres is running on test_old_request_lsn branch')
 
@@ -54,7 +57,7 @@ def test_old_request_lsn(zenith_simple_env: ZenithEnv):
     # Make a lot of updates on a single row, generating a lot of WAL. Trigger
     # garbage collections so that the page server will remove old page versions.
     for i in range(10):
-        pscur.execute(f"do_gc {env.initial_tenant.hex} {timeline} 0")
+        pscur.execute(f"do_gc {env.initial_tenant} {timeline} 0")
         for j in range(100):
             cur.execute('UPDATE foo SET val = val + 1 WHERE id = 1;')
 

test_runner/batch_others/test_pageserver_api.py

@@ -1,15 +1,95 @@
 import json
 from uuid import uuid4, UUID
+import pytest
+import psycopg2
+import requests
 from fixtures.zenith_fixtures import ZenithEnv, ZenithEnvBuilder, ZenithPageserverHttpClient
 from typing import cast
-import pytest, psycopg2
+
+pytest_plugins = ("fixtures.zenith_fixtures")
 
 
-def check_client(client: ZenithPageserverHttpClient, initial_tenant: UUID):
+def test_status_psql(zenith_simple_env: ZenithEnv):
+    env = zenith_simple_env
+    assert env.pageserver.safe_psql('status') == [
+        ('hello world', ),
+    ]
+
+
+def test_branch_list_psql(zenith_simple_env: ZenithEnv):
+    env = zenith_simple_env
+    # Create a branch for us
+    env.zenith_cli(["branch", "test_branch_list_main", "empty"])
+
+    conn = env.pageserver.connect()
+    cur = conn.cursor()
+
+    cur.execute(f'branch_list {env.initial_tenant}')
+    branches = json.loads(cur.fetchone()[0])
+    # Filter out branches created by other tests
+    branches = [x for x in branches if x['name'].startswith('test_branch_list')]
+
+    assert len(branches) == 1
+    assert branches[0]['name'] == 'test_branch_list_main'
+    assert 'timeline_id' in branches[0]
+    assert 'latest_valid_lsn' in branches[0]
+    assert 'ancestor_id' in branches[0]
+    assert 'ancestor_lsn' in branches[0]
+
+    # Create another branch, and start Postgres on it
+    env.zenith_cli(['branch', 'test_branch_list_experimental', 'test_branch_list_main'])
+    env.zenith_cli(['pg', 'create', 'test_branch_list_experimental'])
+
+    cur.execute(f'branch_list {env.initial_tenant}')
+    new_branches = json.loads(cur.fetchone()[0])
+    # Filter out branches created by other tests
+    new_branches = [x for x in new_branches if x['name'].startswith('test_branch_list')]
+    assert len(new_branches) == 2
+    new_branches.sort(key=lambda k: k['name'])
+
+    assert new_branches[0]['name'] == 'test_branch_list_experimental'
+    assert new_branches[0]['timeline_id'] != branches[0]['timeline_id']
+
+    # TODO: do the LSNs have to match here?
+    assert new_branches[1] == branches[0]
+
+    conn.close()
+
+
+def test_tenant_list_psql(zenith_env_builder: ZenithEnvBuilder):
+    # don't use zenith_simple_env, because there might be other tenants there,
+    # left over from other tests.
+    env = zenith_env_builder.init()
+
+    res = env.zenith_cli(["tenant", "list"])
+    res.check_returncode()
+    tenants = sorted(map(lambda t: t.split()[0], res.stdout.splitlines()))
+    assert tenants == [env.initial_tenant]
+
+    conn = env.pageserver.connect()
+    cur = conn.cursor()
+
+    # check same tenant cannot be created twice
+    with pytest.raises(psycopg2.DatabaseError,
+                       match=f'repo for {env.initial_tenant} already exists'):
+        cur.execute(f'tenant_create {env.initial_tenant}')
+
+    # create one more tenant
+    tenant1 = uuid4().hex
+    cur.execute(f'tenant_create {tenant1}')
+
+    cur.execute('tenant_list')
+
+    # compare tenants list
+    new_tenants = sorted(map(lambda t: cast(str, t['id']), json.loads(cur.fetchone()[0])))
+    assert sorted([env.initial_tenant, tenant1]) == new_tenants
+
+
+def check_client(client: ZenithPageserverHttpClient, initial_tenant: str):
     client.check_status()
 
     # check initial tenant is there
-    assert initial_tenant.hex in {t['id'] for t in client.tenant_list()}
+    assert initial_tenant in {t['id'] for t in client.tenant_list()}
 
     # create new tenant and check it is also there
     tenant_id = uuid4()

test_runner/batch_others/test_pageserver_catchup.py

@@ -7,6 +7,8 @@ from multiprocessing import Process, Value
 from fixtures.zenith_fixtures import ZenithEnvBuilder
 from fixtures.log_helper import log
 
+pytest_plugins = ("fixtures.zenith_fixtures")
+
 
 # Test safekeeper sync and pageserver catch up
 # while initial compute node is down and pageserver is lagging behind safekeepers.
@@ -16,7 +18,7 @@ def test_pageserver_catchup_while_compute_down(zenith_env_builder: ZenithEnvBuil
     zenith_env_builder.num_safekeepers = 3
     env = zenith_env_builder.init()
 
-    env.zenith_cli.create_branch("test_pageserver_catchup_while_compute_down", "main")
+    env.zenith_cli(["branch", "test_pageserver_catchup_while_compute_down", "main"])
     pg = env.postgres.create_start('test_pageserver_catchup_while_compute_down')
 
     pg_conn = pg.connect()

test_runner/batch_others/test_pageserver_restart.py

@@ -7,6 +7,8 @@ from multiprocessing import Process, Value
 from fixtures.zenith_fixtures import ZenithEnvBuilder
 from fixtures.log_helper import log
 
+pytest_plugins = ("fixtures.zenith_fixtures")
+
 
 # Test restarting page server, while safekeeper and compute node keep
 # running.
@@ -15,7 +17,7 @@ def test_pageserver_restart(zenith_env_builder: ZenithEnvBuilder):
     zenith_env_builder.num_safekeepers = 1
     env = zenith_env_builder.init()
 
-    env.zenith_cli.create_branch("test_pageserver_restart", "main")
+    env.zenith_cli(["branch", "test_pageserver_restart", "main"])
     pg = env.postgres.create_start('test_pageserver_restart')
 
     pg_conn = pg.connect()

test_runner/batch_others/test_parallel_copy.py

@@ -5,6 +5,8 @@ import subprocess
 from fixtures.zenith_fixtures import ZenithEnv, Postgres
 from fixtures.log_helper import log
 
+pytest_plugins = ("fixtures.zenith_fixtures")
+
 
 async def repeat_bytes(buf, repetitions: int):
     for i in range(repetitions):
@@ -37,7 +39,9 @@ async def parallel_load_same_table(pg: Postgres, n_parallel: int):
 # Load data into one table with COPY TO from 5 parallel connections
 def test_parallel_copy(zenith_simple_env: ZenithEnv, n_parallel=5):
     env = zenith_simple_env
-    env.zenith_cli.create_branch("test_parallel_copy", "empty")
+    # Create a branch for us
+    env.zenith_cli(["branch", "test_parallel_copy", "empty"])
+
     pg = env.postgres.create_start('test_parallel_copy')
     log.info("postgres is running on 'test_parallel_copy' branch")
 

test_runner/batch_others/test_pgbench.py

@@ -1,10 +1,14 @@
 from fixtures.zenith_fixtures import ZenithEnv
 from fixtures.log_helper import log
 
+pytest_plugins = ("fixtures.zenith_fixtures")
+
 
 def test_pgbench(zenith_simple_env: ZenithEnv, pg_bin):
     env = zenith_simple_env
-    env.zenith_cli.create_branch("test_pgbench", "empty")
+    # Create a branch for us
+    env.zenith_cli(["branch", "test_pgbench", "empty"])
+
     pg = env.postgres.create_start('test_pgbench')
     log.info("postgres is running on 'test_pgbench' branch")
 

test_runner/batch_others/test_proxy.py

@@ -1,2 +0,0 @@
-def test_proxy_select_1(static_proxy):
-    static_proxy.safe_psql("select 1;")

test_runner/batch_others/test_readonly_node.py

@@ -2,6 +2,8 @@ import pytest
 from fixtures.log_helper import log
 from fixtures.zenith_fixtures import ZenithEnv
 
+pytest_plugins = ("fixtures.zenith_fixtures")
+
 
 #
 # Create read-only compute nodes, anchored at historical points in time.
@@ -11,7 +13,7 @@ from fixtures.zenith_fixtures import ZenithEnv
 #
 def test_readonly_node(zenith_simple_env: ZenithEnv):
     env = zenith_simple_env
-    env.zenith_cli.create_branch("test_readonly_node", "empty")
+    env.zenith_cli(["branch", "test_readonly_node", "empty"])
 
     pgmain = env.postgres.create_start('test_readonly_node')
     log.info("postgres is running on 'test_readonly_node' branch")
@@ -86,5 +88,4 @@ def test_readonly_node(zenith_simple_env: ZenithEnv):
     # Create node at pre-initdb lsn
     with pytest.raises(Exception, match="invalid basebackup lsn"):
         # compute node startup with invalid LSN should fail
-        env.zenith_cli.pg_start("test_readonly_node_preinitdb",
-                                timeline_spec="test_readonly_node@0/42")
+        env.zenith_cli(["pg", "start", "test_readonly_node_preinitdb", "test_readonly_node@0/42"])

test_runner/batch_others/test_remote_storage.py

@@ -9,6 +9,8 @@ from fixtures.zenith_fixtures import ZenithEnvBuilder
 from fixtures.log_helper import log
 import pytest
 
+pytest_plugins = ("fixtures.zenith_fixtures")
+
 
 #
 # Tests that a piece of data is backed up and restored correctly:
@@ -26,7 +28,6 @@ import pytest
 #   * queries the specific data, ensuring that it matches the one stored before
 #
 # The tests are done for all types of remote storage pageserver supports.
-@pytest.mark.skip(reason="will be fixed with https://github.com/zenithdb/zenith/issues/1193")
 @pytest.mark.parametrize('storage_type', ['local_fs', 'mock_s3'])
 def test_remote_storage_backup_and_restore(zenith_env_builder: ZenithEnvBuilder, storage_type: str):
     zenith_env_builder.rust_log_override = 'debug'
@@ -73,13 +74,8 @@ def test_remote_storage_backup_and_restore(zenith_env_builder: ZenithEnvBuilder,
     ##### Second start, restore the data and ensure it's the same
     env.pageserver.start()
 
-    client = env.pageserver.http_client()
-    client.timeline_attach(UUID(tenant_id), UUID(timeline_id))
-    # FIXME cannot handle duplicate download requests (which might be caused by repeated timeline detail calls)
-    #   subject to fix in https://github.com/zenithdb/zenith/issues/997
-    time.sleep(5)
-
     log.info("waiting for timeline redownload")
+    client = env.pageserver.http_client()
     attempts = 0
     while True:
         timeline_details = client.timeline_detail(UUID(tenant_id), UUID(timeline_id))

test_runner/batch_others/test_restart_compute.py

@@ -4,6 +4,8 @@ from contextlib import closing
 from fixtures.zenith_fixtures import ZenithEnvBuilder
 from fixtures.log_helper import log
 
+pytest_plugins = ("fixtures.zenith_fixtures")
+
 
 #
 # Test restarting and recreating a postgres instance
@@ -15,7 +17,7 @@ def test_restart_compute(zenith_env_builder: ZenithEnvBuilder, with_wal_acceptor
         zenith_env_builder.num_safekeepers = 3
     env = zenith_env_builder.init()
 
-    env.zenith_cli.create_branch("test_restart_compute", "main")
+    env.zenith_cli(["branch", "test_restart_compute", "main"])
 
     pg = env.postgres.create_start('test_restart_compute')
     log.info("postgres is running on 'test_restart_compute' branch")

test_runner/batch_others/test_snapfiles_gc.py

@@ -5,6 +5,8 @@ from fixtures.utils import print_gc_result
 from fixtures.zenith_fixtures import ZenithEnv
 from fixtures.log_helper import log
 
+pytest_plugins = ("fixtures.zenith_fixtures")
+
 
 #
 # Test Garbage Collection of old layer files
@@ -14,7 +16,7 @@ from fixtures.log_helper import log
 #
 def test_layerfiles_gc(zenith_simple_env: ZenithEnv):
     env = zenith_simple_env
-    env.zenith_cli.create_branch("test_layerfiles_gc", "empty")
+    env.zenith_cli(["branch", "test_layerfiles_gc", "empty"])
     pg = env.postgres.create_start('test_layerfiles_gc')
 
     with closing(pg.connect()) as conn:
@@ -48,7 +50,7 @@ def test_layerfiles_gc(zenith_simple_env: ZenithEnv):
                     cur.execute("DELETE FROM foo")
 
                     log.info("Running GC before test")
-                    pscur.execute(f"do_gc {env.initial_tenant.hex} {timeline} 0")
+                    pscur.execute(f"do_gc {env.initial_tenant} {timeline} 0")
                     row = pscur.fetchone()
                     print_gc_result(row)
                     # remember the number of files
@@ -61,7 +63,7 @@ def test_layerfiles_gc(zenith_simple_env: ZenithEnv):
                     # removing the old image and delta layer.
                     log.info("Inserting one row and running GC")
                     cur.execute("INSERT INTO foo VALUES (1)")
-                    pscur.execute(f"do_gc {env.initial_tenant.hex} {timeline} 0")
+                    pscur.execute(f"do_gc {env.initial_tenant} {timeline} 0")
                     row = pscur.fetchone()
                     print_gc_result(row)
                     assert row['layer_relfiles_total'] == layer_relfiles_remain + 2
@@ -75,7 +77,7 @@ def test_layerfiles_gc(zenith_simple_env: ZenithEnv):
                     cur.execute("INSERT INTO foo VALUES (2)")
                     cur.execute("INSERT INTO foo VALUES (3)")
 
-                    pscur.execute(f"do_gc {env.initial_tenant.hex} {timeline} 0")
+                    pscur.execute(f"do_gc {env.initial_tenant} {timeline} 0")
                     row = pscur.fetchone()
                     print_gc_result(row)
                     assert row['layer_relfiles_total'] == layer_relfiles_remain + 2
@@ -87,7 +89,7 @@ def test_layerfiles_gc(zenith_simple_env: ZenithEnv):
                     cur.execute("INSERT INTO foo VALUES (2)")
                     cur.execute("INSERT INTO foo VALUES (3)")
 
-                    pscur.execute(f"do_gc {env.initial_tenant.hex} {timeline} 0")
+                    pscur.execute(f"do_gc {env.initial_tenant} {timeline} 0")
                     row = pscur.fetchone()
                     print_gc_result(row)
                     assert row['layer_relfiles_total'] == layer_relfiles_remain + 2
@@ -96,7 +98,7 @@ def test_layerfiles_gc(zenith_simple_env: ZenithEnv):
 
                     # Run GC again, with no changes in the database. Should not remove anything.
                     log.info("Run GC again, with nothing to do")
-                    pscur.execute(f"do_gc {env.initial_tenant.hex} {timeline} 0")
+                    pscur.execute(f"do_gc {env.initial_tenant} {timeline} 0")
                     row = pscur.fetchone()
                     print_gc_result(row)
                     assert row['layer_relfiles_total'] == layer_relfiles_remain
@@ -109,7 +111,7 @@ def test_layerfiles_gc(zenith_simple_env: ZenithEnv):
                     log.info("Drop table and run GC again")
                     cur.execute("DROP TABLE foo")
 
-                    pscur.execute(f"do_gc {env.initial_tenant.hex} {timeline} 0")
+                    pscur.execute(f"do_gc {env.initial_tenant} {timeline} 0")
                     row = pscur.fetchone()
                     print_gc_result(row)
 

test_runner/batch_others/test_subxacts.py

@@ -1,6 +1,8 @@
 from fixtures.zenith_fixtures import ZenithEnv, check_restored_datadir_content
 from fixtures.log_helper import log
 
+pytest_plugins = ("fixtures.zenith_fixtures")
+
 
 # Test subtransactions
 #
@@ -10,7 +12,8 @@ from fixtures.log_helper import log
 # CLOG.
 def test_subxacts(zenith_simple_env: ZenithEnv, test_output_dir):
     env = zenith_simple_env
-    env.zenith_cli.create_branch("test_subxacts", "empty")
+    # Create a branch for us
+    env.zenith_cli(["branch", "test_subxacts", "empty"])
     pg = env.postgres.create_start('test_subxacts')
 
     log.info("postgres is running on 'test_subxacts' branch")

test_runner/batch_others/test_tenant_relocation.py

@@ -108,13 +108,12 @@ def load(pg: Postgres, stop_event: threading.Event, load_ok_event: threading.Eve
     log.info('load thread stopped')
 
 
-def assert_local(pageserver_http_client: ZenithPageserverHttpClient, tenant: UUID, timeline: str):
-    timeline_detail = pageserver_http_client.timeline_detail(tenant, UUID(timeline))
+def assert_local(pageserver_http_client: ZenithPageserverHttpClient, tenant: str, timeline: str):
+    timeline_detail = pageserver_http_client.timeline_detail(UUID(tenant), UUID(timeline))
     assert timeline_detail.get('type') == "Local", timeline_detail
     return timeline_detail
 
 
-@pytest.mark.skip(reason="will be fixed with https://github.com/zenithdb/zenith/issues/1193")
 @pytest.mark.parametrize('with_load', ['with_load', 'without_load'])
 def test_tenant_relocation(zenith_env_builder: ZenithEnvBuilder,
                            port_distributor: PortDistributor,
@@ -127,10 +126,10 @@ def test_tenant_relocation(zenith_env_builder: ZenithEnvBuilder,
     # create folder for remote storage mock
     remote_storage_mock_path = env.repo_dir / 'local_fs_remote_storage'
 
-    tenant = env.create_tenant(UUID("74ee8b079a0e437eb0afea7d26a07209"))
+    tenant = env.create_tenant("74ee8b079a0e437eb0afea7d26a07209")
     log.info("tenant to relocate %s", tenant)
 
-    env.zenith_cli.create_branch("test_tenant_relocation", "main", tenant_id=tenant)
+    env.zenith_cli(["branch", "test_tenant_relocation", "main", f"--tenantid={tenant}"])
 
     tenant_pg = env.postgres.create_start(
         "test_tenant_relocation",
@@ -167,11 +166,11 @@ def test_tenant_relocation(zenith_env_builder: ZenithEnvBuilder,
     # run checkpoint manually to be sure that data landed in remote storage
     with closing(env.pageserver.connect()) as psconn:
         with psconn.cursor() as pscur:
-            pscur.execute(f"do_gc {tenant.hex} {timeline}")
+            pscur.execute(f"do_gc {tenant} {timeline}")
 
     # ensure upload is completed
     pageserver_http_client = env.pageserver.http_client()
-    timeline_detail = pageserver_http_client.timeline_detail(tenant, UUID(timeline))
+    timeline_detail = pageserver_http_client.timeline_detail(UUID(tenant), UUID(timeline))
     assert timeline_detail['disk_consistent_lsn'] == timeline_detail['timeline_state']['Ready']
 
     log.info("inititalizing new pageserver")
@@ -193,8 +192,8 @@ def test_tenant_relocation(zenith_env_builder: ZenithEnvBuilder,
                                new_pageserver_pg_port,
                                new_pageserver_http_port):
 
-        # call to attach timeline to new pageserver
-        new_pageserver_http_client.timeline_attach(tenant, UUID(timeline))
+        # call to attach timeline to new timeline
+        new_pageserver_http_client.timeline_attach(UUID(tenant), UUID(timeline))
         # FIXME cannot handle duplicate download requests, subject to fix in https://github.com/zenithdb/zenith/issues/997
         time.sleep(5)
         # new pageserver should in sync (modulo wal tail or vacuum activity) with the old one because there was no new writes since checkpoint
@@ -241,7 +240,7 @@ def test_tenant_relocation(zenith_env_builder: ZenithEnvBuilder,
         # detach tenant from old pageserver before we check
         # that all the data is there to be sure that old pageserver
         # is no longer involved, and if it is, we will see the errors
-        pageserver_http_client.timeline_detach(tenant, UUID(timeline))
+        pageserver_http_client.timeline_detach(UUID(tenant), UUID(timeline))
 
         with pg_cur(tenant_pg) as cur:
             # check that data is still there

test_runner/batch_others/test_tenants.py

@@ -15,12 +15,18 @@ def test_tenants_normal_work(zenith_env_builder: ZenithEnvBuilder, with_wal_acce
     tenant_1 = env.create_tenant()
     tenant_2 = env.create_tenant()
 
-    env.zenith_cli.create_branch(f"test_tenants_normal_work_with_wal_acceptors{with_wal_acceptors}",
-                                 "main",
-                                 tenant_id=tenant_1)
-    env.zenith_cli.create_branch(f"test_tenants_normal_work_with_wal_acceptors{with_wal_acceptors}",
-                                 "main",
-                                 tenant_id=tenant_2)
+    env.zenith_cli([
+        "branch",
+        f"test_tenants_normal_work_with_wal_acceptors{with_wal_acceptors}",
+        "main",
+        f"--tenantid={tenant_1}"
+    ])
+    env.zenith_cli([
+        "branch",
+        f"test_tenants_normal_work_with_wal_acceptors{with_wal_acceptors}",
+        "main",
+        f"--tenantid={tenant_2}"
+    ])
 
     pg_tenant1 = env.postgres.create_start(
         f"test_tenants_normal_work_with_wal_acceptors{with_wal_acceptors}",

test_runner/batch_others/test_timeline_size.py

@@ -1,19 +1,17 @@
 from contextlib import closing
 from uuid import UUID
 import psycopg2.extras
-import psycopg2.errors
-from fixtures.zenith_fixtures import ZenithEnv, ZenithEnvBuilder, Postgres
+from fixtures.zenith_fixtures import ZenithEnv
 from fixtures.log_helper import log
-import time
 
 
 def test_timeline_size(zenith_simple_env: ZenithEnv):
     env = zenith_simple_env
     # Branch at the point where only 100 rows were inserted
-    env.zenith_cli.create_branch("test_timeline_size", "empty")
+    env.zenith_cli(["branch", "test_timeline_size", "empty"])
 
     client = env.pageserver.http_client()
-    res = client.branch_detail(env.initial_tenant, "test_timeline_size")
+    res = client.branch_detail(UUID(env.initial_tenant), "test_timeline_size")
     assert res["current_logical_size"] == res["current_logical_size_non_incremental"]
 
     pgmain = env.postgres.create_start("test_timeline_size")
@@ -31,102 +29,9 @@ def test_timeline_size(zenith_simple_env: ZenithEnv):
                     FROM generate_series(1, 10) g
             """)
 
-            res = client.branch_detail(env.initial_tenant, "test_timeline_size")
+            res = client.branch_detail(UUID(env.initial_tenant), "test_timeline_size")
             assert res["current_logical_size"] == res["current_logical_size_non_incremental"]
             cur.execute("TRUNCATE foo")
 
-            res = client.branch_detail(env.initial_tenant, "test_timeline_size")
+            res = client.branch_detail(UUID(env.initial_tenant), "test_timeline_size")
             assert res["current_logical_size"] == res["current_logical_size_non_incremental"]
-
-
-# wait until received_lsn_lag is 0
-def wait_for_pageserver_catchup(pgmain: Postgres, polling_interval=1, timeout=60):
-    started_at = time.time()
-
-    received_lsn_lag = 1
-    while received_lsn_lag > 0:
-        elapsed = time.time() - started_at
-        if elapsed > timeout:
-            raise RuntimeError(
-                f"timed out waiting for pageserver to reach pg_current_wal_flush_lsn()")
-
-        with closing(pgmain.connect()) as conn:
-            with conn.cursor() as cur:
-
-                cur.execute('''
-                    select  pg_size_pretty(pg_cluster_size()),
-                    pg_wal_lsn_diff(pg_current_wal_flush_lsn(),received_lsn) as received_lsn_lag
-                    FROM backpressure_lsns();
-                ''')
-                res = cur.fetchone()
-                log.info(f"pg_cluster_size = {res[0]}, received_lsn_lag = {res[1]}")
-                received_lsn_lag = res[1]
-
-        time.sleep(polling_interval)
-
-
-def test_timeline_size_quota(zenith_env_builder: ZenithEnvBuilder):
-    zenith_env_builder.num_safekeepers = 1
-    env = zenith_env_builder.init()
-    env.zenith_cli.create_branch("test_timeline_size_quota", "main")
-
-    client = env.pageserver.http_client()
-    res = client.branch_detail(env.initial_tenant, "test_timeline_size_quota")
-    assert res["current_logical_size"] == res["current_logical_size_non_incremental"]
-
-    pgmain = env.postgres.create_start(
-        "test_timeline_size_quota",
-        # Set small limit for the test
-        config_lines=['zenith.max_cluster_size=30MB'],
-    )
-    log.info("postgres is running on 'test_timeline_size_quota' branch")
-
-    with closing(pgmain.connect()) as conn:
-        with conn.cursor() as cur:
-            cur.execute("CREATE EXTENSION zenith")  # TODO move it to zenith_fixtures?
-
-            cur.execute("CREATE TABLE foo (t text)")
-
-            wait_for_pageserver_catchup(pgmain)
-
-            # Insert many rows. This query must fail because of space limit
-            try:
-                cur.execute('''
-                    INSERT INTO foo
-                        SELECT 'long string to consume some space' || g
-                        FROM generate_series(1, 100000) g
-                ''')
-
-                wait_for_pageserver_catchup(pgmain)
-
-                cur.execute('''
-                    INSERT INTO foo
-                        SELECT 'long string to consume some space' || g
-                        FROM generate_series(1, 500000) g
-                ''')
-
-                # If we get here, the timeline size limit failed
-                log.error("Query unexpectedly succeeded")
-                assert False
-
-            except psycopg2.errors.DiskFull as err:
-                log.info(f"Query expectedly failed with: {err}")
-
-            # drop table to free space
-            cur.execute('DROP TABLE foo')
-
-            wait_for_pageserver_catchup(pgmain)
-
-            # create it again and insert some rows. This query must succeed
-            cur.execute("CREATE TABLE foo (t text)")
-            cur.execute('''
-                INSERT INTO foo
-                    SELECT 'long string to consume some space' || g
-                    FROM generate_series(1, 10000) g
-            ''')
-
-            wait_for_pageserver_catchup(pgmain)
-
-            cur.execute("SELECT * from pg_size_pretty(pg_cluster_size())")
-            pg_cluster_size = cur.fetchone()
-            log.info(f"pg_cluster_size = {pg_cluster_size}")

test_runner/batch_others/test_twophase.py

@@ -3,13 +3,15 @@ import os
 from fixtures.zenith_fixtures import ZenithEnv
 from fixtures.log_helper import log
 
+pytest_plugins = ("fixtures.zenith_fixtures")
+
 
 #
 # Test branching, when a transaction is in prepared state
 #
 def test_twophase(zenith_simple_env: ZenithEnv):
     env = zenith_simple_env
-    env.zenith_cli.create_branch("test_twophase", "empty")
+    env.zenith_cli(["branch", "test_twophase", "empty"])
 
     pg = env.postgres.create_start('test_twophase', config_lines=['max_prepared_transactions=5'])
     log.info("postgres is running on 'test_twophase' branch")
@@ -56,7 +58,7 @@ def test_twophase(zenith_simple_env: ZenithEnv):
     assert len(twophase_files) == 2
 
     # Create a branch with the transaction in prepared state
-    env.zenith_cli.create_branch("test_twophase_prepared", "test_twophase")
+    env.zenith_cli(["branch", "test_twophase_prepared", "test_twophase"])
 
     # Start compute on the new branch
     pg2 = env.postgres.create_start(
@@ -78,8 +80,8 @@ def test_twophase(zenith_simple_env: ZenithEnv):
     cur2.execute("ROLLBACK PREPARED 'insert_two'")
 
     cur2.execute('SELECT * FROM foo')
-    assert cur2.fetchall() == [('one', ), ('three', )]
+    assert cur2.fetchall() == [('one', ), ('three', )]  # type: ignore[comparison-overlap]
 
     # Only one committed insert is visible on the original branch
     cur.execute('SELECT * FROM foo')
-    assert cur.fetchall() == [('three', )]
+    assert cur.fetchall() == [('three', )]  # type: ignore[comparison-overlap]

test_runner/batch_others/test_vm_bits.py

@@ -1,6 +1,8 @@
 from fixtures.zenith_fixtures import ZenithEnv
 from fixtures.log_helper import log
 
+pytest_plugins = ("fixtures.zenith_fixtures")
+
 
 #
 # Test that the VM bit is cleared correctly at a HEAP_DELETE and
@@ -9,7 +11,8 @@ from fixtures.log_helper import log
 def test_vm_bit_clear(zenith_simple_env: ZenithEnv):
     env = zenith_simple_env
 
-    env.zenith_cli.create_branch("test_vm_bit_clear", "empty")
+    # Create a branch for us
+    env.zenith_cli(["branch", "test_vm_bit_clear", "empty"])
     pg = env.postgres.create_start('test_vm_bit_clear')
 
     log.info("postgres is running on 'test_vm_bit_clear' branch")
@@ -33,7 +36,7 @@ def test_vm_bit_clear(zenith_simple_env: ZenithEnv):
     cur.execute('UPDATE vmtest_update SET id = 5000 WHERE id = 1')
 
     # Branch at this point, to test that later
-    env.zenith_cli.create_branch("test_vm_bit_clear_new", "test_vm_bit_clear")
+    env.zenith_cli(["branch", "test_vm_bit_clear_new", "test_vm_bit_clear"])
 
     # Clear the buffer cache, to force the VM page to be re-fetched from
     # the page server

test_runner/batch_others/test_wal_acceptor.py

@@ -12,11 +12,13 @@ from contextlib import closing
 from dataclasses import dataclass, field
 from multiprocessing import Process, Value
 from pathlib import Path
-from fixtures.zenith_fixtures import PgBin, Postgres, Safekeeper, ZenithEnv, ZenithEnvBuilder, PortDistributor, SafekeeperPort, zenith_binpath, PgProtocol
+from fixtures.zenith_fixtures import PgBin, ZenithEnvBuilder, PortDistributor, SafekeeperPort, zenith_binpath, PgProtocol
 from fixtures.utils import lsn_to_hex, mkdir_if_needed
 from fixtures.log_helper import log
 from typing import List, Optional, Any
 
+pytest_plugins = ("fixtures.zenith_fixtures")
+
 
 # basic test, write something in setup with wal acceptors, ensure that commits
 # succeed and data is written
@@ -24,7 +26,7 @@ def test_normal_work(zenith_env_builder: ZenithEnvBuilder):
     zenith_env_builder.num_safekeepers = 3
     env = zenith_env_builder.init()
 
-    env.zenith_cli.create_branch("test_wal_acceptors_normal_work", "main")
+    env.zenith_cli(["branch", "test_wal_acceptors_normal_work", "main"])
 
     pg = env.postgres.create_start('test_wal_acceptors_normal_work')
 
@@ -60,10 +62,10 @@ def test_many_timelines(zenith_env_builder: ZenithEnvBuilder):
     # start postgres on each timeline
     pgs = []
     for branch in branches:
-        env.zenith_cli.create_branch(branch, "main")
+        env.zenith_cli(["branch", branch, "main"])
         pgs.append(env.postgres.create_start(branch))
 
-    tenant_id = env.initial_tenant
+    tenant_id = uuid.UUID(env.initial_tenant)
 
     def collect_metrics(message: str) -> List[BranchMetrics]:
         with env.pageserver.http_client() as pageserver_http:
@@ -90,8 +92,8 @@ def test_many_timelines(zenith_env_builder: ZenithEnvBuilder):
                     latest_valid_lsn=branch_detail["latest_valid_lsn"],
                 )
                 for sk_m in sk_metrics:
-                    m.flush_lsns.append(sk_m.flush_lsn_inexact[(tenant_id.hex, timeline_id)])
-                    m.commit_lsns.append(sk_m.commit_lsn_inexact[(tenant_id.hex, timeline_id)])
+                    m.flush_lsns.append(sk_m.flush_lsn_inexact[timeline_id])
+                    m.commit_lsns.append(sk_m.commit_lsn_inexact[timeline_id])
 
                 for flush_lsn, commit_lsn in zip(m.flush_lsns, m.commit_lsns):
                     # Invariant. May be < when transaction is in progress.
@@ -183,7 +185,7 @@ def test_restarts(zenith_env_builder: ZenithEnvBuilder):
     zenith_env_builder.num_safekeepers = n_acceptors
     env = zenith_env_builder.init()
 
-    env.zenith_cli.create_branch("test_wal_acceptors_restarts", "main")
+    env.zenith_cli(["branch", "test_wal_acceptors_restarts", "main"])
     pg = env.postgres.create_start('test_wal_acceptors_restarts')
 
     # we rely upon autocommit after each statement
@@ -220,7 +222,7 @@ def test_unavailability(zenith_env_builder: ZenithEnvBuilder):
     zenith_env_builder.num_safekeepers = 2
     env = zenith_env_builder.init()
 
-    env.zenith_cli.create_branch("test_wal_acceptors_unavailability", "main")
+    env.zenith_cli(["branch", "test_wal_acceptors_unavailability", "main"])
     pg = env.postgres.create_start('test_wal_acceptors_unavailability')
 
     # we rely upon autocommit after each statement
@@ -291,7 +293,7 @@ def test_race_conditions(zenith_env_builder: ZenithEnvBuilder, stop_value):
     zenith_env_builder.num_safekeepers = 3
     env = zenith_env_builder.init()
 
-    env.zenith_cli.create_branch("test_wal_acceptors_race_conditions", "main")
+    env.zenith_cli(["branch", "test_wal_acceptors_race_conditions", "main"])
     pg = env.postgres.create_start('test_wal_acceptors_race_conditions')
 
     # we rely upon autocommit after each statement
@@ -319,16 +321,16 @@ class ProposerPostgres(PgProtocol):
     def __init__(self,
                  pgdata_dir: str,
                  pg_bin,
-                 timeline_id: uuid.UUID,
-                 tenant_id: uuid.UUID,
+                 timeline_id: str,
+                 tenant_id: str,
                  listen_addr: str,
                  port: int):
-        super().__init__(host=listen_addr, port=port, username='zenith_admin')
+        super().__init__(host=listen_addr, port=port)
 
         self.pgdata_dir: str = pgdata_dir
         self.pg_bin: PgBin = pg_bin
-        self.timeline_id: uuid.UUID = timeline_id
-        self.tenant_id: uuid.UUID = tenant_id
+        self.timeline_id: str = timeline_id
+        self.tenant_id: str = tenant_id
         self.listen_addr: str = listen_addr
         self.port: int = port
 
@@ -348,8 +350,8 @@ class ProposerPostgres(PgProtocol):
             cfg = [
                 "synchronous_standby_names = 'walproposer'\n",
                 "shared_preload_libraries = 'zenith'\n",
-                f"zenith.zenith_timeline = '{self.timeline_id.hex}'\n",
-                f"zenith.zenith_tenant = '{self.tenant_id.hex}'\n",
+                f"zenith.zenith_timeline = '{self.timeline_id}'\n",
+                f"zenith.zenith_tenant = '{self.tenant_id}'\n",
                 f"zenith.page_server_connstring = ''\n",
                 f"wal_acceptors = '{wal_acceptors}'\n",
                 f"listen_addresses = '{self.listen_addr}'\n",
@@ -406,8 +408,8 @@ def test_sync_safekeepers(zenith_env_builder: ZenithEnvBuilder,
     zenith_env_builder.num_safekeepers = 3
     env = zenith_env_builder.init()
 
-    timeline_id = uuid.uuid4()
-    tenant_id = uuid.uuid4()
+    timeline_id = uuid.uuid4().hex
+    tenant_id = uuid.uuid4().hex
 
     # write config for proposer
     pgdata_dir = os.path.join(env.repo_dir, "proposer_pgdata")
@@ -456,7 +458,7 @@ def test_timeline_status(zenith_env_builder: ZenithEnvBuilder):
     zenith_env_builder.num_safekeepers = 1
     env = zenith_env_builder.init()
 
-    env.zenith_cli.create_branch("test_timeline_status", "main")
+    env.zenith_cli(["branch", "test_timeline_status", "main"])
     pg = env.postgres.create_start('test_timeline_status')
 
     wa = env.safekeepers[0]
@@ -493,15 +495,15 @@ class SafekeeperEnv:
         self.bin_safekeeper = os.path.join(str(zenith_binpath), 'safekeeper')
         self.safekeepers: Optional[List[subprocess.CompletedProcess[Any]]] = None
         self.postgres: Optional[ProposerPostgres] = None
-        self.tenant_id: Optional[uuid.UUID] = None
-        self.timeline_id: Optional[uuid.UUID] = None
+        self.tenant_id: Optional[str] = None
+        self.timeline_id: Optional[str] = None
 
     def init(self) -> "SafekeeperEnv":
         assert self.postgres is None, "postgres is already initialized"
         assert self.safekeepers is None, "safekeepers are already initialized"
 
-        self.timeline_id = uuid.uuid4()
-        self.tenant_id = uuid.uuid4()
+        self.timeline_id = uuid.uuid4().hex
+        self.tenant_id = uuid.uuid4().hex
         mkdir_if_needed(str(self.repo_dir))
 
         # Create config and a Safekeeper object for each safekeeper
@@ -601,92 +603,3 @@ def test_safekeeper_without_pageserver(test_output_dir: str,
         env.postgres.safe_psql("insert into t select generate_series(1, 100)")
         res = env.postgres.safe_psql("select sum(i) from t")[0][0]
         assert res == 5050
-
-
-def test_replace_safekeeper(zenith_env_builder: ZenithEnvBuilder):
-    def safekeepers_guc(env: ZenithEnv, sk_names: List[str]) -> str:
-        return ','.join(
-            [f'localhost:{sk.port.pg}' for sk in env.safekeepers if sk.name in sk_names])
-
-    def execute_payload(pg: Postgres):
-        with closing(pg.connect()) as conn:
-            with conn.cursor() as cur:
-                # we rely upon autocommit after each statement
-                # as waiting for acceptors happens there
-                cur.execute('CREATE TABLE IF NOT EXISTS t(key int, value text)')
-                cur.execute("INSERT INTO t VALUES (0, 'something')")
-                cur.execute('SELECT SUM(key) FROM t')
-                sum_before = cur.fetchone()[0]
-
-                cur.execute("INSERT INTO t SELECT generate_series(1,100000), 'payload'")
-                cur.execute('SELECT SUM(key) FROM t')
-                sum_after = cur.fetchone()[0]
-                assert sum_after == sum_before + 5000050000
-
-    def show_statuses(safekeepers: List[Safekeeper], tenant_id: str, timeline_id: str):
-        for sk in safekeepers:
-            http_cli = sk.http_client()
-            try:
-                status = http_cli.timeline_status(tenant_id, timeline_id)
-                log.info(f"Safekeeper {sk.name} status: {status}")
-            except Exception as e:
-                log.info(f"Safekeeper {sk.name} status error: {e}")
-
-    zenith_env_builder.num_safekeepers = 4
-    env = zenith_env_builder.init()
-    env.zenith_cli.create_branch("test_replace_safekeeper", "main")
-
-    log.info("Use only first 3 safekeepers")
-    env.safekeepers[3].stop()
-    active_safekeepers = ['sk1', 'sk2', 'sk3']
-    pg = env.postgres.create('test_replace_safekeeper')
-    pg.adjust_for_wal_acceptors(safekeepers_guc(env, active_safekeepers))
-    pg.start()
-
-    # learn zenith timeline from compute
-    tenant_id = pg.safe_psql("show zenith.zenith_tenant")[0][0]
-    timeline_id = pg.safe_psql("show zenith.zenith_timeline")[0][0]
-
-    execute_payload(pg)
-    show_statuses(env.safekeepers, tenant_id, timeline_id)
-
-    log.info("Restart all safekeepers to flush everything")
-    env.safekeepers[0].stop(immediate=True)
-    execute_payload(pg)
-    env.safekeepers[0].start()
-    env.safekeepers[1].stop(immediate=True)
-    execute_payload(pg)
-    env.safekeepers[1].start()
-    env.safekeepers[2].stop(immediate=True)
-    execute_payload(pg)
-    env.safekeepers[2].start()
-
-    env.safekeepers[0].stop(immediate=True)
-    env.safekeepers[1].stop(immediate=True)
-    env.safekeepers[2].stop(immediate=True)
-    env.safekeepers[0].start()
-    env.safekeepers[1].start()
-    env.safekeepers[2].start()
-
-    execute_payload(pg)
-    show_statuses(env.safekeepers, tenant_id, timeline_id)
-
-    log.info("Stop sk1 (simulate failure) and use only quorum of sk2 and sk3")
-    env.safekeepers[0].stop(immediate=True)
-    execute_payload(pg)
-    show_statuses(env.safekeepers, tenant_id, timeline_id)
-
-    log.info("Recreate postgres to replace failed sk1 with new sk4")
-    pg.stop_and_destroy().create('test_replace_safekeeper')
-    active_safekeepers = ['sk2', 'sk3', 'sk4']
-    env.safekeepers[3].start()
-    pg.adjust_for_wal_acceptors(safekeepers_guc(env, active_safekeepers))
-    pg.start()
-
-    execute_payload(pg)
-    show_statuses(env.safekeepers, tenant_id, timeline_id)
-
-    log.info("Stop sk2 to require quorum of sk3 and sk4 for normal work")
-    env.safekeepers[1].stop(immediate=True)
-    execute_payload(pg)
-    show_statuses(env.safekeepers, tenant_id, timeline_id)