Build binaries from local code

Now, users can install these extensions themselves if they are owner
of the database they try to install the extension in.

.dockerignore

@@ -3,6 +3,7 @@
 **/.pytest_cache
 
 .git
+.github
 target
 tmp_check
 tmp_install

.git-blame-ignore-revs

@@ -0,0 +1 @@
+4c2bb43775947775401cbb9d774823c5723a91f8

.github/ISSUE_TEMPLATE/bug-template.md

@@ -0,0 +1,23 @@
+---
+name: Bug Template
+about: Used for describing bugs
+title: ''
+labels: t/bug
+assignees: ''
+
+---
+
+## Steps to reproduce
+
+
+## Expected result
+
+
+## Actual result
+
+
+## Environment
+
+
+## Logs, links
+- 

.github/ISSUE_TEMPLATE/epic-template.md

@@ -0,0 +1,25 @@
+---
+name: Epic Template
+about: A set of related tasks contributing towards specific outcome, comprizing of
+  more than 1 week of work.
+title: 'Epic: '
+labels: t/Epic
+assignees: ''
+
+---
+
+## Motivation
+
+
+## DoD
+
+
+## Implementation ideas
+
+
+## Tasks
+- [ ]
+
+
+## Other related tasks and Epics
+- 

.github/PULL_REQUEST_TEMPLATE/release-pr.md

@@ -0,0 +1,20 @@
+## Release 202Y-MM-DD
+
+**NB: this PR must be merged only by 'Create a merge commit'!**
+
+### Checklist when preparing for release
+- [ ] Read or refresh [the release flow guide](https://github.com/neondatabase/cloud/wiki/Release:-general-flow)
+- [ ] Ask in the [cloud Slack channel](https://neondb.slack.com/archives/C033A2WE6BZ) that you are going to rollout the release. Any blockers?
+- [ ] Does this release contain any db migrations? Destructive ones? What is the rollback plan?
+
+<!-- List everything that should be done **before** release, any issues / setting changes / etc -->
+
+### Checklist after release
+- [ ] Based on the merged commits write release notes and open a PR into `website` repo ([example](https://github.com/neondatabase/website/pull/120/files))
+- [ ] Check [#dev-production-stream](https://neondb.slack.com/archives/C03F5SM1N02) Slack channel
+- [ ] Check [stuck projects page](https://console.neon.tech/admin/projects?sort=last_active&order=desc&stuck=true)
+- [ ] Check [recent operation failures](https://console.neon.tech/admin/operations?action=create_timeline%2Cstart_compute%2Cstop_compute%2Csuspend_compute%2Capply_config%2Cdelete_timeline%2Cdelete_tenant%2Ccreate_branch%2Ccheck_availability&sort=updated_at&order=desc&had_retries=some)
+- [ ] Check [cloud SLO dashboard](https://observer.zenith.tech/d/_oWcBMJ7k/cloud-slos?orgId=1)
+- [ ] Check [compute startup metrics dashboard](https://observer.zenith.tech/d/5OkYJEmVz/compute-startup-time)
+
+<!-- List everything that should be done **after** release, any admin UI configuration / Grafana dashboard / alert changes / setting changes / etc -->

.github/actions/allure-report/action.yml

@@ -0,0 +1,219 @@
+name: 'Create Allure report'
+description: 'Create and publish Allure report'
+
+inputs:
+  action:
+    desctiption: 'generate or store'
+    required: true
+  build_type:
+    description: '`build_type` from run-python-test-set action'
+    required: true
+  test_selection:
+    description: '`test_selector` from run-python-test-set action'
+    required: false
+
+runs:
+  using: "composite"
+  steps:
+    - name: Validate input parameters
+      shell: bash -euxo pipefail {0}
+      run: |
+        if [ "${{ inputs.action }}" != "store"] && [ "${{ inputs.action }}" != "generate" ]; then
+          echo 2>&1 "Unknown inputs.action type '${{ inputs.action }}'; allowed 'generate' or 'store' only"
+          exit 1
+        fi
+
+        if [ -z "${{ inputs.test_selection }}" ] && [ "${{ inputs.action }}" == "store" ]; then
+          echo 2>&1 "inputs.test_selection must be set for 'store' action"
+          exit 2
+        fi
+
+    - name: Calculate key
+      id: calculate-key
+      shell: bash -euxo pipefail {0}
+      run: |
+        # TODO: for manually triggered workflows (via workflow_dispatch) we need to have a separate key
+
+        pr_number=$(jq --raw-output .pull_request.number "$GITHUB_EVENT_PATH" || true)
+        if [ "${pr_number}" != "null" ]; then
+          key=pr-${pr_number}
+        elif [ "${GITHUB_REF}" = "refs/heads/main" ]; then
+          # Shortcut for a special branch
+          key=main
+        else
+          key=branch-$(echo ${GITHUB_REF#refs/heads/} | tr -cd "[:alnum:]._-")
+        fi
+        echo "::set-output name=KEY::${key}"
+
+    - uses: actions/setup-java@v3
+      if: ${{ inputs.action == 'generate' }}
+      with:
+        distribution: 'temurin'
+        java-version: '17'
+
+    - name: Install Allure
+      if: ${{ inputs.action == 'generate' }}
+      shell: bash -euxo pipefail {0}
+      run: |
+        if ! which allure; then
+          ALLURE_ZIP=allure-${ALLURE_VERSION}.zip
+          wget -q https://github.com/allure-framework/allure2/releases/download/${ALLURE_VERSION}/${ALLURE_ZIP}
+          echo "${ALLURE_ZIP_MD5}  ${ALLURE_ZIP}" | md5sum -c
+          unzip -q ${ALLURE_ZIP}
+          echo "$(pwd)/allure-${ALLURE_VERSION}/bin" >> $GITHUB_PATH
+          rm -f ${ALLURE_ZIP}
+        fi
+      env:
+        ALLURE_VERSION: 2.19.0
+        ALLURE_ZIP_MD5: ced21401a1a8b9dfb68cee9e4c210464
+
+    - name: Upload Allure results
+      if: ${{ inputs.action == 'store' }}
+      env:
+        REPORT_PREFIX: reports/${{ steps.calculate-key.outputs.KEY }}/${{ inputs.build_type }}
+        RAW_PREFIX: reports-raw/${{ steps.calculate-key.outputs.KEY }}/${{ inputs.build_type }}
+        TEST_OUTPUT: /tmp/test_output
+        BUCKET: neon-github-public-dev
+      shell: bash -euxo pipefail {0}
+      run: |
+        # Add metadata
+        cat <<EOF > $TEST_OUTPUT/allure/results/executor.json
+          {
+            "name": "GitHub Actions",
+            "type": "github",
+            "url": "https://${BUCKET}.s3.amazonaws.com/${REPORT_PREFIX}/latest/index.html",
+            "buildOrder": ${GITHUB_RUN_ID},
+            "buildName": "GitHub Actions Run #${{ github.run_number }}/${GITHUB_RUN_ATTEMPT}",
+            "buildUrl": "${GITHUB_SERVER_URL}/${GITHUB_REPOSITORY}/actions/runs/${GITHUB_RUN_ID}/attempts/${GITHUB_RUN_ATTEMPT}",
+            "reportUrl": "https://${BUCKET}.s3.amazonaws.com/${REPORT_PREFIX}/${GITHUB_RUN_ID}/index.html",
+            "reportName": "Allure Report"
+          }
+        EOF
+        cat <<EOF > $TEST_OUTPUT/allure/results/environment.properties
+          TEST_SELECTION=${{ inputs.test_selection }}
+          BUILD_TYPE=${{ inputs.build_type }}
+        EOF
+
+        ARCHIVE="${GITHUB_RUN_ID}-${{ inputs.test_selection }}-${GITHUB_RUN_ATTEMPT}.tar.zst"
+        ZSTD_NBTHREADS=0
+
+        tar -C ${TEST_OUTPUT}/allure/results -cf ${ARCHIVE} --zstd .
+        aws s3 mv --only-show-errors ${ARCHIVE} "s3://${BUCKET}/${RAW_PREFIX}/${ARCHIVE}"
+
+    # Potentially we could have several running build for the same key (for example for the main branch),  so we use improvised lock for this
+    - name: Acquire Allure lock
+      if: ${{ inputs.action == 'generate' }}
+      shell: bash -euxo pipefail {0}
+      env:
+        LOCK_FILE: reports/${{ steps.calculate-key.outputs.KEY }}/lock.txt
+        BUCKET: neon-github-public-dev
+      run: |
+        LOCK_TIMEOUT=300 # seconds
+
+        for _ in $(seq 1 5); do
+          for i in $(seq 1 ${LOCK_TIMEOUT}); do
+            LOCK_ADDED=$(aws s3api head-object --bucket neon-github-public-dev --key ${LOCK_FILE} | jq --raw-output '.LastModified' || true)
+            # `date --date="..."` is supported only by gnu date (i.e. it doesn't work on BSD/macOS)
+            if [ -z "${LOCK_ADDED}" ] || [ "$(( $(date +%s) - $(date --date="${LOCK_ADDED}" +%s) ))" -gt "${LOCK_TIMEOUT}" ]; then
+              break
+            fi
+            sleep 1
+          done
+          echo "${GITHUB_RUN_ID}-${GITHUB_RUN_ATTEMPT}-${{ inputs.test_selection }}" > lock.txt
+          aws s3 mv --only-show-errors lock.txt "s3://${BUCKET}/${LOCK_FILE}"
+
+          # A double-check that exactly WE have acquired the lock
+          aws s3 cp --only-show-errors "s3://${BUCKET}/${LOCK_FILE}" ./lock.txt
+          if [ "$(cat lock.txt)" = "${GITHUB_RUN_ID}-${GITHUB_RUN_ATTEMPT}-${{ inputs.test_selection }}" ]; then
+            break
+          fi
+        done
+
+    - name: Generate and publish final Allure report
+      if: ${{ inputs.action == 'generate' }}
+      id: generate-report
+      env:
+        REPORT_PREFIX: reports/${{ steps.calculate-key.outputs.KEY }}/${{ inputs.build_type }}
+        RAW_PREFIX: reports-raw/${{ steps.calculate-key.outputs.KEY }}/${{ inputs.build_type }}
+        TEST_OUTPUT: /tmp/test_output
+        BUCKET: neon-github-public-dev
+      shell: bash -euxo pipefail {0}
+      run: |
+        # Get previously uploaded data for this run
+        ZSTD_NBTHREADS=0
+
+        s3_filepaths=$(aws s3api list-objects-v2 --bucket ${BUCKET} --prefix ${RAW_PREFIX}/${GITHUB_RUN_ID}- | jq --raw-output  '.Contents[].Key')
+        if [ -z "$s3_filepaths" ]; then
+          # There's no previously uploaded data for this run
+          exit 0
+        fi
+        for s3_filepath in ${s3_filepaths}; do
+          aws s3 cp --only-show-errors "s3://${BUCKET}/${s3_filepath}" "${TEST_OUTPUT}/allure/"
+
+          archive=${TEST_OUTPUT}/allure/$(basename $s3_filepath)
+          mkdir -p ${archive%.tar.zst}
+          tar -xf ${archive} -C ${archive%.tar.zst}
+          rm -f ${archive}
+        done
+
+        # Get history trend
+        aws s3 cp --recursive --only-show-errors "s3://${BUCKET}/${REPORT_PREFIX}/latest/history" "${TEST_OUTPUT}/allure/latest/history" || true
+
+        # Generate report
+        allure generate --clean --output $TEST_OUTPUT/allure/report $TEST_OUTPUT/allure/*
+
+        # Replace a logo link with a redirect to the latest version of the report
+        sed -i 's|<a href="." class=|<a href="https://'${BUCKET}'.s3.amazonaws.com/'${REPORT_PREFIX}'/latest/index.html" class=|g' $TEST_OUTPUT/allure/report/app.js
+
+        # Upload a history and the final report (in this particular order to not to have duplicated history in 2 places)
+        aws s3 mv --recursive --only-show-errors "${TEST_OUTPUT}/allure/report/history" "s3://${BUCKET}/${REPORT_PREFIX}/latest/history"
+        aws s3 mv --recursive --only-show-errors "${TEST_OUTPUT}/allure/report" "s3://${BUCKET}/${REPORT_PREFIX}/${GITHUB_RUN_ID}"
+
+        REPORT_URL=https://${BUCKET}.s3.amazonaws.com/${REPORT_PREFIX}/${GITHUB_RUN_ID}/index.html
+
+        # Generate redirect
+        cat <<EOF > ./index.html
+          <!DOCTYPE html>
+
+          <meta charset="utf-8">
+          <title>Redirecting to ${REPORT_URL}</title>
+          <meta http-equiv="refresh" content="0; URL=${REPORT_URL}">
+        EOF
+        aws s3 cp --only-show-errors ./index.html "s3://${BUCKET}/${REPORT_PREFIX}/latest/index.html"
+
+        echo "[Allure Report](${REPORT_URL})" >> ${GITHUB_STEP_SUMMARY}
+        echo "::set-output name=REPORT_URL::${REPORT_URL}"
+
+    - name: Release Allure lock
+      if: ${{ inputs.action == 'generate' && always() }}
+      shell: bash -euxo pipefail {0}
+      env:
+        LOCK_FILE: reports/${{ steps.calculate-key.outputs.KEY }}/lock.txt
+        BUCKET: neon-github-public-dev
+      run: |
+        aws s3 cp --only-show-errors "s3://${BUCKET}/${LOCK_FILE}" ./lock.txt || exit 0
+
+        if [ "$(cat lock.txt)" = "${GITHUB_RUN_ID}-${GITHUB_RUN_ATTEMPT}-${{ inputs.test_selection }}" ]; then
+          aws s3 rm "s3://${BUCKET}/${LOCK_FILE}"
+        fi
+
+    - uses: actions/github-script@v6
+      if: ${{ inputs.action == 'generate' && always() }}
+      env:
+        REPORT_URL: ${{ steps.generate-report.outputs.REPORT_URL }}
+        BUILD_TYPE: ${{ inputs.build_type }}
+        SHA: ${{ github.event.pull_request.head.sha || github.sha }}
+      with:
+        script: |
+          const { REPORT_URL, BUILD_TYPE, SHA } = process.env
+
+          result = await github.rest.repos.createCommitStatus({
+            owner: context.repo.owner,
+            repo: context.repo.repo,
+            sha: `${SHA}`,
+            state: 'success',
+            target_url: `${REPORT_URL}`,
+            context: `Allure report / ${BUILD_TYPE}`,
+          })
+
+          console.log(result);

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

@@ -131,8 +131,10 @@ runs:
         # -n4 uses four processes to run tests via pytest-xdist
         # -s is not used to prevent pytest from capturing output, because tests are running
         # in parallel and logs are mixed between different tests
+        mkdir -p $TEST_OUTPUT/allure/results
         "${cov_prefix[@]}" ./scripts/pytest \
           --junitxml=$TEST_OUTPUT/junit.xml \
+          --alluredir=$TEST_OUTPUT/allure/results \
           --tb=short \
           --verbose \
           -m "not remote_cluster" \
@@ -146,6 +148,14 @@ runs:
           fi
         fi
 
+    - name: Upload Allure results
+      if: ${{ always() && (inputs.test_selection == 'batch_others' || inputs.test_selection == 'batch_pg_regress') }}
+      uses: ./.github/actions/allure-report
+      with:
+        action: store
+        build_type: ${{ inputs.build_type }}
+        test_selection: ${{ inputs.test_selection }}
+
     - name: Delete all data but logs
       shell: bash -euxo pipefail {0}
       if: always()

.github/workflows/benchmarking.yml

@@ -106,7 +106,7 @@ jobs:
         mkdir -p perf-report-staging
         # Set --sparse-ordering option of pytest-order plugin to ensure tests are running in order of appears in the file,
         # it's important for test_perf_pgbench.py::test_pgbench_remote_* tests
-        ./scripts/pytest test_runner/performance/ -v -m "remote_cluster" --sparse-ordering --skip-interfering-proc-check --out-dir perf-report-staging --timeout 5400
+        ./scripts/pytest test_runner/performance/ -v -m "remote_cluster" --sparse-ordering --out-dir perf-report-staging --timeout 5400
 
     - name: Submit result
       env:
@@ -186,7 +186,7 @@ jobs:
         mkdir -p perf-report-captest
 
         psql $BENCHMARK_CONNSTR -c "SELECT 1;"
-        ./scripts/pytest test_runner/performance/test_perf_pgbench.py::test_pgbench_remote_init -v -m "remote_cluster" --skip-interfering-proc-check --out-dir perf-report-captest --timeout 21600
+        ./scripts/pytest test_runner/performance/test_perf_pgbench.py::test_pgbench_remote_init -v -m "remote_cluster" --out-dir perf-report-captest --timeout 21600
 
     - name: Benchmark simple-update
       env:
@@ -194,7 +194,7 @@ jobs:
         BENCHMARK_CONNSTR: ${{ secrets[matrix.connstr] }}
       run: |
         psql $BENCHMARK_CONNSTR -c "SELECT 1;"
-        ./scripts/pytest test_runner/performance/test_perf_pgbench.py::test_pgbench_remote_simple_update -v -m "remote_cluster" --skip-interfering-proc-check --out-dir perf-report-captest --timeout 21600
+        ./scripts/pytest test_runner/performance/test_perf_pgbench.py::test_pgbench_remote_simple_update -v -m "remote_cluster" --out-dir perf-report-captest --timeout 21600
 
     - name: Benchmark select-only
       env:
@@ -202,7 +202,7 @@ jobs:
         BENCHMARK_CONNSTR: ${{ secrets[matrix.connstr] }}
       run: |
         psql $BENCHMARK_CONNSTR -c "SELECT 1;"
-        ./scripts/pytest test_runner/performance/test_perf_pgbench.py::test_pgbench_remote_select_only -v -m "remote_cluster" --skip-interfering-proc-check --out-dir perf-report-captest --timeout 21600
+        ./scripts/pytest test_runner/performance/test_perf_pgbench.py::test_pgbench_remote_select_only -v -m "remote_cluster" --out-dir perf-report-captest --timeout 21600
 
     - name: Submit result
       env:

.github/workflows/build_and_test.yml

@@ -95,11 +95,11 @@ jobs:
           if [[ $BUILD_TYPE == "debug" ]]; then
             cov_prefix="scripts/coverage --profraw-prefix=$GITHUB_JOB --dir=/tmp/coverage run"
             CARGO_FEATURES=""
-            CARGO_FLAGS=""
+            CARGO_FLAGS="--locked"
           elif [[ $BUILD_TYPE == "release" ]]; then
             cov_prefix=""
             CARGO_FEATURES="--features profiling"
-            CARGO_FLAGS="--release $CARGO_FEATURES"
+            CARGO_FLAGS="--locked --release $CARGO_FEATURES"
           fi
           echo "cov_prefix=${cov_prefix}" >> $GITHUB_ENV
           echo "CARGO_FEATURES=${CARGO_FEATURES}" >> $GITHUB_ENV
@@ -121,8 +121,8 @@ jobs:
             target/
           # Fall back to older versions of the key, if no cache for current Cargo.lock was found
           key: |
-            v6-${{ runner.os }}-${{ matrix.build_type }}-cargo-${{ matrix.rust_toolchain }}-${{ hashFiles('Cargo.lock') }}
-            v6-${{ runner.os }}-${{ matrix.build_type }}-cargo-${{ matrix.rust_toolchain }}-
+            v7-${{ runner.os }}-${{ matrix.build_type }}-cargo-${{ matrix.rust_toolchain }}-${{ hashFiles('Cargo.lock') }}
+            v7-${{ runner.os }}-${{ matrix.build_type }}-cargo-${{ matrix.rust_toolchain }}-
 
       - name: Cache postgres build
         id: cache_pg
@@ -136,6 +136,10 @@ jobs:
         run: mold -run make postgres -j$(nproc)
         shell: bash -euxo pipefail {0}
 
+      - name: Build neon extensions
+        run: mold -run make neon-pg-ext -j$(nproc)
+        shell: bash -euxo pipefail {0}
+
       - name: Run cargo build
         run: |
           ${cov_prefix} mold -run cargo build $CARGO_FLAGS --features failpoints --bins --tests
@@ -298,6 +302,29 @@ jobs:
       # XXX: no coverage data handling here, since benchmarks are run on release builds,
       # while coverage is currently collected for the debug ones
 
+  merge-allure-report:
+    runs-on: dev
+    container:
+      image: 369495373322.dkr.ecr.eu-central-1.amazonaws.com/rust:pinned
+      options: --init
+    needs: [ other-tests, pg_regress-tests ]
+    if: always()
+    strategy:
+      fail-fast: false
+      matrix:
+        build_type: [ debug, release ]
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v3
+        with:
+          submodules: false
+
+      - name: Merge and  Allure results
+        uses: ./.github/actions/allure-report
+        with:
+          action: generate
+          build_type: ${{ matrix.build_type }}
+
   coverage-report:
     runs-on: dev
     container:
@@ -325,7 +352,7 @@ jobs:
             !~/.cargo/registry/src
             ~/.cargo/git/
             target/
-          key: v5-${{ runner.os }}-${{ matrix.build_type }}-cargo-${{ matrix.rust_toolchain }}-${{ hashFiles('Cargo.lock') }}
+          key: v7-${{ runner.os }}-${{ matrix.build_type }}-cargo-${{ matrix.rust_toolchain }}-${{ hashFiles('Cargo.lock') }}
 
       - name: Get Neon artifact
         uses: ./.github/actions/download
@@ -446,29 +473,9 @@ jobs:
       - name: Kaniko build compute tools
         run: /kaniko/executor --snapshotMode=redo --cache=true --cache-repo 369495373322.dkr.ecr.eu-central-1.amazonaws.com/cache --snapshotMode=redo --context . --dockerfile Dockerfile.compute-tools --destination 369495373322.dkr.ecr.eu-central-1.amazonaws.com/compute-tools:$GITHUB_RUN_ID
 
-  promote-image-compute-tools:
-    runs-on: dev
-    needs: [ compute-tools-image ]
-    if: github.event_name != 'workflow_dispatch'
-    container: amazon/aws-cli
-    strategy:
-      fail-fast: false
-      matrix:
-        name: [ compute-tools ]
-
-    steps:
-      - name: Promote image to latest
-        run:
-          MANIFEST=$(aws ecr batch-get-image --repository-name ${{ matrix.name }} --image-ids imageTag=$GITHUB_RUN_ID --query 'images[].imageManifest' --output text) && aws ecr put-image --repository-name ${{ matrix.name }} --image-tag latest --image-manifest "$MANIFEST"
-
   compute-node-image:
     runs-on: dev
     container: gcr.io/kaniko-project/executor:v1.9.0-debug
-    # note: This image depends on neondatabase/compute-tools:latest (or :thisversion),
-    # which isn't available until after the image is promoted.
-    # Ergo, we must explicitly build and promote compute-tools separately.
-    needs:
-      - promote-image-compute-tools
 
     steps:
       - name: Checkout
@@ -480,19 +487,18 @@ jobs:
       - name: Configure ECR login
         run: echo "{\"credsStore\":\"ecr-login\"}" > /kaniko/.docker/config.json
 
-      - name: Kaniko build compute node
-        working-directory: ./vendor/postgres/
-        run: /kaniko/executor --snapshotMode=redo --cache=true --cache-repo 369495373322.dkr.ecr.eu-central-1.amazonaws.com/cache --snapshotMode=redo --context . --build-arg=COMPUTE_TOOLS_TAG=$GITHUB_RUN_ID --destination 369495373322.dkr.ecr.eu-central-1.amazonaws.com/compute-node:$GITHUB_RUN_ID
+      - name: Kaniko build compute node with extensions
+        run: /kaniko/executor --snapshotMode=redo --cache=true --cache-repo 369495373322.dkr.ecr.eu-central-1.amazonaws.com/cache --snapshotMode=redo --context . --dockerfile Dockerfile.compute-node --destination 369495373322.dkr.ecr.eu-central-1.amazonaws.com/compute-node:$GITHUB_RUN_ID
 
   promote-images:
     runs-on: dev
-    needs: [ neon-image, compute-node-image ]
+    needs: [ neon-image, compute-node-image, compute-tools-image ]
     if: github.event_name != 'workflow_dispatch'
     container: amazon/aws-cli
     strategy:
       fail-fast: false
       matrix:
-        name: [ neon, compute-node ]
+        name: [ neon, compute-node, compute-tools ]
 
     steps:
       - name: Promote image to latest
@@ -509,18 +515,6 @@ jobs:
         run: |
           go install github.com/google/go-containerregistry/cmd/crane@31786c6cbb82d6ec4fb8eb79cd9387905130534e # v0.11.0
           go install github.com/awslabs/amazon-ecr-credential-helper/ecr-login/cli/docker-credential-ecr-login@69c85dc22db6511932bbf119e1a0cc5c90c69a7f # v0.6.0
-          
-#      - name: Get build tag
-#        run: |
-#          if [[ "$GITHUB_REF_NAME" == "main" ]]; then
-#            echo "::set-output name=tag::$(git rev-list --count HEAD)"
-#          elif [[ "$GITHUB_REF_NAME" == "release" ]]; then
-#            echo "::set-output name=tag::release-$(git rev-list --count HEAD)"
-#          else
-#            echo "GITHUB_REF_NAME (value '$GITHUB_REF_NAME') is not set to either 'main' or 'release' "
-#            echo "::set-output name=tag::$GITHUB_RUN_ID"
-#          fi
-#        id: build-tag
 
       - name: Configure ECR login
         run: |
@@ -536,6 +530,9 @@ jobs:
       - name: Pull compute node image from ECR
         run: crane pull 369495373322.dkr.ecr.eu-central-1.amazonaws.com/compute-node:latest compute-node
 
+      - name: Pull rust image from ECR
+        run: crane pull 369495373322.dkr.ecr.eu-central-1.amazonaws.com/rust:pinned rust
+
       - name: Configure docker login
         run: |
           # ECR Credential Helper & Docker Hub don't work together in config, hence reset
@@ -551,6 +548,9 @@ jobs:
       - name: Push compute node image to Docker Hub
         run: crane push compute-node neondatabase/compute-node:${{needs.tag.outputs.build-tag}}
 
+      - name: Push rust image to Docker Hub
+        run: crane push rust neondatabase/rust:pinned
+
       - name: Add latest tag to images
         if: |
           (github.ref_name == 'main' || github.ref_name == 'release') &&

.github/workflows/codestyle.yml

@@ -65,7 +65,7 @@ jobs:
 
       - name: Cache postgres build
         id: cache_pg
-        uses: actions/cache@v2
+        uses: actions/cache@v3
         with:
           path: |
             tmp_install/
@@ -81,6 +81,9 @@ jobs:
         if: steps.cache_pg.outputs.cache-hit != 'true'
         run: make postgres
 
+      - name: Build neon extensions
+        run: make neon-pg-ext
+
       # Plain configure output can contain weird errors like 'error: C compiler cannot create executables'
       # and the real cause will be inside config.log
       - name: Print configure logs in case of failure
@@ -94,20 +97,20 @@ jobs:
 
       - name: Cache cargo deps
         id: cache_cargo
-        uses: actions/cache@v2
+        uses: actions/cache@v3
         with:
           path: |
             ~/.cargo/registry
             !~/.cargo/registry/src
             ~/.cargo/git
             target
-          key: v2-${{ runner.os }}-cargo-${{ hashFiles('./Cargo.lock') }}-rust-${{ matrix.rust_toolchain }}
+          key: v3-${{ runner.os }}-cargo-${{ hashFiles('./Cargo.lock') }}-rust-${{ matrix.rust_toolchain }}
 
       - name: Run cargo clippy
         run: ./run_clippy.sh
 
       - name: Ensure all project builds
-        run: cargo build --all --all-targets
+        run: cargo build --locked --all --all-targets
 
   check-codestyle-python:
     runs-on: [ self-hosted, Linux, k8s-runner ]
@@ -128,8 +131,14 @@ jobs:
       - name: Install Python deps
         run: ./scripts/pysync
 
-      - name: Run yapf to ensure code format
-        run: poetry run yapf --recursive --diff .
+      - name: Run isort to ensure code format
+        run: poetry run isort --diff --check .
+
+      - name: Run black to ensure code format
+        run: poetry run black --diff --check .
+
+      - name: Run flake8 to ensure code format
+        run: poetry run flake8 .
 
       - name: Run mypy to check types
         run: poetry run mypy .

.gitignore

@@ -2,6 +2,7 @@
 /tmp_check
 /tmp_install
 /tmp_check_cli
+/neon_install
 __pycache__/
 test_output/
 .vscode

.yapfignore

@@ -1,10 +0,0 @@
-# This file is only read when `yapf` is run from this directory.
-# Hence we only top-level directories here to avoid confusion.
-# See source code for the exact file format: https://github.com/google/yapf/blob/c6077954245bc3add82dafd853a1c7305a6ebd20/yapf/yapflib/file_resources.py#L40-L43
-vendor/
-target/
-tmp_install/
-__pycache__/
-test_output/
-.neon/
-.git/

Cargo.lock

@@ -1164,20 +1164,14 @@ version = "1.8.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "eabb4a44450da02c90444cf74558da904edde8fb4e9035a9a6a4e15445af0bd7"
 
-[[package]]
-name = "hashbrown"
-version = "0.11.2"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "ab5ef0d4909ef3724cc8cce6ccc8572c5c817592e9285f5464f8e86f8bd3726e"
-dependencies = [
- "ahash",
-]
-
 [[package]]
 name = "hashbrown"
 version = "0.12.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "8a9ee70c43aaf417c914396645a0fa852624801b24ebb7ae78fe8272889ac888"
+dependencies = [
+ "ahash",
+]
 
 [[package]]
 name = "heck"
@@ -1379,7 +1373,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "10a35a97730320ffe8e2d410b5d3b69279b98d2c14bdb8b70ea89ecf7888d41e"
 dependencies = [
  "autocfg",
- "hashbrown 0.12.3",
+ "hashbrown",
 ]
 
 [[package]]
@@ -2111,7 +2105,6 @@ dependencies = [
  "bindgen",
  "byteorder",
  "bytes",
- "chrono",
  "crc32c",
  "env_logger",
  "hex",
@@ -2274,7 +2267,7 @@ dependencies = [
  "clap 3.2.16",
  "futures",
  "git-version",
- "hashbrown 0.11.2",
+ "hashbrown",
  "hex",
  "hmac 0.12.1",
  "hyper",
@@ -2289,7 +2282,7 @@ dependencies = [
  "routerify",
  "rstest",
  "rustls",
- "rustls-pemfile 0.2.1",
+ "rustls-pemfile",
  "scopeguard",
  "serde",
  "serde_json",
@@ -2315,15 +2308,6 @@ dependencies = [
  "memchr",
 ]
 
-[[package]]
-name = "quickcheck"
-version = "1.0.3"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "588f6378e4dd99458b60ec275b4477add41ce4fa9f64dcba6f15adccb19b50d6"
-dependencies = [
- "rand",
-]
-
 [[package]]
 name = "quote"
 version = "1.0.21"
@@ -2508,7 +2492,7 @@ dependencies = [
  "percent-encoding",
  "pin-project-lite",
  "rustls",
- "rustls-pemfile 1.0.1",
+ "rustls-pemfile",
  "serde",
  "serde_json",
  "serde_urlencoded",
@@ -2697,15 +2681,6 @@ dependencies = [
  "webpki",
 ]
 
-[[package]]
-name = "rustls-pemfile"
-version = "0.2.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "5eebeaeb360c87bfb72e84abdb3447159c0eaececf1bef2aecd65a8be949d1c9"
-dependencies = [
- "base64",
-]
-
 [[package]]
 name = "rustls-pemfile"
 version = "1.0.1"
@@ -3242,7 +3217,6 @@ dependencies = [
  "js-sys",
  "libc",
  "num_threads",
- "quickcheck",
  "time-macros",
 ]
 
@@ -3683,7 +3657,7 @@ dependencies = [
  "rand",
  "routerify",
  "rustls",
- "rustls-pemfile 0.2.1",
+ "rustls-pemfile",
  "rustls-split",
  "serde",
  "serde_json",
@@ -3969,6 +3943,7 @@ version = "0.1.0"
 dependencies = [
  "ahash",
  "anyhow",
+ "bstr",
  "bytes",
  "chrono",
  "clap 2.34.0",
@@ -3978,7 +3953,7 @@ dependencies = [
  "futures-task",
  "futures-util",
  "generic-array",
- "hashbrown 0.11.2",
+ "hashbrown",
  "hex",
  "hyper",
  "indexmap",
@@ -3993,6 +3968,7 @@ dependencies = [
  "prost",
  "rand",
  "regex",
+ "regex-automata",
  "regex-syntax",
  "scopeguard",
  "serde",

Dockerfile

@@ -1,18 +1,27 @@
+### Creates a storage Docker image with postgres, pageserver, safekeeper and proxy binaries.
+### The image itself is mainly used as a container for the binaries and for starting e2e tests with custom parameters.
+### By default, the binaries inside the image have some mock parameters and can start, but are not intended to be used
+### inside this image in the real deployments.
+ARG REPOSITORY=369495373322.dkr.ecr.eu-central-1.amazonaws.com
+ARG IMAGE=rust
+ARG TAG=pinned
+
 # Build Postgres
-FROM 369495373322.dkr.ecr.eu-central-1.amazonaws.com/rust:pinned AS pg-build
+FROM $REPOSITORY/$IMAGE:$TAG AS pg-build
 WORKDIR /home/nonroot
 
-COPY vendor/postgres vendor/postgres
-COPY Makefile Makefile
+COPY --chown=nonroot vendor/postgres vendor/postgres
+COPY --chown=nonroot pgxn pgxn
+COPY --chown=nonroot Makefile Makefile
 
 ENV BUILD_TYPE release
 RUN set -e \
-    && mold -run make -j $(nproc) -s postgres \
+    && mold -run make -j $(nproc) -s neon-pg-ext \
     && rm -rf tmp_install/build \
     && tar -C tmp_install -czf /home/nonroot/postgres_install.tar.gz .
 
 # Build zenith binaries
-FROM 369495373322.dkr.ecr.eu-central-1.amazonaws.com/rust:pinned AS build
+FROM $REPOSITORY/$IMAGE:$TAG AS build
 WORKDIR /home/nonroot
 ARG GIT_VERSION=local
 
@@ -32,7 +41,7 @@ COPY . .
 # Show build caching stats to check if it was used in the end.
 # Has to be the part of the same RUN since cachepot daemon is killed in the end of this RUN, losing the compilation stats.
 RUN set -e \
-    && mold -run cargo build --release \
+&& mold -run cargo build --locked --release \
     && cachepot -s
 
 # Build final image
@@ -58,7 +67,18 @@ COPY --from=build --chown=zenith:zenith /home/nonroot/target/release/proxy
 COPY --from=pg-build /home/nonroot/tmp_install/ /usr/local/
 COPY --from=pg-build /home/nonroot/postgres_install.tar.gz /data/
 
+# By default, pageserver uses `.neon/` working directory in WORKDIR, so create one and fill it with the dummy config.
+# Now, when `docker run ... pageserver` is run, it can start without errors, yet will have some default dummy values.
+RUN mkdir -p /data/.neon/ && chown -R zenith:zenith /data/.neon/ \
+    && /usr/local/bin/pageserver -D /data/.neon/ --init \
+       -c "id=1234" \
+       -c "broker_endpoints=['http://etcd:2379']" \
+       -c "pg_distrib_dir='/usr/local'" \
+       -c "listen_pg_addr='0.0.0.0:6400'" \
+       -c "listen_http_addr='0.0.0.0:9898'"
+
 VOLUME ["/data"]
 USER zenith
 EXPOSE 6400
-CMD ["pageserver"]
+EXPOSE 9898
+CMD ["/bin/bash"]

Dockerfile.compute-node

@@ -0,0 +1,111 @@
+ARG TAG=pinned
+
+FROM debian:bullseye-slim AS build-deps
+RUN apt update &&  \
+    apt install -y git autoconf automake libtool build-essential bison flex libreadline-dev zlib1g-dev libxml2-dev \
+    libcurl4-openssl-dev libossp-uuid-dev
+
+# Build Postgres from the neon postgres repository.
+FROM build-deps AS pg-build
+COPY vendor/postgres postgres
+RUN cd postgres && \
+    ./configure CFLAGS='-O2 -g3' --enable-debug --with-uuid=ossp && \
+    make MAKELEVEL=0 -j $(getconf _NPROCESSORS_ONLN) -s install && \
+    make MAKELEVEL=0 -j $(getconf _NPROCESSORS_ONLN) -s -C contrib/ install && \
+    # Install headers
+    make MAKELEVEL=0 -j $(getconf _NPROCESSORS_ONLN) -s -C src/include install && \
+    make MAKELEVEL=0 -j $(getconf _NPROCESSORS_ONLN) -s -C src/interfaces/libpq install
+
+# Build PostGIS from the upstream PostGIS mirror. PostGIS compiles against neon postgres sources without changes.
+# Perhaps we could even use the upstream binaries, compiled against vanilla Postgres, but it would require some
+# investigation to check that it works, and also keeps working in the future. So for now, we compile our own binaries.
+FROM build-deps AS postgis-build
+COPY --from=pg-build /usr/local/pgsql/ /usr/local/pgsql/
+RUN apt update && \
+    apt install -y gdal-bin libgdal-dev libprotobuf-c-dev protobuf-c-compiler xsltproc wget
+
+RUN wget https://download.osgeo.org/postgis/source/postgis-3.2.3.tar.gz && \
+    tar xvzf postgis-3.2.3.tar.gz && \
+    cd postgis-3.2.3 && \
+    ./autogen.sh && \
+    export PATH="/usr/local/pgsql/bin:$PATH" && \
+    ./configure && \
+    make -j $(getconf _NPROCESSORS_ONLN) install && \
+    cd extensions/postgis && \
+    make clean && \
+    make -j $(getconf _NPROCESSORS_ONLN) install && \
+    echo 'trusted = true' >> /usr/local/pgsql/share/extension/postgis.control && \
+    echo 'trusted = true' >> /usr/local/pgsql/share/extension/postgis_raster.control && \
+    echo 'trusted = true' >> /usr/local/pgsql/share/extension/postgis_tiger_geocoder.control && \
+    echo 'trusted = true' >> /usr/local/pgsql/share/extension/postgis_topology.control
+
+# Build plv8
+FROM build-deps AS plv8-build
+COPY --from=postgis-build /usr/local/pgsql/ /usr/local/pgsql/
+RUN apt update && \
+    apt install -y git curl wget make ninja-build build-essential libncurses5 python3-dev pkg-config libc++-dev libc++abi-dev libglib2.0-dev
+
+# https://github.com/plv8/plv8/issues/475
+# Debian bullseye provides binutils 2.35 when >= 2.38 is necessary
+RUN echo "deb http://ftp.debian.org/debian testing main" >> /etc/apt/sources.list && \
+    echo "APT::Default-Release \"stable\";" > /etc/apt/apt.conf.d/default-release && \
+    apt update && \
+    apt install -y --no-install-recommends -t testing binutils
+
+RUN wget https://github.com/plv8/plv8/archive/refs/tags/v3.1.3.tar.gz && \
+    tar xvzf v3.1.3.tar.gz && \
+    cd plv8-3.1.3 && \
+    export PATH="/usr/local/pgsql/bin:$PATH" && \
+    make && \
+    make install && \
+    rm -rf /plv8-* && \
+    echo 'trusted = true' >> /usr/local/pgsql/share/extension/plv8.control
+
+# compile neon extensions
+FROM build-deps AS neon-pg-ext-build
+COPY --from=plv8-build /usr/local/pgsql/ /usr/local/pgsql/
+COPY pgxn/ pgxn/
+
+RUN make -j $(getconf _NPROCESSORS_ONLN) \
+        PG_CONFIG=/usr/local/pgsql/bin/pg_config \
+        -C pgxn/neon \
+        -s install
+
+# Compile and run the Neon-specific `compute_ctl` binary
+FROM 369495373322.dkr.ecr.eu-central-1.amazonaws.com/rust:$TAG AS compute-tools
+USER nonroot
+# Copy entire project to get Cargo.* files with proper dependencies for the whole project
+COPY --chown=nonroot . .
+RUN cd compute_tools && cargo build --locked --release
+
+# Put it all together into the final image
+FROM debian:bullseye-slim
+# Add user postgres
+RUN mkdir /var/db && useradd -m -d /var/db/postgres postgres && \
+    echo "postgres:test_console_pass" | chpasswd && \
+    mkdir /var/db/postgres/compute && mkdir /var/db/postgres/specs && \
+    chown -R postgres:postgres /var/db/postgres && \
+    chmod 0750 /var/db/postgres/compute && \
+    echo '/usr/local/lib' >> /etc/ld.so.conf && /sbin/ldconfig
+
+# TODO: Check if we can make the extension setup more modular versus a linear build
+# currently plv8-build copies the output /usr/local/pgsql from postgis-build, etc#
+COPY --from=neon-pg-ext-build --chown=postgres /usr/local/pgsql /usr/local
+COPY --from=compute-tools --chown=postgres /home/nonroot/target/release/compute_ctl /usr/local/bin/compute_ctl
+
+RUN apt update &&  \
+    apt install -y libreadline-dev libossp-uuid-dev gdal-bin libgdal-dev libprotobuf-c-dev && \
+    rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/*
+
+# Debian bullseye provides GLIBC 2.31 when 2.34 is necessary as we compiled plv8 with that version
+RUN echo "deb http://ftp.debian.org/debian testing main" >> /etc/apt/sources.list && \
+    echo "APT::Default-Release \"stable\";" > /etc/apt/apt.conf.d/default-release && \
+    apt update && \
+    apt install -y --no-install-recommends -t testing binutils && \
+    rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/*
+
+# "temporary" symlink for old control-plane
+RUN ln -s /usr/local/bin/compute_ctl /usr/local/bin/zenith_ctl
+
+USER postgres
+ENTRYPOINT ["/usr/local/bin/compute_ctl"]

Dockerfile.compute-node.legacy

@@ -0,0 +1,87 @@
+#
+# Legacy version of the Dockerfile for the compute node.
+# Used by e2e CI. Building Dockerfile.compute-node will take
+# unreasonable ammount of time without v2 runners.
+#
+# TODO: remove once cloud repo CI is moved to v2 runners.
+#
+
+
+# Allow specifiyng different compute-tools tag and image repo, so we are
+# able to use different images
+ARG REPOSITORY=369495373322.dkr.ecr.eu-central-1.amazonaws.com
+ARG IMAGE=compute-tools
+ARG TAG=latest
+
+#
+# Image with pre-built tools
+#
+FROM $REPOSITORY/$IMAGE:$TAG AS compute-deps
+# Only to get ready compute_ctl binary as deppendency
+
+#
+# Image with Postgres build deps
+#
+FROM debian:buster-slim AS build-deps
+
+RUN apt-get update && apt-get -yq install automake libtool build-essential bison flex libreadline-dev zlib1g-dev libxml2-dev \
+                                          libcurl4-openssl-dev libossp-uuid-dev
+
+#
+# Image with built Postgres
+#
+FROM build-deps AS pg-build
+
+# Add user postgres
+RUN adduser postgres
+RUN mkdir /pg && chown postgres:postgres /pg
+
+# Copy source files
+COPY ./vendor/postgres /pg/
+COPY ./pgxn /pg/
+
+# Build and install Postgres locally
+RUN mkdir /pg/compute_build && cd /pg/compute_build && \
+    ../configure CFLAGS='-O2 -g3' --prefix=$(pwd)/postgres_bin --enable-debug --with-uuid=ossp && \
+    # Install main binaries and contribs
+    make MAKELEVEL=0 -j $(getconf _NPROCESSORS_ONLN) -s install && \
+    make MAKELEVEL=0 -j $(getconf _NPROCESSORS_ONLN) -s -C contrib/ install && \
+    # Install headers
+    make MAKELEVEL=0 -j $(getconf _NPROCESSORS_ONLN) -s -C src/include install
+
+# Install neon contrib
+RUN make MAKELEVEL=0 PG_CONFIG=/pg/compute_build/postgres_bin/bin/pg_config -j $(getconf _NPROCESSORS_ONLN) -C /pg/neon install
+
+USER postgres
+WORKDIR /pg
+
+#
+# Final compute node image to be exported
+#
+FROM debian:buster-slim
+
+# libreadline-dev is required to run psql
+RUN apt-get update && apt-get -yq install libreadline-dev libossp-uuid-dev
+
+# Add user postgres
+RUN mkdir /var/db && useradd -m -d /var/db/postgres postgres && \
+    echo "postgres:test_console_pass" | chpasswd && \
+    mkdir /var/db/postgres/compute && mkdir /var/db/postgres/specs && \
+    chown -R postgres:postgres /var/db/postgres && \
+    chmod 0750 /var/db/postgres/compute
+
+# Copy ready Postgres binaries
+COPY --from=pg-build /pg/compute_build/postgres_bin /usr/local
+
+# Copy binaries from compute-tools
+COPY --from=compute-deps /usr/local/bin/compute_ctl /usr/local/bin/compute_ctl
+
+# XXX: temporary symlink for compatibility with old control-plane
+RUN ln -s /usr/local/bin/compute_ctl /usr/local/bin/zenith_ctl
+
+# Add postgres shared objects to the search path
+RUN echo '/usr/local/lib' >> /etc/ld.so.conf && /sbin/ldconfig
+
+USER postgres
+
+ENTRYPOINT ["/usr/local/bin/compute_ctl"]

Dockerfile.compute-tools

@@ -1,6 +1,10 @@
 # First transient image to build compute_tools binaries
 # NB: keep in sync with rust image version in .github/workflows/build_and_test.yml
-FROM 369495373322.dkr.ecr.eu-central-1.amazonaws.com/rust:pinned AS rust-build
+ARG REPOSITORY=369495373322.dkr.ecr.eu-central-1.amazonaws.com
+ARG IMAGE=rust
+ARG TAG=pinned
+
+FROM $REPOSITORY/$IMAGE:$TAG AS rust-build
 WORKDIR /home/nonroot
 
 # Enable https://github.com/paritytech/cachepot to cache Rust crates' compilation results in Docker builds.
@@ -16,7 +20,7 @@ ARG CACHEPOT_BUCKET=neon-github-dev
 COPY . .
 
 RUN set -e \
-    && mold -run cargo build -p compute_tools --release \
+    && mold -run cargo build -p compute_tools --locked --release \
     && cachepot -s
 
 # Final image that only has one binary

Makefile

@@ -51,7 +51,7 @@ CARGO_CMD_PREFIX += CARGO_TERM_PROGRESS_WHEN=never CI=1
 # Top level Makefile to build Zenith and PostgreSQL
 #
 .PHONY: all
-all: zenith postgres
+all: zenith postgres neon-pg-ext
 
 ### Zenith Rust bits
 #
@@ -87,25 +87,39 @@ postgres: postgres-configure \
 		  postgres-headers # to prevent `make install` conflicts with zenith's `postgres-headers`
 	+@echo "Compiling PostgreSQL"
 	$(MAKE) -C $(POSTGRES_INSTALL_DIR)/build MAKELEVEL=0 install
-	+@echo "Compiling contrib/neon"
-	$(MAKE) -C $(POSTGRES_INSTALL_DIR)/build/contrib/neon install
-	+@echo "Compiling contrib/neon_test_utils"
-	$(MAKE) -C $(POSTGRES_INSTALL_DIR)/build/contrib/neon_test_utils install
+	+@echo "Compiling libpq"
+	$(MAKE) -C $(POSTGRES_INSTALL_DIR)/build/src/interfaces/libpq install
 	+@echo "Compiling pg_buffercache"
 	$(MAKE) -C $(POSTGRES_INSTALL_DIR)/build/contrib/pg_buffercache install
 	+@echo "Compiling pageinspect"
 	$(MAKE) -C $(POSTGRES_INSTALL_DIR)/build/contrib/pageinspect install
 
-
 .PHONY: postgres-clean
 postgres-clean:
 	$(MAKE) -C $(POSTGRES_INSTALL_DIR)/build MAKELEVEL=0 clean
+	$(MAKE) -C $(POSTGRES_INSTALL_DIR)/build/contrib/pg_buffercache clean
+	$(MAKE) -C $(POSTGRES_INSTALL_DIR)/build/contrib/pageinspect clean
+	$(MAKE) -C $(POSTGRES_INSTALL_DIR)/build/src/interfaces/libpq clean
+
+neon-pg-ext: postgres
+	+@echo "Compiling neon"
+	$(MAKE) PG_CONFIG=$(POSTGRES_INSTALL_DIR)/bin/pg_config \
+		-C $(ROOT_PROJECT_DIR)/pgxn/neon install
+	+@echo "Compiling neon_test_utils"
+	$(MAKE) PG_CONFIG=$(POSTGRES_INSTALL_DIR)/bin/pg_config \
+		-C $(ROOT_PROJECT_DIR)/pgxn/neon_test_utils install
+
+.PHONY: neon-pg-ext-clean
+	$(MAKE) -C $(ROOT_PROJECT_DIR)/pgxn/neon clean
+	$(MAKE) -C $(ROOT_PROJECT_DIR)/pgxn/neon_test_utils clean
 
 # This doesn't remove the effects of 'configure'.
 .PHONY: clean
 clean:
 	cd $(POSTGRES_INSTALL_DIR)/build && $(MAKE) clean
 	$(CARGO_CMD_PREFIX) cargo clean
+	cd pgxn/neon && $(MAKE) clean
+	cd pgxn/neon_test_utils && $(MAKE) clean
 
 # This removes everything
 .PHONY: distclean

docs/glossary.md

@@ -92,6 +92,7 @@ The layer map tracks what layers exist in a timeline.
 ### Layered repository
 
 Neon repository implementation that keeps data in layers.
+
 ### LSN
 
 The Log Sequence Number (LSN) is a unique identifier of the WAL record[] in the WAL log.
@@ -125,6 +126,26 @@ TODO: use this name consistently in remote storage code. Now `disk_consistent_ls
 * `ancestor_lsn` - LSN of the branch point (the LSN at which this branch was created)
 
 TODO: add table that describes mapping between PostgreSQL (compute), safekeeper and pageserver LSNs.
+
+### Logical size
+
+The pageserver tracks the "logical size" of a timeline. It is the
+total size of all relations in all Postgres databases on the
+timeline. It includes all user and system tables, including their FSM
+and VM forks. But it does not include SLRUs, twophase files or any
+other such data or metadata that lives outside relations.
+
+The logical size is calculated by the pageserver, and is sent to
+PostgreSQL via feedback messages to the safekeepers. PostgreSQL uses
+the logical size to enforce the size limit in the free tier. The
+logical size is also shown to users in the web console.
+
+The logical size is not affected by branches or the physical layout of
+layer files in the pageserver. If you have a database with 1 GB
+logical size and you create a branch of it, both branches will have 1
+GB logical size, even though the branch is copy-on-write and won't
+consume any extra physical disk space until you make changes to it.
+
 ### Page (block)
 
 The basic structure used to store relation data. All pages are of the same size.

docs/sourcetree.md

@@ -112,11 +112,13 @@ Run `poetry shell` to activate the virtual environment.
 Alternatively, use `poetry run` to run a single command in the venv, e.g. `poetry run pytest`.
 
 ### Obligatory checks
-We force code formatting via `yapf` and type hints via `mypy`.
-Run the following commands in the repository's root (next to `setup.cfg`):
+We force code formatting via `black`, `isort` and type hints via `mypy`.
+Run the following commands in the repository's root (next to `pyproject.toml`):
 
 ```bash
-poetry run yapf -ri .  # All code is reformatted
+poetry run isort .  # Imports are reformatted
+poetry run black .  # All code is reformatted
+poetry run flake8 .  # Python linter
 poetry run mypy .  # Ensure there are no typing errors
 ```
 
@@ -125,7 +127,7 @@ Otherwise it will not find its configuration.
 
 Also consider:
 
-* Running `flake8` (or a linter of your choice, e.g. `pycodestyle`) and fixing possible defects, if any.
+* Running `pycodestyle` (or a linter of your choice) and fixing possible defects, if any.
 * Adding more type hints to your code to avoid `Any`.
 
 ### Changing dependencies

libs/postgres_ffi/Cargo.toml

@@ -4,7 +4,6 @@ version = "0.1.0"
 edition = "2021"
 
 [dependencies]
-chrono = "0.4.19"
 rand = "0.8.3"
 regex = "1.4.5"
 bytes = "1.0.1"

libs/postgres_ffi/src/nonrelfile_utils.rs

@@ -1,8 +1,8 @@
 //!
 //! Common utilities for dealing with PostgreSQL non-relation files.
 //!
-use crate::transaction_id_precedes;
 use super::pg_constants;
+use crate::transaction_id_precedes;
 use bytes::BytesMut;
 use log::*;
 

libs/postgres_ffi/src/waldecoder.rs

@@ -8,9 +8,9 @@
 //! to look deeper into the WAL records to also understand which blocks they modify, the code
 //! for that is in pageserver/src/walrecord.rs
 //!
+use super::bindings::{XLogLongPageHeaderData, XLogPageHeaderData, XLogRecord, XLOG_PAGE_MAGIC};
 use super::pg_constants;
 use super::xlog_utils::*;
-use super::bindings::{XLogLongPageHeaderData, XLogPageHeaderData, XLogRecord, XLOG_PAGE_MAGIC};
 use bytes::{Buf, BufMut, Bytes, BytesMut};
 use crc32c::*;
 use log::*;
@@ -170,6 +170,7 @@ impl WalStreamDecoder {
                 }
                 State::SkippingEverything { .. } => {}
             }
+            // now read page contents
             match &mut self.state {
                 State::WaitingForRecord => {
                     // need to have at least the xl_tot_len field
@@ -194,8 +195,8 @@ impl WalStreamDecoder {
                         return Ok(Some(self.complete_record(recordbuf)?));
                     } else {
                         // Need to assemble the record from pieces. Remember the size of the
-                        // record, and loop back. On next iteration, we will reach the 'else'
-                        // branch below, and copy the part of the record that was on this page
+                        // record, and loop back. On next iterations, we will reach the branch
+                        // below, and copy the part of the record that was on this or next page(s)
                         // to 'recordbuf'.  Subsequent iterations will skip page headers, and
                         // append the continuations from the next pages to 'recordbuf'.
                         self.state = State::ReassemblingRecord {

libs/utils/Cargo.toml

@@ -39,7 +39,7 @@ bytes = "1.0.1"
 hex-literal = "0.3"
 tempfile = "3.2"
 criterion = "0.3"
-rustls-pemfile = "0.2.1"
+rustls-pemfile = "1"
 
 [[bench]]
 name = "benchmarks"

libs/utils/src/lib.rs

@@ -8,6 +8,9 @@ pub mod lsn;
 /// SeqWait allows waiting for a future sequence number to arrive
 pub mod seqwait;
 
+/// A simple Read-Copy-Update implementation.
+pub mod simple_rcu;
+
 /// append only ordered map implemented with a Vec
 pub mod vec_map;
 

libs/utils/src/postgres_backend.rs

@@ -163,14 +163,9 @@ pub fn is_socket_read_timed_out(error: &anyhow::Error) -> bool {
     false
 }
 
-// Truncate 0 from C string in Bytes and stringify it (returns slice, no allocations)
-// PG protocol strings are always C strings.
-fn cstr_to_str(b: &Bytes) -> Result<&str> {
-    let without_null = if b.last() == Some(&0) {
-        &b[..b.len() - 1]
-    } else {
-        &b[..]
-    };
+// Cast a byte slice to a string slice, dropping null terminator if there's one.
+fn cstr_to_str(bytes: &[u8]) -> Result<&str> {
+    let without_null = bytes.strip_suffix(&[0]).unwrap_or(bytes);
     std::str::from_utf8(without_null).map_err(|e| e.into())
 }
 
@@ -423,9 +418,9 @@ impl PostgresBackend {
                 self.state = ProtoState::Established;
             }
 
-            FeMessage::Query(m) => {
+            FeMessage::Query(body) => {
                 // remove null terminator
-                let query_string = cstr_to_str(&m.body)?;
+                let query_string = cstr_to_str(&body)?;
 
                 trace!("got query {:?}", query_string);
                 // xxx distinguish fatal and recoverable errors?

libs/utils/src/pq_proto.rs

@@ -25,8 +25,10 @@ pub const TEXT_OID: Oid = 25;
 #[derive(Debug)]
 pub enum FeMessage {
     StartupPacket(FeStartupPacket),
-    Query(FeQueryMessage), // Simple query
-    Parse(FeParseMessage), // Extended query protocol
+    // Simple query.
+    Query(Bytes),
+    // Extended query protocol.
+    Parse(FeParseMessage),
     Describe(FeDescribeMessage),
     Bind(FeBindMessage),
     Execute(FeExecuteMessage),
@@ -69,11 +71,6 @@ impl Distribution<CancelKeyData> for Standard {
     }
 }
 
-#[derive(Debug)]
-pub struct FeQueryMessage {
-    pub body: Bytes,
-}
-
 // We only support the simple case of Parse on unnamed prepared statement and
 // no params
 #[derive(Debug)]
@@ -89,7 +86,7 @@ pub struct FeDescribeMessage {
 
 // we only support unnamed prepared stmt and portal
 #[derive(Debug)]
-pub struct FeBindMessage {}
+pub struct FeBindMessage;
 
 // we only support unnamed prepared stmt or portal
 #[derive(Debug)]
@@ -100,7 +97,7 @@ pub struct FeExecuteMessage {
 
 // we only support unnamed prepared stmt and portal
 #[derive(Debug)]
-pub struct FeCloseMessage {}
+pub struct FeCloseMessage;
 
 /// Retry a read on EINTR
 ///
@@ -163,22 +160,20 @@ impl FeMessage {
                 Err(e) if e.kind() == io::ErrorKind::UnexpectedEof => return Ok(None),
                 Err(e) => return Err(e.into()),
             };
-            let len = retry_read!(stream.read_u32().await)?;
 
-            // The message length includes itself, so it better be at least 4
-            let bodylen = len
+            // The message length includes itself, so it better be at least 4.
+            let len = retry_read!(stream.read_u32().await)?
                 .checked_sub(4)
-                .context("invalid message length: parsing u32")?;
+                .context("invalid message length")?;
 
-            // Read message body
-            let mut body_buf: Vec<u8> = vec![0; bodylen as usize];
-            stream.read_exact(&mut body_buf).await?;
+            let body = {
+                let mut buffer = vec![0u8; len as usize];
+                stream.read_exact(&mut buffer).await?;
+                Bytes::from(buffer)
+            };
 
-            let body = Bytes::from(body_buf);
-
-            // Parse it
             match tag {
-                b'Q' => Ok(Some(FeMessage::Query(FeQueryMessage { body }))),
+                b'Q' => Ok(Some(FeMessage::Query(body))),
                 b'P' => Ok(Some(FeParseMessage::parse(body)?)),
                 b'D' => Ok(Some(FeDescribeMessage::parse(body)?)),
                 b'E' => Ok(Some(FeExecuteMessage::parse(body)?)),
@@ -302,124 +297,71 @@ impl FeStartupPacket {
 }
 
 impl FeParseMessage {
-    pub fn parse(mut buf: Bytes) -> anyhow::Result<FeMessage> {
-        let _pstmt_name = read_null_terminated(&mut buf)?;
-        let query_string = read_null_terminated(&mut buf)?;
-        let nparams = buf.get_i16();
-
+    fn parse(mut buf: Bytes) -> anyhow::Result<FeMessage> {
         // FIXME: the rust-postgres driver uses a named prepared statement
         // for copy_out(). We're not prepared to handle that correctly. For
         // now, just ignore the statement name, assuming that the client never
         // uses more than one prepared statement at a time.
-        /*
-        if !pstmt_name.is_empty() {
-            return Err(io::Error::new(
-                io::ErrorKind::InvalidInput,
-                "named prepared statements not implemented in Parse",
-            ));
-        }
-         */
 
-        if nparams != 0 {
-            bail!("query params not implemented");
-        }
+        let _pstmt_name = read_cstr(&mut buf)?;
+        let query_string = read_cstr(&mut buf)?;
+        let nparams = buf.get_i16();
+
+        ensure!(nparams == 0, "query params not implemented");
 
         Ok(FeMessage::Parse(FeParseMessage { query_string }))
     }
 }
 
 impl FeDescribeMessage {
-    pub fn parse(mut buf: Bytes) -> anyhow::Result<FeMessage> {
+    fn parse(mut buf: Bytes) -> anyhow::Result<FeMessage> {
         let kind = buf.get_u8();
-        let _pstmt_name = read_null_terminated(&mut buf)?;
+        let _pstmt_name = read_cstr(&mut buf)?;
 
         // FIXME: see FeParseMessage::parse
-        /*
-        if !pstmt_name.is_empty() {
-            return Err(io::Error::new(
-                io::ErrorKind::InvalidInput,
-                "named prepared statements not implemented in Describe",
-            ));
-        }
-        */
-
-        if kind != b'S' {
-            bail!("only prepared statmement Describe is implemented");
-        }
+        ensure!(
+            kind == b'S',
+            "only prepared statemement Describe is implemented"
+        );
 
         Ok(FeMessage::Describe(FeDescribeMessage { kind }))
     }
 }
 
 impl FeExecuteMessage {
-    pub fn parse(mut buf: Bytes) -> anyhow::Result<FeMessage> {
-        let portal_name = read_null_terminated(&mut buf)?;
+    fn parse(mut buf: Bytes) -> anyhow::Result<FeMessage> {
+        let portal_name = read_cstr(&mut buf)?;
         let maxrows = buf.get_i32();
 
-        if !portal_name.is_empty() {
-            bail!("named portals not implemented");
-        }
-
-        if maxrows != 0 {
-            bail!("row limit in Execute message not supported");
-        }
+        ensure!(portal_name.is_empty(), "named portals not implemented");
+        ensure!(maxrows == 0, "row limit in Execute message not implemented");
 
         Ok(FeMessage::Execute(FeExecuteMessage { maxrows }))
     }
 }
 
 impl FeBindMessage {
-    pub fn parse(mut buf: Bytes) -> anyhow::Result<FeMessage> {
-        let portal_name = read_null_terminated(&mut buf)?;
-        let _pstmt_name = read_null_terminated(&mut buf)?;
-
-        if !portal_name.is_empty() {
-            bail!("named portals not implemented");
-        }
+    fn parse(mut buf: Bytes) -> anyhow::Result<FeMessage> {
+        let portal_name = read_cstr(&mut buf)?;
+        let _pstmt_name = read_cstr(&mut buf)?;
 
         // FIXME: see FeParseMessage::parse
-        /*
-        if !pstmt_name.is_empty() {
-            return Err(io::Error::new(
-                io::ErrorKind::InvalidInput,
-                "named prepared statements not implemented",
-            ));
-        }
-        */
+        ensure!(portal_name.is_empty(), "named portals not implemented");
 
-        Ok(FeMessage::Bind(FeBindMessage {}))
+        Ok(FeMessage::Bind(FeBindMessage))
     }
 }
 
 impl FeCloseMessage {
-    pub fn parse(mut buf: Bytes) -> anyhow::Result<FeMessage> {
+    fn parse(mut buf: Bytes) -> anyhow::Result<FeMessage> {
         let _kind = buf.get_u8();
-        let _pstmt_or_portal_name = read_null_terminated(&mut buf)?;
+        let _pstmt_or_portal_name = read_cstr(&mut buf)?;
 
         // FIXME: we do nothing with Close
-
-        Ok(FeMessage::Close(FeCloseMessage {}))
+        Ok(FeMessage::Close(FeCloseMessage))
     }
 }
 
-fn read_null_terminated(buf: &mut Bytes) -> anyhow::Result<Bytes> {
-    let mut result = BytesMut::new();
-
-    loop {
-        if !buf.has_remaining() {
-            bail!("no null-terminator in string");
-        }
-
-        let byte = buf.get_u8();
-
-        if byte == 0 {
-            break;
-        }
-        result.put_u8(byte);
-    }
-    Ok(result.freeze())
-}
-
 // Backend
 
 #[derive(Debug)]
@@ -441,7 +383,7 @@ pub enum BeMessage<'a> {
     // None means column is NULL
     DataRow(&'a [Option<&'a [u8]>]),
     ErrorResponse(&'a str),
-    // single byte - used in response to SSLRequest/GSSENCRequest
+    /// Single byte - used in response to SSLRequest/GSSENCRequest.
     EncryptionResponse(bool),
     NoData,
     ParameterDescription,
@@ -554,49 +496,22 @@ pub static SINGLE_COL_ROWDESC: BeMessage = BeMessage::RowDescription(&[RowDescri
     formatcode: 0,
 }]);
 
-// Safe usize -> i32|i16 conversion, from rust-postgres
-trait FromUsize: Sized {
-    fn from_usize(x: usize) -> Result<Self, io::Error>;
-}
-
-macro_rules! from_usize {
-    ($t:ty) => {
-        impl FromUsize for $t {
-            #[inline]
-            fn from_usize(x: usize) -> io::Result<$t> {
-                if x > <$t>::max_value() as usize {
-                    Err(io::Error::new(
-                        io::ErrorKind::InvalidInput,
-                        "value too large to transmit",
-                    ))
-                } else {
-                    Ok(x as $t)
-                }
-            }
-        }
-    };
-}
-
-from_usize!(i32);
-
 /// Call f() to write body of the message and prepend it with 4-byte len as
 /// prescribed by the protocol.
-fn write_body<F>(buf: &mut BytesMut, f: F) -> io::Result<()>
-where
-    F: FnOnce(&mut BytesMut) -> io::Result<()>,
-{
+fn write_body<R>(buf: &mut BytesMut, f: impl FnOnce(&mut BytesMut) -> R) -> R {
     let base = buf.len();
     buf.extend_from_slice(&[0; 4]);
 
-    f(buf)?;
+    let res = f(buf);
 
-    let size = i32::from_usize(buf.len() - base)?;
+    let size = i32::try_from(buf.len() - base).expect("message too big to transmit");
     (&mut buf[base..]).put_slice(&size.to_be_bytes());
-    Ok(())
+
+    res
 }
 
 /// Safe write of s into buf as cstring (String in the protocol).
-pub fn write_cstr(s: &[u8], buf: &mut BytesMut) -> Result<(), io::Error> {
+fn write_cstr(s: &[u8], buf: &mut BytesMut) -> Result<(), io::Error> {
     if s.contains(&0) {
         return Err(io::Error::new(
             io::ErrorKind::InvalidInput,
@@ -608,15 +523,11 @@ pub fn write_cstr(s: &[u8], buf: &mut BytesMut) -> Result<(), io::Error> {
     Ok(())
 }
 
-// Truncate 0 from C string in Bytes and stringify it (returns slice, no allocations)
-// PG protocol strings are always C strings.
-fn cstr_to_str(b: &Bytes) -> Result<&str> {
-    let without_null = if b.last() == Some(&0) {
-        &b[..b.len() - 1]
-    } else {
-        &b[..]
-    };
-    std::str::from_utf8(without_null).map_err(|e| e.into())
+fn read_cstr(buf: &mut Bytes) -> anyhow::Result<Bytes> {
+    let pos = buf.iter().position(|x| *x == 0);
+    let result = buf.split_to(pos.context("missing terminator")?);
+    buf.advance(1); // drop the null terminator
+    Ok(result)
 }
 
 impl<'a> BeMessage<'a> {
@@ -631,18 +542,14 @@ impl<'a> BeMessage<'a> {
                 buf.put_u8(b'R');
                 write_body(buf, |buf| {
                     buf.put_i32(0); // Specifies that the authentication was successful.
-                    Ok::<_, io::Error>(())
-                })
-                .unwrap(); // write into BytesMut can't fail
+                });
             }
 
             BeMessage::AuthenticationCleartextPassword => {
                 buf.put_u8(b'R');
                 write_body(buf, |buf| {
                     buf.put_i32(3); // Specifies that clear text password is required.
-                    Ok::<_, io::Error>(())
-                })
-                .unwrap(); // write into BytesMut can't fail
+                });
             }
 
             BeMessage::AuthenticationMD5Password(salt) => {
@@ -650,9 +557,7 @@ impl<'a> BeMessage<'a> {
                 write_body(buf, |buf| {
                     buf.put_i32(5); // Specifies that an MD5-encrypted password is required.
                     buf.put_slice(&salt[..]);
-                    Ok::<_, io::Error>(())
-                })
-                .unwrap(); // write into BytesMut can't fail
+                });
             }
 
             BeMessage::AuthenticationSasl(msg) => {
@@ -677,8 +582,7 @@ impl<'a> BeMessage<'a> {
                         }
                     }
                     Ok::<_, io::Error>(())
-                })
-                .unwrap()
+                })?;
             }
 
             BeMessage::BackendKeyData(key_data) => {
@@ -686,77 +590,64 @@ impl<'a> BeMessage<'a> {
                 write_body(buf, |buf| {
                     buf.put_i32(key_data.backend_pid);
                     buf.put_i32(key_data.cancel_key);
-                    Ok(())
-                })
-                .unwrap();
+                });
             }
 
             BeMessage::BindComplete => {
                 buf.put_u8(b'2');
-                write_body(buf, |_| Ok::<(), io::Error>(())).unwrap();
+                write_body(buf, |_| {});
             }
 
             BeMessage::CloseComplete => {
                 buf.put_u8(b'3');
-                write_body(buf, |_| Ok::<(), io::Error>(())).unwrap();
+                write_body(buf, |_| {});
             }
 
             BeMessage::CommandComplete(cmd) => {
                 buf.put_u8(b'C');
-                write_body(buf, |buf| {
-                    write_cstr(cmd, buf)?;
-                    Ok::<_, io::Error>(())
-                })?;
+                write_body(buf, |buf| write_cstr(cmd, buf))?;
             }
 
             BeMessage::CopyData(data) => {
                 buf.put_u8(b'd');
                 write_body(buf, |buf| {
                     buf.put_slice(data);
-                    Ok::<_, io::Error>(())
-                })
-                .unwrap();
+                });
             }
 
             BeMessage::CopyDone => {
                 buf.put_u8(b'c');
-                write_body(buf, |_| Ok::<(), io::Error>(())).unwrap();
+                write_body(buf, |_| {});
             }
 
             BeMessage::CopyFail => {
                 buf.put_u8(b'f');
-                write_body(buf, |_| Ok::<(), io::Error>(())).unwrap();
+                write_body(buf, |_| {});
             }
 
             BeMessage::CopyInResponse => {
                 buf.put_u8(b'G');
                 write_body(buf, |buf| {
-                    buf.put_u8(1); /* copy_is_binary */
-                    buf.put_i16(0); /* numAttributes */
-                    Ok::<_, io::Error>(())
-                })
-                .unwrap();
+                    buf.put_u8(1); // copy_is_binary
+                    buf.put_i16(0); // numAttributes
+                });
             }
 
             BeMessage::CopyOutResponse => {
                 buf.put_u8(b'H');
                 write_body(buf, |buf| {
-                    buf.put_u8(0); /* copy_is_binary */
-                    buf.put_i16(0); /* numAttributes */
-                    Ok::<_, io::Error>(())
-                })
-                .unwrap();
+                    buf.put_u8(0); // copy_is_binary
+                    buf.put_i16(0); // numAttributes
+                });
             }
 
             BeMessage::CopyBothResponse => {
                 buf.put_u8(b'W');
                 write_body(buf, |buf| {
                     // doesn't matter, used only for replication
-                    buf.put_u8(0); /* copy_is_binary */
-                    buf.put_i16(0); /* numAttributes */
-                    Ok::<_, io::Error>(())
-                })
-                .unwrap();
+                    buf.put_u8(0); // copy_is_binary
+                    buf.put_i16(0); // numAttributes
+                });
             }
 
             BeMessage::DataRow(vals) => {
@@ -771,9 +662,7 @@ impl<'a> BeMessage<'a> {
                             buf.put_i32(-1);
                         }
                     }
-                    Ok::<_, io::Error>(())
-                })
-                .unwrap();
+                });
             }
 
             // ErrorResponse is a zero-terminated array of zero-terminated fields.
@@ -788,18 +677,17 @@ impl<'a> BeMessage<'a> {
                 buf.put_u8(b'E');
                 write_body(buf, |buf| {
                     buf.put_u8(b'S'); // severity
-                    write_cstr(&Bytes::from("ERROR"), buf)?;
+                    buf.put_slice(b"ERROR\0");
 
                     buf.put_u8(b'C'); // SQLSTATE error code
-                    write_cstr(&Bytes::from("CXX000"), buf)?;
+                    buf.put_slice(b"CXX000\0");
 
                     buf.put_u8(b'M'); // the message
                     write_cstr(error_msg.as_bytes(), buf)?;
 
                     buf.put_u8(0); // terminator
                     Ok::<_, io::Error>(())
-                })
-                .unwrap();
+                })?;
             }
 
             // NoticeResponse has the same format as ErrorResponse. From doc: "The frontend should display the
@@ -812,23 +700,22 @@ impl<'a> BeMessage<'a> {
                 buf.put_u8(b'N');
                 write_body(buf, |buf| {
                     buf.put_u8(b'S'); // severity
-                    write_cstr(&Bytes::from("NOTICE"), buf)?;
+                    buf.put_slice(b"NOTICE\0");
 
                     buf.put_u8(b'C'); // SQLSTATE error code
-                    write_cstr(&Bytes::from("CXX000"), buf)?;
+                    buf.put_slice(b"CXX000\0");
 
                     buf.put_u8(b'M'); // the message
                     write_cstr(error_msg.as_bytes(), buf)?;
 
                     buf.put_u8(0); // terminator
                     Ok::<_, io::Error>(())
-                })
-                .unwrap();
+                })?;
             }
 
             BeMessage::NoData => {
                 buf.put_u8(b'n');
-                write_body(buf, |_| Ok::<(), io::Error>(())).unwrap();
+                write_body(buf, |_| {});
             }
 
             BeMessage::EncryptionResponse(should_negotiate) => {
@@ -853,9 +740,7 @@ impl<'a> BeMessage<'a> {
                 buf.put_u8(b'S');
                 write_body(buf, |buf| {
                     buf.put_slice(&buffer[..cnt]);
-                    Ok::<_, io::Error>(())
-                })
-                .unwrap();
+                });
             }
 
             BeMessage::ParameterDescription => {
@@ -863,23 +748,19 @@ impl<'a> BeMessage<'a> {
                 write_body(buf, |buf| {
                     // we don't support params, so always 0
                     buf.put_i16(0);
-                    Ok::<_, io::Error>(())
-                })
-                .unwrap();
+                });
             }
 
             BeMessage::ParseComplete => {
                 buf.put_u8(b'1');
-                write_body(buf, |_| Ok::<(), io::Error>(())).unwrap();
+                write_body(buf, |_| {});
             }
 
             BeMessage::ReadyForQuery => {
                 buf.put_u8(b'Z');
                 write_body(buf, |buf| {
                     buf.put_u8(b'I');
-                    Ok::<_, io::Error>(())
-                })
-                .unwrap();
+                });
             }
 
             BeMessage::RowDescription(rows) => {
@@ -907,9 +788,7 @@ impl<'a> BeMessage<'a> {
                     buf.put_u64(body.wal_end);
                     buf.put_i64(body.timestamp);
                     buf.put_slice(body.data);
-                    Ok::<_, io::Error>(())
-                })
-                .unwrap();
+                });
             }
 
             BeMessage::KeepAlive(req) => {
@@ -918,10 +797,8 @@ impl<'a> BeMessage<'a> {
                     buf.put_u8(b'k');
                     buf.put_u64(req.sent_ptr);
                     buf.put_i64(req.timestamp);
-                    buf.put_u8(if req.request_reply { 1u8 } else { 0u8 });
-                    Ok::<_, io::Error>(())
-                })
-                .unwrap();
+                    buf.put_u8(if req.request_reply { 1 } else { 0 });
+                });
             }
         }
         Ok(())
@@ -968,17 +845,17 @@ impl ReplicationFeedback {
     // value itself
     pub fn serialize(&self, buf: &mut BytesMut) -> Result<()> {
         buf.put_u8(REPLICATION_FEEDBACK_FIELDS_NUMBER); // # of keys
-        write_cstr(&Bytes::from("current_timeline_size"), buf)?;
+        buf.put_slice(b"current_timeline_size\0");
         buf.put_i32(8);
         buf.put_u64(self.current_timeline_size);
 
-        write_cstr(&Bytes::from("ps_writelsn"), buf)?;
+        buf.put_slice(b"ps_writelsn\0");
         buf.put_i32(8);
         buf.put_u64(self.ps_writelsn);
-        write_cstr(&Bytes::from("ps_flushlsn"), buf)?;
+        buf.put_slice(b"ps_flushlsn\0");
         buf.put_i32(8);
         buf.put_u64(self.ps_flushlsn);
-        write_cstr(&Bytes::from("ps_applylsn"), buf)?;
+        buf.put_slice(b"ps_applylsn\0");
         buf.put_i32(8);
         buf.put_u64(self.ps_applylsn);
 
@@ -988,7 +865,7 @@ impl ReplicationFeedback {
             .expect("failed to serialize pg_replytime earlier than PG_EPOCH")
             .as_micros() as i64;
 
-        write_cstr(&Bytes::from("ps_replytime"), buf)?;
+        buf.put_slice(b"ps_replytime\0");
         buf.put_i32(8);
         buf.put_i64(timestamp);
         Ok(())
@@ -998,33 +875,30 @@ impl ReplicationFeedback {
     pub fn parse(mut buf: Bytes) -> ReplicationFeedback {
         let mut zf = ReplicationFeedback::empty();
         let nfields = buf.get_u8();
-        let mut i = 0;
-        while i < nfields {
-            i += 1;
-            let key_cstr = read_null_terminated(&mut buf).unwrap();
-            let key = cstr_to_str(&key_cstr).unwrap();
-            match key {
-                "current_timeline_size" => {
+        for _ in 0..nfields {
+            let key = read_cstr(&mut buf).unwrap();
+            match key.as_ref() {
+                b"current_timeline_size" => {
                     let len = buf.get_i32();
                     assert_eq!(len, 8);
                     zf.current_timeline_size = buf.get_u64();
                 }
-                "ps_writelsn" => {
+                b"ps_writelsn" => {
                     let len = buf.get_i32();
                     assert_eq!(len, 8);
                     zf.ps_writelsn = buf.get_u64();
                 }
-                "ps_flushlsn" => {
+                b"ps_flushlsn" => {
                     let len = buf.get_i32();
                     assert_eq!(len, 8);
                     zf.ps_flushlsn = buf.get_u64();
                 }
-                "ps_applylsn" => {
+                b"ps_applylsn" => {
                     let len = buf.get_i32();
                     assert_eq!(len, 8);
                     zf.ps_applylsn = buf.get_u64();
                 }
-                "ps_replytime" => {
+                b"ps_replytime" => {
                     let len = buf.get_i32();
                     assert_eq!(len, 8);
                     let raw_time = buf.get_i64();
@@ -1037,8 +911,8 @@ impl ReplicationFeedback {
                 _ => {
                     let len = buf.get_i32();
                     warn!(
-                        "ReplicationFeedback parse. unknown key {} of len {}. Skip it.",
-                        key, len
+                        "ReplicationFeedback parse. unknown key {} of len {len}. Skip it.",
+                        String::from_utf8_lossy(key.as_ref())
                     );
                     buf.advance(len as usize);
                 }
@@ -1084,7 +958,7 @@ mod tests {
             *first = REPLICATION_FEEDBACK_FIELDS_NUMBER + 1;
         }
 
-        write_cstr(&Bytes::from("new_field_one"), &mut data).unwrap();
+        data.put_slice(b"new_field_one\0");
         data.put_i32(8);
         data.put_u64(42);
 

libs/utils/src/seqwait_async.rs

@@ -1,8 +1,8 @@
-///
-/// Async version of 'seqwait.rs'
-///
-/// NOTE: This is currently unused. If you need this, you'll need to uncomment this in lib.rs.
-///
+//!
+//! Async version of 'seqwait.rs'
+//!
+//! NOTE: This is currently unused. If you need this, you'll need to uncomment this in lib.rs.
+//!
 
 #![warn(missing_docs)]
 

libs/utils/src/simple_rcu.rs

@@ -0,0 +1,217 @@
+//!
+//! RCU stands for Read-Copy-Update. It's a synchronization mechanism somewhat
+//! similar to a lock, but it allows readers to "hold on" to an old value of RCU
+//! without blocking writers, and allows writing a new values without blocking
+//! readers. When you update the new value, the new value is immediately visible
+//! to new readers, but the update waits until all existing readers have
+//! finishe, so that no one sees the old value anymore.
+//!
+//! This implementation isn't wait-free; it uses an RwLock that is held for a
+//! short duration when the value is read or updated.
+//!
+#![warn(missing_docs)]
+
+use std::ops::Deref;
+use std::sync::mpsc::{sync_channel, Receiver, SyncSender};
+use std::sync::{Arc, Weak};
+use std::sync::{Mutex, RwLock, RwLockWriteGuard};
+
+///
+/// Rcu allows multiple readers to read and hold onto a value without blocking
+/// (for very long).  Storing to the Rcu updates the value, making new readers
+/// immediately see the new value, but it also waits for all current readers to
+/// finish.
+///
+pub struct Rcu<V> {
+    inner: RwLock<RcuInner<V>>,
+}
+
+struct RcuInner<V> {
+    current_cell: Arc<RcuCell<V>>,
+    old_cells: Vec<Weak<RcuCell<V>>>,
+}
+
+///
+/// RcuCell holds one value. It can be the latest one, or an old one.
+///
+struct RcuCell<V> {
+    value: V,
+
+    /// A dummy channel. We never send anything to this channel. The point is
+    /// that when the RcuCell is dropped, any cloned Senders will be notified
+    /// that the channel is closed. Updaters can use this to wait out until the
+    /// RcuCell has been dropped, i.e. until the old value is no longer in use.
+    ///
+    /// We never do anything with the receiver, we just need to hold onto it so
+    /// that the Senders will be notified when it's dropped. But because it's
+    /// not Sync, we need a Mutex on it.
+    watch: (SyncSender<()>, Mutex<Receiver<()>>),
+}
+
+impl<V> RcuCell<V> {
+    fn new(value: V) -> Self {
+        let (watch_sender, watch_receiver) = sync_channel(0);
+        RcuCell {
+            value,
+            watch: (watch_sender, Mutex::new(watch_receiver)),
+        }
+    }
+}
+
+impl<V> Rcu<V> {
+    /// Create a new `Rcu`, initialized to `starting_val`
+    pub fn new(starting_val: V) -> Self {
+        let inner = RcuInner {
+            current_cell: Arc::new(RcuCell::new(starting_val)),
+            old_cells: Vec::new(),
+        };
+        Self {
+            inner: RwLock::new(inner),
+        }
+    }
+
+    ///
+    /// Read current value. Any store() calls will block until the returned
+    /// guard object is dropped.
+    ///
+    pub fn read(&self) -> RcuReadGuard<V> {
+        let current_cell = Arc::clone(&self.inner.read().unwrap().current_cell);
+        RcuReadGuard { cell: current_cell }
+    }
+
+    ///
+    /// Lock the current value for updating. Returns a guard object that can be
+    /// used to read the current value, and to store a new value.
+    ///
+    /// Note: holding the write-guard blocks concurrent readers, so you should
+    /// finish the update and drop the guard quickly!
+    ///
+    pub fn write(&self) -> RcuWriteGuard<'_, V> {
+        let inner = self.inner.write().unwrap();
+        RcuWriteGuard { inner }
+    }
+}
+
+///
+/// Read guard returned by `read`
+///
+pub struct RcuReadGuard<V> {
+    cell: Arc<RcuCell<V>>,
+}
+
+impl<V> Deref for RcuReadGuard<V> {
+    type Target = V;
+
+    fn deref(&self) -> &V {
+        &self.cell.value
+    }
+}
+
+///
+/// Read guard returned by `read`
+///
+pub struct RcuWriteGuard<'a, V> {
+    inner: RwLockWriteGuard<'a, RcuInner<V>>,
+}
+
+impl<'a, V> Deref for RcuWriteGuard<'a, V> {
+    type Target = V;
+
+    fn deref(&self) -> &V {
+        &self.inner.current_cell.value
+    }
+}
+
+impl<'a, V> RcuWriteGuard<'a, V> {
+    ///
+    /// Store a new value. The new value will be written to the Rcu immediately,
+    /// and will be immediately seen by any `read` calls that start afterwards.
+    /// But if there are any readers still holding onto the old value, or any
+    /// even older values, this will await until they have been released.
+    ///
+    /// This will drop the write-guard before it starts waiting for the reads to
+    /// finish, so a new write operation can begin before this functio returns.
+    ///
+    pub fn store(mut self, new_val: V) {
+        let new_cell = Arc::new(RcuCell::new(new_val));
+
+        let mut watches = Vec::new();
+        {
+            let old = std::mem::replace(&mut self.inner.current_cell, new_cell);
+            self.inner.old_cells.push(Arc::downgrade(&old));
+
+            // cleanup old cells that no longer have any readers, and collect
+            // the watches for any that do.
+            self.inner.old_cells.retain(|weak| {
+                if let Some(cell) = weak.upgrade() {
+                    watches.push(cell.watch.0.clone());
+                    true
+                } else {
+                    false
+                }
+            });
+        }
+        drop(self);
+
+        // after all the old_cells are no longer in use, we're done
+        for w in watches.iter_mut() {
+            // This will block until the Receiver is closed. That happens then
+            // the RcuCell is dropped.
+            #[allow(clippy::single_match)]
+            match w.send(()) {
+                Ok(_) => panic!("send() unexpectedly succeeded on dummy channel"),
+                Err(_) => {
+                    // closed, which means that the cell has been dropped, and
+                    // its value is no longer in use
+                }
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::sync::{Arc, Mutex};
+    use std::thread::{sleep, spawn};
+    use std::time::Duration;
+
+    #[test]
+    fn basic() {
+        let rcu = Arc::new(Rcu::new(1));
+        let log = Arc::new(Mutex::new(Vec::new()));
+
+        let a = rcu.read();
+        assert_eq!(*a, 1);
+        log.lock().unwrap().push("one");
+
+        let (rcu_clone, log_clone) = (Arc::clone(&rcu), Arc::clone(&log));
+        let thread = spawn(move || {
+            log_clone.lock().unwrap().push("store two start");
+            let write_guard = rcu_clone.write();
+            assert_eq!(*write_guard, 1);
+            write_guard.store(2);
+            log_clone.lock().unwrap().push("store two done");
+        });
+        // without this sleep the test can pass on accident if the writer is slow
+        sleep(Duration::from_secs(1));
+
+        // new read should see the new value
+        let b = rcu.read();
+        assert_eq!(*b, 2);
+
+        // old guard still sees the old value
+        assert_eq!(*a, 1);
+
+        // Release the old guard. This lets the store in the thread to finish.
+        log.lock().unwrap().push("release a");
+        drop(a);
+
+        thread.join().unwrap();
+
+        assert_eq!(
+            log.lock().unwrap().as_slice(),
+            &["one", "store two start", "release a", "store two done",]
+        );
+    }
+}

pageserver/Cargo.toml

@@ -15,7 +15,7 @@ failpoints = ["fail/failpoints"]
 chrono = "0.4.19"
 rand = "0.8.3"
 regex = "1.4.5"
-bytes = { version = "1.0.1", features = ['serde'] }
+bytes = "1.0.1"
 byteorder = "1.4.3"
 futures = "0.3.13"
 hex = "0.4.3"

pageserver/src/basebackup.rs

@@ -22,8 +22,8 @@ use std::time::SystemTime;
 use tar::{Builder, EntryType, Header};
 use tracing::*;
 
+use crate::layered_repository::Timeline;
 use crate::reltag::{RelTag, SlruKind};
-use crate::DatadirTimeline;
 
 use postgres_ffi::v14::pg_constants;
 use postgres_ffi::v14::xlog_utils::{generate_wal_segment, normalize_lsn, XLogFileName};
@@ -36,13 +36,12 @@ use utils::lsn::Lsn;
 /// This is short-living object only for the time of tarball creation,
 /// created mostly to avoid passing a lot of parameters between various functions
 /// used for constructing tarball.
-pub struct Basebackup<'a, W, T>
+pub struct Basebackup<'a, W>
 where
     W: Write,
-    T: DatadirTimeline,
 {
     ar: Builder<AbortableWrite<W>>,
-    timeline: &'a Arc<T>,
+    timeline: &'a Arc<Timeline>,
     pub lsn: Lsn,
     prev_record_lsn: Lsn,
     full_backup: bool,
@@ -57,18 +56,17 @@ where
 //  * When working without safekeepers. In this situation it is important to match the lsn
 //    we are taking basebackup on with the lsn that is used in pageserver's walreceiver
 //    to start the replication.
-impl<'a, W, T> Basebackup<'a, W, T>
+impl<'a, W> Basebackup<'a, W>
 where
     W: Write,
-    T: DatadirTimeline,
 {
     pub fn new(
         write: W,
-        timeline: &'a Arc<T>,
+        timeline: &'a Arc<Timeline>,
         req_lsn: Option<Lsn>,
         prev_lsn: Option<Lsn>,
         full_backup: bool,
-    ) -> Result<Basebackup<'a, W, T>> {
+    ) -> Result<Basebackup<'a, W>> {
         // Compute postgres doesn't have any previous WAL files, but the first
         // record that it's going to write needs to include the LSN of the
         // previous record (xl_prev). We include prev_record_lsn in the
@@ -404,10 +402,9 @@ where
     }
 }
 
-impl<'a, W, T> Drop for Basebackup<'a, W, T>
+impl<'a, W> Drop for Basebackup<'a, W>
 where
     W: Write,
-    T: DatadirTimeline,
 {
     /// If the basebackup was not finished, prevent the Archive::drop() from
     /// writing the end-of-archive marker.

pageserver/src/http/models.rs

@@ -129,9 +129,9 @@ pub struct LocalTimelineInfo {
     pub latest_gc_cutoff_lsn: Lsn,
     #[serde_as(as = "DisplayFromStr")]
     pub disk_consistent_lsn: Lsn,
-    pub current_logical_size: Option<usize>, // is None when timeline is Unloaded
-    pub current_physical_size: Option<u64>,  // is None when timeline is Unloaded
-    pub current_logical_size_non_incremental: Option<usize>,
+    pub current_logical_size: Option<u64>, // is None when timeline is Unloaded
+    pub current_physical_size: Option<u64>, // is None when timeline is Unloaded
+    pub current_logical_size_non_incremental: Option<u64>,
     pub current_physical_size_non_incremental: Option<u64>,
     pub timeline_state: LocalTimelineState,
 
@@ -150,6 +150,9 @@ pub struct RemoteTimelineInfo {
     pub awaits_download: bool,
 }
 
+///
+/// This represents the output of the "timeline_detail" API call.
+///
 #[serde_as]
 #[derive(Debug, Serialize, Deserialize, Clone)]
 pub struct TimelineInfo {

pageserver/src/http/routes.rs

@@ -11,10 +11,8 @@ use super::models::{
     StatusResponse, TenantConfigRequest, TenantCreateRequest, TenantCreateResponse, TenantInfo,
     TimelineCreateRequest,
 };
-use crate::layered_repository::{metadata::TimelineMetadata, LayeredTimeline};
-use crate::pgdatadir_mapping::DatadirTimeline;
+use crate::layered_repository::{metadata::TimelineMetadata, Timeline};
 use crate::repository::{LocalTimelineState, RepositoryTimeline};
-use crate::repository::{Repository, Timeline};
 use crate::storage_sync;
 use crate::storage_sync::index::{RemoteIndex, RemoteTimeline};
 use crate::tenant_config::TenantConfOpt;
@@ -85,7 +83,7 @@ fn get_config(request: &Request<Body>) -> &'static PageServerConf {
 // Helper functions to construct a LocalTimelineInfo struct for a timeline
 
 fn local_timeline_info_from_loaded_timeline(
-    timeline: &LayeredTimeline,
+    timeline: &Timeline,
     include_non_incremental_logical_size: bool,
     include_non_incremental_physical_size: bool,
 ) -> anyhow::Result<LocalTimelineInfo> {
@@ -160,7 +158,7 @@ fn local_timeline_info_from_unloaded_timeline(metadata: &TimelineMetadata) -> Lo
 }
 
 fn local_timeline_info_from_repo_timeline(
-    repo_timeline: &RepositoryTimeline<LayeredTimeline>,
+    repo_timeline: &RepositoryTimeline<Timeline>,
     include_non_incremental_logical_size: bool,
     include_non_incremental_physical_size: bool,
 ) -> anyhow::Result<LocalTimelineInfo> {
@@ -208,7 +206,6 @@ async fn status_handler(request: Request<Body>) -> Result<Response<Body>, ApiErr
 async fn timeline_create_handler(mut request: Request<Body>) -> Result<Response<Body>, ApiError> {
     let tenant_id: ZTenantId = parse_request_param(&request, "tenant_id")?;
     let request_data: TimelineCreateRequest = json_request(&mut request).await?;
-
     check_permission(&request, Some(tenant_id))?;
 
     let new_timeline_info = tokio::task::spawn_blocking(move || {
@@ -246,11 +243,12 @@ async fn timeline_create_handler(mut request: Request<Body>) -> Result<Response<
 
 async fn timeline_list_handler(request: Request<Body>) -> Result<Response<Body>, ApiError> {
     let tenant_id: ZTenantId = parse_request_param(&request, "tenant_id")?;
-    check_permission(&request, Some(tenant_id))?;
     let include_non_incremental_logical_size =
         query_param_present(&request, "include-non-incremental-logical-size");
     let include_non_incremental_physical_size =
         query_param_present(&request, "include-non-incremental-physical-size");
+    check_permission(&request, Some(tenant_id))?;
+
     let local_timeline_infos = tokio::task::spawn_blocking(move || {
         let _enter = info_span!("timeline_list", tenant = %tenant_id).entered();
         list_local_timelines(
@@ -301,13 +299,12 @@ fn query_param_present(request: &Request<Body>, param: &str) -> bool {
 
 async fn timeline_detail_handler(request: Request<Body>) -> Result<Response<Body>, ApiError> {
     let tenant_id: ZTenantId = parse_request_param(&request, "tenant_id")?;
-    check_permission(&request, Some(tenant_id))?;
-
     let timeline_id: ZTimelineId = parse_request_param(&request, "timeline_id")?;
     let include_non_incremental_logical_size =
         query_param_present(&request, "include-non-incremental-logical-size");
     let include_non_incremental_physical_size =
         query_param_present(&request, "include-non-incremental-physical-size");
+    check_permission(&request, Some(tenant_id))?;
 
     let (local_timeline_info, remote_timeline_info) = async {
         // any error here will render local timeline as None
@@ -371,7 +368,7 @@ async fn tenant_attach_handler(request: Request<Body>) -> Result<Response<Body>,
     let tenant_id: ZTenantId = parse_request_param(&request, "tenant_id")?;
     check_permission(&request, Some(tenant_id))?;
 
-    info!("Handling tenant attach {}", tenant_id,);
+    info!("Handling tenant attach {}", tenant_id);
 
     tokio::task::spawn_blocking(move || {
         if tenant_mgr::get_tenant_state(tenant_id).is_some() {
@@ -480,9 +477,8 @@ async fn gather_tenant_timelines_index_parts(
 
 async fn timeline_delete_handler(request: Request<Body>) -> Result<Response<Body>, ApiError> {
     let tenant_id: ZTenantId = parse_request_param(&request, "tenant_id")?;
-    check_permission(&request, Some(tenant_id))?;
-
     let timeline_id: ZTimelineId = parse_request_param(&request, "timeline_id")?;
+    check_permission(&request, Some(tenant_id))?;
 
     let state = get_state(&request);
     tokio::task::spawn_blocking(move || {
@@ -521,7 +517,6 @@ async fn tenant_detach_handler(request: Request<Body>) -> Result<Response<Body>,
 }
 
 async fn tenant_list_handler(request: Request<Body>) -> Result<Response<Body>, ApiError> {
-    // check for management permission
     check_permission(&request, None)?;
 
     let state = get_state(&request);
@@ -589,7 +584,6 @@ async fn tenant_status(request: Request<Body>) -> Result<Response<Body>, ApiErro
 }
 
 async fn tenant_create_handler(mut request: Request<Body>) -> Result<Response<Body>, ApiError> {
-    // check for management permission
     check_permission(&request, None)?;
 
     let request_data: TenantCreateRequest = json_request(&mut request).await?;
@@ -658,7 +652,6 @@ async fn tenant_create_handler(mut request: Request<Body>) -> Result<Response<Bo
 async fn tenant_config_handler(mut request: Request<Body>) -> Result<Response<Body>, ApiError> {
     let request_data: TenantConfigRequest = json_request(&mut request).await?;
     let tenant_id = request_data.tenant_id;
-    // check for management permission
     check_permission(&request, Some(tenant_id))?;
 
     let mut tenant_conf: TenantConfOpt = Default::default();

pageserver/src/import_datadir.rs

@@ -11,6 +11,7 @@ use bytes::Bytes;
 use tracing::*;
 use walkdir::WalkDir;
 
+use crate::layered_repository::Timeline;
 use crate::pgdatadir_mapping::*;
 use crate::reltag::{RelTag, SlruKind};
 use crate::walingest::WalIngest;
@@ -39,9 +40,9 @@ pub fn get_lsn_from_controlfile(path: &Path) -> Result<Lsn> {
 /// This is currently only used to import a cluster freshly created by initdb.
 /// The code that deals with the checkpoint would not work right if the
 /// cluster was not shut down cleanly.
-pub fn import_timeline_from_postgres_datadir<T: DatadirTimeline>(
+pub fn import_timeline_from_postgres_datadir(
     path: &Path,
-    tline: &T,
+    tline: &Timeline,
     lsn: Lsn,
 ) -> Result<()> {
     let mut pg_control: Option<ControlFileData> = None;
@@ -99,8 +100,8 @@ pub fn import_timeline_from_postgres_datadir<T: DatadirTimeline>(
 }
 
 // subroutine of import_timeline_from_postgres_datadir(), to load one relation file.
-fn import_rel<T: DatadirTimeline, Reader: Read>(
-    modification: &mut DatadirModification<T>,
+fn import_rel<Reader: Read>(
+    modification: &mut DatadirModification,
     path: &Path,
     spcoid: Oid,
     dboid: Oid,
@@ -178,8 +179,8 @@ fn import_rel<T: DatadirTimeline, Reader: Read>(
 
 /// Import an SLRU segment file
 ///
-fn import_slru<T: DatadirTimeline, Reader: Read>(
-    modification: &mut DatadirModification<T>,
+fn import_slru<Reader: Read>(
+    modification: &mut DatadirModification,
     slru: SlruKind,
     path: &Path,
     mut reader: Reader,
@@ -234,12 +235,7 @@ fn import_slru<T: DatadirTimeline, Reader: Read>(
 
 /// Scan PostgreSQL WAL files in given directory and load all records between
 /// 'startpoint' and 'endpoint' into the repository.
-fn import_wal<T: DatadirTimeline>(
-    walpath: &Path,
-    tline: &T,
-    startpoint: Lsn,
-    endpoint: Lsn,
-) -> Result<()> {
+fn import_wal(walpath: &Path, tline: &Timeline, startpoint: Lsn, endpoint: Lsn) -> Result<()> {
     let mut waldecoder = WalStreamDecoder::new(startpoint);
 
     let mut segno = startpoint.segment_number(pg_constants::WAL_SEGMENT_SIZE);
@@ -305,12 +301,12 @@ fn import_wal<T: DatadirTimeline>(
     Ok(())
 }
 
-pub fn import_basebackup_from_tar<T: DatadirTimeline, Reader: Read>(
-    tline: &T,
+pub fn import_basebackup_from_tar<Reader: Read>(
+    tline: &Timeline,
     reader: Reader,
     base_lsn: Lsn,
 ) -> Result<()> {
-    info!("importing base at {}", base_lsn);
+    info!("importing base at {base_lsn}");
     let mut modification = tline.begin_modification(base_lsn);
     modification.init_empty()?;
 
@@ -335,7 +331,11 @@ pub fn import_basebackup_from_tar<T: DatadirTimeline, Reader: Read>(
                 debug!("directory {:?}", file_path);
             }
             _ => {
-                panic!("tar::EntryType::?? {}", file_path.display());
+                bail!(
+                    "entry {} in backup tar archive is of unexpected type: {:?}",
+                    file_path.display(),
+                    header.entry_type()
+                );
             }
         }
     }
@@ -347,8 +347,8 @@ pub fn import_basebackup_from_tar<T: DatadirTimeline, Reader: Read>(
     Ok(())
 }
 
-pub fn import_wal_from_tar<T: DatadirTimeline, Reader: Read>(
-    tline: &T,
+pub fn import_wal_from_tar<Reader: Read>(
+    tline: &Timeline,
     reader: Reader,
     start_lsn: Lsn,
     end_lsn: Lsn,
@@ -388,7 +388,11 @@ pub fn import_wal_from_tar<T: DatadirTimeline, Reader: Read>(
                     continue;
                 }
                 _ => {
-                    panic!("tar::EntryType::?? {}", file_path.display());
+                    bail!(
+                        "entry {} in WAL tar archive is of unexpected type: {:?}",
+                        file_path.display(),
+                        header.entry_type()
+                    );
                 }
             }
         };
@@ -428,14 +432,12 @@ pub fn import_wal_from_tar<T: DatadirTimeline, Reader: Read>(
     Ok(())
 }
 
-pub fn import_file<T: DatadirTimeline, Reader: Read>(
-    modification: &mut DatadirModification<T>,
+fn import_file<Reader: Read>(
+    modification: &mut DatadirModification,
     file_path: &Path,
     reader: Reader,
     len: usize,
 ) -> Result<Option<ControlFileData>> {
-    debug!("looking at {:?}", file_path);
-
     if file_path.starts_with("global") {
         let spcnode = pg_constants::GLOBALTABLESPACE_OID;
         let dbnode = 0;
@@ -557,7 +559,10 @@ pub fn import_file<T: DatadirTimeline, Reader: Read>(
         // this to import arbitrary postgres databases.
         bail!("Importing pg_tblspc is not implemented");
     } else {
-        debug!("ignored");
+        debug!(
+            "ignoring unrecognized file \"{}\" in tar archive",
+            file_path.display()
+        );
     }
 
     Ok(None)

pageserver/src/layered_repository.rs

@@ -13,6 +13,7 @@
 
 use anyhow::{bail, ensure, Context, Result};
 use tracing::*;
+use utils::zid::ZTenantTimelineId;
 
 use std::cmp::min;
 use std::collections::hash_map::Entry;
@@ -31,7 +32,8 @@ use crate::config::PageServerConf;
 use crate::storage_sync::index::RemoteIndex;
 use crate::tenant_config::{TenantConf, TenantConfOpt};
 
-use crate::repository::{GcResult, Repository, RepositoryTimeline, Timeline};
+use crate::repository::{GcResult, RepositoryTimeline};
+use crate::tenant_mgr::LocalTimelineUpdate;
 use crate::thread_mgr;
 use crate::walredo::WalRedoManager;
 use crate::CheckpointConfig;
@@ -61,13 +63,13 @@ mod timeline;
 use storage_layer::Layer;
 use timeline::LayeredTimelineEntry;
 
-pub use timeline::LayeredTimeline;
+pub use timeline::Timeline;
 
 // re-export this function so that page_cache.rs can use it.
 pub use crate::layered_repository::ephemeral_file::writeback as writeback_ephemeral_file;
 
 // re-export for use in storage_sync.rs
-pub use crate::layered_repository::timeline::save_metadata;
+pub use crate::layered_repository::metadata::save_metadata;
 
 // re-export for use in walreceiver
 pub use crate::layered_repository::timeline::WalReceiverInfo;
@@ -78,7 +80,7 @@ pub const TIMELINES_SEGMENT_NAME: &str = "timelines";
 ///
 /// Repository consists of multiple timelines. Keep them in a hash table.
 ///
-pub struct LayeredRepository {
+pub struct Repository {
     // Global pageserver config parameters
     pub conf: &'static PageServerConf,
 
@@ -119,17 +121,22 @@ pub struct LayeredRepository {
     upload_layers: bool,
 }
 
-/// Public interface
-impl Repository for LayeredRepository {
-    type Timeline = LayeredTimeline;
-
-    fn get_timeline(&self, timelineid: ZTimelineId) -> Option<RepositoryTimeline<Self::Timeline>> {
-        let timelines = self.timelines.lock().unwrap();
-        self.get_timeline_internal(timelineid, &timelines)
+/// A repository corresponds to one .neon directory. One repository holds multiple
+/// timelines, forked off from the same initial call to 'initdb'.
+impl Repository {
+    /// Get Timeline handle for given zenith timeline ID.
+    /// This function is idempotent. It doesn't change internal state in any way.
+    pub fn get_timeline(&self, timelineid: ZTimelineId) -> Option<RepositoryTimeline<Timeline>> {
+        self.timelines
+            .lock()
+            .unwrap()
+            .get(&timelineid)
+            .cloned()
             .map(RepositoryTimeline::from)
     }
 
-    fn get_timeline_load(&self, timelineid: ZTimelineId) -> Result<Arc<LayeredTimeline>> {
+    /// Get Timeline handle for locally available timeline. Load it into memory if it is not loaded.
+    pub fn get_timeline_load(&self, timelineid: ZTimelineId) -> Result<Arc<Timeline>> {
         let mut timelines = self.timelines.lock().unwrap();
         match self.get_timeline_load_internal(timelineid, &mut timelines)? {
             Some(local_loaded_timeline) => Ok(local_loaded_timeline),
@@ -140,7 +147,9 @@ impl Repository for LayeredRepository {
         }
     }
 
-    fn list_timelines(&self) -> Vec<(ZTimelineId, RepositoryTimeline<Self::Timeline>)> {
+    /// Lists timelines the repository contains.
+    /// Up to repository's implementation to omit certain timelines that ar not considered ready for use.
+    pub fn list_timelines(&self) -> Vec<(ZTimelineId, RepositoryTimeline<Timeline>)> {
         self.timelines
             .lock()
             .unwrap()
@@ -154,11 +163,13 @@ impl Repository for LayeredRepository {
             .collect()
     }
 
-    fn create_empty_timeline(
+    /// Create a new, empty timeline. The caller is responsible for loading data into it
+    /// Initdb lsn is provided for timeline impl to be able to perform checks for some operations against it.
+    pub fn create_empty_timeline(
         &self,
         timeline_id: ZTimelineId,
         initdb_lsn: Lsn,
-    ) -> Result<Arc<LayeredTimeline>> {
+    ) -> Result<Arc<Timeline>> {
         let mut timelines = self.timelines.lock().unwrap();
         let vacant_timeline_entry = match timelines.entry(timeline_id) {
             Entry::Occupied(_) => bail!("Timeline already exists"),
@@ -174,9 +185,9 @@ impl Repository for LayeredRepository {
         crashsafe_dir::create_dir_all(timeline_path)?;
 
         let metadata = TimelineMetadata::new(Lsn(0), None, None, Lsn(0), initdb_lsn, initdb_lsn);
-        timeline::save_metadata(self.conf, timeline_id, self.tenant_id, &metadata, true)?;
+        save_metadata(self.conf, timeline_id, self.tenant_id, &metadata, true)?;
 
-        let timeline = LayeredTimeline::new(
+        let timeline = Timeline::new(
             self.conf,
             Arc::clone(&self.tenant_conf),
             metadata,
@@ -192,11 +203,16 @@ impl Repository for LayeredRepository {
         let timeline = Arc::new(timeline);
         vacant_timeline_entry.insert(LayeredTimelineEntry::Loaded(Arc::clone(&timeline)));
 
+        crate::tenant_mgr::try_send_timeline_update(LocalTimelineUpdate::Attach {
+            id: ZTenantTimelineId::new(self.tenant_id(), timeline_id),
+            timeline: Arc::clone(&timeline),
+        });
+
         Ok(timeline)
     }
 
     /// Branch a timeline
-    fn branch_timeline(
+    pub fn branch_timeline(
         &self,
         src: ZTimelineId,
         dst: ZTimelineId,
@@ -238,7 +254,8 @@ impl Repository for LayeredRepository {
         src_timeline
             .check_lsn_is_in_scope(start_lsn, &latest_gc_cutoff_lsn)
             .context(format!(
-                "invalid branch start lsn: less than latest GC cutoff {latest_gc_cutoff_lsn}"
+                "invalid branch start lsn: less than latest GC cutoff {}",
+                *latest_gc_cutoff_lsn
             ))?;
         {
             let gc_info = src_timeline.gc_info.read().unwrap();
@@ -274,11 +291,11 @@ impl Repository for LayeredRepository {
             dst_prev,
             Some(src),
             start_lsn,
-            *src_timeline.latest_gc_cutoff_lsn.read().unwrap(),
+            *src_timeline.latest_gc_cutoff_lsn.read(),
             src_timeline.initdb_lsn,
         );
         crashsafe_dir::create_dir_all(self.conf.timeline_path(&dst, &self.tenant_id))?;
-        timeline::save_metadata(self.conf, dst, self.tenant_id, &metadata, true)?;
+        save_metadata(self.conf, dst, self.tenant_id, &metadata, true)?;
         timelines.insert(dst, LayeredTimelineEntry::Unloaded { id: dst, metadata });
 
         info!("branched timeline {} from {} at {}", dst, src, start_lsn);
@@ -286,10 +303,16 @@ impl Repository for LayeredRepository {
         Ok(())
     }
 
-    /// Public entry point to GC. All the logic is in the private
-    /// gc_iteration_internal function, this public facade just wraps it for
-    /// metrics collection.
-    fn gc_iteration(
+    /// perform one garbage collection iteration, removing old data files from disk.
+    /// this function is periodically called by gc thread.
+    /// also it can be explicitly requested through page server api 'do_gc' command.
+    ///
+    /// 'timelineid' specifies the timeline to GC, or None for all.
+    /// `horizon` specifies delta from last lsn to preserve all object versions (pitr interval).
+    /// `checkpoint_before_gc` parameter is used to force compaction of storage before GC
+    /// to make tests more deterministic.
+    /// TODO Do we still need it or we can call checkpoint explicitly in tests where needed?
+    pub fn gc_iteration(
         &self,
         target_timeline_id: Option<ZTimelineId>,
         horizon: u64,
@@ -307,7 +330,11 @@ impl Repository for LayeredRepository {
             })
     }
 
-    fn compaction_iteration(&self) -> Result<()> {
+    /// Perform one compaction iteration.
+    /// This function is periodically called by compactor thread.
+    /// Also it can be explicitly requested per timeline through page server
+    /// api's 'compact' command.
+    pub fn compaction_iteration(&self) -> Result<()> {
         // Scan through the hashmap and collect a list of all the timelines,
         // while holding the lock. Then drop the lock and actually perform the
         // compactions.  We don't want to block everything else while the
@@ -335,12 +362,11 @@ impl Repository for LayeredRepository {
         Ok(())
     }
 
-    ///
     /// Flush all in-memory data to disk.
     ///
-    /// Used at shutdown.
+    /// Used at graceful shutdown.
     ///
-    fn checkpoint(&self) -> Result<()> {
+    pub fn checkpoint(&self) -> Result<()> {
         // Scan through the hashmap and collect a list of all the timelines,
         // while holding the lock. Then drop the lock and actually perform the
         // checkpoints. We don't want to block everything else while the
@@ -370,7 +396,8 @@ impl Repository for LayeredRepository {
         Ok(())
     }
 
-    fn delete_timeline(&self, timeline_id: ZTimelineId) -> anyhow::Result<()> {
+    /// Removes timeline-related in-memory data
+    pub fn delete_timeline(&self, timeline_id: ZTimelineId) -> anyhow::Result<()> {
         // in order to be retriable detach needs to be idempotent
         // (or at least to a point that each time the detach is called it can make progress)
         let mut timelines = self.timelines.lock().unwrap();
@@ -407,7 +434,9 @@ impl Repository for LayeredRepository {
         Ok(())
     }
 
-    fn attach_timeline(&self, timeline_id: ZTimelineId) -> Result<()> {
+    /// Updates timeline based on the `TimelineSyncStatusUpdate`, received from the remote storage synchronization.
+    /// See [`crate::remote_storage`] for more details about the synchronization.
+    pub fn attach_timeline(&self, timeline_id: ZTimelineId) -> Result<()> {
         debug!("attach timeline_id: {}", timeline_id,);
         match self.timelines.lock().unwrap().entry(timeline_id) {
             Entry::Occupied(_) => bail!("We completed a download for a timeline that already exists in repository. This is a bug."),
@@ -421,13 +450,14 @@ impl Repository for LayeredRepository {
         Ok(())
     }
 
-    fn get_remote_index(&self) -> &RemoteIndex {
+    /// Allows to retrieve remote timeline index from the tenant. Used in walreceiver to grab remote consistent lsn.
+    pub fn get_remote_index(&self) -> &RemoteIndex {
         &self.remote_index
     }
 }
 
 /// Private functions
-impl LayeredRepository {
+impl Repository {
     pub fn get_checkpoint_distance(&self) -> u64 {
         let tenant_conf = self.tenant_conf.read().unwrap();
         tenant_conf
@@ -517,49 +547,37 @@ impl LayeredRepository {
 
         tenant_conf.update(&new_tenant_conf);
 
-        LayeredRepository::persist_tenant_config(self.conf, self.tenant_id, *tenant_conf)?;
+        Repository::persist_tenant_config(self.conf, self.tenant_id, *tenant_conf)?;
         Ok(())
     }
 
-    // Implementation of the public `get_timeline` function.
-    // Differences from the public:
-    //  * interface in that the caller must already hold the mutex on the 'timelines' hashmap.
-    fn get_timeline_internal(
-        &self,
-        timelineid: ZTimelineId,
-        timelines: &HashMap<ZTimelineId, LayeredTimelineEntry>,
-    ) -> Option<LayeredTimelineEntry> {
-        timelines.get(&timelineid).cloned()
-    }
-
     // Implementation of the public `get_timeline_load` function.
     // Differences from the public:
     //  * interface in that the caller must already hold the mutex on the 'timelines' hashmap.
     fn get_timeline_load_internal(
         &self,
-        timelineid: ZTimelineId,
+        timeline_id: ZTimelineId,
         timelines: &mut HashMap<ZTimelineId, LayeredTimelineEntry>,
-    ) -> anyhow::Result<Option<Arc<LayeredTimeline>>> {
-        match timelines.get(&timelineid) {
+    ) -> anyhow::Result<Option<Arc<Timeline>>> {
+        match timelines.get(&timeline_id) {
             Some(entry) => match entry {
                 LayeredTimelineEntry::Loaded(local_timeline) => {
-                    debug!("timeline {} found loaded into memory", &timelineid);
+                    debug!("timeline {timeline_id} found loaded into memory");
                     return Ok(Some(Arc::clone(local_timeline)));
                 }
                 LayeredTimelineEntry::Unloaded { .. } => {}
             },
             None => {
-                debug!("timeline {} not found", &timelineid);
+                debug!("timeline {timeline_id} not found");
                 return Ok(None);
             }
         };
         debug!(
-            "timeline {} found on a local disk, but not loaded into the memory, loading",
-            &timelineid
+            "timeline {timeline_id} found on a local disk, but not loaded into the memory, loading"
         );
-        let timeline = self.load_local_timeline(timelineid, timelines)?;
+        let timeline = self.load_local_timeline(timeline_id, timelines)?;
         let was_loaded = timelines.insert(
-            timelineid,
+            timeline_id,
             LayeredTimelineEntry::Loaded(Arc::clone(&timeline)),
         );
         ensure!(
@@ -574,7 +592,7 @@ impl LayeredRepository {
         &self,
         timeline_id: ZTimelineId,
         timelines: &mut HashMap<ZTimelineId, LayeredTimelineEntry>,
-    ) -> anyhow::Result<Arc<LayeredTimeline>> {
+    ) -> anyhow::Result<Arc<Timeline>> {
         let metadata = load_metadata(self.conf, timeline_id, self.tenant_id)
             .context("failed to load metadata")?;
         let disk_consistent_lsn = metadata.disk_consistent_lsn();
@@ -591,7 +609,7 @@ impl LayeredRepository {
             .map(LayeredTimelineEntry::Loaded);
         let _enter = info_span!("loading local timeline").entered();
 
-        let timeline = LayeredTimeline::new(
+        let timeline = Timeline::new(
             self.conf,
             Arc::clone(&self.tenant_conf),
             metadata,
@@ -605,7 +623,14 @@ impl LayeredRepository {
             .load_layer_map(disk_consistent_lsn)
             .context("failed to load layermap")?;
 
-        Ok(Arc::new(timeline))
+        let timeline = Arc::new(timeline);
+
+        crate::tenant_mgr::try_send_timeline_update(LocalTimelineUpdate::Attach {
+            id: ZTenantTimelineId::new(self.tenant_id(), timeline_id),
+            timeline: Arc::clone(&timeline),
+        });
+
+        Ok(timeline)
     }
 
     pub fn new(
@@ -615,8 +640,8 @@ impl LayeredRepository {
         tenant_id: ZTenantId,
         remote_index: RemoteIndex,
         upload_layers: bool,
-    ) -> LayeredRepository {
-        LayeredRepository {
+    ) -> Repository {
+        Repository {
             tenant_id,
             file_lock: RwLock::new(()),
             conf,
@@ -632,9 +657,9 @@ impl LayeredRepository {
     /// Locate and load config
     pub fn load_tenant_config(
         conf: &'static PageServerConf,
-        tenantid: ZTenantId,
+        tenant_id: ZTenantId,
     ) -> anyhow::Result<TenantConfOpt> {
-        let target_config_path = TenantConf::path(conf, tenantid);
+        let target_config_path = TenantConf::path(conf, tenant_id);
 
         info!("load tenantconf from {}", target_config_path.display());
 
@@ -669,11 +694,11 @@ impl LayeredRepository {
 
     pub fn persist_tenant_config(
         conf: &'static PageServerConf,
-        tenantid: ZTenantId,
+        tenant_id: ZTenantId,
         tenant_conf: TenantConfOpt,
     ) -> anyhow::Result<()> {
         let _enter = info_span!("saving tenantconf").entered();
-        let target_config_path = TenantConf::path(conf, tenantid);
+        let target_config_path = TenantConf::path(conf, tenant_id);
         info!("save tenantconf to {}", target_config_path.display());
 
         let mut conf_content = r#"# This file contains a specific per-tenant's config.
@@ -810,7 +835,7 @@ impl LayeredRepository {
         // compaction (both require `layer_removal_cs` lock),
         // but the GC iteration can run concurrently with branch creation.
         //
-        // See comments in [`LayeredRepository::branch_timeline`] for more information
+        // See comments in [`Repository::branch_timeline`] for more information
         // about why branch creation task can run concurrently with timeline's GC iteration.
         for timeline in gc_timelines {
             if thread_mgr::is_shutdown_requested() {
@@ -886,22 +911,525 @@ pub fn load_metadata(
     })
 }
 
-///
-/// 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)]
-pub mod tests {
+pub mod repo_harness {
+    use bytes::{Bytes, BytesMut};
+    use once_cell::sync::Lazy;
+    use std::sync::{Arc, RwLock, RwLockReadGuard, RwLockWriteGuard};
+    use std::{fs, path::PathBuf};
+    use utils::lsn::Lsn;
+
+    use crate::storage_sync::index::RemoteIndex;
+    use crate::{
+        config::PageServerConf,
+        layered_repository::Repository,
+        repository::Key,
+        walrecord::ZenithWalRecord,
+        walredo::{WalRedoError, WalRedoManager},
+    };
+
+    use super::*;
+    use crate::tenant_config::{TenantConf, TenantConfOpt};
+    use hex_literal::hex;
+    use utils::zid::{ZTenantId, ZTimelineId};
+
+    pub const TIMELINE_ID: ZTimelineId =
+        ZTimelineId::from_array(hex!("11223344556677881122334455667788"));
+    pub const NEW_TIMELINE_ID: ZTimelineId =
+        ZTimelineId::from_array(hex!("AA223344556677881122334455667788"));
+
+    /// Convenience function to create a page image with given string as the only content
+    #[allow(non_snake_case)]
+    pub fn TEST_IMG(s: &str) -> Bytes {
+        let mut buf = BytesMut::new();
+        buf.extend_from_slice(s.as_bytes());
+        buf.resize(64, 0);
+
+        buf.freeze()
+    }
+
+    static LOCK: Lazy<RwLock<()>> = Lazy::new(|| RwLock::new(()));
+
+    impl From<TenantConf> for TenantConfOpt {
+        fn from(tenant_conf: TenantConf) -> Self {
+            Self {
+                checkpoint_distance: Some(tenant_conf.checkpoint_distance),
+                checkpoint_timeout: Some(tenant_conf.checkpoint_timeout),
+                compaction_target_size: Some(tenant_conf.compaction_target_size),
+                compaction_period: Some(tenant_conf.compaction_period),
+                compaction_threshold: Some(tenant_conf.compaction_threshold),
+                gc_horizon: Some(tenant_conf.gc_horizon),
+                gc_period: Some(tenant_conf.gc_period),
+                image_creation_threshold: Some(tenant_conf.image_creation_threshold),
+                pitr_interval: Some(tenant_conf.pitr_interval),
+                walreceiver_connect_timeout: Some(tenant_conf.walreceiver_connect_timeout),
+                lagging_wal_timeout: Some(tenant_conf.lagging_wal_timeout),
+                max_lsn_wal_lag: Some(tenant_conf.max_lsn_wal_lag),
+            }
+        }
+    }
+
+    pub struct RepoHarness<'a> {
+        pub conf: &'static PageServerConf,
+        pub tenant_conf: TenantConf,
+        pub tenant_id: ZTenantId,
+
+        pub lock_guard: (
+            Option<RwLockReadGuard<'a, ()>>,
+            Option<RwLockWriteGuard<'a, ()>>,
+        ),
+    }
+
+    impl<'a> RepoHarness<'a> {
+        pub fn create(test_name: &'static str) -> Result<Self> {
+            Self::create_internal(test_name, false)
+        }
+        pub fn create_exclusive(test_name: &'static str) -> Result<Self> {
+            Self::create_internal(test_name, true)
+        }
+        fn create_internal(test_name: &'static str, exclusive: bool) -> Result<Self> {
+            let lock_guard = if exclusive {
+                (None, Some(LOCK.write().unwrap()))
+            } else {
+                (Some(LOCK.read().unwrap()), None)
+            };
+
+            let repo_dir = PageServerConf::test_repo_dir(test_name);
+            let _ = fs::remove_dir_all(&repo_dir);
+            fs::create_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 tenant_conf = TenantConf::dummy_conf();
+
+            let tenant_id = ZTenantId::generate();
+            fs::create_dir_all(conf.tenant_path(&tenant_id))?;
+            fs::create_dir_all(conf.timelines_path(&tenant_id))?;
+
+            Ok(Self {
+                conf,
+                tenant_conf,
+                tenant_id,
+                lock_guard,
+            })
+        }
+
+        pub fn load(&self) -> Repository {
+            self.try_load().expect("failed to load test repo")
+        }
+
+        pub fn try_load(&self) -> Result<Repository> {
+            let walredo_mgr = Arc::new(TestRedoManager);
+
+            let repo = Repository::new(
+                self.conf,
+                TenantConfOpt::from(self.tenant_conf),
+                walredo_mgr,
+                self.tenant_id,
+                RemoteIndex::default(),
+                false,
+            );
+            // populate repo with locally available timelines
+            for timeline_dir_entry in fs::read_dir(self.conf.timelines_path(&self.tenant_id))
+                .expect("should be able to read timelines dir")
+            {
+                let timeline_dir_entry = timeline_dir_entry.unwrap();
+                let timeline_id: ZTimelineId = timeline_dir_entry
+                    .path()
+                    .file_name()
+                    .unwrap()
+                    .to_string_lossy()
+                    .parse()
+                    .unwrap();
+
+                repo.attach_timeline(timeline_id)?;
+            }
+
+            Ok(repo)
+        }
+
+        pub fn timeline_path(&self, timeline_id: &ZTimelineId) -> PathBuf {
+            self.conf.timeline_path(timeline_id, &self.tenant_id)
+        }
+    }
+
+    // Mock WAL redo manager that doesn't do much
+    pub struct TestRedoManager;
+
+    impl WalRedoManager for TestRedoManager {
+        fn request_redo(
+            &self,
+            key: Key,
+            lsn: Lsn,
+            base_img: Option<Bytes>,
+            records: Vec<(Lsn, ZenithWalRecord)>,
+        ) -> Result<Bytes, WalRedoError> {
+            let s = format!(
+                "redo for {} to get to {}, with {} and {} records",
+                key,
+                lsn,
+                if base_img.is_some() {
+                    "base image"
+                } else {
+                    "no base image"
+                },
+                records.len()
+            );
+            println!("{}", s);
+
+            Ok(TEST_IMG(&s))
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
     use super::metadata::METADATA_FILE_NAME;
     use super::*;
     use crate::keyspace::KeySpaceAccum;
-    use crate::repository::repo_harness::*;
+    use crate::layered_repository::repo_harness::*;
     use crate::repository::{Key, Value};
+    use bytes::BytesMut;
+    use hex_literal::hex;
+    use once_cell::sync::Lazy;
     use rand::{thread_rng, Rng};
 
+    static TEST_KEY: Lazy<Key> =
+        Lazy::new(|| Key::from_slice(&hex!("112222222233333333444444445500000001")));
+
+    #[test]
+    fn test_basic() -> Result<()> {
+        let repo = RepoHarness::create("test_basic")?.load();
+        let tline = repo.create_empty_timeline(TIMELINE_ID, Lsn(0))?;
+
+        let writer = tline.writer();
+        writer.put(*TEST_KEY, Lsn(0x10), &Value::Image(TEST_IMG("foo at 0x10")))?;
+        writer.finish_write(Lsn(0x10));
+        drop(writer);
+
+        let writer = tline.writer();
+        writer.put(*TEST_KEY, Lsn(0x20), &Value::Image(TEST_IMG("foo at 0x20")))?;
+        writer.finish_write(Lsn(0x20));
+        drop(writer);
+
+        assert_eq!(tline.get(*TEST_KEY, Lsn(0x10))?, TEST_IMG("foo at 0x10"));
+        assert_eq!(tline.get(*TEST_KEY, Lsn(0x1f))?, TEST_IMG("foo at 0x10"));
+        assert_eq!(tline.get(*TEST_KEY, Lsn(0x20))?, TEST_IMG("foo at 0x20"));
+
+        Ok(())
+    }
+
+    #[test]
+    fn no_duplicate_timelines() -> Result<()> {
+        let repo = RepoHarness::create("no_duplicate_timelines")?.load();
+        let _ = repo.create_empty_timeline(TIMELINE_ID, Lsn(0))?;
+
+        match repo.create_empty_timeline(TIMELINE_ID, Lsn(0)) {
+            Ok(_) => panic!("duplicate timeline creation should fail"),
+            Err(e) => assert_eq!(e.to_string(), "Timeline already exists"),
+        }
+
+        Ok(())
+    }
+
+    /// Convenience function to create a page image with given string as the only content
+    pub fn test_value(s: &str) -> Value {
+        let mut buf = BytesMut::new();
+        buf.extend_from_slice(s.as_bytes());
+        Value::Image(buf.freeze())
+    }
+
+    ///
+    /// Test branch creation
+    ///
+    #[test]
+    fn test_branch() -> Result<()> {
+        let repo = RepoHarness::create("test_branch")?.load();
+        let tline = repo.create_empty_timeline(TIMELINE_ID, Lsn(0))?;
+        let writer = tline.writer();
+        use std::str::from_utf8;
+
+        #[allow(non_snake_case)]
+        let TEST_KEY_A: Key = Key::from_hex("112222222233333333444444445500000001").unwrap();
+        #[allow(non_snake_case)]
+        let TEST_KEY_B: Key = Key::from_hex("112222222233333333444444445500000002").unwrap();
+
+        // Insert a value on the timeline
+        writer.put(TEST_KEY_A, Lsn(0x20), &test_value("foo at 0x20"))?;
+        writer.put(TEST_KEY_B, Lsn(0x20), &test_value("foobar at 0x20"))?;
+        writer.finish_write(Lsn(0x20));
+
+        writer.put(TEST_KEY_A, Lsn(0x30), &test_value("foo at 0x30"))?;
+        writer.finish_write(Lsn(0x30));
+        writer.put(TEST_KEY_A, Lsn(0x40), &test_value("foo at 0x40"))?;
+        writer.finish_write(Lsn(0x40));
+
+        //assert_current_logical_size(&tline, Lsn(0x40));
+
+        // Branch the history, modify relation differently on the new timeline
+        repo.branch_timeline(TIMELINE_ID, NEW_TIMELINE_ID, Some(Lsn(0x30)))?;
+        let newtline = repo
+            .get_timeline_load(NEW_TIMELINE_ID)
+            .expect("Should have a local timeline");
+        let new_writer = newtline.writer();
+        new_writer.put(TEST_KEY_A, Lsn(0x40), &test_value("bar at 0x40"))?;
+        new_writer.finish_write(Lsn(0x40));
+
+        // Check page contents on both branches
+        assert_eq!(
+            from_utf8(&tline.get(TEST_KEY_A, Lsn(0x40))?)?,
+            "foo at 0x40"
+        );
+        assert_eq!(
+            from_utf8(&newtline.get(TEST_KEY_A, Lsn(0x40))?)?,
+            "bar at 0x40"
+        );
+        assert_eq!(
+            from_utf8(&newtline.get(TEST_KEY_B, Lsn(0x40))?)?,
+            "foobar at 0x20"
+        );
+
+        //assert_current_logical_size(&tline, Lsn(0x40));
+
+        Ok(())
+    }
+
+    fn make_some_layers(tline: &Timeline, start_lsn: Lsn) -> Result<()> {
+        let mut lsn = start_lsn;
+        #[allow(non_snake_case)]
+        {
+            let writer = tline.writer();
+            // Create a relation on the timeline
+            writer.put(
+                *TEST_KEY,
+                lsn,
+                &Value::Image(TEST_IMG(&format!("foo at {}", lsn))),
+            )?;
+            writer.finish_write(lsn);
+            lsn += 0x10;
+            writer.put(
+                *TEST_KEY,
+                lsn,
+                &Value::Image(TEST_IMG(&format!("foo at {}", lsn))),
+            )?;
+            writer.finish_write(lsn);
+            lsn += 0x10;
+        }
+        tline.checkpoint(CheckpointConfig::Forced)?;
+        {
+            let writer = tline.writer();
+            writer.put(
+                *TEST_KEY,
+                lsn,
+                &Value::Image(TEST_IMG(&format!("foo at {}", lsn))),
+            )?;
+            writer.finish_write(lsn);
+            lsn += 0x10;
+            writer.put(
+                *TEST_KEY,
+                lsn,
+                &Value::Image(TEST_IMG(&format!("foo at {}", lsn))),
+            )?;
+            writer.finish_write(lsn);
+        }
+        tline.checkpoint(CheckpointConfig::Forced)
+    }
+
+    #[test]
+    fn test_prohibit_branch_creation_on_garbage_collected_data() -> Result<()> {
+        let repo =
+            RepoHarness::create("test_prohibit_branch_creation_on_garbage_collected_data")?.load();
+        let tline = repo.create_empty_timeline(TIMELINE_ID, Lsn(0))?;
+        make_some_layers(tline.as_ref(), Lsn(0x20))?;
+
+        // this removes layers before lsn 40 (50 minus 10), so there are two remaining layers, image and delta for 31-50
+        // FIXME: this doesn't actually remove any layer currently, given how the checkpointing
+        // and compaction works. But it does set the 'cutoff' point so that the cross check
+        // below should fail.
+        repo.gc_iteration(Some(TIMELINE_ID), 0x10, Duration::ZERO, false)?;
+
+        // try to branch at lsn 25, should fail because we already garbage collected the data
+        match repo.branch_timeline(TIMELINE_ID, NEW_TIMELINE_ID, Some(Lsn(0x25))) {
+            Ok(_) => panic!("branching should have failed"),
+            Err(err) => {
+                assert!(err.to_string().contains("invalid branch start lsn"));
+                assert!(err
+                    .source()
+                    .unwrap()
+                    .to_string()
+                    .contains("we might've already garbage collected needed data"))
+            }
+        }
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_prohibit_branch_creation_on_pre_initdb_lsn() -> Result<()> {
+        let repo = RepoHarness::create("test_prohibit_branch_creation_on_pre_initdb_lsn")?.load();
+
+        repo.create_empty_timeline(TIMELINE_ID, Lsn(0x50))?;
+        // try to branch at lsn 0x25, should fail because initdb lsn is 0x50
+        match repo.branch_timeline(TIMELINE_ID, NEW_TIMELINE_ID, Some(Lsn(0x25))) {
+            Ok(_) => panic!("branching should have failed"),
+            Err(err) => {
+                assert!(&err.to_string().contains("invalid branch start lsn"));
+                assert!(&err
+                    .source()
+                    .unwrap()
+                    .to_string()
+                    .contains("is earlier than latest GC horizon"));
+            }
+        }
+
+        Ok(())
+    }
+
+    /*
+    // FIXME: This currently fails to error out. Calling GC doesn't currently
+    // remove the old value, we'd need to work a little harder
+    #[test]
+    fn test_prohibit_get_for_garbage_collected_data() -> Result<()> {
+        let repo =
+            RepoHarness::create("test_prohibit_get_for_garbage_collected_data")?
+            .load();
+
+        let tline = repo.create_empty_timeline(TIMELINE_ID, Lsn(0))?;
+        make_some_layers(tline.as_ref(), Lsn(0x20))?;
+
+        repo.gc_iteration(Some(TIMELINE_ID), 0x10, Duration::ZERO, false)?;
+        let latest_gc_cutoff_lsn = tline.get_latest_gc_cutoff_lsn();
+        assert!(*latest_gc_cutoff_lsn > Lsn(0x25));
+        match tline.get(*TEST_KEY, Lsn(0x25)) {
+            Ok(_) => panic!("request for page should have failed"),
+            Err(err) => assert!(err.to_string().contains("not found at")),
+        }
+        Ok(())
+    }
+     */
+
+    #[test]
+    fn test_retain_data_in_parent_which_is_needed_for_child() -> Result<()> {
+        let repo =
+            RepoHarness::create("test_retain_data_in_parent_which_is_needed_for_child")?.load();
+        let tline = repo.create_empty_timeline(TIMELINE_ID, Lsn(0))?;
+        make_some_layers(tline.as_ref(), Lsn(0x20))?;
+
+        repo.branch_timeline(TIMELINE_ID, NEW_TIMELINE_ID, Some(Lsn(0x40)))?;
+        let newtline = repo
+            .get_timeline_load(NEW_TIMELINE_ID)
+            .expect("Should have a local timeline");
+        // this removes layers before lsn 40 (50 minus 10), so there are two remaining layers, image and delta for 31-50
+        repo.gc_iteration(Some(TIMELINE_ID), 0x10, Duration::ZERO, false)?;
+        assert!(newtline.get(*TEST_KEY, Lsn(0x25)).is_ok());
+
+        Ok(())
+    }
+    #[test]
+    fn test_parent_keeps_data_forever_after_branching() -> Result<()> {
+        let repo = RepoHarness::create("test_parent_keeps_data_forever_after_branching")?.load();
+        let tline = repo.create_empty_timeline(TIMELINE_ID, Lsn(0))?;
+        make_some_layers(tline.as_ref(), Lsn(0x20))?;
+
+        repo.branch_timeline(TIMELINE_ID, NEW_TIMELINE_ID, Some(Lsn(0x40)))?;
+        let newtline = repo
+            .get_timeline_load(NEW_TIMELINE_ID)
+            .expect("Should have a local timeline");
+
+        make_some_layers(newtline.as_ref(), Lsn(0x60))?;
+
+        // run gc on parent
+        repo.gc_iteration(Some(TIMELINE_ID), 0x10, Duration::ZERO, false)?;
+
+        // Check that the data is still accessible on the branch.
+        assert_eq!(
+            newtline.get(*TEST_KEY, Lsn(0x50))?,
+            TEST_IMG(&format!("foo at {}", Lsn(0x40)))
+        );
+
+        Ok(())
+    }
+
+    #[test]
+    fn timeline_load() -> Result<()> {
+        const TEST_NAME: &str = "timeline_load";
+        let harness = RepoHarness::create(TEST_NAME)?;
+        {
+            let repo = harness.load();
+            let tline = repo.create_empty_timeline(TIMELINE_ID, Lsn(0x8000))?;
+            make_some_layers(tline.as_ref(), Lsn(0x8000))?;
+            tline.checkpoint(CheckpointConfig::Forced)?;
+        }
+
+        let repo = harness.load();
+        let tline = repo
+            .get_timeline(TIMELINE_ID)
+            .expect("cannot load timeline");
+        assert!(matches!(tline, RepositoryTimeline::Unloaded { .. }));
+
+        assert!(repo.get_timeline_load(TIMELINE_ID).is_ok());
+
+        let tline = repo
+            .get_timeline(TIMELINE_ID)
+            .expect("cannot load timeline");
+        assert!(matches!(tline, RepositoryTimeline::Loaded(_)));
+
+        Ok(())
+    }
+
+    #[test]
+    fn timeline_load_with_ancestor() -> Result<()> {
+        const TEST_NAME: &str = "timeline_load_with_ancestor";
+        let harness = RepoHarness::create(TEST_NAME)?;
+        // create two timelines
+        {
+            let repo = harness.load();
+            let tline = repo.create_empty_timeline(TIMELINE_ID, Lsn(0))?;
+
+            make_some_layers(tline.as_ref(), Lsn(0x20))?;
+            tline.checkpoint(CheckpointConfig::Forced)?;
+
+            repo.branch_timeline(TIMELINE_ID, NEW_TIMELINE_ID, Some(Lsn(0x40)))?;
+
+            let newtline = repo
+                .get_timeline_load(NEW_TIMELINE_ID)
+                .expect("Should have a local timeline");
+
+            make_some_layers(newtline.as_ref(), Lsn(0x60))?;
+            tline.checkpoint(CheckpointConfig::Forced)?;
+        }
+
+        // check that both of them are initially unloaded
+        let repo = harness.load();
+        {
+            let tline = repo.get_timeline(TIMELINE_ID).expect("cannot get timeline");
+            assert!(matches!(tline, RepositoryTimeline::Unloaded { .. }));
+
+            let tline = repo
+                .get_timeline(NEW_TIMELINE_ID)
+                .expect("cannot get timeline");
+            assert!(matches!(tline, RepositoryTimeline::Unloaded { .. }));
+        }
+        // load only child timeline
+        let _ = repo
+            .get_timeline_load(NEW_TIMELINE_ID)
+            .expect("cannot load timeline");
+
+        // check that both, child and ancestor are loaded
+        let tline = repo
+            .get_timeline(NEW_TIMELINE_ID)
+            .expect("cannot get timeline");
+        assert!(matches!(tline, RepositoryTimeline::Loaded(_)));
+
+        let tline = repo.get_timeline(TIMELINE_ID).expect("cannot get timeline");
+        assert!(matches!(tline, RepositoryTimeline::Loaded(_)));
+
+        Ok(())
+    }
+
     #[test]
     fn corrupt_metadata() -> Result<()> {
         const TEST_NAME: &str = "corrupt_metadata";
@@ -940,22 +1468,13 @@ pub mod tests {
         Ok(())
     }
 
-    // Target file size in the unit tests. In production, the target
-    // file size is much larger, maybe 1 GB. But a small size makes it
-    // much faster to exercise all the logic for creating the files,
-    // garbage collection, compaction etc.
-    pub const TEST_FILE_SIZE: u64 = 4 * 1024 * 1024;
-
     #[test]
     fn test_images() -> Result<()> {
         let repo = RepoHarness::create("test_images")?.load();
         let tline = repo.create_empty_timeline(TIMELINE_ID, Lsn(0))?;
 
-        #[allow(non_snake_case)]
-        let TEST_KEY: Key = Key::from_hex("112222222233333333444444445500000001").unwrap();
-
         let writer = tline.writer();
-        writer.put(TEST_KEY, Lsn(0x10), &Value::Image(TEST_IMG("foo at 0x10")))?;
+        writer.put(*TEST_KEY, Lsn(0x10), &Value::Image(TEST_IMG("foo at 0x10")))?;
         writer.finish_write(Lsn(0x10));
         drop(writer);
 
@@ -963,7 +1482,7 @@ pub mod tests {
         tline.compact()?;
 
         let writer = tline.writer();
-        writer.put(TEST_KEY, Lsn(0x20), &Value::Image(TEST_IMG("foo at 0x20")))?;
+        writer.put(*TEST_KEY, Lsn(0x20), &Value::Image(TEST_IMG("foo at 0x20")))?;
         writer.finish_write(Lsn(0x20));
         drop(writer);
 
@@ -971,7 +1490,7 @@ pub mod tests {
         tline.compact()?;
 
         let writer = tline.writer();
-        writer.put(TEST_KEY, Lsn(0x30), &Value::Image(TEST_IMG("foo at 0x30")))?;
+        writer.put(*TEST_KEY, Lsn(0x30), &Value::Image(TEST_IMG("foo at 0x30")))?;
         writer.finish_write(Lsn(0x30));
         drop(writer);
 
@@ -979,18 +1498,18 @@ pub mod tests {
         tline.compact()?;
 
         let writer = tline.writer();
-        writer.put(TEST_KEY, Lsn(0x40), &Value::Image(TEST_IMG("foo at 0x40")))?;
+        writer.put(*TEST_KEY, Lsn(0x40), &Value::Image(TEST_IMG("foo at 0x40")))?;
         writer.finish_write(Lsn(0x40));
         drop(writer);
 
         tline.checkpoint(CheckpointConfig::Forced)?;
         tline.compact()?;
 
-        assert_eq!(tline.get(TEST_KEY, Lsn(0x10))?, TEST_IMG("foo at 0x10"));
-        assert_eq!(tline.get(TEST_KEY, Lsn(0x1f))?, TEST_IMG("foo at 0x10"));
-        assert_eq!(tline.get(TEST_KEY, Lsn(0x20))?, TEST_IMG("foo at 0x20"));
-        assert_eq!(tline.get(TEST_KEY, Lsn(0x30))?, TEST_IMG("foo at 0x30"));
-        assert_eq!(tline.get(TEST_KEY, Lsn(0x40))?, TEST_IMG("foo at 0x40"));
+        assert_eq!(tline.get(*TEST_KEY, Lsn(0x10))?, TEST_IMG("foo at 0x10"));
+        assert_eq!(tline.get(*TEST_KEY, Lsn(0x1f))?, TEST_IMG("foo at 0x10"));
+        assert_eq!(tline.get(*TEST_KEY, Lsn(0x20))?, TEST_IMG("foo at 0x20"));
+        assert_eq!(tline.get(*TEST_KEY, Lsn(0x30))?, TEST_IMG("foo at 0x30"));
+        assert_eq!(tline.get(*TEST_KEY, Lsn(0x40))?, TEST_IMG("foo at 0x40"));
 
         Ok(())
     }

pageserver/src/layered_repository/filename.rs

@@ -10,7 +10,7 @@ use std::path::PathBuf;
 
 use utils::lsn::Lsn;
 
-// Note: LayeredTimeline::load_layer_map() relies on this sort order
+// Note: Timeline::load_layer_map() relies on this sort order
 #[derive(Debug, PartialEq, Eq, Clone)]
 pub struct DeltaFileName {
     pub key_range: Range<Key>,

pageserver/src/layered_repository/metadata.rs

@@ -1,4 +1,4 @@
-//! Every image of a certain timeline from [`crate::layered_repository::LayeredRepository`]
+//! Every image of a certain timeline from [`crate::layered_repository::Repository`]
 //! has a metadata that needs to be stored persistently.
 //!
 //! Later, the file gets is used in [`crate::remote_storage::storage_sync`] as a part of
@@ -6,10 +6,13 @@
 //!
 //! The module contains all structs and related helper methods related to timeline metadata.
 
+use std::fs::{File, OpenOptions};
+use std::io::Write;
 use std::path::PathBuf;
 
-use anyhow::ensure;
+use anyhow::{bail, ensure, Context};
 use serde::{Deserialize, Serialize};
+use tracing::info_span;
 use utils::{
     bin_ser::BeSer,
     lsn::Lsn,
@@ -17,6 +20,7 @@ use utils::{
 };
 
 use crate::config::PageServerConf;
+use crate::virtual_file::VirtualFile;
 use crate::STORAGE_FORMAT_VERSION;
 
 /// We assume that a write of up to METADATA_MAX_SIZE bytes is atomic.
@@ -30,7 +34,7 @@ pub const METADATA_FILE_NAME: &str = "metadata";
 
 /// Metadata stored on disk for each timeline
 ///
-/// The fields correspond to the values we hold in memory, in LayeredTimeline.
+/// The fields correspond to the values we hold in memory, in Timeline.
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct TimelineMetadata {
     hdr: TimelineMetadataHeader,
@@ -65,17 +69,6 @@ struct TimelineMetadataBody {
     initdb_lsn: Lsn,
 }
 
-/// Points to a place in pageserver's local directory,
-/// where certain timeline's metadata file should be located.
-pub fn metadata_path(
-    conf: &'static PageServerConf,
-    timelineid: ZTimelineId,
-    tenantid: ZTenantId,
-) -> PathBuf {
-    conf.timeline_path(&timelineid, &tenantid)
-        .join(METADATA_FILE_NAME)
-}
-
 impl TimelineMetadata {
     pub fn new(
         disk_consistent_lsn: Lsn,
@@ -173,11 +166,57 @@ impl TimelineMetadata {
     }
 }
 
+/// Points to a place in pageserver's local directory,
+/// where certain timeline's metadata file should be located.
+pub fn metadata_path(
+    conf: &'static PageServerConf,
+    timelineid: ZTimelineId,
+    tenantid: ZTenantId,
+) -> PathBuf {
+    conf.timeline_path(&timelineid, &tenantid)
+        .join(METADATA_FILE_NAME)
+}
+
+/// Save timeline metadata to file
+pub fn save_metadata(
+    conf: &'static PageServerConf,
+    timelineid: ZTimelineId,
+    tenantid: ZTenantId,
+    data: &TimelineMetadata,
+    first_save: bool,
+) -> anyhow::Result<()> {
+    let _enter = info_span!("saving metadata").entered();
+    let path = metadata_path(conf, timelineid, tenantid);
+    // use OpenOptions to ensure file presence is consistent with first_save
+    let mut file = VirtualFile::open_with_options(
+        &path,
+        OpenOptions::new().write(true).create_new(first_save),
+    )?;
+
+    let metadata_bytes = data.to_bytes().context("Failed to get metadata bytes")?;
+
+    if file.write(&metadata_bytes)? != metadata_bytes.len() {
+        bail!("Could not write all the metadata bytes in a single call");
+    }
+    file.sync_all()?;
+
+    // fsync the parent directory to ensure the directory entry is durable
+    if first_save {
+        let timeline_dir = File::open(
+            &path
+                .parent()
+                .expect("Metadata should always have a parent dir"),
+        )?;
+        timeline_dir.sync_all()?;
+    }
+
+    Ok(())
+}
+
 #[cfg(test)]
 mod tests {
-    use crate::repository::repo_harness::TIMELINE_ID;
-
     use super::*;
+    use crate::layered_repository::repo_harness::TIMELINE_ID;
 
     #[test]
     fn metadata_serializes_correctly() {

pageserver/src/layered_repository/timeline.rs

@@ -9,14 +9,12 @@ use once_cell::sync::Lazy;
 use tracing::*;
 
 use std::cmp::{max, min, Ordering};
-use std::collections::{hash_map::Entry, HashMap, HashSet};
+use std::collections::{HashMap, HashSet};
 use std::fs;
-use std::fs::{File, OpenOptions};
-use std::io::Write;
 use std::ops::{Deref, Range};
 use std::path::PathBuf;
-use std::sync::atomic::{self, AtomicBool, AtomicIsize, Ordering as AtomicOrdering};
-use std::sync::{Arc, Mutex, MutexGuard, RwLock, RwLockReadGuard, TryLockError};
+use std::sync::atomic::{self, AtomicBool, AtomicI64, Ordering as AtomicOrdering};
+use std::sync::{Arc, Mutex, MutexGuard, RwLock, TryLockError};
 use std::time::{Duration, Instant, SystemTime};
 
 use metrics::{
@@ -32,7 +30,7 @@ use crate::layered_repository::{
     image_layer::{ImageLayer, ImageLayerWriter},
     inmemory_layer::InMemoryLayer,
     layer_map::{LayerMap, SearchResult},
-    metadata::{metadata_path, TimelineMetadata, METADATA_FILE_NAME},
+    metadata::{save_metadata, TimelineMetadata, METADATA_FILE_NAME},
     par_fsync,
     storage_layer::{Layer, ValueReconstructResult, ValueReconstructState},
 };
@@ -43,19 +41,18 @@ use crate::pgdatadir_mapping::BlockNumber;
 use crate::pgdatadir_mapping::LsnForTimestamp;
 use crate::reltag::RelTag;
 use crate::tenant_config::TenantConfOpt;
-use crate::DatadirTimeline;
 
 use postgres_ffi::v14::xlog_utils::to_pg_timestamp;
 use utils::{
     lsn::{AtomicLsn, Lsn, RecordLsn},
     seqwait::SeqWait,
+    simple_rcu::{Rcu, RcuReadGuard},
     zid::{ZTenantId, ZTimelineId},
 };
 
-use crate::repository::{GcResult, RepositoryTimeline, Timeline, TimelineWriter};
+use crate::repository::{GcResult, RepositoryTimeline};
 use crate::repository::{Key, Value};
 use crate::thread_mgr;
-use crate::virtual_file::VirtualFile;
 use crate::walreceiver::IS_WAL_RECEIVER;
 use crate::walredo::WalRedoManager;
 use crate::CheckpointConfig;
@@ -140,6 +137,15 @@ static CURRENT_PHYSICAL_SIZE: Lazy<UIntGaugeVec> = Lazy::new(|| {
     .expect("failed to define a metric")
 });
 
+static CURRENT_LOGICAL_SIZE: Lazy<IntGaugeVec> = Lazy::new(|| {
+    register_int_gauge_vec!(
+        "pageserver_current_logical_size",
+        "Current logical size grouped by timeline",
+        &["tenant_id", "timeline_id"]
+    )
+    .expect("failed to define a metric")
+});
+
 // Metrics for cloud upload. These metrics reflect data uploaded to cloud storage,
 // or in testing they estimate how much we would upload if we did.
 static NUM_PERSISTENT_FILES_CREATED: Lazy<IntCounter> = Lazy::new(|| {
@@ -160,7 +166,7 @@ static PERSISTENT_BYTES_WRITTEN: Lazy<IntCounter> = Lazy::new(|| {
 
 #[derive(Clone)]
 pub enum LayeredTimelineEntry {
-    Loaded(Arc<LayeredTimeline>),
+    Loaded(Arc<Timeline>),
     Unloaded {
         id: ZTimelineId,
         metadata: TimelineMetadata,
@@ -191,7 +197,7 @@ impl LayeredTimelineEntry {
         }
     }
 
-    fn ensure_loaded(&self) -> anyhow::Result<&Arc<LayeredTimeline>> {
+    fn ensure_loaded(&self) -> anyhow::Result<&Arc<Timeline>> {
         match self {
             LayeredTimelineEntry::Loaded(timeline) => Ok(timeline),
             LayeredTimelineEntry::Unloaded { .. } => {
@@ -213,7 +219,7 @@ impl LayeredTimelineEntry {
     }
 }
 
-impl From<LayeredTimelineEntry> for RepositoryTimeline<LayeredTimeline> {
+impl From<LayeredTimelineEntry> for RepositoryTimeline<Timeline> {
     fn from(entry: LayeredTimelineEntry) -> Self {
         match entry {
             LayeredTimelineEntry::Loaded(timeline) => RepositoryTimeline::Loaded(timeline as _),
@@ -235,6 +241,8 @@ struct TimelineMetrics {
     pub last_record_gauge: IntGauge,
     pub wait_lsn_time_histo: Histogram,
     pub current_physical_size_gauge: UIntGauge,
+    /// copy of LayeredTimeline.current_logical_size
+    pub current_logical_size_gauge: IntGauge,
 }
 
 impl TimelineMetrics {
@@ -272,6 +280,9 @@ impl TimelineMetrics {
         let current_physical_size_gauge = CURRENT_PHYSICAL_SIZE
             .get_metric_with_label_values(&[&tenant_id, &timeline_id])
             .unwrap();
+        let current_logical_size_gauge = CURRENT_LOGICAL_SIZE
+            .get_metric_with_label_values(&[&tenant_id, &timeline_id])
+            .unwrap();
 
         TimelineMetrics {
             reconstruct_time_histo,
@@ -284,11 +295,12 @@ impl TimelineMetrics {
             last_record_gauge,
             wait_lsn_time_histo,
             current_physical_size_gauge,
+            current_logical_size_gauge,
         }
     }
 }
 
-pub struct LayeredTimeline {
+pub struct Timeline {
     conf: &'static PageServerConf,
     tenant_conf: Arc<RwLock<TenantConfOpt>>,
 
@@ -340,8 +352,8 @@ pub struct LayeredTimeline {
     upload_layers: AtomicBool,
 
     /// Ensures layers aren't frozen by checkpointer between
-    /// [`LayeredTimeline::get_layer_for_write`] and layer reads.
-    /// Locked automatically by [`LayeredTimelineWriter`] and checkpointer.
+    /// [`Timeline::get_layer_for_write`] and layer reads.
+    /// Locked automatically by [`TimelineWriter`] and checkpointer.
     /// Must always be acquired before the layer map/individual layer lock
     /// to avoid deadlock.
     write_lock: Mutex<()>,
@@ -351,12 +363,12 @@ pub struct LayeredTimeline {
 
     /// Layer removal lock.
     /// A lock to ensure that no layer of the timeline is removed concurrently by other threads.
-    /// This lock is acquired in [`LayeredTimeline::gc`], [`LayeredTimeline::compact`],
-    /// and [`LayeredRepository::delete_timeline`].
+    /// This lock is acquired in [`Timeline::gc`], [`Timeline::compact`],
+    /// and [`Repository::delete_timeline`].
     layer_removal_cs: Mutex<()>,
 
     // Needed to ensure that we can't create a branch at a point that was already garbage collected
-    pub latest_gc_cutoff_lsn: RwLock<Lsn>,
+    pub latest_gc_cutoff_lsn: Rcu<Lsn>,
 
     // List of child timelines and their branch points. This is needed to avoid
     // garbage collecting data that is still needed by the child timelines.
@@ -377,7 +389,32 @@ pub struct LayeredTimeline {
     repartition_threshold: u64,
 
     /// Current logical size of the "datadir", at the last LSN.
-    current_logical_size: AtomicIsize,
+    ///
+    /// Size shouldn't ever be negative, but this is signed for two reasons:
+    ///
+    /// 1. If we initialized the "baseline" size lazily, while we already
+    /// process incoming WAL, the incoming WAL records could decrement the
+    /// variable and temporarily make it negative. (This is just future-proofing;
+    /// the initialization is currently not done lazily.)
+    ///
+    /// 2. If there is a bug and we e.g. forget to increment it in some cases
+    /// when size grows, but remember to decrement it when it shrinks again, the
+    /// variable could go negative. In that case, it seems better to at least
+    /// try to keep tracking it, rather than clamp or overflow it. Note that
+    /// get_current_logical_size() will clamp the returned value to zero if it's
+    /// negative, and log an error. Could set it permanently to zero or some
+    /// special value to indicate "broken" instead, but this will do for now.
+    ///
+    /// Note that we also expose a copy of this value as a prometheus metric,
+    /// see `current_logical_size_gauge`. Use the `update_current_logical_size`
+    /// and `set_current_logical_size` functions to modify this, they will
+    /// also keep the prometheus metric in sync.
+    current_logical_size: AtomicI64,
+    // TODO we don't have a good, API to ensure on a compilation level
+    // that the timeline passes all initialization.
+    // Hence we ensure that we init at least once for every timeline
+    // and keep this flag to avoid potentually long recomputes.
+    logical_size_initialized: AtomicBool,
 
     /// Information about the last processed message by the WAL receiver,
     /// or None if WAL receiver has not received anything for this timeline
@@ -385,7 +422,7 @@ pub struct LayeredTimeline {
     pub last_received_wal: Mutex<Option<WalReceiverInfo>>,
 
     /// Relation size cache
-    rel_size_cache: RwLock<HashMap<RelTag, (Lsn, BlockNumber)>>,
+    pub rel_size_cache: RwLock<HashMap<RelTag, (Lsn, BlockNumber)>>,
 }
 
 pub struct WalReceiverInfo {
@@ -394,46 +431,6 @@ pub struct WalReceiverInfo {
     pub last_received_msg_ts: u128,
 }
 
-/// Inherit all the functions from DatadirTimeline, to provide the
-/// functionality to store PostgreSQL relations, SLRUs, etc. in a
-/// LayeredTimeline.
-impl DatadirTimeline for LayeredTimeline {
-    fn get_cached_rel_size(&self, tag: &RelTag, lsn: Lsn) -> Option<BlockNumber> {
-        let rel_size_cache = self.rel_size_cache.read().unwrap();
-        if let Some((cached_lsn, nblocks)) = rel_size_cache.get(tag) {
-            if lsn >= *cached_lsn {
-                return Some(*nblocks);
-            }
-        }
-        None
-    }
-
-    fn update_cached_rel_size(&self, tag: RelTag, lsn: Lsn, nblocks: BlockNumber) {
-        let mut rel_size_cache = self.rel_size_cache.write().unwrap();
-        match rel_size_cache.entry(tag) {
-            Entry::Occupied(mut entry) => {
-                let cached_lsn = entry.get_mut();
-                if lsn >= cached_lsn.0 {
-                    *cached_lsn = (lsn, nblocks);
-                }
-            }
-            Entry::Vacant(entry) => {
-                entry.insert((lsn, nblocks));
-            }
-        }
-    }
-
-    fn set_cached_rel_size(&self, tag: RelTag, lsn: Lsn, nblocks: BlockNumber) {
-        let mut rel_size_cache = self.rel_size_cache.write().unwrap();
-        rel_size_cache.insert(tag, (lsn, nblocks));
-    }
-
-    fn remove_cached_rel_size(&self, tag: &RelTag) {
-        let mut rel_size_cache = self.rel_size_cache.write().unwrap();
-        rel_size_cache.remove(tag);
-    }
-}
-
 ///
 /// Information about how much history needs to be retained, needed by
 /// Garbage Collection.
@@ -464,45 +461,37 @@ pub struct GcInfo {
 }
 
 /// Public interface functions
-impl Timeline for LayeredTimeline {
-    fn get_ancestor_lsn(&self) -> Lsn {
+impl Timeline {
+    //------------------------------------------------------------------------------
+    // Public GET functions
+    //------------------------------------------------------------------------------
+
+    /// Get the LSN where this branch was created
+    pub fn get_ancestor_lsn(&self) -> Lsn {
         self.ancestor_lsn
     }
 
-    fn get_ancestor_timeline_id(&self) -> Option<ZTimelineId> {
+    /// Get the ancestor's timeline id
+    pub fn get_ancestor_timeline_id(&self) -> Option<ZTimelineId> {
         self.ancestor_timeline
             .as_ref()
             .map(LayeredTimelineEntry::timeline_id)
     }
 
-    /// Wait until WAL has been received up to the given LSN.
-    fn wait_lsn(&self, lsn: Lsn) -> anyhow::Result<()> {
-        // This should never be called from the WAL receiver thread, because that could lead
-        // to a deadlock.
-        ensure!(
-            !IS_WAL_RECEIVER.with(|c| c.get()),
-            "wait_lsn called by WAL receiver thread"
-        );
-
-        self.metrics.wait_lsn_time_histo.observe_closure_duration(
-            || self.last_record_lsn
-                .wait_for_timeout(lsn, self.conf.wait_lsn_timeout)
-                .with_context(|| {
-                    format!(
-                        "Timed out while waiting for WAL record at LSN {} to arrive, last_record_lsn {} disk consistent LSN={}",
-                        lsn, self.get_last_record_lsn(), self.get_disk_consistent_lsn()
-                    )
-                }))?;
-
-        Ok(())
+    /// Lock and get timeline's GC cuttof
+    pub fn get_latest_gc_cutoff_lsn(&self) -> RcuReadGuard<Lsn> {
+        self.latest_gc_cutoff_lsn.read()
     }
 
-    fn get_latest_gc_cutoff_lsn(&self) -> RwLockReadGuard<Lsn> {
-        self.latest_gc_cutoff_lsn.read().unwrap()
-    }
-
-    /// Look up the value with the given a key
-    fn get(&self, key: Key, lsn: Lsn) -> Result<Bytes> {
+    /// Look up given page version.
+    ///
+    /// NOTE: It is considered an error to 'get' a key that doesn't exist. The abstraction
+    /// above this needs to store suitable metadata to track what data exists with
+    /// what keys, in separate metadata entries. If a non-existent key is requested,
+    /// the Repository implementation may incorrectly return a value from an ancestor
+    /// branch, for example, or waste a lot of cycles chasing the non-existing key.
+    ///
+    pub fn get(&self, key: Key, lsn: Lsn) -> Result<Bytes> {
         // Check the page cache. We will get back the most recent page with lsn <= `lsn`.
         // The cached image can be returned directly if there is no WAL between the cached image
         // and requested LSN. The cached image can also be used to reduce the amount of WAL needed
@@ -531,68 +520,31 @@ impl Timeline for LayeredTimeline {
             .observe_closure_duration(|| self.reconstruct_value(key, lsn, reconstruct_state))
     }
 
-    /// Public entry point for checkpoint(). All the logic is in the private
-    /// checkpoint_internal function, this public facade just wraps it for
-    /// metrics collection.
-    fn checkpoint(&self, cconf: CheckpointConfig) -> anyhow::Result<()> {
-        match cconf {
-            CheckpointConfig::Flush => {
-                self.freeze_inmem_layer(false);
-                self.flush_frozen_layers(true)
-            }
-            CheckpointConfig::Forced => {
-                self.freeze_inmem_layer(false);
-                self.flush_frozen_layers(true)?;
-                self.compact()
-            }
-        }
-    }
-
-    ///
-    /// 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<()> {
-        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,
-        );
-        Ok(())
-    }
-
-    fn get_last_record_lsn(&self) -> Lsn {
+    /// Get last or prev record separately. Same as get_last_record_rlsn().last/prev.
+    pub fn get_last_record_lsn(&self) -> Lsn {
         self.last_record_lsn.load().last
     }
 
-    fn get_prev_record_lsn(&self) -> Lsn {
+    pub fn get_prev_record_lsn(&self) -> Lsn {
         self.last_record_lsn.load().prev
     }
 
-    fn get_last_record_rlsn(&self) -> RecordLsn {
+    /// Atomically get both last and prev.
+    pub fn get_last_record_rlsn(&self) -> RecordLsn {
         self.last_record_lsn.load()
     }
 
-    fn get_disk_consistent_lsn(&self) -> Lsn {
+    pub fn get_disk_consistent_lsn(&self) -> Lsn {
         self.disk_consistent_lsn.load()
     }
 
-    fn writer<'a>(&'a self) -> Box<dyn TimelineWriter + 'a> {
-        Box::new(LayeredTimelineWriter {
-            tl: self,
-            _write_guard: self.write_lock.lock().unwrap(),
-        })
-    }
-
-    fn get_physical_size(&self) -> u64 {
+    /// Get the physical size of the timeline at the latest LSN
+    pub fn get_physical_size(&self) -> u64 {
         self.metrics.current_physical_size_gauge.get()
     }
 
-    fn get_physical_size_non_incremental(&self) -> anyhow::Result<u64> {
+    /// Get the physical size of the timeline at the latest LSN non incrementally
+    pub fn get_physical_size_non_incremental(&self) -> anyhow::Result<u64> {
         let timeline_path = self.conf.timeline_path(&self.timeline_id, &self.tenant_id);
         // total size of layer files in the current timeline directory
         let mut total_physical_size = 0;
@@ -611,9 +563,88 @@ impl Timeline for LayeredTimeline {
 
         Ok(total_physical_size)
     }
+
+    ///
+    /// Wait until WAL has been received and processed up to this LSN.
+    ///
+    /// You should call this before any of the other get_* or list_* functions. Calling
+    /// those functions with an LSN that has been processed yet is an error.
+    ///
+    pub fn wait_lsn(&self, lsn: Lsn) -> anyhow::Result<()> {
+        // This should never be called from the WAL receiver thread, because that could lead
+        // to a deadlock.
+        ensure!(
+            !IS_WAL_RECEIVER.with(|c| c.get()),
+            "wait_lsn called by WAL receiver thread"
+        );
+
+        self.metrics.wait_lsn_time_histo.observe_closure_duration(
+            || self.last_record_lsn
+                .wait_for_timeout(lsn, self.conf.wait_lsn_timeout)
+                .with_context(|| {
+                    format!(
+                        "Timed out while waiting for WAL record at LSN {} to arrive, last_record_lsn {} disk consistent LSN={}",
+                        lsn, self.get_last_record_lsn(), self.get_disk_consistent_lsn()
+                    )
+                }))?;
+
+        Ok(())
+    }
+
+    /// Check that it is valid to request operations with that lsn.
+    pub fn check_lsn_is_in_scope(
+        &self,
+        lsn: Lsn,
+        latest_gc_cutoff_lsn: &RcuReadGuard<Lsn>,
+    ) -> Result<()> {
+        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,
+        );
+        Ok(())
+    }
+
+    //------------------------------------------------------------------------------
+    // Public PUT functions, to update the repository with new page versions.
+    //
+    // These are called by the WAL receiver to digest WAL records.
+    //------------------------------------------------------------------------------
+
+    /// Flush to disk all data that was written with the put_* functions
+    ///
+    /// NOTE: This has nothing to do with checkpoint in PostgreSQL. We don't
+    /// know anything about them here in the repository.
+    pub fn checkpoint(&self, cconf: CheckpointConfig) -> anyhow::Result<()> {
+        match cconf {
+            CheckpointConfig::Flush => {
+                self.freeze_inmem_layer(false);
+                self.flush_frozen_layers(true)
+            }
+            CheckpointConfig::Forced => {
+                self.freeze_inmem_layer(false);
+                self.flush_frozen_layers(true)?;
+                self.compact()
+            }
+        }
+    }
+
+    /// Mutate the timeline with a [`TimelineWriter`].
+    ///
+    /// FIXME: This ought to return &'a TimelineWriter, where TimelineWriter
+    /// is a generic type in this trait. But that doesn't currently work in
+    /// Rust: https://rust-lang.github.io/rfcs/1598-generic_associated_types.html
+    pub fn writer(&self) -> TimelineWriter<'_> {
+        TimelineWriter {
+            tl: self,
+            _write_guard: self.write_lock.lock().unwrap(),
+        }
+    }
 }
 
-impl LayeredTimeline {
+// Private functions
+impl Timeline {
     fn get_checkpoint_distance(&self) -> u64 {
         let tenant_conf = self.tenant_conf.read().unwrap();
         tenant_conf
@@ -662,8 +693,8 @@ impl LayeredTimeline {
         tenant_id: ZTenantId,
         walredo_mgr: Arc<dyn WalRedoManager + Send + Sync>,
         upload_layers: bool,
-    ) -> LayeredTimeline {
-        let mut result = LayeredTimeline {
+    ) -> Timeline {
+        let mut result = Timeline {
             conf,
             tenant_conf,
             timeline_id,
@@ -699,10 +730,11 @@ impl LayeredTimeline {
                 pitr_cutoff: Lsn(0),
             }),
 
-            latest_gc_cutoff_lsn: RwLock::new(metadata.latest_gc_cutoff_lsn()),
+            latest_gc_cutoff_lsn: Rcu::new(metadata.latest_gc_cutoff_lsn()),
             initdb_lsn: metadata.initdb_lsn(),
 
-            current_logical_size: AtomicIsize::new(0),
+            current_logical_size: AtomicI64::new(0),
+            logical_size_initialized: AtomicBool::new(false),
             partitioning: Mutex::new((KeyPartitioning::new(), Lsn(0))),
             repartition_threshold: 0,
 
@@ -807,6 +839,10 @@ impl LayeredTimeline {
     ///
     /// This can be a slow operation.
     pub fn init_logical_size(&self) -> Result<()> {
+        if self.logical_size_initialized.load(AtomicOrdering::Acquire) {
+            return Ok(());
+        }
+
         // Try a fast-path first:
         // Copy logical size from ancestor timeline if there has been no changes on this
         // branch, and no changes on the ancestor branch since the branch point.
@@ -819,8 +855,7 @@ impl LayeredTimeline {
             //
             // Logical size 0 means that it was not initialized, so don't believe that.
             if ancestor_logical_size != 0 && ancestor.get_last_record_lsn() == self.ancestor_lsn {
-                self.current_logical_size
-                    .store(ancestor_logical_size as isize, AtomicOrdering::SeqCst);
+                self.set_current_logical_size(ancestor_logical_size);
                 debug!(
                     "logical size copied from ancestor: {}",
                     ancestor_logical_size
@@ -834,8 +869,7 @@ impl LayeredTimeline {
         // Have to calculate it the hard way
         let last_lsn = self.get_last_record_lsn();
         let logical_size = self.get_current_logical_size_non_incremental(last_lsn)?;
-        self.current_logical_size
-            .store(logical_size as isize, AtomicOrdering::SeqCst);
+        self.set_current_logical_size(logical_size);
         debug!("calculated logical size the hard way: {}", logical_size);
 
         timer.stop_and_record();
@@ -844,10 +878,10 @@ impl LayeredTimeline {
 
     /// Retrieve current logical size of the timeline
     ///
-    /// NOTE: counted incrementally, includes ancestors,
-    pub fn get_current_logical_size(&self) -> usize {
+    /// NOTE: counted incrementally, includes ancestors.
+    pub fn get_current_logical_size(&self) -> u64 {
         let current_logical_size = self.current_logical_size.load(AtomicOrdering::Acquire);
-        match usize::try_from(current_logical_size) {
+        match u64::try_from(current_logical_size) {
             Ok(sz) => sz,
             Err(_) => {
                 error!(
@@ -859,6 +893,36 @@ impl LayeredTimeline {
         }
     }
 
+    /// Update current logical size, adding `delta' to the old value.
+    fn update_current_logical_size(&self, delta: i64) {
+        let new_size = self
+            .current_logical_size
+            .fetch_add(delta, AtomicOrdering::SeqCst);
+
+        // Also set the value in the prometheus gauge. Note that
+        // there is a race condition here: if this is is called by two
+        // threads concurrently, the prometheus gauge might be set to
+        // one value while current_logical_size is set to the
+        // other. Currently, only initialization and the WAL receiver
+        // updates the logical size, and they don't run concurrently,
+        // so it cannot happen. And even if it did, it wouldn't be
+        // very serious, the metrics would just be slightly off until
+        // the next update.
+        self.metrics.current_logical_size_gauge.set(new_size);
+    }
+
+    /// Set current logical size.
+    fn set_current_logical_size(&self, new_size: u64) {
+        self.current_logical_size
+            .store(new_size as i64, AtomicOrdering::SeqCst);
+        self.logical_size_initialized
+            .store(true, AtomicOrdering::SeqCst);
+
+        // Also set the value in the prometheus gauge. Same race condition
+        // here as in `update_current_logical_size`.
+        self.metrics.current_logical_size_gauge.set(new_size as i64);
+    }
+
     ///
     /// Get a handle to a Layer for reading.
     ///
@@ -1014,7 +1078,7 @@ impl LayeredTimeline {
         Some((lsn, img))
     }
 
-    fn get_ancestor_timeline(&self) -> Result<Arc<LayeredTimeline>> {
+    fn get_ancestor_timeline(&self) -> Result<Arc<Timeline>> {
         let ancestor = self
             .ancestor_timeline
             .as_ref()
@@ -1135,7 +1199,7 @@ impl LayeredTimeline {
     /// Also flush after a period of time without new data -- it helps
     /// safekeepers to regard pageserver as caught up and suspend activity.
     ///
-    pub fn check_checkpoint_distance(self: &Arc<LayeredTimeline>) -> Result<()> {
+    pub fn check_checkpoint_distance(self: &Arc<Timeline>) -> Result<()> {
         let last_lsn = self.get_last_record_lsn();
         let layers = self.layers.read().unwrap();
         if let Some(open_layer) = &layers.open_layer {
@@ -1314,7 +1378,7 @@ impl LayeredTimeline {
                 ondisk_prev_record_lsn,
                 ancestor_timelineid,
                 self.ancestor_lsn,
-                *self.latest_gc_cutoff_lsn.read().unwrap(),
+                *self.latest_gc_cutoff_lsn.read(),
                 self.initdb_lsn,
             );
 
@@ -1969,9 +2033,21 @@ impl LayeredTimeline {
 
         let _enter = info_span!("garbage collection", timeline = %self.timeline_id, tenant = %self.tenant_id, cutoff = %new_gc_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() = new_gc_cutoff;
+        // We need to ensure that no one tries to read page versions or create
+        // branches at a point before latest_gc_cutoff_lsn. See branch_timeline()
+        // for details. This will block until the old value is no longer in use.
+        //
+        // The GC cutoff should only ever move forwards.
+        {
+            let write_guard = self.latest_gc_cutoff_lsn.write();
+            ensure!(
+                *write_guard <= new_gc_cutoff,
+                "Cannot move GC cutoff LSN backwards (was {}, new {})",
+                *write_guard,
+                new_gc_cutoff
+            );
+            write_guard.store(new_gc_cutoff);
+        }
 
         info!("GC starting");
 
@@ -2210,39 +2286,50 @@ fn layer_traversal_error(
     Err(msg_iter.fold(err, |err, msg| err.context(msg)))
 }
 
-struct LayeredTimelineWriter<'a> {
-    tl: &'a LayeredTimeline,
+/// Various functions to mutate the timeline.
+// TODO Currently, Deref is used to allow easy access to read methods from this trait.
+// This is probably considered a bad practice in Rust and should be fixed eventually,
+// but will cause large code changes.
+pub struct TimelineWriter<'a> {
+    tl: &'a Timeline,
     _write_guard: MutexGuard<'a, ()>,
 }
 
-impl Deref for LayeredTimelineWriter<'_> {
-    type Target = dyn Timeline;
+impl Deref for TimelineWriter<'_> {
+    type Target = Timeline;
 
     fn deref(&self) -> &Self::Target {
         self.tl
     }
 }
 
-impl<'a> TimelineWriter<'_> for LayeredTimelineWriter<'a> {
-    fn put(&self, key: Key, lsn: Lsn, value: &Value) -> Result<()> {
+impl<'a> TimelineWriter<'a> {
+    /// Put a new page version that can be constructed from a WAL record
+    ///
+    /// This will implicitly extend the relation, if the page is beyond the
+    /// current end-of-file.
+    pub fn put(&self, key: Key, lsn: Lsn, value: &Value) -> Result<()> {
         self.tl.put_value(key, lsn, value)
     }
 
-    fn delete(&self, key_range: Range<Key>, lsn: Lsn) -> Result<()> {
+    pub fn delete(&self, key_range: Range<Key>, lsn: Lsn) -> Result<()> {
         self.tl.put_tombstone(key_range, lsn)
     }
 
-    ///
+    /// Track the end of the latest digested WAL record.
     /// Remember the (end of) last valid WAL record remembered in the timeline.
     ///
-    fn finish_write(&self, new_lsn: Lsn) {
+    /// Call this after you have finished writing all the WAL up to 'lsn'.
+    ///
+    /// 'lsn' must be aligned. This wakes up any wait_lsn() callers waiting for
+    /// the 'lsn' or anything older. The previous last record LSN is stored alongside
+    /// the latest and can be read.
+    pub fn finish_write(&self, new_lsn: Lsn) {
         self.tl.finish_write(new_lsn);
     }
 
-    fn update_current_logical_size(&self, delta: isize) {
-        self.tl
-            .current_logical_size
-            .fetch_add(delta, AtomicOrdering::SeqCst);
+    pub fn update_current_logical_size(&self, delta: i64) {
+        self.tl.update_current_logical_size(delta)
     }
 }
 
@@ -2265,39 +2352,3 @@ fn rename_to_backup(path: PathBuf) -> anyhow::Result<()> {
 
     bail!("couldn't find an unused backup number for {:?}", path)
 }
-
-/// Save timeline metadata to file
-pub fn save_metadata(
-    conf: &'static PageServerConf,
-    timelineid: ZTimelineId,
-    tenantid: ZTenantId,
-    data: &TimelineMetadata,
-    first_save: bool,
-) -> Result<()> {
-    let _enter = info_span!("saving metadata").entered();
-    let path = metadata_path(conf, timelineid, tenantid);
-    // use OpenOptions to ensure file presence is consistent with first_save
-    let mut file = VirtualFile::open_with_options(
-        &path,
-        OpenOptions::new().write(true).create_new(first_save),
-    )?;
-
-    let metadata_bytes = data.to_bytes().context("Failed to get metadata bytes")?;
-
-    if file.write(&metadata_bytes)? != metadata_bytes.len() {
-        bail!("Could not write all the metadata bytes in a single call");
-    }
-    file.sync_all()?;
-
-    // fsync the parent directory to ensure the directory entry is durable
-    if first_save {
-        let timeline_dir = File::open(
-            &path
-                .parent()
-                .expect("Metadata should always have a parent dir"),
-        )?;
-        timeline_dir.sync_all()?;
-    }
-
-    Ok(())
-}

pageserver/src/lib.rs

@@ -28,8 +28,6 @@ use tracing::info;
 use crate::thread_mgr::ThreadKind;
 use metrics::{register_int_gauge_vec, IntGaugeVec};
 
-use pgdatadir_mapping::DatadirTimeline;
-
 /// Current storage format version
 ///
 /// This is embedded in the metadata file, and also in the header of all the

pageserver/src/page_service.rs

@@ -17,24 +17,24 @@ use std::io::{self, Read};
 use std::net::TcpListener;
 use std::str;
 use std::str::FromStr;
-use std::sync::{Arc, RwLockReadGuard};
+use std::sync::Arc;
 use tracing::*;
 use utils::{
     auth::{self, Claims, JwtAuth, Scope},
     lsn::Lsn,
     postgres_backend::{self, is_socket_read_timed_out, AuthType, PostgresBackend},
     pq_proto::{BeMessage, FeMessage, RowDescriptor, SINGLE_COL_ROWDESC},
+    simple_rcu::RcuReadGuard,
     zid::{ZTenantId, ZTimelineId},
 };
 
 use crate::basebackup;
 use crate::config::{PageServerConf, ProfilingConfig};
 use crate::import_datadir::{import_basebackup_from_tar, import_wal_from_tar};
-use crate::pgdatadir_mapping::{DatadirTimeline, LsnForTimestamp};
+use crate::layered_repository::Timeline;
+use crate::pgdatadir_mapping::LsnForTimestamp;
 use crate::profiling::profpoint_start;
 use crate::reltag::RelTag;
-use crate::repository::Repository;
-use crate::repository::Timeline;
 use crate::tenant_mgr;
 use crate::thread_mgr;
 use crate::thread_mgr::ThreadKind;
@@ -495,22 +495,22 @@ impl PageServerHandler {
                             PagestreamFeMessage::Exists(req) => SMGR_QUERY_TIME
                                 .with_label_values(&["get_rel_exists", &tenant_id, &timeline_id])
                                 .observe_closure_duration(|| {
-                                    self.handle_get_rel_exists_request(timeline.as_ref(), &req)
+                                    self.handle_get_rel_exists_request(&timeline, &req)
                                 }),
                             PagestreamFeMessage::Nblocks(req) => SMGR_QUERY_TIME
                                 .with_label_values(&["get_rel_size", &tenant_id, &timeline_id])
                                 .observe_closure_duration(|| {
-                                    self.handle_get_nblocks_request(timeline.as_ref(), &req)
+                                    self.handle_get_nblocks_request(&timeline, &req)
                                 }),
                             PagestreamFeMessage::GetPage(req) => SMGR_QUERY_TIME
                                 .with_label_values(&["get_page_at_lsn", &tenant_id, &timeline_id])
                                 .observe_closure_duration(|| {
-                                    self.handle_get_page_at_lsn_request(timeline.as_ref(), &req)
+                                    self.handle_get_page_at_lsn_request(&timeline, &req)
                                 }),
                             PagestreamFeMessage::DbSize(req) => SMGR_QUERY_TIME
                                 .with_label_values(&["get_db_size", &tenant_id, &timeline_id])
                                 .observe_closure_duration(|| {
-                                    self.handle_db_size_request(timeline.as_ref(), &req)
+                                    self.handle_db_size_request(&timeline, &req)
                                 }),
                         };
 
@@ -636,11 +636,11 @@ 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<T: DatadirTimeline>(
-        timeline: &T,
+    fn wait_or_get_last_lsn(
+        timeline: &Timeline,
         mut lsn: Lsn,
         latest: bool,
-        latest_gc_cutoff_lsn: &RwLockReadGuard<Lsn>,
+        latest_gc_cutoff_lsn: &RcuReadGuard<Lsn>,
     ) -> Result<Lsn> {
         if latest {
             // Latest page version was requested. If LSN is given, it is a hint
@@ -684,9 +684,9 @@ impl PageServerHandler {
         Ok(lsn)
     }
 
-    fn handle_get_rel_exists_request<T: DatadirTimeline>(
+    fn handle_get_rel_exists_request(
         &self,
-        timeline: &T,
+        timeline: &Timeline,
         req: &PagestreamExistsRequest,
     ) -> Result<PagestreamBeMessage> {
         let _enter = info_span!("get_rel_exists", rel = %req.rel, req_lsn = %req.lsn).entered();
@@ -701,9 +701,9 @@ impl PageServerHandler {
         }))
     }
 
-    fn handle_get_nblocks_request<T: DatadirTimeline>(
+    fn handle_get_nblocks_request(
         &self,
-        timeline: &T,
+        timeline: &Timeline,
         req: &PagestreamNblocksRequest,
     ) -> Result<PagestreamBeMessage> {
         let _enter = info_span!("get_nblocks", rel = %req.rel, req_lsn = %req.lsn).entered();
@@ -717,9 +717,9 @@ impl PageServerHandler {
         }))
     }
 
-    fn handle_db_size_request<T: DatadirTimeline>(
+    fn handle_db_size_request(
         &self,
-        timeline: &T,
+        timeline: &Timeline,
         req: &PagestreamDbSizeRequest,
     ) -> Result<PagestreamBeMessage> {
         let _enter = info_span!("get_db_size", dbnode = %req.dbnode, req_lsn = %req.lsn).entered();
@@ -735,9 +735,9 @@ impl PageServerHandler {
         }))
     }
 
-    fn handle_get_page_at_lsn_request<T: DatadirTimeline>(
+    fn handle_get_page_at_lsn_request(
         &self,
-        timeline: &T,
+        timeline: &Timeline,
         req: &PagestreamGetPageRequest,
     ) -> Result<PagestreamBeMessage> {
         let _enter = info_span!("get_page", rel = %req.rel, blkno = &req.blkno, req_lsn = %req.lsn)
@@ -745,7 +745,7 @@ impl PageServerHandler {
         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
+        // Add a 1s delay to some requests. The delay helps the requests to
         // hit the race condition from github issue #1047 more easily.
         use rand::Rng;
         if rand::thread_rng().gen::<u8>() < 5 {

pageserver/src/pgdatadir_mapping.rs

@@ -7,8 +7,8 @@
 //! Clarify that)
 //!
 use crate::keyspace::{KeySpace, KeySpaceAccum};
+use crate::layered_repository::Timeline;
 use crate::reltag::{RelTag, SlruKind};
-use crate::repository::Timeline;
 use crate::repository::*;
 use crate::walrecord::ZenithWalRecord;
 use anyhow::{bail, ensure, Result};
@@ -18,7 +18,7 @@ use postgres_ffi::v14::xlog_utils::TimestampTz;
 use postgres_ffi::BLCKSZ;
 use postgres_ffi::{Oid, TransactionId};
 use serde::{Deserialize, Serialize};
-use std::collections::{HashMap, HashSet};
+use std::collections::{hash_map, HashMap, HashSet};
 use std::ops::Range;
 use tracing::{debug, trace, warn};
 use utils::{bin_ser::BeSer, lsn::Lsn};
@@ -35,23 +35,13 @@ pub enum LsnForTimestamp {
 }
 
 ///
-/// This trait provides all the functionality to store PostgreSQL relations, SLRUs,
+/// This impl provides all the functionality to store PostgreSQL relations, SLRUs,
 /// and other special kinds of files, in a versioned key-value store. The
-/// Timeline trait provides the key-value store.
+/// Timeline struct provides the key-value store.
 ///
-/// This is a trait, so that we can easily include all these functions in a Timeline
-/// implementation. You're not expected to have different implementations of this trait,
-/// rather, this provides an interface and implementation, over Timeline.
-///
-/// If you wanted to store other kinds of data in the Neon repository, e.g.
-/// flat files or MySQL, you would create a new trait like this, with all the
-/// functions that make sense for the kind of data you're storing. For flat files,
-/// for example, you might have a function like "fn read(path, offset, size)".
-/// We might also have that situation in the future, to support multiple PostgreSQL
-/// versions, if there are big changes in how the data is organized in the data
-/// directory, or if new special files are introduced.
-///
-pub trait DatadirTimeline: Timeline {
+/// This is a separate impl, so that we can easily include all these functions in a Timeline
+/// implementation, and might be moved into a separate struct later.
+impl Timeline {
     /// Start ingesting a WAL record, or other atomic modification of
     /// the timeline.
     ///
@@ -75,7 +65,7 @@ pub trait DatadirTimeline: Timeline {
     /// functions of the timeline until you finish! And if you update the
     /// same page twice, the last update wins.
     ///
-    fn begin_modification(&self, lsn: Lsn) -> DatadirModification<Self>
+    pub fn begin_modification(&self, lsn: Lsn) -> DatadirModification
     where
         Self: Sized,
     {
@@ -93,7 +83,7 @@ pub trait DatadirTimeline: Timeline {
     //------------------------------------------------------------------------------
 
     /// Look up given page version.
-    fn get_rel_page_at_lsn(&self, tag: RelTag, blknum: BlockNumber, lsn: Lsn) -> Result<Bytes> {
+    pub fn get_rel_page_at_lsn(&self, tag: RelTag, blknum: BlockNumber, lsn: Lsn) -> Result<Bytes> {
         ensure!(tag.relnode != 0, "invalid relnode");
 
         let nblocks = self.get_rel_size(tag, lsn)?;
@@ -110,7 +100,7 @@ pub trait DatadirTimeline: Timeline {
     }
 
     // Get size of a database in blocks
-    fn get_db_size(&self, spcnode: Oid, dbnode: Oid, lsn: Lsn) -> Result<usize> {
+    pub fn get_db_size(&self, spcnode: Oid, dbnode: Oid, lsn: Lsn) -> Result<usize> {
         let mut total_blocks = 0;
 
         let rels = self.list_rels(spcnode, dbnode, lsn)?;
@@ -123,7 +113,7 @@ pub trait DatadirTimeline: Timeline {
     }
 
     /// Get size of a relation file
-    fn get_rel_size(&self, tag: RelTag, lsn: Lsn) -> Result<BlockNumber> {
+    pub fn get_rel_size(&self, tag: RelTag, lsn: Lsn) -> Result<BlockNumber> {
         ensure!(tag.relnode != 0, "invalid relnode");
 
         if let Some(nblocks) = self.get_cached_rel_size(&tag, lsn) {
@@ -151,7 +141,7 @@ pub trait DatadirTimeline: Timeline {
     }
 
     /// Does relation exist?
-    fn get_rel_exists(&self, tag: RelTag, lsn: Lsn) -> Result<bool> {
+    pub fn get_rel_exists(&self, tag: RelTag, lsn: Lsn) -> Result<bool> {
         ensure!(tag.relnode != 0, "invalid relnode");
 
         // first try to lookup relation in cache
@@ -169,7 +159,7 @@ pub trait DatadirTimeline: Timeline {
     }
 
     /// Get a list of all existing relations in given tablespace and database.
-    fn list_rels(&self, spcnode: Oid, dbnode: Oid, lsn: Lsn) -> Result<HashSet<RelTag>> {
+    pub fn list_rels(&self, spcnode: Oid, dbnode: Oid, lsn: Lsn) -> Result<HashSet<RelTag>> {
         // fetch directory listing
         let key = rel_dir_to_key(spcnode, dbnode);
         let buf = self.get(key, lsn)?;
@@ -187,7 +177,7 @@ pub trait DatadirTimeline: Timeline {
     }
 
     /// Look up given SLRU page version.
-    fn get_slru_page_at_lsn(
+    pub fn get_slru_page_at_lsn(
         &self,
         kind: SlruKind,
         segno: u32,
@@ -199,14 +189,19 @@ pub trait DatadirTimeline: Timeline {
     }
 
     /// Get size of an SLRU segment
-    fn get_slru_segment_size(&self, kind: SlruKind, segno: u32, lsn: Lsn) -> Result<BlockNumber> {
+    pub fn get_slru_segment_size(
+        &self,
+        kind: SlruKind,
+        segno: u32,
+        lsn: Lsn,
+    ) -> Result<BlockNumber> {
         let key = slru_segment_size_to_key(kind, segno);
         let mut buf = self.get(key, lsn)?;
         Ok(buf.get_u32_le())
     }
 
     /// Get size of an SLRU segment
-    fn get_slru_segment_exists(&self, kind: SlruKind, segno: u32, lsn: Lsn) -> Result<bool> {
+    pub fn get_slru_segment_exists(&self, kind: SlruKind, segno: u32, lsn: Lsn) -> Result<bool> {
         // fetch directory listing
         let key = slru_dir_to_key(kind);
         let buf = self.get(key, lsn)?;
@@ -223,7 +218,7 @@ pub trait DatadirTimeline: Timeline {
     /// so it's not well defined which LSN you get if there were multiple commits
     /// "in flight" at that point in time.
     ///
-    fn find_lsn_for_timestamp(&self, search_timestamp: TimestampTz) -> Result<LsnForTimestamp> {
+    pub fn find_lsn_for_timestamp(&self, search_timestamp: TimestampTz) -> Result<LsnForTimestamp> {
         let gc_cutoff_lsn_guard = self.get_latest_gc_cutoff_lsn();
         let min_lsn = *gc_cutoff_lsn_guard;
         let max_lsn = self.get_last_record_lsn();
@@ -286,7 +281,7 @@ pub trait DatadirTimeline: Timeline {
     /// Additionally, sets 'found_smaller'/'found_Larger, if encounters any commits
     /// with a smaller/larger timestamp.
     ///
-    fn is_latest_commit_timestamp_ge_than(
+    pub fn is_latest_commit_timestamp_ge_than(
         &self,
         search_timestamp: TimestampTz,
         probe_lsn: Lsn,
@@ -317,7 +312,7 @@ pub trait DatadirTimeline: Timeline {
     }
 
     /// Get a list of SLRU segments
-    fn list_slru_segments(&self, kind: SlruKind, lsn: Lsn) -> Result<HashSet<u32>> {
+    pub fn list_slru_segments(&self, kind: SlruKind, lsn: Lsn) -> Result<HashSet<u32>> {
         // fetch directory entry
         let key = slru_dir_to_key(kind);
 
@@ -327,14 +322,14 @@ pub trait DatadirTimeline: Timeline {
         Ok(dir.segments)
     }
 
-    fn get_relmap_file(&self, spcnode: Oid, dbnode: Oid, lsn: Lsn) -> Result<Bytes> {
+    pub fn get_relmap_file(&self, spcnode: Oid, dbnode: Oid, lsn: Lsn) -> Result<Bytes> {
         let key = relmap_file_key(spcnode, dbnode);
 
         let buf = self.get(key, lsn)?;
         Ok(buf)
     }
 
-    fn list_dbdirs(&self, lsn: Lsn) -> Result<HashMap<(Oid, Oid), bool>> {
+    pub fn list_dbdirs(&self, lsn: Lsn) -> Result<HashMap<(Oid, Oid), bool>> {
         // fetch directory entry
         let buf = self.get(DBDIR_KEY, lsn)?;
         let dir = DbDirectory::des(&buf)?;
@@ -342,13 +337,13 @@ pub trait DatadirTimeline: Timeline {
         Ok(dir.dbdirs)
     }
 
-    fn get_twophase_file(&self, xid: TransactionId, lsn: Lsn) -> Result<Bytes> {
+    pub fn get_twophase_file(&self, xid: TransactionId, lsn: Lsn) -> Result<Bytes> {
         let key = twophase_file_key(xid);
         let buf = self.get(key, lsn)?;
         Ok(buf)
     }
 
-    fn list_twophase_files(&self, lsn: Lsn) -> Result<HashSet<TransactionId>> {
+    pub fn list_twophase_files(&self, lsn: Lsn) -> Result<HashSet<TransactionId>> {
         // fetch directory entry
         let buf = self.get(TWOPHASEDIR_KEY, lsn)?;
         let dir = TwoPhaseDirectory::des(&buf)?;
@@ -356,11 +351,11 @@ pub trait DatadirTimeline: Timeline {
         Ok(dir.xids)
     }
 
-    fn get_control_file(&self, lsn: Lsn) -> Result<Bytes> {
+    pub fn get_control_file(&self, lsn: Lsn) -> Result<Bytes> {
         self.get(CONTROLFILE_KEY, lsn)
     }
 
-    fn get_checkpoint(&self, lsn: Lsn) -> Result<Bytes> {
+    pub fn get_checkpoint(&self, lsn: Lsn) -> Result<Bytes> {
         self.get(CHECKPOINT_KEY, lsn)
     }
 
@@ -369,29 +364,29 @@ pub trait DatadirTimeline: Timeline {
     ///
     /// Only relation blocks are counted currently. That excludes metadata,
     /// SLRUs, twophase files etc.
-    fn get_current_logical_size_non_incremental(&self, lsn: Lsn) -> Result<usize> {
+    pub fn get_current_logical_size_non_incremental(&self, lsn: Lsn) -> Result<u64> {
         // Fetch list of database dirs and iterate them
         let buf = self.get(DBDIR_KEY, lsn)?;
         let dbdir = DbDirectory::des(&buf)?;
 
-        let mut total_size: usize = 0;
+        let mut total_size: u64 = 0;
         for (spcnode, dbnode) in dbdir.dbdirs.keys() {
             for rel in self.list_rels(*spcnode, *dbnode, lsn)? {
                 let relsize_key = rel_size_to_key(rel);
                 let mut buf = self.get(relsize_key, lsn)?;
                 let relsize = buf.get_u32_le();
 
-                total_size += relsize as usize;
+                total_size += relsize as u64;
             }
         }
-        Ok(total_size * BLCKSZ as usize)
+        Ok(total_size * BLCKSZ as u64)
     }
 
     ///
     /// Get a KeySpace that covers all the Keys that are in use at the given LSN.
     /// Anything that's not listed maybe removed from the underlying storage (from
     /// that LSN forwards).
-    fn collect_keyspace(&self, lsn: Lsn) -> Result<KeySpace> {
+    pub fn collect_keyspace(&self, lsn: Lsn) -> Result<KeySpace> {
         // Iterate through key ranges, greedily packing them into partitions
         let mut result = KeySpaceAccum::new();
 
@@ -465,27 +460,54 @@ pub trait DatadirTimeline: Timeline {
     }
 
     /// Get cached size of relation if it not updated after specified LSN
-    fn get_cached_rel_size(&self, tag: &RelTag, lsn: Lsn) -> Option<BlockNumber>;
+    pub fn get_cached_rel_size(&self, tag: &RelTag, lsn: Lsn) -> Option<BlockNumber> {
+        let rel_size_cache = self.rel_size_cache.read().unwrap();
+        if let Some((cached_lsn, nblocks)) = rel_size_cache.get(tag) {
+            if lsn >= *cached_lsn {
+                return Some(*nblocks);
+            }
+        }
+        None
+    }
 
     /// Update cached relation size if there is no more recent update
-    fn update_cached_rel_size(&self, tag: RelTag, lsn: Lsn, nblocks: BlockNumber);
+    pub fn update_cached_rel_size(&self, tag: RelTag, lsn: Lsn, nblocks: BlockNumber) {
+        let mut rel_size_cache = self.rel_size_cache.write().unwrap();
+        match rel_size_cache.entry(tag) {
+            hash_map::Entry::Occupied(mut entry) => {
+                let cached_lsn = entry.get_mut();
+                if lsn >= cached_lsn.0 {
+                    *cached_lsn = (lsn, nblocks);
+                }
+            }
+            hash_map::Entry::Vacant(entry) => {
+                entry.insert((lsn, nblocks));
+            }
+        }
+    }
 
     /// Store cached relation size
-    fn set_cached_rel_size(&self, tag: RelTag, lsn: Lsn, nblocks: BlockNumber);
+    pub fn set_cached_rel_size(&self, tag: RelTag, lsn: Lsn, nblocks: BlockNumber) {
+        let mut rel_size_cache = self.rel_size_cache.write().unwrap();
+        rel_size_cache.insert(tag, (lsn, nblocks));
+    }
 
     /// Remove cached relation size
-    fn remove_cached_rel_size(&self, tag: &RelTag);
+    pub fn remove_cached_rel_size(&self, tag: &RelTag) {
+        let mut rel_size_cache = self.rel_size_cache.write().unwrap();
+        rel_size_cache.remove(tag);
+    }
 }
 
 /// DatadirModification represents an operation to ingest an atomic set of
 /// updates to the repository. It is created by the 'begin_record'
 /// function. It is called for each WAL record, so that all the modifications
 /// by a one WAL record appear atomic.
-pub struct DatadirModification<'a, T: DatadirTimeline> {
+pub struct DatadirModification<'a> {
     /// The timeline this modification applies to. You can access this to
     /// read the state, but note that any pending updates are *not* reflected
     /// in the state in 'tline' yet.
-    pub tline: &'a T,
+    pub tline: &'a Timeline,
 
     /// Lsn assigned by begin_modification
     pub lsn: Lsn,
@@ -495,10 +517,10 @@ pub struct DatadirModification<'a, T: DatadirTimeline> {
     // underlying key-value store by the 'finish' function.
     pending_updates: HashMap<Key, Value>,
     pending_deletions: Vec<Range<Key>>,
-    pending_nblocks: isize,
+    pending_nblocks: i64,
 }
 
-impl<'a, T: DatadirTimeline> DatadirModification<'a, T> {
+impl<'a> DatadirModification<'a> {
     /// Initialize a completely new repository.
     ///
     /// This inserts the directory metadata entries that are assumed to
@@ -654,7 +676,7 @@ impl<'a, T: DatadirTimeline> DatadirModification<'a, T> {
         }
 
         // Update logical database size.
-        self.pending_nblocks -= total_blocks as isize;
+        self.pending_nblocks -= total_blocks as i64;
 
         // Delete all relations and metadata files for the spcnode/dnode
         self.delete(dbdir_key_range(spcnode, dbnode));
@@ -697,7 +719,7 @@ impl<'a, T: DatadirTimeline> DatadirModification<'a, T> {
         let buf = nblocks.to_le_bytes();
         self.put(size_key, Value::Image(Bytes::from(buf.to_vec())));
 
-        self.pending_nblocks += nblocks as isize;
+        self.pending_nblocks += nblocks as i64;
 
         // Update relation size cache
         self.tline.set_cached_rel_size(rel, self.lsn, nblocks);
@@ -727,7 +749,7 @@ impl<'a, T: DatadirTimeline> DatadirModification<'a, T> {
             self.tline.set_cached_rel_size(rel, self.lsn, nblocks);
 
             // Update logical database size.
-            self.pending_nblocks -= old_size as isize - nblocks as isize;
+            self.pending_nblocks -= old_size as i64 - nblocks as i64;
         }
         Ok(())
     }
@@ -749,7 +771,7 @@ impl<'a, T: DatadirTimeline> DatadirModification<'a, T> {
             // Update relation size cache
             self.tline.set_cached_rel_size(rel, self.lsn, nblocks);
 
-            self.pending_nblocks += nblocks as isize - old_size as isize;
+            self.pending_nblocks += nblocks as i64 - old_size as i64;
         }
         Ok(())
     }
@@ -772,7 +794,7 @@ impl<'a, T: DatadirTimeline> DatadirModification<'a, T> {
         // update logical size
         let size_key = rel_size_to_key(rel);
         let old_size = self.get(size_key)?.get_u32_le();
-        self.pending_nblocks -= old_size as isize;
+        self.pending_nblocks -= old_size as i64;
 
         // Remove enty from relation size cache
         self.tline.remove_cached_rel_size(&rel);
@@ -914,7 +936,7 @@ impl<'a, T: DatadirTimeline> DatadirModification<'a, T> {
         result?;
 
         if pending_nblocks != 0 {
-            writer.update_current_logical_size(pending_nblocks * BLCKSZ as isize);
+            writer.update_current_logical_size(pending_nblocks * BLCKSZ as i64);
             self.pending_nblocks = 0;
         }
 
@@ -942,7 +964,7 @@ impl<'a, T: DatadirTimeline> DatadirModification<'a, T> {
         writer.finish_write(lsn);
 
         if pending_nblocks != 0 {
-            writer.update_current_logical_size(pending_nblocks * BLCKSZ as isize);
+            writer.update_current_logical_size(pending_nblocks * BLCKSZ as i64);
         }
 
         Ok(())
@@ -1368,10 +1390,10 @@ fn is_slru_block_key(key: Key) -> bool {
 //
 
 #[cfg(test)]
-pub fn create_test_timeline<R: Repository>(
-    repo: R,
+pub fn create_test_timeline(
+    repo: &crate::layered_repository::Repository,
     timeline_id: utils::zid::ZTimelineId,
-) -> Result<std::sync::Arc<R::Timeline>> {
+) -> Result<std::sync::Arc<Timeline>> {
     let tline = repo.create_empty_timeline(timeline_id, Lsn(8))?;
     let mut m = tline.begin_modification(Lsn(8));
     m.init_empty()?;

pageserver/src/repository.rs

@@ -1,19 +1,13 @@
 use crate::layered_repository::metadata::TimelineMetadata;
-use crate::storage_sync::index::RemoteIndex;
 use crate::walrecord::ZenithWalRecord;
-use crate::CheckpointConfig;
 use anyhow::{bail, Result};
 use byteorder::{ByteOrder, BE};
 use bytes::Bytes;
 use serde::{Deserialize, Serialize};
 use std::fmt;
 use std::ops::{AddAssign, Range};
-use std::sync::{Arc, RwLockReadGuard};
+use std::sync::Arc;
 use std::time::Duration;
-use utils::{
-    lsn::{Lsn, RecordLsn},
-    zid::ZTimelineId,
-};
 
 #[derive(Debug, Clone, Copy, Hash, PartialEq, Eq, Ord, PartialOrd, Serialize, Deserialize)]
 /// Key used in the Repository kv-store.
@@ -181,78 +175,6 @@ impl Value {
     }
 }
 
-///
-/// A repository corresponds to one .neon directory. One repository holds multiple
-/// timelines, forked off from the same initial call to 'initdb'.
-pub trait Repository: Send + Sync {
-    type Timeline: crate::DatadirTimeline;
-
-    /// Updates timeline based on the `TimelineSyncStatusUpdate`, received from the remote storage synchronization.
-    /// See [`crate::remote_storage`] for more details about the synchronization.
-    fn attach_timeline(&self, timeline_id: ZTimelineId) -> Result<()>;
-
-    /// Get Timeline handle for given zenith timeline ID.
-    /// This function is idempotent. It doesn't change internal state in any way.
-    fn get_timeline(&self, timelineid: ZTimelineId) -> Option<RepositoryTimeline<Self::Timeline>>;
-
-    /// Get Timeline handle for locally available timeline. Load it into memory if it is not loaded.
-    fn get_timeline_load(&self, timelineid: ZTimelineId) -> Result<Arc<Self::Timeline>>;
-
-    /// Lists timelines the repository contains.
-    /// Up to repository's implementation to omit certain timelines that ar not considered ready for use.
-    fn list_timelines(&self) -> Vec<(ZTimelineId, RepositoryTimeline<Self::Timeline>)>;
-
-    /// Create a new, empty timeline. The caller is responsible for loading data into it
-    /// Initdb lsn is provided for timeline impl to be able to perform checks for some operations against it.
-    fn create_empty_timeline(
-        &self,
-        timeline_id: ZTimelineId,
-        initdb_lsn: Lsn,
-    ) -> Result<Arc<Self::Timeline>>;
-
-    /// Branch a timeline
-    fn branch_timeline(
-        &self,
-        src: ZTimelineId,
-        dst: ZTimelineId,
-        start_lsn: Option<Lsn>,
-    ) -> Result<()>;
-
-    /// Flush all data to disk.
-    ///
-    /// this is used at graceful shutdown.
-    fn checkpoint(&self) -> Result<()>;
-
-    /// perform one garbage collection iteration, removing old data files from disk.
-    /// this function is periodically called by gc thread.
-    /// also it can be explicitly requested through page server api 'do_gc' command.
-    ///
-    /// 'timelineid' specifies the timeline to GC, or None for all.
-    /// `horizon` specifies delta from last lsn to preserve all object versions (pitr interval).
-    /// `checkpoint_before_gc` parameter is used to force compaction of storage before GC
-    /// to make tests more deterministic.
-    /// TODO Do we still need it or we can call checkpoint explicitly in tests where needed?
-    fn gc_iteration(
-        &self,
-        timelineid: Option<ZTimelineId>,
-        horizon: u64,
-        pitr: Duration,
-        checkpoint_before_gc: bool,
-    ) -> Result<GcResult>;
-
-    /// Perform one compaction iteration.
-    /// This function is periodically called by compactor thread.
-    /// Also it can be explicitly requested per timeline through page server
-    /// api's 'compact' command.
-    fn compaction_iteration(&self) -> Result<()>;
-
-    /// removes timeline-related in-memory data
-    fn delete_timeline(&self, timeline_id: ZTimelineId) -> anyhow::Result<()>;
-
-    /// Allows to retrieve remote timeline index from the repo. Used in walreceiver to grab remote consistent lsn.
-    fn get_remote_index(&self) -> &RemoteIndex;
-}
-
 /// A timeline, that belongs to the current repository.
 pub enum RepositoryTimeline<T> {
     /// Timeline, with its files present locally in pageserver's working directory.
@@ -304,621 +226,3 @@ impl AddAssign for GcResult {
         self.elapsed += other.elapsed;
     }
 }
-
-pub trait Timeline: Send + Sync {
-    //------------------------------------------------------------------------------
-    // Public GET functions
-    //------------------------------------------------------------------------------
-
-    ///
-    /// Wait until WAL has been received and processed up to this LSN.
-    ///
-    /// You should call this before any of the other get_* or list_* functions. Calling
-    /// those functions with an LSN that has been processed yet is an error.
-    ///
-    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.
-    ///
-    /// NOTE: It is considered an error to 'get' a key that doesn't exist. The abstraction
-    /// above this needs to store suitable metadata to track what data exists with
-    /// what keys, in separate metadata entries. If a non-existent key is requested,
-    /// the Repository implementation may incorrectly return a value from an ancestor
-    /// branch, for example, or waste a lot of cycles chasing the non-existing key.
-    ///
-    fn get(&self, key: Key, lsn: Lsn) -> Result<Bytes>;
-
-    /// Get the ancestor's timeline id
-    fn get_ancestor_timeline_id(&self) -> Option<ZTimelineId>;
-
-    /// Get the LSN where this branch was created
-    fn get_ancestor_lsn(&self) -> Lsn;
-
-    //------------------------------------------------------------------------------
-    // Public PUT functions, to update the repository with new page versions.
-    //
-    // These are called by the WAL receiver to digest WAL records.
-    //------------------------------------------------------------------------------
-    /// 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_disk_consistent_lsn(&self) -> Lsn;
-
-    /// Mutate the timeline with a [`TimelineWriter`].
-    ///
-    /// FIXME: This ought to return &'a TimelineWriter, where TimelineWriter
-    /// is a generic type in this trait. But that doesn't currently work in
-    /// Rust: https://rust-lang.github.io/rfcs/1598-generic_associated_types.html
-    fn writer<'a>(&'a self) -> Box<dyn TimelineWriter + 'a>;
-
-    ///
-    /// Flush to disk all data that was written with the put_* functions
-    ///
-    /// NOTE: This has nothing to do with checkpoint in PostgreSQL. We don't
-    /// know anything about them here in the repository.
-    fn checkpoint(&self, cconf: CheckpointConfig) -> Result<()>;
-
-    ///
-    /// 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<()>;
-
-    /// Get the physical size of the timeline at the latest LSN
-    fn get_physical_size(&self) -> u64;
-    /// Get the physical size of the timeline at the latest LSN non incrementally
-    fn get_physical_size_non_incremental(&self) -> Result<u64>;
-}
-
-/// Various functions to mutate the timeline.
-// TODO Currently, Deref is used to allow easy access to read methods from this trait.
-// This is probably considered a bad practice in Rust and should be fixed eventually,
-// but will cause large code changes.
-pub trait TimelineWriter<'a> {
-    /// Put a new page version that can be constructed from a WAL record
-    ///
-    /// This will implicitly extend the relation, if the page is beyond the
-    /// current end-of-file.
-    fn put(&self, key: Key, lsn: Lsn, value: &Value) -> Result<()>;
-
-    fn delete(&self, key_range: Range<Key>, lsn: Lsn) -> Result<()>;
-
-    /// Track the end of the latest digested WAL record.
-    ///
-    /// Call this after you have finished writing all the WAL up to 'lsn'.
-    ///
-    /// 'lsn' must be aligned. This wakes up any wait_lsn() callers waiting for
-    /// the 'lsn' or anything older. The previous last record LSN is stored alongside
-    /// the latest and can be read.
-    fn finish_write(&self, lsn: Lsn);
-
-    fn update_current_logical_size(&self, delta: isize);
-}
-
-#[cfg(test)]
-pub mod repo_harness {
-    use bytes::BytesMut;
-    use once_cell::sync::Lazy;
-    use std::sync::{Arc, RwLock, RwLockReadGuard, RwLockWriteGuard};
-    use std::{fs, path::PathBuf};
-
-    use crate::{
-        config::PageServerConf,
-        layered_repository::LayeredRepository,
-        walredo::{WalRedoError, WalRedoManager},
-    };
-
-    use super::*;
-    use crate::tenant_config::{TenantConf, TenantConfOpt};
-    use hex_literal::hex;
-    use utils::zid::ZTenantId;
-
-    pub const TIMELINE_ID: ZTimelineId =
-        ZTimelineId::from_array(hex!("11223344556677881122334455667788"));
-    pub const NEW_TIMELINE_ID: ZTimelineId =
-        ZTimelineId::from_array(hex!("AA223344556677881122334455667788"));
-
-    /// Convenience function to create a page image with given string as the only content
-    #[allow(non_snake_case)]
-    pub fn TEST_IMG(s: &str) -> Bytes {
-        let mut buf = BytesMut::new();
-        buf.extend_from_slice(s.as_bytes());
-        buf.resize(64, 0);
-
-        buf.freeze()
-    }
-
-    static LOCK: Lazy<RwLock<()>> = Lazy::new(|| RwLock::new(()));
-
-    impl From<TenantConf> for TenantConfOpt {
-        fn from(tenant_conf: TenantConf) -> Self {
-            Self {
-                checkpoint_distance: Some(tenant_conf.checkpoint_distance),
-                checkpoint_timeout: Some(tenant_conf.checkpoint_timeout),
-                compaction_target_size: Some(tenant_conf.compaction_target_size),
-                compaction_period: Some(tenant_conf.compaction_period),
-                compaction_threshold: Some(tenant_conf.compaction_threshold),
-                gc_horizon: Some(tenant_conf.gc_horizon),
-                gc_period: Some(tenant_conf.gc_period),
-                image_creation_threshold: Some(tenant_conf.image_creation_threshold),
-                pitr_interval: Some(tenant_conf.pitr_interval),
-                walreceiver_connect_timeout: Some(tenant_conf.walreceiver_connect_timeout),
-                lagging_wal_timeout: Some(tenant_conf.lagging_wal_timeout),
-                max_lsn_wal_lag: Some(tenant_conf.max_lsn_wal_lag),
-            }
-        }
-    }
-
-    pub struct RepoHarness<'a> {
-        pub conf: &'static PageServerConf,
-        pub tenant_conf: TenantConf,
-        pub tenant_id: ZTenantId,
-
-        pub lock_guard: (
-            Option<RwLockReadGuard<'a, ()>>,
-            Option<RwLockWriteGuard<'a, ()>>,
-        ),
-    }
-
-    impl<'a> RepoHarness<'a> {
-        pub fn create(test_name: &'static str) -> Result<Self> {
-            Self::create_internal(test_name, false)
-        }
-        pub fn create_exclusive(test_name: &'static str) -> Result<Self> {
-            Self::create_internal(test_name, true)
-        }
-        fn create_internal(test_name: &'static str, exclusive: bool) -> Result<Self> {
-            let lock_guard = if exclusive {
-                (None, Some(LOCK.write().unwrap()))
-            } else {
-                (Some(LOCK.read().unwrap()), None)
-            };
-
-            let repo_dir = PageServerConf::test_repo_dir(test_name);
-            let _ = fs::remove_dir_all(&repo_dir);
-            fs::create_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 tenant_conf = TenantConf::dummy_conf();
-
-            let tenant_id = ZTenantId::generate();
-            fs::create_dir_all(conf.tenant_path(&tenant_id))?;
-            fs::create_dir_all(conf.timelines_path(&tenant_id))?;
-
-            Ok(Self {
-                conf,
-                tenant_conf,
-                tenant_id,
-                lock_guard,
-            })
-        }
-
-        pub fn load(&self) -> LayeredRepository {
-            self.try_load().expect("failed to load test repo")
-        }
-
-        pub fn try_load(&self) -> Result<LayeredRepository> {
-            let walredo_mgr = Arc::new(TestRedoManager);
-
-            let repo = LayeredRepository::new(
-                self.conf,
-                TenantConfOpt::from(self.tenant_conf),
-                walredo_mgr,
-                self.tenant_id,
-                RemoteIndex::default(),
-                false,
-            );
-            // populate repo with locally available timelines
-            for timeline_dir_entry in fs::read_dir(self.conf.timelines_path(&self.tenant_id))
-                .expect("should be able to read timelines dir")
-            {
-                let timeline_dir_entry = timeline_dir_entry.unwrap();
-                let timeline_id: ZTimelineId = timeline_dir_entry
-                    .path()
-                    .file_name()
-                    .unwrap()
-                    .to_string_lossy()
-                    .parse()
-                    .unwrap();
-
-                repo.attach_timeline(timeline_id)?;
-            }
-
-            Ok(repo)
-        }
-
-        pub fn timeline_path(&self, timeline_id: &ZTimelineId) -> PathBuf {
-            self.conf.timeline_path(timeline_id, &self.tenant_id)
-        }
-    }
-
-    // Mock WAL redo manager that doesn't do much
-    pub struct TestRedoManager;
-
-    impl WalRedoManager for TestRedoManager {
-        fn request_redo(
-            &self,
-            key: Key,
-            lsn: Lsn,
-            base_img: Option<Bytes>,
-            records: Vec<(Lsn, ZenithWalRecord)>,
-        ) -> Result<Bytes, WalRedoError> {
-            let s = format!(
-                "redo for {} to get to {}, with {} and {} records",
-                key,
-                lsn,
-                if base_img.is_some() {
-                    "base image"
-                } else {
-                    "no base image"
-                },
-                records.len()
-            );
-            println!("{}", s);
-
-            Ok(TEST_IMG(&s))
-        }
-    }
-}
-
-///
-/// Tests that should work the same with any Repository/Timeline implementation.
-///
-#[allow(clippy::bool_assert_comparison)]
-#[cfg(test)]
-mod tests {
-    use super::repo_harness::*;
-    use super::*;
-    //use postgres_ffi::{pg_constants, xlog_utils::SIZEOF_CHECKPOINT};
-    //use std::sync::Arc;
-    use bytes::BytesMut;
-    use hex_literal::hex;
-    use once_cell::sync::Lazy;
-
-    static TEST_KEY: Lazy<Key> =
-        Lazy::new(|| Key::from_slice(&hex!("112222222233333333444444445500000001")));
-
-    #[test]
-    fn test_basic() -> Result<()> {
-        let repo = RepoHarness::create("test_basic")?.load();
-        let tline = repo.create_empty_timeline(TIMELINE_ID, Lsn(0))?;
-
-        let writer = tline.writer();
-        writer.put(*TEST_KEY, Lsn(0x10), &Value::Image(TEST_IMG("foo at 0x10")))?;
-        writer.finish_write(Lsn(0x10));
-        drop(writer);
-
-        let writer = tline.writer();
-        writer.put(*TEST_KEY, Lsn(0x20), &Value::Image(TEST_IMG("foo at 0x20")))?;
-        writer.finish_write(Lsn(0x20));
-        drop(writer);
-
-        assert_eq!(tline.get(*TEST_KEY, Lsn(0x10))?, TEST_IMG("foo at 0x10"));
-        assert_eq!(tline.get(*TEST_KEY, Lsn(0x1f))?, TEST_IMG("foo at 0x10"));
-        assert_eq!(tline.get(*TEST_KEY, Lsn(0x20))?, TEST_IMG("foo at 0x20"));
-
-        Ok(())
-    }
-
-    #[test]
-    fn no_duplicate_timelines() -> Result<()> {
-        let repo = RepoHarness::create("no_duplicate_timelines")?.load();
-        let _ = repo.create_empty_timeline(TIMELINE_ID, Lsn(0))?;
-
-        match repo.create_empty_timeline(TIMELINE_ID, Lsn(0)) {
-            Ok(_) => panic!("duplicate timeline creation should fail"),
-            Err(e) => assert_eq!(e.to_string(), "Timeline already exists"),
-        }
-
-        Ok(())
-    }
-
-    /// Convenience function to create a page image with given string as the only content
-    pub fn test_value(s: &str) -> Value {
-        let mut buf = BytesMut::new();
-        buf.extend_from_slice(s.as_bytes());
-        Value::Image(buf.freeze())
-    }
-
-    ///
-    /// Test branch creation
-    ///
-    #[test]
-    fn test_branch() -> Result<()> {
-        let repo = RepoHarness::create("test_branch")?.load();
-        let tline = repo.create_empty_timeline(TIMELINE_ID, Lsn(0))?;
-        let writer = tline.writer();
-        use std::str::from_utf8;
-
-        #[allow(non_snake_case)]
-        let TEST_KEY_A: Key = Key::from_hex("112222222233333333444444445500000001").unwrap();
-        #[allow(non_snake_case)]
-        let TEST_KEY_B: Key = Key::from_hex("112222222233333333444444445500000002").unwrap();
-
-        // Insert a value on the timeline
-        writer.put(TEST_KEY_A, Lsn(0x20), &test_value("foo at 0x20"))?;
-        writer.put(TEST_KEY_B, Lsn(0x20), &test_value("foobar at 0x20"))?;
-        writer.finish_write(Lsn(0x20));
-
-        writer.put(TEST_KEY_A, Lsn(0x30), &test_value("foo at 0x30"))?;
-        writer.finish_write(Lsn(0x30));
-        writer.put(TEST_KEY_A, Lsn(0x40), &test_value("foo at 0x40"))?;
-        writer.finish_write(Lsn(0x40));
-
-        //assert_current_logical_size(&tline, Lsn(0x40));
-
-        // Branch the history, modify relation differently on the new timeline
-        repo.branch_timeline(TIMELINE_ID, NEW_TIMELINE_ID, Some(Lsn(0x30)))?;
-        let newtline = repo
-            .get_timeline_load(NEW_TIMELINE_ID)
-            .expect("Should have a local timeline");
-        let new_writer = newtline.writer();
-        new_writer.put(TEST_KEY_A, Lsn(0x40), &test_value("bar at 0x40"))?;
-        new_writer.finish_write(Lsn(0x40));
-
-        // Check page contents on both branches
-        assert_eq!(
-            from_utf8(&tline.get(TEST_KEY_A, Lsn(0x40))?)?,
-            "foo at 0x40"
-        );
-        assert_eq!(
-            from_utf8(&newtline.get(TEST_KEY_A, Lsn(0x40))?)?,
-            "bar at 0x40"
-        );
-        assert_eq!(
-            from_utf8(&newtline.get(TEST_KEY_B, Lsn(0x40))?)?,
-            "foobar at 0x20"
-        );
-
-        //assert_current_logical_size(&tline, Lsn(0x40));
-
-        Ok(())
-    }
-
-    fn make_some_layers<T: Timeline>(tline: &T, start_lsn: Lsn) -> Result<()> {
-        let mut lsn = start_lsn;
-        #[allow(non_snake_case)]
-        {
-            let writer = tline.writer();
-            // Create a relation on the timeline
-            writer.put(
-                *TEST_KEY,
-                lsn,
-                &Value::Image(TEST_IMG(&format!("foo at {}", lsn))),
-            )?;
-            writer.finish_write(lsn);
-            lsn += 0x10;
-            writer.put(
-                *TEST_KEY,
-                lsn,
-                &Value::Image(TEST_IMG(&format!("foo at {}", lsn))),
-            )?;
-            writer.finish_write(lsn);
-            lsn += 0x10;
-        }
-        tline.checkpoint(CheckpointConfig::Forced)?;
-        {
-            let writer = tline.writer();
-            writer.put(
-                *TEST_KEY,
-                lsn,
-                &Value::Image(TEST_IMG(&format!("foo at {}", lsn))),
-            )?;
-            writer.finish_write(lsn);
-            lsn += 0x10;
-            writer.put(
-                *TEST_KEY,
-                lsn,
-                &Value::Image(TEST_IMG(&format!("foo at {}", lsn))),
-            )?;
-            writer.finish_write(lsn);
-        }
-        tline.checkpoint(CheckpointConfig::Forced)
-    }
-
-    #[test]
-    fn test_prohibit_branch_creation_on_garbage_collected_data() -> Result<()> {
-        let repo =
-            RepoHarness::create("test_prohibit_branch_creation_on_garbage_collected_data")?.load();
-        let tline = repo.create_empty_timeline(TIMELINE_ID, Lsn(0))?;
-        make_some_layers(tline.as_ref(), Lsn(0x20))?;
-
-        // this removes layers before lsn 40 (50 minus 10), so there are two remaining layers, image and delta for 31-50
-        // FIXME: this doesn't actually remove any layer currently, given how the checkpointing
-        // and compaction works. But it does set the 'cutoff' point so that the cross check
-        // below should fail.
-        repo.gc_iteration(Some(TIMELINE_ID), 0x10, Duration::ZERO, false)?;
-
-        // try to branch at lsn 25, should fail because we already garbage collected the data
-        match repo.branch_timeline(TIMELINE_ID, NEW_TIMELINE_ID, Some(Lsn(0x25))) {
-            Ok(_) => panic!("branching should have failed"),
-            Err(err) => {
-                assert!(err.to_string().contains("invalid branch start lsn"));
-                assert!(err
-                    .source()
-                    .unwrap()
-                    .to_string()
-                    .contains("we might've already garbage collected needed data"))
-            }
-        }
-
-        Ok(())
-    }
-
-    #[test]
-    fn test_prohibit_branch_creation_on_pre_initdb_lsn() -> Result<()> {
-        let repo = RepoHarness::create("test_prohibit_branch_creation_on_pre_initdb_lsn")?.load();
-
-        repo.create_empty_timeline(TIMELINE_ID, Lsn(0x50))?;
-        // try to branch at lsn 0x25, should fail because initdb lsn is 0x50
-        match repo.branch_timeline(TIMELINE_ID, NEW_TIMELINE_ID, Some(Lsn(0x25))) {
-            Ok(_) => panic!("branching should have failed"),
-            Err(err) => {
-                assert!(&err.to_string().contains("invalid branch start lsn"));
-                assert!(&err
-                    .source()
-                    .unwrap()
-                    .to_string()
-                    .contains("is earlier than latest GC horizon"));
-            }
-        }
-
-        Ok(())
-    }
-
-    /*
-    // FIXME: This currently fails to error out. Calling GC doesn't currently
-    // remove the old value, we'd need to work a little harder
-    #[test]
-    fn test_prohibit_get_for_garbage_collected_data() -> Result<()> {
-        let repo =
-            RepoHarness::create("test_prohibit_get_for_garbage_collected_data")?
-            .load();
-
-        let tline = repo.create_empty_timeline(TIMELINE_ID, Lsn(0))?;
-        make_some_layers(tline.as_ref(), Lsn(0x20))?;
-
-        repo.gc_iteration(Some(TIMELINE_ID), 0x10, Duration::ZERO, false)?;
-        let latest_gc_cutoff_lsn = tline.get_latest_gc_cutoff_lsn();
-        assert!(*latest_gc_cutoff_lsn > Lsn(0x25));
-        match tline.get(*TEST_KEY, Lsn(0x25)) {
-            Ok(_) => panic!("request for page should have failed"),
-            Err(err) => assert!(err.to_string().contains("not found at")),
-        }
-        Ok(())
-    }
-     */
-
-    #[test]
-    fn test_retain_data_in_parent_which_is_needed_for_child() -> Result<()> {
-        let repo =
-            RepoHarness::create("test_retain_data_in_parent_which_is_needed_for_child")?.load();
-        let tline = repo.create_empty_timeline(TIMELINE_ID, Lsn(0))?;
-        make_some_layers(tline.as_ref(), Lsn(0x20))?;
-
-        repo.branch_timeline(TIMELINE_ID, NEW_TIMELINE_ID, Some(Lsn(0x40)))?;
-        let newtline = repo
-            .get_timeline_load(NEW_TIMELINE_ID)
-            .expect("Should have a local timeline");
-        // this removes layers before lsn 40 (50 minus 10), so there are two remaining layers, image and delta for 31-50
-        repo.gc_iteration(Some(TIMELINE_ID), 0x10, Duration::ZERO, false)?;
-        assert!(newtline.get(*TEST_KEY, Lsn(0x25)).is_ok());
-
-        Ok(())
-    }
-    #[test]
-    fn test_parent_keeps_data_forever_after_branching() -> Result<()> {
-        let repo = RepoHarness::create("test_parent_keeps_data_forever_after_branching")?.load();
-        let tline = repo.create_empty_timeline(TIMELINE_ID, Lsn(0))?;
-        make_some_layers(tline.as_ref(), Lsn(0x20))?;
-
-        repo.branch_timeline(TIMELINE_ID, NEW_TIMELINE_ID, Some(Lsn(0x40)))?;
-        let newtline = repo
-            .get_timeline_load(NEW_TIMELINE_ID)
-            .expect("Should have a local timeline");
-
-        make_some_layers(newtline.as_ref(), Lsn(0x60))?;
-
-        // run gc on parent
-        repo.gc_iteration(Some(TIMELINE_ID), 0x10, Duration::ZERO, false)?;
-
-        // Check that the data is still accessible on the branch.
-        assert_eq!(
-            newtline.get(*TEST_KEY, Lsn(0x50))?,
-            TEST_IMG(&format!("foo at {}", Lsn(0x40)))
-        );
-
-        Ok(())
-    }
-
-    #[test]
-    fn timeline_load() -> Result<()> {
-        const TEST_NAME: &str = "timeline_load";
-        let harness = RepoHarness::create(TEST_NAME)?;
-        {
-            let repo = harness.load();
-            let tline = repo.create_empty_timeline(TIMELINE_ID, Lsn(0x8000))?;
-            make_some_layers(tline.as_ref(), Lsn(0x8000))?;
-            tline.checkpoint(CheckpointConfig::Forced)?;
-        }
-
-        let repo = harness.load();
-        let tline = repo
-            .get_timeline(TIMELINE_ID)
-            .expect("cannot load timeline");
-        assert!(matches!(tline, RepositoryTimeline::Unloaded { .. }));
-
-        assert!(repo.get_timeline_load(TIMELINE_ID).is_ok());
-
-        let tline = repo
-            .get_timeline(TIMELINE_ID)
-            .expect("cannot load timeline");
-        assert!(matches!(tline, RepositoryTimeline::Loaded(_)));
-
-        Ok(())
-    }
-
-    #[test]
-    fn timeline_load_with_ancestor() -> Result<()> {
-        const TEST_NAME: &str = "timeline_load_with_ancestor";
-        let harness = RepoHarness::create(TEST_NAME)?;
-        // create two timelines
-        {
-            let repo = harness.load();
-            let tline = repo.create_empty_timeline(TIMELINE_ID, Lsn(0))?;
-
-            make_some_layers(tline.as_ref(), Lsn(0x20))?;
-            tline.checkpoint(CheckpointConfig::Forced)?;
-
-            repo.branch_timeline(TIMELINE_ID, NEW_TIMELINE_ID, Some(Lsn(0x40)))?;
-
-            let newtline = repo
-                .get_timeline_load(NEW_TIMELINE_ID)
-                .expect("Should have a local timeline");
-
-            make_some_layers(newtline.as_ref(), Lsn(0x60))?;
-            tline.checkpoint(CheckpointConfig::Forced)?;
-        }
-
-        // check that both of them are initially unloaded
-        let repo = harness.load();
-        {
-            let tline = repo.get_timeline(TIMELINE_ID).expect("cannot get timeline");
-            assert!(matches!(tline, RepositoryTimeline::Unloaded { .. }));
-
-            let tline = repo
-                .get_timeline(NEW_TIMELINE_ID)
-                .expect("cannot get timeline");
-            assert!(matches!(tline, RepositoryTimeline::Unloaded { .. }));
-        }
-        // load only child timeline
-        let _ = repo
-            .get_timeline_load(NEW_TIMELINE_ID)
-            .expect("cannot load timeline");
-
-        // check that both, child and ancestor are loaded
-        let tline = repo
-            .get_timeline(NEW_TIMELINE_ID)
-            .expect("cannot get timeline");
-        assert!(matches!(tline, RepositoryTimeline::Loaded(_)));
-
-        let tline = repo.get_timeline(TIMELINE_ID).expect("cannot get timeline");
-        assert!(matches!(tline, RepositoryTimeline::Loaded(_)));
-
-        Ok(())
-    }
-}

pageserver/src/storage_sync.rs

@@ -1642,7 +1642,7 @@ fn register_sync_status(
 mod test_utils {
     use utils::lsn::Lsn;
 
-    use crate::repository::repo_harness::RepoHarness;
+    use crate::layered_repository::repo_harness::RepoHarness;
 
     use super::*;
 
@@ -1687,7 +1687,7 @@ mod test_utils {
 #[cfg(test)]
 mod tests {
     use super::test_utils::dummy_metadata;
-    use crate::repository::repo_harness::TIMELINE_ID;
+    use crate::layered_repository::repo_harness::TIMELINE_ID;
     use hex_literal::hex;
     use utils::lsn::Lsn;
 

pageserver/src/storage_sync/delete.rs

@@ -111,7 +111,7 @@ mod tests {
     use utils::lsn::Lsn;
 
     use crate::{
-        repository::repo_harness::{RepoHarness, TIMELINE_ID},
+        layered_repository::repo_harness::{RepoHarness, TIMELINE_ID},
         storage_sync::test_utils::{create_local_timeline, dummy_metadata},
     };
     use remote_storage::LocalFs;

pageserver/src/storage_sync/download.rs

@@ -445,7 +445,7 @@ mod tests {
     use utils::lsn::Lsn;
 
     use crate::{
-        repository::repo_harness::{RepoHarness, TIMELINE_ID},
+        layered_repository::repo_harness::{RepoHarness, TIMELINE_ID},
         storage_sync::{
             index::RelativePath,
             test_utils::{create_local_timeline, dummy_metadata},

pageserver/src/storage_sync/index.rs

@@ -210,7 +210,7 @@ impl RemoteTimelineIndex {
 }
 
 /// Restored index part data about the timeline, stored in the remote index.
-#[derive(Debug, PartialEq, Eq, Clone)]
+#[derive(Debug, Clone)]
 pub struct RemoteTimeline {
     timeline_layers: HashSet<PathBuf>,
     missing_layers: HashSet<PathBuf>,
@@ -341,7 +341,7 @@ mod tests {
     use std::collections::BTreeSet;
 
     use super::*;
-    use crate::repository::repo_harness::{RepoHarness, TIMELINE_ID};
+    use crate::layered_repository::repo_harness::{RepoHarness, TIMELINE_ID};
 
     #[test]
     fn index_part_conversion() {

pageserver/src/storage_sync/upload.rs

@@ -248,7 +248,7 @@ mod tests {
     use utils::lsn::Lsn;
 
     use crate::{
-        repository::repo_harness::{RepoHarness, TIMELINE_ID},
+        layered_repository::repo_harness::{RepoHarness, TIMELINE_ID},
         storage_sync::{
             index::RelativePath,
             test_utils::{create_local_timeline, dummy_metadata},

pageserver/src/tenant_mgr.rs

@@ -3,8 +3,8 @@
 
 use crate::config::PageServerConf;
 use crate::http::models::TenantInfo;
-use crate::layered_repository::{load_metadata, LayeredRepository, LayeredTimeline};
-use crate::repository::Repository;
+use crate::layered_repository::{load_metadata, Repository, Timeline};
+use crate::repository::RepositoryTimeline;
 use crate::storage_sync::index::{RemoteIndex, RemoteTimelineIndex};
 use crate::storage_sync::{self, LocalTimelineInitStatus, SyncStartupData};
 use crate::tenant_config::TenantConfOpt;
@@ -21,6 +21,7 @@ use tokio::sync::mpsc;
 use tracing::*;
 use utils::lsn::Lsn;
 
+pub use tenants_state::try_send_timeline_update;
 use utils::zid::{ZTenantId, ZTenantTimelineId, ZTimelineId};
 
 mod tenants_state {
@@ -68,7 +69,7 @@ mod tenants_state {
         Ok(())
     }
 
-    pub(super) fn try_send_timeline_update(update: LocalTimelineUpdate) {
+    pub fn try_send_timeline_update(update: LocalTimelineUpdate) {
         match TIMELINE_UPDATE_SENDER
             .read()
             .expect("Failed to read() timeline_update_sender lock, it got poisoned")
@@ -94,13 +95,7 @@ mod tenants_state {
 struct Tenant {
     state: TenantState,
     /// Contains in-memory state, including the timeline that might not yet flushed on disk or loaded form disk.
-    repo: Arc<LayeredRepository>,
-    /// Timelines, located locally in the pageserver's datadir.
-    /// Timelines can entirely be removed entirely by the `detach` operation only.
-    ///
-    /// Local timelines have more metadata that's loaded into memory,
-    /// that is located in the `repo.timelines` field, [`crate::layered_repository::LayeredTimelineEntry`].
-    local_timelines: HashMap<ZTimelineId, Arc<LayeredTimeline>>,
+    repo: Arc<Repository>,
 }
 
 #[derive(Debug, Serialize, Deserialize, Clone, Copy, PartialEq, Eq)]
@@ -177,15 +172,15 @@ pub enum LocalTimelineUpdate {
     },
     Attach {
         id: ZTenantTimelineId,
-        datadir: Arc<LayeredTimeline>,
+        timeline: Arc<Timeline>,
     },
 }
 
 impl std::fmt::Debug for LocalTimelineUpdate {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         match self {
-            Self::Detach { id, .. } => f.debug_tuple("Remove").field(id).finish(),
-            Self::Attach { id, .. } => f.debug_tuple("Add").field(id).finish(),
+            Self::Detach { id, .. } => f.debug_tuple("Detach").field(id).finish(),
+            Self::Attach { id, .. } => f.debug_tuple("Attach").field(id).finish(),
         }
     }
 }
@@ -289,7 +284,6 @@ pub fn create_tenant_repository(
             v.insert(Tenant {
                 state: TenantState::Idle,
                 repo,
-                local_timelines: HashMap::new(),
             });
             Ok(Some(tenant_id))
         }
@@ -365,7 +359,7 @@ pub fn set_tenant_state(tenant_id: ZTenantId, new_state: TenantState) -> anyhow:
     Ok(())
 }
 
-pub fn get_repository_for_tenant(tenant_id: ZTenantId) -> anyhow::Result<Arc<LayeredRepository>> {
+pub fn get_repository_for_tenant(tenant_id: ZTenantId) -> anyhow::Result<Arc<Repository>> {
     let m = tenants_state::read_tenants();
     let tenant = m
         .get(&tenant_id)
@@ -379,21 +373,15 @@ pub fn get_repository_for_tenant(tenant_id: ZTenantId) -> anyhow::Result<Arc<Lay
 pub fn get_local_timeline_with_load(
     tenant_id: ZTenantId,
     timeline_id: ZTimelineId,
-) -> anyhow::Result<Arc<LayeredTimeline>> {
-    let mut m = tenants_state::write_tenants();
-    let tenant = m
-        .get_mut(&tenant_id)
-        .with_context(|| format!("Tenant {tenant_id} not found"))?;
-
-    if let Some(page_tline) = tenant.local_timelines.get(&timeline_id) {
-        Ok(Arc::clone(page_tline))
-    } else {
-        let page_tline = load_local_timeline(&tenant.repo, timeline_id)
-            .with_context(|| format!("Failed to load local timeline for tenant {tenant_id}"))?;
-        tenant
-            .local_timelines
-            .insert(timeline_id, Arc::clone(&page_tline));
-        Ok(page_tline)
+) -> anyhow::Result<Arc<Timeline>> {
+    let repository = get_repository_for_tenant(tenant_id)?;
+    match repository.get_timeline(timeline_id) {
+        Some(RepositoryTimeline::Loaded(loaded_timeline)) => {
+            loaded_timeline.init_logical_size()?;
+            Ok(loaded_timeline)
+        }
+        _ => load_local_timeline(&repository, timeline_id)
+            .with_context(|| format!("Failed to load local timeline for tenant {tenant_id}")),
     }
 }
 
@@ -420,10 +408,7 @@ pub fn delete_timeline(tenant_id: ZTenantId, timeline_id: ZTimelineId) -> anyhow
     thread_mgr::shutdown_threads(None, None, Some(timeline_id));
     debug!("thread shutdown completed");
     match tenants_state::write_tenants().get_mut(&tenant_id) {
-        Some(tenant) => {
-            tenant.repo.delete_timeline(timeline_id)?;
-            tenant.local_timelines.remove(&timeline_id);
-        }
+        Some(tenant) => tenant.repo.delete_timeline(timeline_id)?,
         None => anyhow::bail!("Tenant {tenant_id} not found in local tenant state"),
     }
 
@@ -435,23 +420,21 @@ pub fn detach_tenant(conf: &'static PageServerConf, tenant_id: ZTenantId) -> any
     // shutdown the tenant and timeline threads: gc, compaction, page service threads)
     thread_mgr::shutdown_threads(None, Some(tenant_id), None);
 
-    // FIXME should we protect somehow from starting new threads/walreceivers when tenant is in stopping state?
-    // send stop signal to wal receiver and collect join handles while holding the lock
-    let walreceiver_join_handles = {
-        let tenants = tenants_state::write_tenants();
-        let tenant = tenants.get(&tenant_id).context("tenant not found")?;
-        let mut walreceiver_join_handles = Vec::with_capacity(tenant.local_timelines.len());
-        for timeline_id in tenant.local_timelines.keys() {
+    let mut walreceiver_join_handles = Vec::new();
+    let removed_tenant = {
+        let mut tenants_accessor = tenants_state::write_tenants();
+        tenants_accessor.remove(&tenant_id)
+    };
+    if let Some(tenant) = removed_tenant {
+        for (timeline_id, _) in tenant.repo.list_timelines() {
             let (sender, receiver) = std::sync::mpsc::channel::<()>();
             tenants_state::try_send_timeline_update(LocalTimelineUpdate::Detach {
-                id: ZTenantTimelineId::new(tenant_id, *timeline_id),
+                id: ZTenantTimelineId::new(tenant_id, timeline_id),
                 join_confirmation_sender: sender,
             });
-            walreceiver_join_handles.push((*timeline_id, receiver));
+            walreceiver_join_handles.push((timeline_id, receiver));
         }
-        // drop the tenants lock
-        walreceiver_join_handles
-    };
+    }
 
     // wait for wal receivers to stop without holding the lock, because walreceiver
     // will attempt to change tenant state which is protected by the same global tenants lock.
@@ -484,19 +467,13 @@ pub fn detach_tenant(conf: &'static PageServerConf, tenant_id: ZTenantId) -> any
 }
 
 fn load_local_timeline(
-    repo: &LayeredRepository,
+    repo: &Repository,
     timeline_id: ZTimelineId,
-) -> anyhow::Result<Arc<LayeredTimeline>> {
+) -> anyhow::Result<Arc<Timeline>> {
     let inmem_timeline = repo.get_timeline_load(timeline_id).with_context(|| {
         format!("Inmem timeline {timeline_id} not found in tenant's repository")
     })?;
     inmem_timeline.init_logical_size()?;
-
-    tenants_state::try_send_timeline_update(LocalTimelineUpdate::Attach {
-        id: ZTenantTimelineId::new(repo.tenant_id(), timeline_id),
-        datadir: Arc::clone(&inmem_timeline),
-    });
-
     Ok(inmem_timeline)
 }
 
@@ -588,37 +565,24 @@ fn init_local_repository(
 }
 
 fn attach_downloaded_tenant(
-    repo: &LayeredRepository,
+    repo: &Repository,
     downloaded_timelines: HashSet<ZTimelineId>,
 ) -> anyhow::Result<()> {
-    let mut registration_queue = Vec::with_capacity(downloaded_timelines.len());
-
-    // first need to register the in-mem representations, to avoid missing ancestors during the local disk data registration
-    for timeline_id in downloaded_timelines {
+    // first, register timeline metadata to ensure ancestors will be found later during layer load
+    for &timeline_id in &downloaded_timelines {
         repo.attach_timeline(timeline_id).with_context(|| {
             format!("Failed to load timeline {timeline_id} into in-memory repository")
         })?;
-        registration_queue.push(timeline_id);
     }
 
-    for timeline_id in registration_queue {
-        let tenant_id = repo.tenant_id();
-        match tenants_state::write_tenants().get_mut(&tenant_id) {
-            Some(tenant) => match tenant.local_timelines.entry(timeline_id) {
-                Entry::Occupied(_) => {
-                    anyhow::bail!("Local timeline {timeline_id} already registered")
-                }
-                Entry::Vacant(v) => {
-                    v.insert(load_local_timeline(repo, timeline_id).with_context(|| {
-                        format!("Failed to register add local timeline for tenant {tenant_id}")
-                    })?);
-                }
-            },
-            None => anyhow::bail!(
-                "Tenant {} not found in local tenant state",
-                repo.tenant_id()
-            ),
-        }
+    // and then load its layers in memory
+    for timeline_id in downloaded_timelines {
+        let _ = load_local_timeline(repo, timeline_id).with_context(|| {
+            format!(
+                "Failed to register add local timeline for tenant {}",
+                repo.tenant_id(),
+            )
+        })?;
     }
 
     Ok(())
@@ -630,14 +594,14 @@ fn load_local_repo(
     conf: &'static PageServerConf,
     tenant_id: ZTenantId,
     remote_index: &RemoteIndex,
-) -> anyhow::Result<Arc<LayeredRepository>> {
+) -> anyhow::Result<Arc<Repository>> {
     let mut m = tenants_state::write_tenants();
     let tenant = m.entry(tenant_id).or_insert_with(|| {
         // Set up a WAL redo manager, for applying WAL records.
         let walredo_mgr = PostgresRedoManager::new(conf, tenant_id);
 
         // Set up an object repository, for actual data storage.
-        let repo: Arc<LayeredRepository> = Arc::new(LayeredRepository::new(
+        let repo: Arc<Repository> = Arc::new(Repository::new(
             conf,
             TenantConfOpt::default(),
             Arc::new(walredo_mgr),
@@ -648,12 +612,11 @@ fn load_local_repo(
         Tenant {
             state: TenantState::Idle,
             repo,
-            local_timelines: HashMap::new(),
         }
     });
 
     // Restore tenant config
-    let tenant_conf = LayeredRepository::load_tenant_config(conf, tenant_id)?;
+    let tenant_conf = Repository::load_tenant_config(conf, tenant_id)?;
     tenant.repo.update_tenant_config(tenant_conf)?;
 
     Ok(Arc::clone(&tenant.repo))

pageserver/src/tenant_tasks.rs

@@ -5,7 +5,6 @@ use std::collections::HashMap;
 use std::ops::ControlFlow;
 use std::time::Duration;
 
-use crate::repository::Repository;
 use crate::tenant_mgr::TenantState;
 use crate::thread_mgr::ThreadKind;
 use crate::{tenant_mgr, thread_mgr};

pageserver/src/timelines.rs

@@ -20,15 +20,14 @@ use utils::{
 
 use crate::import_datadir;
 use crate::tenant_mgr;
+use crate::CheckpointConfig;
 use crate::{
-    config::PageServerConf, repository::Repository, storage_sync::index::RemoteIndex,
-    tenant_config::TenantConfOpt,
+    config::PageServerConf, storage_sync::index::RemoteIndex, tenant_config::TenantConfOpt,
 };
 use crate::{
-    layered_repository::{LayeredRepository, LayeredTimeline},
+    layered_repository::{Repository, Timeline},
     walredo::WalRedoManager,
 };
-use crate::{repository::Timeline, CheckpointConfig};
 
 #[derive(Debug, Clone, Copy)]
 pub struct PointInTime {
@@ -42,7 +41,7 @@ pub fn create_repo(
     tenant_id: ZTenantId,
     wal_redo_manager: Arc<dyn WalRedoManager + Send + Sync>,
     remote_index: RemoteIndex,
-) -> Result<Arc<LayeredRepository>> {
+) -> Result<Arc<Repository>> {
     let repo_dir = conf.tenant_path(&tenant_id);
     ensure!(
         !repo_dir.exists(),
@@ -57,9 +56,9 @@ pub fn create_repo(
     info!("created directory structure in {}", repo_dir.display());
 
     // Save tenant's config
-    LayeredRepository::persist_tenant_config(conf, tenant_id, tenant_conf)?;
+    Repository::persist_tenant_config(conf, tenant_id, tenant_conf)?;
 
-    Ok(Arc::new(LayeredRepository::new(
+    Ok(Arc::new(Repository::new(
         conf,
         tenant_conf,
         wal_redo_manager,
@@ -104,11 +103,11 @@ fn run_initdb(conf: &'static PageServerConf, initdbpath: &Path) -> Result<()> {
 // - run initdb to init temporary instance and get bootstrap data
 // - after initialization complete, remove the temp dir.
 //
-fn bootstrap_timeline<R: Repository>(
+fn bootstrap_timeline(
     conf: &'static PageServerConf,
     tenantid: ZTenantId,
     tli: ZTimelineId,
-    repo: &R,
+    repo: &Repository,
 ) -> Result<()> {
     let initdb_path = conf
         .tenant_path(&tenantid)
@@ -160,7 +159,7 @@ pub(crate) fn create_timeline(
     new_timeline_id: Option<ZTimelineId>,
     ancestor_timeline_id: Option<ZTimelineId>,
     mut ancestor_start_lsn: Option<Lsn>,
-) -> Result<Option<(ZTimelineId, Arc<LayeredTimeline>)>> {
+) -> Result<Option<(ZTimelineId, Arc<Timeline>)>> {
     let new_timeline_id = new_timeline_id.unwrap_or_else(ZTimelineId::generate);
     let repo = tenant_mgr::get_repository_for_tenant(tenant_id)?;
 

pageserver/src/walingest.rs

@@ -30,6 +30,7 @@ use anyhow::Result;
 use bytes::{Buf, Bytes, BytesMut};
 use tracing::*;
 
+use crate::layered_repository::Timeline;
 use crate::pgdatadir_mapping::*;
 use crate::reltag::{RelTag, SlruKind};
 use crate::walrecord::*;
@@ -43,15 +44,15 @@ use utils::lsn::Lsn;
 
 static ZERO_PAGE: Bytes = Bytes::from_static(&[0u8; 8192]);
 
-pub struct WalIngest<'a, T: DatadirTimeline> {
-    timeline: &'a T,
+pub struct WalIngest<'a> {
+    timeline: &'a Timeline,
 
     checkpoint: CheckPoint,
     checkpoint_modified: bool,
 }
 
-impl<'a, T: DatadirTimeline> WalIngest<'a, T> {
-    pub fn new(timeline: &T, startpoint: Lsn) -> Result<WalIngest<T>> {
+impl<'a> WalIngest<'a> {
+    pub fn new(timeline: &Timeline, startpoint: Lsn) -> Result<WalIngest> {
         // Fetch the latest checkpoint into memory, so that we can compare with it
         // quickly in `ingest_record` and update it when it changes.
         let checkpoint_bytes = timeline.get_checkpoint(startpoint)?;
@@ -77,7 +78,7 @@ impl<'a, T: DatadirTimeline> WalIngest<'a, T> {
         &mut self,
         recdata: Bytes,
         lsn: Lsn,
-        modification: &mut DatadirModification<T>,
+        modification: &mut DatadirModification,
         decoded: &mut DecodedWALRecord,
     ) -> Result<()> {
         modification.lsn = lsn;
@@ -266,7 +267,7 @@ impl<'a, T: DatadirTimeline> WalIngest<'a, T> {
 
     fn ingest_decoded_block(
         &mut self,
-        modification: &mut DatadirModification<T>,
+        modification: &mut DatadirModification,
         lsn: Lsn,
         decoded: &DecodedWALRecord,
         blk: &DecodedBkpBlock,
@@ -326,7 +327,7 @@ impl<'a, T: DatadirTimeline> WalIngest<'a, T> {
     fn ingest_heapam_record(
         &mut self,
         buf: &mut Bytes,
-        modification: &mut DatadirModification<T>,
+        modification: &mut DatadirModification,
         decoded: &mut DecodedWALRecord,
     ) -> Result<()> {
         // Handle VM bit updates that are implicitly part of heap records.
@@ -470,7 +471,7 @@ impl<'a, T: DatadirTimeline> WalIngest<'a, T> {
     /// Subroutine of ingest_record(), to handle an XLOG_DBASE_CREATE record.
     fn ingest_xlog_dbase_create(
         &mut self,
-        modification: &mut DatadirModification<T>,
+        modification: &mut DatadirModification,
         rec: &XlCreateDatabase,
     ) -> Result<()> {
         let db_id = rec.db_id;
@@ -537,7 +538,7 @@ impl<'a, T: DatadirTimeline> WalIngest<'a, T> {
 
     fn ingest_xlog_smgr_create(
         &mut self,
-        modification: &mut DatadirModification<T>,
+        modification: &mut DatadirModification,
         rec: &XlSmgrCreate,
     ) -> Result<()> {
         let rel = RelTag {
@@ -555,7 +556,7 @@ impl<'a, T: DatadirTimeline> WalIngest<'a, T> {
     /// This is the same logic as in PostgreSQL's smgr_redo() function.
     fn ingest_xlog_smgr_truncate(
         &mut self,
-        modification: &mut DatadirModification<T>,
+        modification: &mut DatadirModification,
         rec: &XlSmgrTruncate,
     ) -> Result<()> {
         let spcnode = rec.rnode.spcnode;
@@ -620,7 +621,7 @@ impl<'a, T: DatadirTimeline> WalIngest<'a, T> {
     ///
     fn ingest_xact_record(
         &mut self,
-        modification: &mut DatadirModification<T>,
+        modification: &mut DatadirModification,
         parsed: &XlXactParsedRecord,
         is_commit: bool,
     ) -> Result<()> {
@@ -689,7 +690,7 @@ impl<'a, T: DatadirTimeline> WalIngest<'a, T> {
 
     fn ingest_clog_truncate_record(
         &mut self,
-        modification: &mut DatadirModification<T>,
+        modification: &mut DatadirModification,
         xlrec: &XlClogTruncate,
     ) -> Result<()> {
         info!(
@@ -747,7 +748,7 @@ impl<'a, T: DatadirTimeline> WalIngest<'a, T> {
 
     fn ingest_multixact_create_record(
         &mut self,
-        modification: &mut DatadirModification<T>,
+        modification: &mut DatadirModification,
         xlrec: &XlMultiXactCreate,
     ) -> Result<()> {
         // Create WAL record for updating the multixact-offsets page
@@ -826,7 +827,7 @@ impl<'a, T: DatadirTimeline> WalIngest<'a, T> {
 
     fn ingest_multixact_truncate_record(
         &mut self,
-        modification: &mut DatadirModification<T>,
+        modification: &mut DatadirModification,
         xlrec: &XlMultiXactTruncate,
     ) -> Result<()> {
         self.checkpoint.oldestMulti = xlrec.end_trunc_off;
@@ -860,7 +861,7 @@ impl<'a, T: DatadirTimeline> WalIngest<'a, T> {
 
     fn ingest_relmap_page(
         &mut self,
-        modification: &mut DatadirModification<T>,
+        modification: &mut DatadirModification,
         xlrec: &XlRelmapUpdate,
         decoded: &DecodedWALRecord,
     ) -> Result<()> {
@@ -876,7 +877,7 @@ impl<'a, T: DatadirTimeline> WalIngest<'a, T> {
 
     fn put_rel_creation(
         &mut self,
-        modification: &mut DatadirModification<T>,
+        modification: &mut DatadirModification,
         rel: RelTag,
     ) -> Result<()> {
         modification.put_rel_creation(rel, 0)?;
@@ -885,7 +886,7 @@ impl<'a, T: DatadirTimeline> WalIngest<'a, T> {
 
     fn put_rel_page_image(
         &mut self,
-        modification: &mut DatadirModification<T>,
+        modification: &mut DatadirModification,
         rel: RelTag,
         blknum: BlockNumber,
         img: Bytes,
@@ -897,7 +898,7 @@ impl<'a, T: DatadirTimeline> WalIngest<'a, T> {
 
     fn put_rel_wal_record(
         &mut self,
-        modification: &mut DatadirModification<T>,
+        modification: &mut DatadirModification,
         rel: RelTag,
         blknum: BlockNumber,
         rec: ZenithWalRecord,
@@ -909,7 +910,7 @@ impl<'a, T: DatadirTimeline> WalIngest<'a, T> {
 
     fn put_rel_truncation(
         &mut self,
-        modification: &mut DatadirModification<T>,
+        modification: &mut DatadirModification,
         rel: RelTag,
         nblocks: BlockNumber,
     ) -> Result<()> {
@@ -917,11 +918,7 @@ impl<'a, T: DatadirTimeline> WalIngest<'a, T> {
         Ok(())
     }
 
-    fn put_rel_drop(
-        &mut self,
-        modification: &mut DatadirModification<T>,
-        rel: RelTag,
-    ) -> Result<()> {
+    fn put_rel_drop(&mut self, modification: &mut DatadirModification, rel: RelTag) -> Result<()> {
         modification.put_rel_drop(rel)?;
         Ok(())
     }
@@ -937,7 +934,7 @@ impl<'a, T: DatadirTimeline> WalIngest<'a, T> {
 
     fn handle_rel_extend(
         &mut self,
-        modification: &mut DatadirModification<T>,
+        modification: &mut DatadirModification,
         rel: RelTag,
         blknum: BlockNumber,
     ) -> Result<()> {
@@ -968,7 +965,7 @@ impl<'a, T: DatadirTimeline> WalIngest<'a, T> {
 
     fn put_slru_page_image(
         &mut self,
-        modification: &mut DatadirModification<T>,
+        modification: &mut DatadirModification,
         kind: SlruKind,
         segno: u32,
         blknum: BlockNumber,
@@ -981,7 +978,7 @@ impl<'a, T: DatadirTimeline> WalIngest<'a, T> {
 
     fn handle_slru_extend(
         &mut self,
-        modification: &mut DatadirModification<T>,
+        modification: &mut DatadirModification,
         kind: SlruKind,
         segno: u32,
         blknum: BlockNumber,
@@ -1032,9 +1029,9 @@ impl<'a, T: DatadirTimeline> WalIngest<'a, T> {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::layered_repository::repo_harness::*;
+    use crate::layered_repository::Timeline;
     use crate::pgdatadir_mapping::create_test_timeline;
-    use crate::repository::repo_harness::*;
-    use crate::repository::Timeline;
     use postgres_ffi::v14::xlog_utils::SIZEOF_CHECKPOINT;
     use postgres_ffi::RELSEG_SIZE;
 
@@ -1046,13 +1043,13 @@ mod tests {
         forknum: 0,
     };
 
-    fn assert_current_logical_size<T: Timeline>(_timeline: &T, _lsn: Lsn) {
+    fn assert_current_logical_size(_timeline: &Timeline, _lsn: Lsn) {
         // TODO
     }
 
     static ZERO_CHECKPOINT: Bytes = Bytes::from_static(&[0u8; SIZEOF_CHECKPOINT]);
 
-    fn init_walingest_test<T: DatadirTimeline>(tline: &T) -> Result<WalIngest<T>> {
+    fn init_walingest_test(tline: &Timeline) -> Result<WalIngest> {
         let mut m = tline.begin_modification(Lsn(0x10));
         m.put_checkpoint(ZERO_CHECKPOINT.clone())?;
         m.put_relmap_file(0, 111, Bytes::from(""))?; // dummy relmapper file
@@ -1065,7 +1062,7 @@ mod tests {
     #[test]
     fn test_relsize() -> Result<()> {
         let repo = RepoHarness::create("test_relsize")?.load();
-        let tline = create_test_timeline(repo, TIMELINE_ID)?;
+        let tline = create_test_timeline(&repo, TIMELINE_ID)?;
         let mut walingest = init_walingest_test(&*tline)?;
 
         let mut m = tline.begin_modification(Lsn(0x20));
@@ -1193,7 +1190,7 @@ mod tests {
     #[test]
     fn test_drop_extend() -> Result<()> {
         let repo = RepoHarness::create("test_drop_extend")?.load();
-        let tline = create_test_timeline(repo, TIMELINE_ID)?;
+        let tline = create_test_timeline(&repo, TIMELINE_ID)?;
         let mut walingest = init_walingest_test(&*tline)?;
 
         let mut m = tline.begin_modification(Lsn(0x20));
@@ -1233,7 +1230,7 @@ mod tests {
     #[test]
     fn test_truncate_extend() -> Result<()> {
         let repo = RepoHarness::create("test_truncate_extend")?.load();
-        let tline = create_test_timeline(repo, TIMELINE_ID)?;
+        let tline = create_test_timeline(&repo, TIMELINE_ID)?;
         let mut walingest = init_walingest_test(&*tline)?;
 
         // Create a 20 MB relation (the size is arbitrary)
@@ -1321,7 +1318,7 @@ mod tests {
     #[test]
     fn test_large_rel() -> Result<()> {
         let repo = RepoHarness::create("test_large_rel")?.load();
-        let tline = create_test_timeline(repo, TIMELINE_ID)?;
+        let tline = create_test_timeline(&repo, TIMELINE_ID)?;
         let mut walingest = init_walingest_test(&*tline)?;
 
         let mut lsn = 0x10;

pageserver/src/walreceiver.rs

@@ -269,7 +269,7 @@ async fn wal_receiver_main_thread_loop_step<'a>(
                     }
                 }
                 // Timeline got attached, retrieve all necessary information to start its broker loop and maintain this loop endlessly.
-                LocalTimelineUpdate::Attach { id, datadir } => {
+                LocalTimelineUpdate::Attach { id, timeline } => {
                     let timeline_connection_managers = local_timeline_wal_receivers
                         .entry(id.tenant_id)
                         .or_default();
@@ -305,7 +305,7 @@ async fn wal_receiver_main_thread_loop_step<'a>(
                             id,
                             broker_prefix.to_owned(),
                             etcd_client.clone(),
-                            datadir,
+                            timeline,
                             wal_connect_timeout,
                             lagging_wal_timeout,
                             max_lsn_wal_lag,

pageserver/src/walreceiver/connection_manager.rs

@@ -16,7 +16,7 @@ use std::{
     time::Duration,
 };
 
-use crate::{layered_repository::LayeredTimeline, repository::Timeline};
+use crate::layered_repository::Timeline;
 use anyhow::Context;
 use chrono::{NaiveDateTime, Utc};
 use etcd_broker::{
@@ -39,7 +39,7 @@ pub(super) fn spawn_connection_manager_task(
     id: ZTenantTimelineId,
     broker_loop_prefix: String,
     mut client: Client,
-    local_timeline: Arc<LayeredTimeline>,
+    local_timeline: Arc<Timeline>,
     wal_connect_timeout: Duration,
     lagging_wal_timeout: Duration,
     max_lsn_wal_lag: NonZeroU64,
@@ -242,7 +242,7 @@ const WALCONNECTION_RETRY_BACKOFF_MULTIPLIER: f64 = 1.5;
 struct WalreceiverState {
     id: ZTenantTimelineId,
     /// Use pageserver data about the timeline to filter out some of the safekeepers.
-    local_timeline: Arc<LayeredTimeline>,
+    local_timeline: Arc<Timeline>,
     /// The timeout on the connection to safekeeper for WAL streaming.
     wal_connect_timeout: Duration,
     /// The timeout to use to determine when the current connection is "stale" and reconnect to the other one.
@@ -300,7 +300,7 @@ struct EtcdSkTimeline {
 impl WalreceiverState {
     fn new(
         id: ZTenantTimelineId,
-        local_timeline: Arc<LayeredTimeline>,
+        local_timeline: Arc<Timeline>,
         wal_connect_timeout: Duration,
         lagging_wal_timeout: Duration,
         max_lsn_wal_lag: NonZeroU64,
@@ -735,12 +735,8 @@ fn wal_stream_connection_string(
 
 #[cfg(test)]
 mod tests {
-    use crate::repository::{
-        repo_harness::{RepoHarness, TIMELINE_ID},
-        Repository,
-    };
-
     use super::*;
+    use crate::layered_repository::repo_harness::{RepoHarness, TIMELINE_ID};
 
     #[test]
     fn no_connection_no_candidate() -> anyhow::Result<()> {

pageserver/src/walreceiver/walreceiver_connection.rs

@@ -20,11 +20,7 @@ use tracing::{debug, error, info, info_span, trace, warn, Instrument};
 
 use super::TaskEvent;
 use crate::{
-    layered_repository::WalReceiverInfo,
-    pgdatadir_mapping::DatadirTimeline,
-    repository::{Repository, Timeline},
-    tenant_mgr,
-    walingest::WalIngest,
+    layered_repository::WalReceiverInfo, tenant_mgr, walingest::WalIngest,
     walrecord::DecodedWALRecord,
 };
 use postgres_ffi::v14::waldecoder::WalStreamDecoder;
@@ -67,7 +63,7 @@ pub async fn handle_walreceiver_connection(
     )
     .await
     .context("Timed out while waiting for walreceiver connection to open")?
-    .context("Failed to open walreceiver conection")?;
+    .context("Failed to open walreceiver connection")?;
 
     info!("connected!");
     let mut connection_status = WalConnectionStatus {

pageserver/src/walredo.rs

@@ -89,15 +89,52 @@ pub trait WalRedoManager: Send + Sync {
 // for access to the postgres process ('wait') since there is only one for
 // each tenant.
 
+/// Time buckets are small because we want to be able to measure the
+/// smallest redo processing times. These buckets allow us to measure down
+/// to 5us, which equates to 200'000 pages/sec, which equates to 1.6GB/sec.
+/// This is much better than the previous 5ms aka 200 pages/sec aka 1.6MB/sec.
+macro_rules! redo_histogram_time_buckets {
+    () => {
+        vec![
+            0.000_005, 0.000_010, 0.000_025, 0.000_050, 0.000_100, 0.000_250, 0.000_500, 0.001_000,
+            0.002_500, 0.005_000, 0.010_000, 0.025_000, 0.050_000,
+        ]
+    };
+}
+
+/// While we're at it, also measure the amount of records replayed in each
+/// operation. We have a global 'total replayed' counter, but that's not
+/// as useful as 'what is the skew for how many records we replay in one
+/// operation'.
+macro_rules! redo_histogram_count_buckets {
+    () => {
+        vec![0.0, 1.0, 2.0, 5.0, 10.0, 25.0, 50.0, 100.0, 250.0, 500.0]
+    };
+}
+
 static WAL_REDO_TIME: Lazy<Histogram> = Lazy::new(|| {
-    register_histogram!("pageserver_wal_redo_seconds", "Time spent on WAL redo")
-        .expect("failed to define a metric")
+    register_histogram!(
+        "pageserver_wal_redo_seconds",
+        "Time spent on WAL redo",
+        redo_histogram_time_buckets!()
+    )
+    .expect("failed to define a metric")
 });
 
 static WAL_REDO_WAIT_TIME: Lazy<Histogram> = Lazy::new(|| {
     register_histogram!(
         "pageserver_wal_redo_wait_seconds",
-        "Time spent waiting for access to the WAL redo process"
+        "Time spent waiting for access to the WAL redo process",
+        redo_histogram_time_buckets!(),
+    )
+    .expect("failed to define a metric")
+});
+
+static WAL_REDO_RECORDS_HISTOGRAM: Lazy<Histogram> = Lazy::new(|| {
+    register_histogram!(
+        "pageserver_wal_redo_records_histogram",
+        "Histogram of number of records replayed per redo",
+        redo_histogram_count_buckets!(),
     )
     .expect("failed to define a metric")
 });
@@ -262,7 +299,10 @@ impl PostgresRedoManager {
 
         let end_time = Instant::now();
         let duration = end_time.duration_since(lock_time);
+
         WAL_REDO_TIME.observe(duration.as_secs_f64());
+        WAL_REDO_RECORDS_HISTOGRAM.observe(records.len() as f64);
+
         debug!(
             "postgres applied {} WAL records in {} us to reconstruct page image at LSN {}",
             records.len(),

pgxn/neon/Makefile

@@ -0,0 +1,26 @@
+# pgxs/neon/Makefile
+
+
+MODULE_big = neon
+OBJS = \
+	$(WIN32RES) \
+	inmem_smgr.o \
+	libpagestore.o \
+	libpqwalproposer.o \
+	pagestore_smgr.o \
+	relsize_cache.o \
+	neon.o \
+	walproposer.o \
+	walproposer_utils.o
+
+PG_CPPFLAGS = -I$(libpq_srcdir)
+SHLIB_LINK_INTERNAL = $(libpq)
+
+EXTENSION = neon
+DATA = neon--1.0.sql
+PGFILEDESC = "neon - cloud storage for PostgreSQL"
+
+
+PG_CONFIG = pg_config
+PGXS := $(shell $(PG_CONFIG) --pgxs)
+include $(PGXS)

pgxn/neon/inmem_smgr.c

@@ -0,0 +1,286 @@
+/*-------------------------------------------------------------------------
+ *
+ * inmem_smgr.c
+ *
+ * This is an implementation of the SMGR interface, used in the WAL redo
+ * process (see src/backend/tcop/zenith_wal_redo.c). It has no persistent
+ * storage, the pages that are written out are kept in a small number of
+ * in-memory buffers.
+ *
+ * Normally, replaying a WAL record only needs to access a handful of
+ * buffers, which fit in the normal buffer cache, so this is just for
+ * "overflow" storage when the buffer cache is not large enough.
+ *
+ *
+ * Portions Copyright (c) 1996-2021, PostgreSQL Global Development Group
+ * Portions Copyright (c) 1994, Regents of the University of California
+ *
+ * IDENTIFICATION
+ *	  contrib/neon/inmem_smgr.c
+ *
+ *-------------------------------------------------------------------------
+ */
+#include "postgres.h"
+
+#include "access/xlog.h"
+#include "pagestore_client.h"
+#include "storage/block.h"
+#include "storage/buf_internals.h"
+#include "storage/relfilenode.h"
+#include "storage/smgr.h"
+
+/* Size of the in-memory smgr */
+#define MAX_PAGES 64
+
+/* If more than WARN_PAGES are used, print a warning in the log */
+#define WARN_PAGES 32
+
+static BufferTag page_tag[MAX_PAGES];
+static char page_body[MAX_PAGES][BLCKSZ];
+static int	used_pages;
+
+static int
+locate_page(SMgrRelation reln, ForkNumber forknum, BlockNumber blkno)
+{
+	/* We only hold a small number of pages, so linear search */
+	for (int i = 0; i < used_pages; i++)
+	{
+		if (RelFileNodeEquals(reln->smgr_rnode.node, page_tag[i].rnode)
+			&& forknum == page_tag[i].forkNum
+			&& blkno == page_tag[i].blockNum)
+		{
+			return i;
+		}
+	}
+	return -1;
+}
+
+/*
+ *	inmem_init() -- Initialize private state
+ */
+void
+inmem_init(void)
+{
+	used_pages = 0;
+}
+
+/*
+ *	inmem_exists() -- Does the physical file exist?
+ */
+bool
+inmem_exists(SMgrRelation reln, ForkNumber forknum)
+{
+	for (int i = 0; i < used_pages; i++)
+	{
+		if (RelFileNodeEquals(reln->smgr_rnode.node, page_tag[i].rnode)
+			&& forknum == page_tag[i].forkNum)
+		{
+			return true;
+		}
+	}
+	return false;
+}
+
+/*
+ *	inmem_create() -- Create a new relation on zenithd storage
+ *
+ * If isRedo is true, it's okay for the relation to exist already.
+ */
+void
+inmem_create(SMgrRelation reln, ForkNumber forknum, bool isRedo)
+{
+}
+
+/*
+ *	inmem_unlink() -- Unlink a relation.
+ */
+void
+inmem_unlink(RelFileNodeBackend rnode, ForkNumber forknum, bool isRedo)
+{
+}
+
+/*
+ *	inmem_extend() -- Add a block to the specified relation.
+ *
+ *		The semantics are nearly the same as mdwrite(): write at the
+ *		specified position.  However, this is to be used for the case of
+ *		extending a relation (i.e., blocknum is at or beyond the current
+ *		EOF).  Note that we assume writing a block beyond current EOF
+ *		causes intervening file space to become filled with zeroes.
+ */
+void
+inmem_extend(SMgrRelation reln, ForkNumber forknum, BlockNumber blkno,
+			 char *buffer, bool skipFsync)
+{
+	/* same as smgwrite() for us */
+	inmem_write(reln, forknum, blkno, buffer, skipFsync);
+}
+
+/*
+ *  inmem_open() -- Initialize newly-opened relation.
+ */
+void
+inmem_open(SMgrRelation reln)
+{
+}
+
+/*
+ *	inmem_close() -- Close the specified relation, if it isn't closed already.
+ */
+void
+inmem_close(SMgrRelation reln, ForkNumber forknum)
+{
+}
+
+/*
+ *	inmem_prefetch() -- Initiate asynchronous read of the specified block of a relation
+ */
+bool
+inmem_prefetch(SMgrRelation reln, ForkNumber forknum, BlockNumber blocknum)
+{
+	return true;
+}
+
+/*
+ * inmem_writeback() -- Tell the kernel to write pages back to storage.
+ */
+void
+inmem_writeback(SMgrRelation reln, ForkNumber forknum,
+				BlockNumber blocknum, BlockNumber nblocks)
+{
+}
+
+/*
+ *	inmem_read() -- Read the specified block from a relation.
+ */
+void
+inmem_read(SMgrRelation reln, ForkNumber forknum, BlockNumber blkno,
+		   char *buffer)
+{
+	int			pg;
+
+	pg = locate_page(reln, forknum, blkno);
+	if (pg < 0)
+		memset(buffer, 0, BLCKSZ);
+	else
+		memcpy(buffer, page_body[pg], BLCKSZ);
+}
+
+/*
+ *	inmem_write() -- Write the supplied block at the appropriate location.
+ *
+ *		This is to be used only for updating already-existing blocks of a
+ *		relation (ie, those before the current EOF).  To extend a relation,
+ *		use mdextend().
+ */
+void
+inmem_write(SMgrRelation reln, ForkNumber forknum, BlockNumber blocknum,
+			char *buffer, bool skipFsync)
+{
+	int			pg;
+
+	pg = locate_page(reln, forknum, blocknum);
+	if (pg < 0)
+	{
+		/*
+		 * We assume the buffer cache is large enough to hold all the buffers
+		 * needed for most operations. Overflowing to this "in-mem smgr" in rare
+		 * cases is OK. But if we find that we're using more than WARN_PAGES,
+		 * print a warning so that we get alerted and get to investigate why
+		 * we're accessing so many buffers.
+		 */
+		elog(used_pages >= WARN_PAGES ? WARNING : DEBUG1,
+			 "inmem_write() called for %u/%u/%u.%u blk %u: used_pages %u",
+			 reln->smgr_rnode.node.spcNode,
+			 reln->smgr_rnode.node.dbNode,
+			 reln->smgr_rnode.node.relNode,
+			 forknum,
+			 blocknum,
+			 used_pages);
+		if (used_pages == MAX_PAGES)
+			elog(ERROR, "Inmem storage overflow");
+
+		pg = used_pages;
+		used_pages++;
+		INIT_BUFFERTAG(page_tag[pg], reln->smgr_rnode.node, forknum, blocknum);
+	}  else {
+		elog(DEBUG1, "inmem_write() called for %u/%u/%u.%u blk %u: found at %u",
+			 reln->smgr_rnode.node.spcNode,
+			 reln->smgr_rnode.node.dbNode,
+			 reln->smgr_rnode.node.relNode,
+			 forknum,
+			 blocknum,
+			 used_pages);
+	}
+	memcpy(page_body[pg], buffer, BLCKSZ);
+}
+
+/*
+ *	inmem_nblocks() -- Get the number of blocks stored in a relation.
+ */
+BlockNumber
+inmem_nblocks(SMgrRelation reln, ForkNumber forknum)
+{
+	/*
+	 * It's not clear why a WAL redo function would call smgrnblocks().
+	 * During recovery, at least before reaching consistency, the size of a
+	 * relation could be arbitrarily small, if it was truncated after the
+	 * record being replayed, or arbitrarily large if it was extended
+	 * afterwards. But one place where it's called is in
+	 * XLogReadBufferExtended(): it extends the relation, if it's smaller than
+	 * the requested page. That's a waste of time in the WAL redo
+	 * process. Pretend that all relations are maximally sized to avoid it.
+	 */
+	return MaxBlockNumber;
+}
+
+/*
+ *	inmem_truncate() -- Truncate relation to specified number of blocks.
+ */
+void
+inmem_truncate(SMgrRelation reln, ForkNumber forknum, BlockNumber nblocks)
+{
+}
+
+/*
+ *	inmem_immedsync() -- Immediately sync a relation to stable storage.
+ */
+void
+inmem_immedsync(SMgrRelation reln, ForkNumber forknum)
+{
+}
+
+static const struct f_smgr inmem_smgr =
+{
+	.smgr_init = inmem_init,
+	.smgr_shutdown = NULL,
+	.smgr_open = inmem_open,
+	.smgr_close = inmem_close,
+	.smgr_create = inmem_create,
+	.smgr_exists = inmem_exists,
+	.smgr_unlink = inmem_unlink,
+	.smgr_extend = inmem_extend,
+	.smgr_prefetch = inmem_prefetch,
+	.smgr_read = inmem_read,
+	.smgr_write = inmem_write,
+	.smgr_writeback = inmem_writeback,
+	.smgr_nblocks = inmem_nblocks,
+	.smgr_truncate = inmem_truncate,
+	.smgr_immedsync = inmem_immedsync,
+};
+
+const f_smgr *
+smgr_inmem(BackendId backend, RelFileNode rnode)
+{
+	Assert(InRecovery);
+	if (backend != InvalidBackendId)
+		return smgr_standard(backend, rnode);
+	else
+		return &inmem_smgr;
+}
+
+void
+smgr_init_inmem()
+{
+	inmem_init();
+}

pgxn/neon/libpagestore.c

@@ -0,0 +1,432 @@
+/*-------------------------------------------------------------------------
+ *
+ * libpagestore.c
+ *	  Handles network communications with the remote pagestore.
+ *
+ * Portions Copyright (c) 1996-2021, PostgreSQL Global Development Group
+ * Portions Copyright (c) 1994, Regents of the University of California
+ *
+ *
+ * IDENTIFICATION
+ *	 contrib/neon/libpqpagestore.c
+ *
+ *-------------------------------------------------------------------------
+ */
+#include "postgres.h"
+
+#include "pagestore_client.h"
+#include "fmgr.h"
+#include "access/xlog.h"
+
+#include "libpq-fe.h"
+#include "libpq/pqformat.h"
+#include "libpq/libpq.h"
+
+#include "miscadmin.h"
+#include "pgstat.h"
+#include "utils/guc.h"
+
+#include "neon.h"
+#include "walproposer.h"
+#include "walproposer_utils.h"
+
+
+#define PageStoreTrace DEBUG5
+
+#define NEON_TAG "[NEON_SMGR] "
+#define neon_log(tag, fmt, ...) ereport(tag, \
+		(errmsg(NEON_TAG fmt, ## __VA_ARGS__), \
+		 errhidestmt(true), errhidecontext(true)))
+
+bool		connected = false;
+PGconn	   *pageserver_conn = NULL;
+
+char	   *page_server_connstring_raw;
+
+static ZenithResponse *pageserver_call(ZenithRequest *request);
+page_server_api api = {
+	.request = pageserver_call
+};
+
+static void
+pageserver_connect()
+{
+	char	   *query;
+	int			ret;
+
+	Assert(!connected);
+
+	pageserver_conn = PQconnectdb(page_server_connstring);
+
+	if (PQstatus(pageserver_conn) == CONNECTION_BAD)
+	{
+		char	   *msg = pchomp(PQerrorMessage(pageserver_conn));
+
+		PQfinish(pageserver_conn);
+		pageserver_conn = NULL;
+		ereport(ERROR,
+				(errcode(ERRCODE_SQLCLIENT_UNABLE_TO_ESTABLISH_SQLCONNECTION),
+				 errmsg(NEON_TAG "could not establish connection to pageserver"),
+				 errdetail_internal("%s", msg)));
+	}
+
+	query = psprintf("pagestream %s %s", zenith_tenant, zenith_timeline);
+	ret = PQsendQuery(pageserver_conn, query);
+	if (ret != 1)
+	{
+		PQfinish(pageserver_conn);
+		pageserver_conn = NULL;
+		neon_log(ERROR, "could not send pagestream command to pageserver");
+	}
+
+	while (PQisBusy(pageserver_conn))
+	{
+		int			wc;
+
+		/* Sleep until there's something to do */
+		wc = WaitLatchOrSocket(MyLatch,
+							   WL_LATCH_SET | WL_SOCKET_READABLE |
+							   WL_EXIT_ON_PM_DEATH,
+							   PQsocket(pageserver_conn),
+							   -1L, PG_WAIT_EXTENSION);
+		ResetLatch(MyLatch);
+
+		CHECK_FOR_INTERRUPTS();
+
+		/* Data available in socket? */
+		if (wc & WL_SOCKET_READABLE)
+		{
+			if (!PQconsumeInput(pageserver_conn))
+			{
+				char	   *msg = pchomp(PQerrorMessage(pageserver_conn));
+
+				PQfinish(pageserver_conn);
+				pageserver_conn = NULL;
+
+				neon_log(ERROR, "could not complete handshake with pageserver: %s",
+						 msg);
+			}
+		}
+	}
+
+	neon_log(LOG, "libpagestore: connected to '%s'", page_server_connstring_raw);
+
+	connected = true;
+}
+
+/*
+ * A wrapper around PQgetCopyData that checks for interrupts while sleeping.
+ */
+static int
+call_PQgetCopyData(PGconn *conn, char **buffer)
+{
+	int			ret;
+
+retry:
+	ret = PQgetCopyData(conn, buffer, 1 /* async */ );
+
+	if (ret == 0)
+	{
+		int			wc;
+
+		/* Sleep until there's something to do */
+		wc = WaitLatchOrSocket(MyLatch,
+							   WL_LATCH_SET | WL_SOCKET_READABLE |
+							   WL_EXIT_ON_PM_DEATH,
+							   PQsocket(conn),
+							   -1L, PG_WAIT_EXTENSION);
+		ResetLatch(MyLatch);
+
+		CHECK_FOR_INTERRUPTS();
+
+		/* Data available in socket? */
+		if (wc & WL_SOCKET_READABLE)
+		{
+			if (!PQconsumeInput(conn))
+				neon_log(ERROR, "could not get response from pageserver: %s",
+						 PQerrorMessage(conn));
+		}
+
+		goto retry;
+	}
+
+	return ret;
+}
+
+
+static ZenithResponse *
+pageserver_call(ZenithRequest *request)
+{
+	StringInfoData req_buff;
+	StringInfoData resp_buff;
+	ZenithResponse *resp;
+
+	PG_TRY();
+	{
+		/* If the connection was lost for some reason, reconnect */
+		if (connected && PQstatus(pageserver_conn) == CONNECTION_BAD)
+		{
+			PQfinish(pageserver_conn);
+			pageserver_conn = NULL;
+			connected = false;
+		}
+
+		if (!connected)
+			pageserver_connect();
+
+		req_buff = zm_pack_request(request);
+
+		/*
+		 * Send request.
+		 *
+		 * In principle, this could block if the output buffer is full, and we
+		 * should use async mode and check for interrupts while waiting. In
+		 * practice, our requests are small enough to always fit in the output
+		 * and TCP buffer.
+		 */
+		if (PQputCopyData(pageserver_conn, req_buff.data, req_buff.len) <= 0 || PQflush(pageserver_conn))
+		{
+			neon_log(ERROR, "failed to send page request: %s",
+					 PQerrorMessage(pageserver_conn));
+		}
+		pfree(req_buff.data);
+
+		if (message_level_is_interesting(PageStoreTrace))
+		{
+			char	   *msg = zm_to_string((ZenithMessage *) request);
+
+			neon_log(PageStoreTrace, "sent request: %s", msg);
+			pfree(msg);
+		}
+
+		/* read response */
+		resp_buff.len = call_PQgetCopyData(pageserver_conn, &resp_buff.data);
+		resp_buff.cursor = 0;
+
+		if (resp_buff.len == -1)
+			neon_log(ERROR, "end of COPY");
+		else if (resp_buff.len == -2)
+			neon_log(ERROR, "could not read COPY data: %s", PQerrorMessage(pageserver_conn));
+
+		resp = zm_unpack_response(&resp_buff);
+		PQfreemem(resp_buff.data);
+
+		if (message_level_is_interesting(PageStoreTrace))
+		{
+			char	   *msg = zm_to_string((ZenithMessage *) resp);
+
+			neon_log(PageStoreTrace, "got response: %s", msg);
+			pfree(msg);
+		}
+	}
+	PG_CATCH();
+	{
+		/*
+		 * If anything goes wrong while we were sending a request, it's not
+		 * clear what state the connection is in. For example, if we sent the
+		 * request but didn't receive a response yet, we might receive the
+		 * response some time later after we have already sent a new unrelated
+		 * request. Close the connection to avoid getting confused.
+		 */
+		if (connected)
+		{
+			neon_log(LOG, "dropping connection to page server due to error");
+			PQfinish(pageserver_conn);
+			pageserver_conn = NULL;
+			connected = false;
+		}
+		PG_RE_THROW();
+	}
+	PG_END_TRY();
+
+	return (ZenithResponse *) resp;
+}
+
+
+static bool
+check_zenith_id(char **newval, void **extra, GucSource source)
+{
+	uint8		zid[16];
+
+	return **newval == '\0' || HexDecodeString(zid, *newval, 16);
+}
+
+static char *
+substitute_pageserver_password(const char *page_server_connstring_raw)
+{
+	char	   *host = NULL;
+	char	   *port = NULL;
+	char	   *user = NULL;
+	char	   *auth_token = NULL;
+	char	   *err = NULL;
+	char	   *page_server_connstring = NULL;
+	PQconninfoOption *conn_options;
+	PQconninfoOption *conn_option;
+	MemoryContext oldcontext;
+
+	/*
+	 * Here we substitute password in connection string with an environment
+	 * variable. To simplify things we construct a connection string back with
+	 * only known options. In particular: host port user and password. We do
+	 * not currently use other options and constructing full connstring in an
+	 * URI shape is quite messy.
+	 */
+
+	if (page_server_connstring_raw == NULL || page_server_connstring_raw[0] == '\0')
+		return NULL;
+
+	/* extract the auth token from the connection string */
+	conn_options = PQconninfoParse(page_server_connstring_raw, &err);
+	if (conn_options == NULL)
+	{
+		/* The error string is malloc'd, so we must free it explicitly */
+		char	   *errcopy = err ? pstrdup(err) : "out of memory";
+
+		PQfreemem(err);
+		ereport(ERROR,
+				(errcode(ERRCODE_SYNTAX_ERROR),
+				 errmsg("invalid connection string syntax: %s", errcopy)));
+	}
+
+	/*
+	 * Trying to populate pageserver connection string with auth token from
+	 * environment. We are looking for password in with placeholder value like
+	 * $ENV_VAR_NAME, so if password field is present and starts with $ we try
+	 * to fetch environment variable value and fail loudly if it is not set.
+	 */
+	for (conn_option = conn_options; conn_option->keyword != NULL; conn_option++)
+	{
+		if (strcmp(conn_option->keyword, "host") == 0)
+		{
+			if (conn_option->val != NULL && conn_option->val[0] != '\0')
+				host = conn_option->val;
+		}
+		else if (strcmp(conn_option->keyword, "port") == 0)
+		{
+			if (conn_option->val != NULL && conn_option->val[0] != '\0')
+				port = conn_option->val;
+		}
+		else if (strcmp(conn_option->keyword, "user") == 0)
+		{
+			if (conn_option->val != NULL && conn_option->val[0] != '\0')
+				user = conn_option->val;
+		}
+		else if (strcmp(conn_option->keyword, "password") == 0)
+		{
+			if (conn_option->val != NULL && conn_option->val[0] != '\0')
+			{
+				/* ensure that this is a template */
+				if (strncmp(conn_option->val, "$", 1) != 0)
+					ereport(ERROR,
+							(errcode(ERRCODE_CONNECTION_EXCEPTION),
+							 errmsg("expected placeholder value in pageserver password starting from $ but found: %s", &conn_option->val[1])));
+
+				neon_log(LOG, "found auth token placeholder in pageserver conn string '%s'", &conn_option->val[1]);
+				auth_token = getenv(&conn_option->val[1]);
+				if (!auth_token)
+				{
+					ereport(ERROR,
+							(errcode(ERRCODE_CONNECTION_EXCEPTION),
+							 errmsg("cannot get auth token, environment variable %s is not set", &conn_option->val[1])));
+				}
+				else
+				{
+					neon_log(LOG, "using auth token from environment passed via env");
+				}
+			}
+		}
+	}
+
+	/*
+	 * allocate connection string in TopMemoryContext to make sure it is not
+	 * freed
+	 */
+	oldcontext = CurrentMemoryContext;
+	MemoryContextSwitchTo(TopMemoryContext);
+	page_server_connstring = psprintf("postgresql://%s:%s@%s:%s", user, auth_token ? auth_token : "", host, port);
+	MemoryContextSwitchTo(oldcontext);
+
+	PQconninfoFree(conn_options);
+	return page_server_connstring;
+}
+
+/*
+ * Module initialization function
+ */
+void
+pg_init_libpagestore(void)
+{
+	DefineCustomStringVariable("neon.pageserver_connstring",
+							   "connection string to the page server",
+							   NULL,
+							   &page_server_connstring_raw,
+							   "",
+							   PGC_POSTMASTER,
+							   0,	/* no flags required */
+							   NULL, NULL, NULL);
+
+	DefineCustomStringVariable("neon.timeline_id",
+							   "Zenith timelineid the server is running on",
+							   NULL,
+							   &zenith_timeline,
+							   "",
+							   PGC_POSTMASTER,
+							   0,	/* no flags required */
+							   check_zenith_id, NULL, NULL);
+
+	DefineCustomStringVariable("neon.tenant_id",
+							   "Neon tenantid the server is running on",
+							   NULL,
+							   &zenith_tenant,
+							   "",
+							   PGC_POSTMASTER,
+							   0,	/* no flags required */
+							   check_zenith_id, NULL, NULL);
+
+	DefineCustomBoolVariable("neon.wal_redo",
+							 "start in wal-redo mode",
+							 NULL,
+							 &wal_redo,
+							 false,
+							 PGC_POSTMASTER,
+							 0,
+							 NULL, NULL, NULL);
+
+	DefineCustomIntVariable("neon.max_cluster_size",
+							"cluster size limit",
+							NULL,
+							&max_cluster_size,
+							-1, -1, INT_MAX,
+							PGC_SIGHUP,
+							GUC_UNIT_MB,
+							NULL, NULL, NULL);
+
+	relsize_hash_init();
+
+	if (page_server != NULL)
+		neon_log(ERROR, "libpagestore already loaded");
+
+	neon_log(PageStoreTrace, "libpagestore already loaded");
+	page_server = &api;
+
+	/* substitute password in pageserver_connstring */
+	page_server_connstring = substitute_pageserver_password(page_server_connstring_raw);
+
+	/* Is there more correct way to pass CustomGUC to postgres code? */
+	zenith_timeline_walproposer = zenith_timeline;
+	zenith_tenant_walproposer = zenith_tenant;
+
+	if (wal_redo)
+	{
+		neon_log(PageStoreTrace, "set inmem_smgr hook");
+		smgr_hook = smgr_inmem;
+		smgr_init_hook = smgr_init_inmem;
+	}
+	else if (page_server_connstring && page_server_connstring[0])
+	{
+		neon_log(PageStoreTrace, "set neon_smgr hook");
+		smgr_hook = smgr_zenith;
+		smgr_init_hook = smgr_init_zenith;
+		dbsize_hook = zenith_dbsize;
+	}
+}

pgxn/neon/libpqwalproposer.c

@@ -0,0 +1,413 @@
+#include "postgres.h"
+
+#include "libpq-fe.h"
+#include "neon.h"
+#include "walproposer.h"
+
+/* Header in walproposer.h -- Wrapper struct to abstract away the libpq connection */
+struct WalProposerConn
+{
+	PGconn* pg_conn;
+	bool    is_nonblocking; /* whether the connection is non-blocking */
+	char   *recvbuf;	/* last received data from libpqprop_async_read */
+};
+
+/* Prototypes for exported functions */
+static char*							libpqprop_error_message(WalProposerConn* conn);
+static WalProposerConnStatusType		libpqprop_status(WalProposerConn* conn);
+static WalProposerConn*					libpqprop_connect_start(char* conninfo);
+static WalProposerConnectPollStatusType	libpqprop_connect_poll(WalProposerConn* conn);
+static bool								libpqprop_send_query(WalProposerConn* conn, char* query);
+static WalProposerExecStatusType		libpqprop_get_query_result(WalProposerConn* conn);
+static pgsocket							libpqprop_socket(WalProposerConn* conn);
+static int								libpqprop_flush(WalProposerConn* conn);
+static void								libpqprop_finish(WalProposerConn* conn);
+static PGAsyncReadResult				libpqprop_async_read(WalProposerConn* conn, char** buf, int* amount);
+static PGAsyncWriteResult				libpqprop_async_write(WalProposerConn* conn, void const* buf, size_t size);
+static bool                             libpqprop_blocking_write(WalProposerConn* conn, void const* buf, size_t size);
+
+static WalProposerFunctionsType PQWalProposerFunctions = {
+	libpqprop_error_message,
+	libpqprop_status,
+	libpqprop_connect_start,
+	libpqprop_connect_poll,
+	libpqprop_send_query,
+	libpqprop_get_query_result,
+	libpqprop_socket,
+	libpqprop_flush,
+	libpqprop_finish,
+	libpqprop_async_read,
+	libpqprop_async_write,
+	libpqprop_blocking_write,
+};
+
+/* Module initialization */
+void
+pg_init_libpqwalproposer(void)
+{
+	if (WalProposerFunctions != NULL)
+		elog(ERROR, "libpqwalproposer already loaded");
+	WalProposerFunctions = &PQWalProposerFunctions;
+}
+
+/* Helper function */
+static bool
+ensure_nonblocking_status(WalProposerConn* conn, bool is_nonblocking)
+{
+	/* If we're already correctly blocking or nonblocking, all good */
+	if (is_nonblocking == conn->is_nonblocking)
+		return true;
+
+	/* Otherwise, set it appropriately */
+	if (PQsetnonblocking(conn->pg_conn, is_nonblocking) == -1)
+		return false;
+
+	conn->is_nonblocking = is_nonblocking;
+	return true;
+}
+
+/* Exported function definitions */
+static char*
+libpqprop_error_message(WalProposerConn* conn)
+{
+	return PQerrorMessage(conn->pg_conn);
+}
+
+static WalProposerConnStatusType
+libpqprop_status(WalProposerConn* conn)
+{
+	switch (PQstatus(conn->pg_conn))
+	{
+		case CONNECTION_OK:
+			return WP_CONNECTION_OK;
+		case CONNECTION_BAD:
+			return WP_CONNECTION_BAD;
+		default:
+			return WP_CONNECTION_IN_PROGRESS;
+	}
+}
+
+static WalProposerConn*
+libpqprop_connect_start(char* conninfo)
+{
+	WalProposerConn*	conn;
+	PGconn*				pg_conn;
+
+	pg_conn = PQconnectStart(conninfo);
+	/*
+	 * Allocation of a PQconn can fail, and will return NULL. We want to fully replicate the
+	 * behavior of PQconnectStart here.
+	 */
+	if (!pg_conn)
+		return NULL;
+
+	/*
+	 * And in theory this allocation can fail as well, but it's incredibly unlikely if we just
+	 * successfully allocated a PGconn.
+	 *
+	 * palloc will exit on failure though, so there's not much we could do if it *did* fail.
+	 */
+	conn = palloc(sizeof(WalProposerConn));
+	conn->pg_conn = pg_conn;
+	conn->is_nonblocking = false; /* connections always start in blocking mode */
+	conn->recvbuf = NULL;
+	return conn;
+}
+
+static WalProposerConnectPollStatusType
+libpqprop_connect_poll(WalProposerConn* conn)
+{
+	WalProposerConnectPollStatusType return_val;
+
+	switch (PQconnectPoll(conn->pg_conn))
+	{
+		case PGRES_POLLING_FAILED:
+			return_val = WP_CONN_POLLING_FAILED;
+			break;
+		case PGRES_POLLING_READING:
+			return_val = WP_CONN_POLLING_READING;
+			break;
+		case PGRES_POLLING_WRITING:
+			return_val = WP_CONN_POLLING_WRITING;
+			break;
+		case PGRES_POLLING_OK:
+			return_val = WP_CONN_POLLING_OK;
+			break;
+
+		/* There's a comment at its source about this constant being unused. We'll expect it's never
+		 * returned. */
+		case PGRES_POLLING_ACTIVE:
+			elog(FATAL, "Unexpected PGRES_POLLING_ACTIVE returned from PQconnectPoll");
+			/* This return is never actually reached, but it's here to make the compiler happy */
+			return WP_CONN_POLLING_FAILED;
+
+		default:
+			Assert(false);
+			return_val = WP_CONN_POLLING_FAILED; /* keep the compiler quiet */
+	}
+
+	return return_val;
+}
+
+static bool
+libpqprop_send_query(WalProposerConn* conn, char* query)
+{
+	/* We need to be in blocking mode for sending the query to run without
+	 * requiring a call to PQflush */
+	if (!ensure_nonblocking_status(conn, false))
+		return false;
+
+	/* PQsendQuery returns 1 on success, 0 on failure */
+	if (!PQsendQuery(conn->pg_conn, query))
+		return false;
+
+	return true;
+}
+
+static WalProposerExecStatusType
+libpqprop_get_query_result(WalProposerConn* conn)
+{
+	PGresult* result;
+	WalProposerExecStatusType return_val;
+
+	/* Marker variable if we need to log an unexpected success result */
+	char* unexpected_success = NULL;
+
+	/* Consume any input that we might be missing */
+	if (!PQconsumeInput(conn->pg_conn))
+		return WP_EXEC_FAILED;
+
+	if (PQisBusy(conn->pg_conn))
+		return WP_EXEC_NEEDS_INPUT;
+
+
+	result = PQgetResult(conn->pg_conn);
+	/* PQgetResult returns NULL only if getting the result was successful & there's no more of the
+	 * result to get. */
+	if (!result)
+	{
+		elog(WARNING, "[libpqwalproposer] Unexpected successful end of command results");
+		return WP_EXEC_UNEXPECTED_SUCCESS;
+	}
+
+	/* Helper macro to reduce boilerplate */
+	#define UNEXPECTED_SUCCESS(msg) \
+		return_val = WP_EXEC_UNEXPECTED_SUCCESS; \
+		unexpected_success = msg; \
+		break;
+
+
+	switch (PQresultStatus(result))
+	{
+		/* "true" success case */
+		case PGRES_COPY_BOTH:
+			return_val = WP_EXEC_SUCCESS_COPYBOTH;
+			break;
+
+		/* Unexpected success case */
+		case PGRES_EMPTY_QUERY:
+			UNEXPECTED_SUCCESS("empty query return");
+		case PGRES_COMMAND_OK:
+			UNEXPECTED_SUCCESS("data-less command end");
+		case PGRES_TUPLES_OK:
+			UNEXPECTED_SUCCESS("tuples return");
+		case PGRES_COPY_OUT:
+			UNEXPECTED_SUCCESS("'Copy Out' response");
+		case PGRES_COPY_IN:
+			UNEXPECTED_SUCCESS("'Copy In' response");
+		case PGRES_SINGLE_TUPLE:
+			UNEXPECTED_SUCCESS("single tuple return");
+		case PGRES_PIPELINE_SYNC:
+			UNEXPECTED_SUCCESS("pipeline sync point");
+
+		/* Failure cases */
+		case PGRES_BAD_RESPONSE:
+		case PGRES_NONFATAL_ERROR:
+		case PGRES_FATAL_ERROR:
+		case PGRES_PIPELINE_ABORTED:
+			return_val = WP_EXEC_FAILED;
+			break;
+
+		default:
+			Assert(false);
+			return_val = WP_EXEC_FAILED; /* keep the compiler quiet */
+	}
+
+	if (unexpected_success)
+		elog(WARNING, "[libpqwalproposer] Unexpected successful %s", unexpected_success);
+
+	return return_val;
+}
+
+static pgsocket
+libpqprop_socket(WalProposerConn* conn)
+{
+	return PQsocket(conn->pg_conn);
+}
+
+static int
+libpqprop_flush(WalProposerConn* conn)
+{
+	return (PQflush(conn->pg_conn));
+}
+
+static void
+libpqprop_finish(WalProposerConn* conn)
+{
+	if (conn->recvbuf != NULL)
+		PQfreemem(conn->recvbuf);
+	PQfinish(conn->pg_conn);
+	pfree(conn);
+}
+
+/*
+ * Receive a message from the safekeeper.
+ *
+ * On success, the data is placed in *buf. It is valid until the next call
+ * to this function.
+ */
+static PGAsyncReadResult
+libpqprop_async_read(WalProposerConn* conn, char** buf, int* amount)
+{
+	int result;
+
+	if (conn->recvbuf != NULL)
+	{
+		PQfreemem(conn->recvbuf);
+		conn->recvbuf = NULL;
+	}
+
+	/* Call PQconsumeInput so that we have the data we need */
+	if (!PQconsumeInput(conn->pg_conn))
+	{
+		*amount = 0;
+		*buf = NULL;
+		return PG_ASYNC_READ_FAIL;
+	}
+
+	/* The docs for PQgetCopyData list the return values as:
+	 *      0 if the copy is still in progress, but no "complete row" is
+	 *        available
+	 *     -1 if the copy is done
+	 *     -2 if an error occured
+	 *  (> 0) if it was successful; that value is the amount transferred.
+	 *
+	 * The protocol we use between walproposer and safekeeper means that we
+	 * *usually* wouldn't expect to see that the copy is done, but this can
+	 * sometimes be triggered by the server returning an ErrorResponse (which
+	 * also happens to have the effect that the copy is done).
+	 */
+	switch (result = PQgetCopyData(conn->pg_conn, &conn->recvbuf, true))
+	{
+		case 0:
+			*amount = 0;
+			*buf = NULL;
+			return PG_ASYNC_READ_TRY_AGAIN;
+		case -1:
+		{
+			/*
+			 * If we get -1, it's probably because of a server error; the
+			 * safekeeper won't normally send a CopyDone message.
+			 *
+			 * We can check PQgetResult to make sure that the server failed;
+			 * it'll always result in PGRES_FATAL_ERROR
+			 */
+			ExecStatusType status = PQresultStatus(PQgetResult(conn->pg_conn));
+
+			if (status != PGRES_FATAL_ERROR)
+				elog(FATAL, "unexpected result status %d after failed PQgetCopyData", status);
+
+			/* If there was actually an error, it'll be properly reported by
+			 * calls to PQerrorMessage -- we don't have to do anything else */
+			*amount = 0;
+			*buf = NULL;
+			return PG_ASYNC_READ_FAIL;
+		}
+		case -2:
+			*amount = 0;
+			*buf = NULL;
+			return PG_ASYNC_READ_FAIL;
+		default:
+			/* Positive values indicate the size of the returned result */
+			*amount = result;
+			*buf = conn->recvbuf;
+			return PG_ASYNC_READ_SUCCESS;
+	}
+}
+
+static PGAsyncWriteResult
+libpqprop_async_write(WalProposerConn* conn, void const* buf, size_t size)
+{
+	int result;
+
+	/* If we aren't in non-blocking mode, switch to it. */
+	if (!ensure_nonblocking_status(conn, true))
+		return PG_ASYNC_WRITE_FAIL;
+
+	/* The docs for PQputcopyData list the return values as:
+	 *   1 if the data was queued,
+	 *   0 if it was not queued because of full buffers, or
+	 *  -1 if an error occured
+	 */
+	result = PQputCopyData(conn->pg_conn, buf, size);
+
+	/* We won't get a result of zero because walproposer always empties the
+	 * connection's buffers before sending more */
+	Assert(result != 0);
+
+	switch (result)
+	{
+		case 1:
+			/* good -- continue */
+			break;
+		case -1:
+			return PG_ASYNC_WRITE_FAIL;
+		default:
+			elog(FATAL, "invalid return %d from PQputCopyData", result);
+	}
+
+	/* After queueing the data, we still need to flush to get it to send.
+	 * This might take multiple tries, but we don't want to wait around
+	 * until it's done.
+	 *
+	 * PQflush has the following returns (directly quoting the docs):
+	 *   0 if sucessful,
+	 *   1 if it was unable to send all the data in the send queue yet
+	 *  -1 if it failed for some reason
+	 */
+	switch (result = PQflush(conn->pg_conn)) {
+		case 0:
+			return PG_ASYNC_WRITE_SUCCESS;
+		case 1:
+			return PG_ASYNC_WRITE_TRY_FLUSH;
+		case -1:
+			return PG_ASYNC_WRITE_FAIL;
+		default:
+			elog(FATAL, "invalid return %d from PQflush", result);
+	}
+}
+
+static bool
+libpqprop_blocking_write(WalProposerConn* conn, void const* buf, size_t size)
+{
+	int result;
+
+	/* If we are in non-blocking mode, switch out of it. */
+	if (!ensure_nonblocking_status(conn, false))
+		return false;
+
+	/* Ths function is very similar to libpqprop_async_write. For more
+	 * information, refer to the comments there */
+	if ((result = PQputCopyData(conn->pg_conn, buf, size)) == -1)
+		return false;
+
+	Assert(result == 1);
+
+	/* Because the connection is non-blocking, flushing returns 0 or -1 */
+
+	if ((result = PQflush(conn->pg_conn)) == -1)
+		return false;
+
+	Assert(result == 0);
+	return true;
+}

pgxn/neon/neon--1.0.sql

@@ -0,0 +1,17 @@
+\echo Use "CREATE EXTENSION neon" to load this file. \quit
+
+CREATE FUNCTION pg_cluster_size()
+RETURNS bigint
+AS 'MODULE_PATHNAME', 'pg_cluster_size'
+LANGUAGE C STRICT
+PARALLEL UNSAFE;
+
+CREATE FUNCTION backpressure_lsns(
+    OUT received_lsn pg_lsn,
+    OUT disk_consistent_lsn pg_lsn,
+    OUT remote_consistent_lsn pg_lsn
+)
+RETURNS record
+AS 'MODULE_PATHNAME', 'backpressure_lsns'
+LANGUAGE C STRICT
+PARALLEL UNSAFE;

pgxn/neon/neon.c

@@ -0,0 +1,82 @@
+/*-------------------------------------------------------------------------
+ *
+ * neon.c
+ *	  Utility functions to expose neon specific information to user
+ *
+ * IDENTIFICATION
+ *	 contrib/neon/neon.c
+ *
+ *-------------------------------------------------------------------------
+ */
+#include "postgres.h"
+#include "fmgr.h"
+
+#include "access/xact.h"
+#include "access/xlog.h"
+#include "storage/buf_internals.h"
+#include "storage/bufmgr.h"
+#include "catalog/pg_type.h"
+#include "replication/walsender.h"
+#include "funcapi.h"
+#include "access/htup_details.h"
+#include "utils/pg_lsn.h"
+#include "utils/guc.h"
+
+#include "neon.h"
+#include "walproposer.h"
+
+PG_MODULE_MAGIC;
+void		_PG_init(void);
+
+
+void		_PG_init(void)
+{
+	pg_init_libpagestore();
+	pg_init_libpqwalproposer();
+	pg_init_walproposer();
+
+	EmitWarningsOnPlaceholders("neon");
+}
+
+PG_FUNCTION_INFO_V1(pg_cluster_size);
+PG_FUNCTION_INFO_V1(backpressure_lsns);
+
+Datum
+pg_cluster_size(PG_FUNCTION_ARGS)
+{
+	int64		size;
+
+	size = GetZenithCurrentClusterSize();
+
+	if (size == 0)
+		PG_RETURN_NULL();
+
+	PG_RETURN_INT64(size);
+}
+
+
+Datum
+backpressure_lsns(PG_FUNCTION_ARGS)
+{
+	XLogRecPtr writePtr;
+	XLogRecPtr flushPtr;
+	XLogRecPtr applyPtr;
+	Datum		values[3];
+	bool		nulls[3];
+	TupleDesc	tupdesc;
+
+	replication_feedback_get_lsns(&writePtr, &flushPtr, &applyPtr);
+
+	tupdesc = CreateTemplateTupleDesc(3);
+	TupleDescInitEntry(tupdesc, (AttrNumber) 1, "received_lsn", PG_LSNOID, -1, 0);
+	TupleDescInitEntry(tupdesc, (AttrNumber) 2, "disk_consistent_lsn", PG_LSNOID, -1, 0);
+	TupleDescInitEntry(tupdesc, (AttrNumber) 3, "remote_consistent_lsn", PG_LSNOID, -1, 0);
+	tupdesc = BlessTupleDesc(tupdesc);
+
+	MemSet(nulls, 0, sizeof(nulls));
+	values[0] = LSNGetDatum(writePtr);
+	values[1] = LSNGetDatum(flushPtr);
+	values[2] = LSNGetDatum(applyPtr);
+
+	PG_RETURN_DATUM(HeapTupleGetDatum(heap_form_tuple(tupdesc, values, nulls)));
+}

pgxn/neon/neon.control

@@ -0,0 +1,4 @@
+# neon extension
+comment = 'cloud storage for PostgreSQL'
+default_version = '1.0'
+module_pathname = '$libdir/neon'

pgxn/neon/neon.h

@@ -0,0 +1,19 @@
+/*-------------------------------------------------------------------------
+ *
+ * neon.h
+ *	  Functions used in the initialization of this extension.
+ *
+ * IDENTIFICATION
+ *	 contrib/neon/neon.h
+ *
+ *-------------------------------------------------------------------------
+ */
+
+#ifndef NEON_H
+#define NEON_H
+
+extern void pg_init_libpagestore(void);
+extern void pg_init_libpqwalproposer(void);
+extern void pg_init_walproposer(void);
+
+#endif /* NEON_H */

pgxn/neon/pagestore_client.h

@@ -0,0 +1,221 @@
+/*-------------------------------------------------------------------------
+ *
+ * pagestore_client.h
+ *
+ *
+ * Portions Copyright (c) 1996-2021, PostgreSQL Global Development Group
+ * Portions Copyright (c) 1994, Regents of the University of California
+ *
+ * contrib/neon/pagestore_client.h
+ *
+ *-------------------------------------------------------------------------
+ */
+#ifndef pageserver_h
+#define pageserver_h
+
+#include "postgres.h"
+
+#include "access/xlogdefs.h"
+#include "storage/relfilenode.h"
+#include "storage/block.h"
+#include "storage/smgr.h"
+#include "lib/stringinfo.h"
+#include "libpq/pqformat.h"
+#include "utils/memutils.h"
+
+#include "pg_config.h"
+
+typedef enum
+{
+	/* pagestore_client -> pagestore */
+	T_ZenithExistsRequest = 0,
+	T_ZenithNblocksRequest,
+	T_ZenithGetPageRequest,
+	T_ZenithDbSizeRequest,
+
+	/* pagestore -> pagestore_client */
+	T_ZenithExistsResponse = 100,
+	T_ZenithNblocksResponse,
+	T_ZenithGetPageResponse,
+	T_ZenithErrorResponse,
+	T_ZenithDbSizeResponse,
+} ZenithMessageTag;
+
+
+
+/* base struct for c-style inheritance */
+typedef struct
+{
+	ZenithMessageTag tag;
+} ZenithMessage;
+
+#define messageTag(m)		(((const ZenithMessage *)(m))->tag)
+
+/*
+ * supertype of all the Zenith*Request structs below
+ *
+ * If 'latest' is true, we are requesting the latest page version, and 'lsn'
+ * is just a hint to the server that we know there are no versions of the page
+ * (or relation size, for exists/nblocks requests) later than the 'lsn'.
+ */
+typedef struct
+{
+	ZenithMessageTag tag;
+	bool		latest;			/* if true, request latest page version */
+	XLogRecPtr	lsn;			/* request page version @ this LSN */
+} ZenithRequest;
+
+typedef struct
+{
+	ZenithRequest req;
+	RelFileNode rnode;
+	ForkNumber	forknum;
+} ZenithExistsRequest;
+
+typedef struct
+{
+	ZenithRequest req;
+	RelFileNode rnode;
+	ForkNumber	forknum;
+} ZenithNblocksRequest;
+
+
+typedef struct
+{
+	ZenithRequest req;
+	Oid dbNode;
+} ZenithDbSizeRequest;
+
+
+typedef struct
+{
+	ZenithRequest req;
+	RelFileNode rnode;
+	ForkNumber	forknum;
+	BlockNumber blkno;
+} ZenithGetPageRequest;
+
+/* supertype of all the Zenith*Response structs below */
+typedef struct
+{
+	ZenithMessageTag tag;
+} ZenithResponse;
+
+typedef struct
+{
+	ZenithMessageTag tag;
+	bool		exists;
+} ZenithExistsResponse;
+
+typedef struct
+{
+	ZenithMessageTag tag;
+	uint32		n_blocks;
+} ZenithNblocksResponse;
+
+typedef struct
+{
+	ZenithMessageTag tag;
+	char		page[FLEXIBLE_ARRAY_MEMBER];
+} ZenithGetPageResponse;
+
+typedef struct
+{
+	ZenithMessageTag tag;
+	int64		db_size;
+} ZenithDbSizeResponse;
+
+typedef struct
+{
+	ZenithMessageTag tag;
+	char		message[FLEXIBLE_ARRAY_MEMBER]; /* null-terminated error message */
+} ZenithErrorResponse;
+
+extern StringInfoData zm_pack_request(ZenithRequest *msg);
+extern ZenithResponse *zm_unpack_response(StringInfo s);
+extern char *zm_to_string(ZenithMessage *msg);
+
+/*
+ * API
+ */
+
+typedef struct
+{
+	ZenithResponse *(*request) (ZenithRequest *request);
+} page_server_api;
+
+extern page_server_api *page_server;
+
+extern char *page_server_connstring;
+extern char *zenith_timeline;
+extern char *zenith_tenant;
+extern bool wal_redo;
+extern int32 max_cluster_size;
+
+extern const f_smgr *smgr_zenith(BackendId backend, RelFileNode rnode);
+extern void smgr_init_zenith(void);
+
+extern const f_smgr *smgr_inmem(BackendId backend, RelFileNode rnode);
+extern void smgr_init_inmem(void);
+extern void smgr_shutdown_inmem(void);
+
+/* zenith storage manager functionality */
+
+extern void zenith_init(void);
+extern void zenith_open(SMgrRelation reln);
+extern void zenith_close(SMgrRelation reln, ForkNumber forknum);
+extern void zenith_create(SMgrRelation reln, ForkNumber forknum, bool isRedo);
+extern bool zenith_exists(SMgrRelation reln, ForkNumber forknum);
+extern void zenith_unlink(RelFileNodeBackend rnode, ForkNumber forknum, bool isRedo);
+extern void zenith_extend(SMgrRelation reln, ForkNumber forknum,
+						  BlockNumber blocknum, char *buffer, bool skipFsync);
+extern bool zenith_prefetch(SMgrRelation reln, ForkNumber forknum,
+							BlockNumber blocknum);
+extern void zenith_read(SMgrRelation reln, ForkNumber forknum, BlockNumber blocknum,
+						char *buffer);
+
+extern void zenith_read_at_lsn(RelFileNode rnode, ForkNumber forkNum, BlockNumber blkno,
+			XLogRecPtr request_lsn, bool request_latest, char *buffer);
+
+extern void zenith_write(SMgrRelation reln, ForkNumber forknum,
+						 BlockNumber blocknum, char *buffer, bool skipFsync);
+extern void zenith_writeback(SMgrRelation reln, ForkNumber forknum,
+							 BlockNumber blocknum, BlockNumber nblocks);
+extern BlockNumber zenith_nblocks(SMgrRelation reln, ForkNumber forknum);
+extern const int64 zenith_dbsize(Oid dbNode);
+extern void zenith_truncate(SMgrRelation reln, ForkNumber forknum,
+							BlockNumber nblocks);
+extern void zenith_immedsync(SMgrRelation reln, ForkNumber forknum);
+
+/* zenith wal-redo storage manager functionality */
+
+extern void inmem_init(void);
+extern void inmem_open(SMgrRelation reln);
+extern void inmem_close(SMgrRelation reln, ForkNumber forknum);
+extern void inmem_create(SMgrRelation reln, ForkNumber forknum, bool isRedo);
+extern bool inmem_exists(SMgrRelation reln, ForkNumber forknum);
+extern void inmem_unlink(RelFileNodeBackend rnode, ForkNumber forknum, bool isRedo);
+extern void inmem_extend(SMgrRelation reln, ForkNumber forknum,
+						 BlockNumber blocknum, char *buffer, bool skipFsync);
+extern bool inmem_prefetch(SMgrRelation reln, ForkNumber forknum,
+						   BlockNumber blocknum);
+extern void inmem_read(SMgrRelation reln, ForkNumber forknum, BlockNumber blocknum,
+					   char *buffer);
+extern void inmem_write(SMgrRelation reln, ForkNumber forknum,
+						BlockNumber blocknum, char *buffer, bool skipFsync);
+extern void inmem_writeback(SMgrRelation reln, ForkNumber forknum,
+							BlockNumber blocknum, BlockNumber nblocks);
+extern BlockNumber inmem_nblocks(SMgrRelation reln, ForkNumber forknum);
+extern void inmem_truncate(SMgrRelation reln, ForkNumber forknum,
+						   BlockNumber nblocks);
+extern void inmem_immedsync(SMgrRelation reln, ForkNumber forknum);
+
+
+/* utils for zenith relsize cache */
+extern void relsize_hash_init(void);
+extern bool get_cached_relsize(RelFileNode rnode, ForkNumber forknum, BlockNumber* size);
+extern void set_cached_relsize(RelFileNode rnode, ForkNumber forknum, BlockNumber size);
+extern void update_cached_relsize(RelFileNode rnode, ForkNumber forknum, BlockNumber size);
+extern void forget_cached_relsize(RelFileNode rnode, ForkNumber forknum);
+
+#endif

pgxn/neon/relsize_cache.c

@@ -0,0 +1,167 @@
+/*-------------------------------------------------------------------------
+ *
+ * relsize_cache.c
+ *      Relation size cache for better zentih performance.
+ *
+ * Portions Copyright (c) 1996-2021, PostgreSQL Global Development Group
+ * Portions Copyright (c) 1994, Regents of the University of California
+ *
+ *
+ * IDENTIFICATION
+ *	  contrib/neon/relsize_cache.c
+ *
+ *-------------------------------------------------------------------------
+ */
+#include "postgres.h"
+
+#include "pagestore_client.h"
+#include "storage/relfilenode.h"
+#include "storage/smgr.h"
+#include "storage/lwlock.h"
+#include "storage/ipc.h"
+#include "storage/shmem.h"
+#include "catalog/pg_tablespace_d.h"
+#include "utils/dynahash.h"
+#include "utils/guc.h"
+
+
+typedef struct
+{
+	RelFileNode rnode;
+	ForkNumber	forknum;
+} RelTag;
+
+typedef struct
+{
+	RelTag		tag;
+	BlockNumber size;
+} RelSizeEntry;
+
+static HTAB *relsize_hash;
+static LWLockId relsize_lock;
+static int	relsize_hash_size;
+static shmem_startup_hook_type prev_shmem_startup_hook = NULL;
+
+/*
+ * Size of a cache entry is 20 bytes. So this default will take about 1.2 MB,
+ * which seems reasonable.
+ */
+#define DEFAULT_RELSIZE_HASH_SIZE (64 * 1024)
+
+static void
+zenith_smgr_shmem_startup(void)
+{
+	static HASHCTL info;
+
+	if (prev_shmem_startup_hook)
+		prev_shmem_startup_hook();
+
+	LWLockAcquire(AddinShmemInitLock, LW_EXCLUSIVE);
+	relsize_lock = (LWLockId) GetNamedLWLockTranche("neon_relsize");
+	info.keysize = sizeof(RelTag);
+	info.entrysize = sizeof(RelSizeEntry);
+	relsize_hash = ShmemInitHash("neon_relsize",
+								 relsize_hash_size, relsize_hash_size,
+								 &info,
+								 HASH_ELEM | HASH_BLOBS);
+	LWLockRelease(AddinShmemInitLock);
+}
+
+bool
+get_cached_relsize(RelFileNode rnode, ForkNumber forknum, BlockNumber *size)
+{
+	bool		found = false;
+
+	if (relsize_hash_size > 0)
+	{
+		RelTag		tag;
+		RelSizeEntry *entry;
+
+		tag.rnode = rnode;
+		tag.forknum = forknum;
+		LWLockAcquire(relsize_lock, LW_SHARED);
+		entry = hash_search(relsize_hash, &tag, HASH_FIND, NULL);
+		if (entry != NULL)
+		{
+			*size = entry->size;
+			found = true;
+		}
+		LWLockRelease(relsize_lock);
+	}
+	return found;
+}
+
+void
+set_cached_relsize(RelFileNode rnode, ForkNumber forknum, BlockNumber size)
+{
+	if (relsize_hash_size > 0)
+	{
+		RelTag		tag;
+		RelSizeEntry *entry;
+
+		tag.rnode = rnode;
+		tag.forknum = forknum;
+		LWLockAcquire(relsize_lock, LW_EXCLUSIVE);
+		entry = hash_search(relsize_hash, &tag, HASH_ENTER, NULL);
+		entry->size = size;
+		LWLockRelease(relsize_lock);
+	}
+}
+
+void
+update_cached_relsize(RelFileNode rnode, ForkNumber forknum, BlockNumber size)
+{
+	if (relsize_hash_size > 0)
+	{
+		RelTag		tag;
+		RelSizeEntry *entry;
+		bool		found;
+
+		tag.rnode = rnode;
+		tag.forknum = forknum;
+		LWLockAcquire(relsize_lock, LW_EXCLUSIVE);
+		entry = hash_search(relsize_hash, &tag, HASH_ENTER, &found);
+		if (!found || entry->size < size)
+			entry->size = size;
+		LWLockRelease(relsize_lock);
+	}
+}
+
+void
+forget_cached_relsize(RelFileNode rnode, ForkNumber forknum)
+{
+	if (relsize_hash_size > 0)
+	{
+		RelTag		tag;
+
+		tag.rnode = rnode;
+		tag.forknum = forknum;
+		LWLockAcquire(relsize_lock, LW_EXCLUSIVE);
+		hash_search(relsize_hash, &tag, HASH_REMOVE, NULL);
+		LWLockRelease(relsize_lock);
+	}
+}
+
+void
+relsize_hash_init(void)
+{
+	DefineCustomIntVariable("neon.relsize_hash_size",
+							"Sets the maximum number of cached relation sizes for neon",
+							NULL,
+							&relsize_hash_size,
+							DEFAULT_RELSIZE_HASH_SIZE,
+							0,
+							INT_MAX,
+							PGC_POSTMASTER,
+							0,
+							NULL, NULL, NULL);
+
+	if (relsize_hash_size > 0)
+	{
+		RequestAddinShmemSpace(hash_estimate_size(relsize_hash_size, sizeof(RelSizeEntry)));
+		RequestNamedLWLockTranche("neon_relsize", 1);
+
+		prev_shmem_startup_hook = shmem_startup_hook;
+		shmem_startup_hook = zenith_smgr_shmem_startup;
+	}
+}

pgxn/neon/walproposer.h

@@ -0,0 +1,540 @@
+#ifndef __NEON_WALPROPOSER_H__
+#define __NEON_WALPROPOSER_H__
+
+#include "access/xlogdefs.h"
+#include "postgres.h"
+#include "port.h"
+#include "access/xlog_internal.h"
+#include "access/transam.h"
+#include "nodes/replnodes.h"
+#include "utils/uuid.h"
+#include "replication/walreceiver.h"
+
+#define SK_MAGIC              0xCafeCeefu
+#define SK_PROTOCOL_VERSION   2
+
+#define MAX_SAFEKEEPERS        32
+#define MAX_SEND_SIZE         (XLOG_BLCKSZ * 16) /* max size of a single WAL message */
+#define XLOG_HDR_SIZE         (1+8*3)  /* 'w' + startPos + walEnd + timestamp */
+#define XLOG_HDR_START_POS    1        /* offset of start position in wal sender message header */
+#define XLOG_HDR_END_POS      (1+8)    /* offset of end position in wal sender message header */
+
+/*
+ * In the spirit of WL_SOCKET_READABLE and others, this corresponds to no events having occured,
+ * because all WL_* events are given flags equal to some (1 << i), starting from i = 0
+ */
+#define WL_NO_EVENTS 0
+
+extern char* wal_acceptors_list;
+extern int   wal_acceptor_reconnect_timeout;
+extern int   wal_acceptor_connect_timeout;
+extern bool  am_wal_proposer;
+
+struct WalProposerConn; /* Defined in libpqwalproposer */
+typedef struct WalProposerConn WalProposerConn;
+
+struct WalMessage;
+typedef struct WalMessage WalMessage;
+
+extern char *zenith_timeline_walproposer;
+extern char *zenith_tenant_walproposer;
+
+/* Possible return values from ReadPGAsync */
+typedef enum
+{
+	/* The full read was successful. buf now points to the data */
+	PG_ASYNC_READ_SUCCESS,
+	/* The read is ongoing. Wait until the connection is read-ready, then try
+	 * again. */
+	PG_ASYNC_READ_TRY_AGAIN,
+	/* Reading failed. Check PQerrorMessage(conn) */
+	PG_ASYNC_READ_FAIL,
+} PGAsyncReadResult;
+
+/* Possible return values from WritePGAsync */
+typedef enum
+{
+	/* The write fully completed */
+	PG_ASYNC_WRITE_SUCCESS,
+	/* The write started, but you'll need to call PQflush some more times
+	 * to finish it off. We just tried, so it's best to wait until the
+	 * connection is read- or write-ready to try again.
+	 *
+	 * If it becomes read-ready, call PQconsumeInput and flush again. If it
+	 * becomes write-ready, just call PQflush.
+	 */
+	PG_ASYNC_WRITE_TRY_FLUSH,
+	/* Writing failed. Check PQerrorMessage(conn) */
+	PG_ASYNC_WRITE_FAIL,
+} PGAsyncWriteResult;
+
+/*
+ * WAL safekeeper state, which is used to wait for some event.
+ *
+ * States are listed here in the order that they're executed.
+ *
+ * Most states, upon failure, will move back to SS_OFFLINE by calls to
+ * ResetConnection or ShutdownConnection.
+ */
+typedef enum
+{
+	/*
+	 * Does not have an active connection and will stay that way until
+	 * further notice.
+	 *
+	 * Moves to SS_CONNECTING_WRITE by calls to ResetConnection.
+	 */
+	SS_OFFLINE,
+
+	/*
+	 * Connecting states. "_READ" waits for the socket to be available for
+	 * reading, "_WRITE" waits for writing. There's no difference in the code
+	 * they execute when polled, but we have this distinction in order to
+	 * recreate the event set in HackyRemoveWalProposerEvent.
+	 *
+	 * After the connection is made, "START_WAL_PUSH" query is sent.
+	 */
+	SS_CONNECTING_WRITE,
+	SS_CONNECTING_READ,
+
+	/*
+	 * Waiting for the result of the "START_WAL_PUSH" command.
+	 *
+	 * After we get a successful result, sends handshake to safekeeper.
+	 */
+	SS_WAIT_EXEC_RESULT,
+
+	/*
+	 * Executing the receiving half of the handshake. After receiving, moves to
+	 * SS_VOTING.
+	 */
+	SS_HANDSHAKE_RECV,
+
+	/*
+	 * Waiting to participate in voting, but a quorum hasn't yet been reached.
+	 * This is an idle state - we do not expect AdvancePollState to be called.
+	 *
+	 * Moved externally by execution of SS_HANDSHAKE_RECV, when we received a
+	 * quorum of handshakes.
+	 */
+	SS_VOTING,
+
+	/*
+	 * Already sent voting information, waiting to receive confirmation from the
+	 * node. After receiving, moves to SS_IDLE, if the quorum isn't reached yet.
+	 */
+	SS_WAIT_VERDICT,
+
+	/* Need to flush ProposerElected message. */
+	SS_SEND_ELECTED_FLUSH,
+
+	/*
+	 * Waiting for quorum to send WAL. Idle state. If the socket becomes
+	 * read-ready, the connection has been closed.
+	 *
+	 * Moves to SS_ACTIVE only by call to StartStreaming.
+	 */
+	SS_IDLE,
+
+	/*
+	 * Active phase, when we acquired quorum and have WAL to send or feedback
+	 * to read.
+	 */
+	SS_ACTIVE,
+} SafekeeperState;
+
+/* Consensus logical timestamp. */
+typedef uint64 term_t;
+
+/* neon storage node id */
+typedef uint64 NNodeId;
+
+/*
+ * Proposer <-> Acceptor messaging.
+ */
+
+/* Initial Proposer -> Acceptor message */
+typedef struct ProposerGreeting
+{
+	uint64	   tag;				  /* message tag */
+	uint32	   protocolVersion;	  /* proposer-safekeeper protocol version */
+	uint32	   pgVersion;
+	pg_uuid_t  proposerId;
+	uint64	   systemId;		  /* Postgres system identifier */
+	uint8	   ztimelineid[16];	  /* Zenith timeline id */
+	uint8	   ztenantid[16];
+	TimeLineID timeline;
+	uint32	   walSegSize;
+} ProposerGreeting;
+
+typedef struct AcceptorProposerMessage
+{
+	uint64 tag;
+} AcceptorProposerMessage;
+
+/*
+ * Acceptor -> Proposer initial response: the highest term acceptor voted for.
+ */
+typedef struct AcceptorGreeting
+{
+	AcceptorProposerMessage apm;
+	term_t		term;
+	NNodeId		nodeId;
+} AcceptorGreeting;
+
+/*
+ * Proposer -> Acceptor vote request.
+ */
+typedef struct VoteRequest
+{
+	uint64		tag;
+	term_t		term;
+	pg_uuid_t   proposerId; /* for monitoring/debugging */
+} VoteRequest;
+
+/* Element of term switching chain. */
+typedef struct TermSwitchEntry
+{
+	term_t term;
+	XLogRecPtr lsn;
+} TermSwitchEntry;
+
+typedef struct TermHistory
+{
+	uint32 n_entries;
+	TermSwitchEntry *entries;
+} TermHistory;
+
+/* Vote itself, sent from safekeeper to proposer */
+typedef struct VoteResponse {
+	AcceptorProposerMessage apm;
+	term_t term;
+	uint64 voteGiven;
+	/*
+	 * Safekeeper flush_lsn (end of WAL) + history of term switches allow
+     * proposer to choose the most advanced one.
+	 */
+	XLogRecPtr flushLsn;
+	XLogRecPtr truncateLsn;  /* minimal LSN which may be needed for recovery of some safekeeper */
+	TermHistory termHistory;
+	XLogRecPtr timelineStartLsn; /* timeline globally starts at this LSN */
+} VoteResponse;
+
+/*
+ * Proposer -> Acceptor message announcing proposer is elected and communicating
+ * epoch history to it.
+ */
+typedef struct ProposerElected
+{
+	uint64 tag;
+	term_t term;
+	/* proposer will send since this point */
+	XLogRecPtr startStreamingAt;
+	/* history of term switches up to this proposer */
+	TermHistory *termHistory;
+	/* timeline globally starts at this LSN */
+	XLogRecPtr timelineStartLsn;
+} ProposerElected;
+
+/*
+ * Header of request with WAL message sent from proposer to safekeeper.
+ */
+typedef struct AppendRequestHeader
+{
+	uint64 tag;
+	term_t term; /* term of the proposer */
+	/*
+	 * LSN since which current proposer appends WAL (begin_lsn of its first
+	 * record); determines epoch switch point.
+	 */
+	XLogRecPtr epochStartLsn;
+	XLogRecPtr beginLsn;    /* start position of message in WAL */
+	XLogRecPtr endLsn;      /* end position of message in WAL */
+	XLogRecPtr commitLsn;   /* LSN committed by quorum of safekeepers */
+	/*
+	 *  minimal LSN which may be needed for recovery of some safekeeper (end lsn
+	 *  + 1 of last chunk streamed to everyone)
+	 */
+    XLogRecPtr truncateLsn;
+    pg_uuid_t  proposerId; /* for monitoring/debugging */
+} AppendRequestHeader;
+
+/*
+ * Hot standby feedback received from replica
+ */
+typedef struct HotStandbyFeedback
+{
+	TimestampTz       ts;
+	FullTransactionId xmin;
+	FullTransactionId catalog_xmin;
+} HotStandbyFeedback;
+
+
+typedef	struct ReplicationFeedback
+{
+	// current size of the timeline on pageserver
+	uint64 currentClusterSize;
+	// standby_status_update fields that safekeeper received from pageserver
+	XLogRecPtr ps_writelsn;
+	XLogRecPtr ps_flushlsn;
+	XLogRecPtr ps_applylsn;
+	TimestampTz ps_replytime;
+} ReplicationFeedback;
+
+
+typedef struct WalproposerShmemState
+{
+	slock_t		mutex;
+	ReplicationFeedback feedback;
+	term_t		mineLastElectedTerm;
+} WalproposerShmemState;
+
+/*
+ * Report safekeeper state to proposer
+ */
+typedef struct AppendResponse
+{
+	AcceptorProposerMessage apm;
+	/*
+	 * Current term of the safekeeper; if it is higher than proposer's, the
+	 * compute is out of date.
+	 */
+	term_t     term;
+	// TODO: add comment
+	XLogRecPtr flushLsn;
+	// Safekeeper reports back his awareness about which WAL is committed, as
+	// this is a criterion for walproposer --sync mode exit
+	XLogRecPtr commitLsn;
+	HotStandbyFeedback hs;
+	// Feedback recieved from pageserver includes standby_status_update fields
+	// and custom zenith feedback.
+	// This part of the message is extensible.
+	ReplicationFeedback rf;
+} AppendResponse;
+
+// ReplicationFeedback is extensible part of the message that is parsed separately
+// Other fields are fixed part
+#define APPENDRESPONSE_FIXEDPART_SIZE offsetof(AppendResponse, rf)
+
+
+/*
+ * Descriptor of safekeeper
+ */
+typedef struct Safekeeper
+{
+	char const*        host;
+	char const*        port;
+	char               conninfo[MAXCONNINFO]; /* connection info for connecting/reconnecting */
+
+	/*
+	 * postgres protocol connection to the WAL acceptor
+	 *
+	 * Equals NULL only when state = SS_OFFLINE. Nonblocking is set once we
+	 * reach SS_ACTIVE; not before.
+	 */
+	WalProposerConn*   conn;
+	/*
+	 * Temporary buffer for the message being sent to the safekeeper.
+	 */
+	StringInfoData outbuf;
+	/*
+	 * WAL reader, allocated for each safekeeper.
+	 */
+	XLogReaderState* xlogreader;
+
+	/*
+	 * Streaming will start here; must be record boundary.
+	 */
+	XLogRecPtr startStreamingAt;
+
+	bool                flushWrite;     /* set to true if we need to call AsyncFlush, to flush pending messages */
+	XLogRecPtr          streamingAt;    /* current streaming position */
+	AppendRequestHeader appendRequest;  /* request for sending to safekeeper */
+
+	int                 eventPos;       /* position in wait event set. Equal to -1 if no event */
+	SafekeeperState     state;          /* safekeeper state machine state */
+	TimestampTz         startedConnAt;  /* when connection attempt started */
+	AcceptorGreeting    greetResponse;  /* acceptor greeting */
+	VoteResponse        voteResponse;   /* the vote */
+	AppendResponse      appendResponse; /* feedback for master */
+} Safekeeper;
+
+
+extern PGDLLIMPORT void WalProposerMain(Datum main_arg);
+void       WalProposerBroadcast(XLogRecPtr startpos, XLogRecPtr endpos);
+void       WalProposerPoll(void);
+void       WalProposerRegister(void);
+void ParseReplicationFeedbackMessage(StringInfo reply_message,
+								ReplicationFeedback *rf);
+extern void StartProposerReplication(StartReplicationCmd *cmd);
+
+Size WalproposerShmemSize(void);
+bool WalproposerShmemInit(void);
+void replication_feedback_set(ReplicationFeedback *rf);
+void replication_feedback_get_lsns(XLogRecPtr *writeLsn, XLogRecPtr *flushLsn, XLogRecPtr *applyLsn);
+
+/* libpqwalproposer hooks & helper type */
+
+/* Re-exported PostgresPollingStatusType */
+typedef enum
+{
+	WP_CONN_POLLING_FAILED = 0,
+	WP_CONN_POLLING_READING,
+	WP_CONN_POLLING_WRITING,
+	WP_CONN_POLLING_OK,
+	/*
+	 * 'libpq-fe.h' still has PGRES_POLLING_ACTIVE, but says it's unused.
+	 * We've removed it here to avoid clutter.
+	 */
+} WalProposerConnectPollStatusType;
+
+/* Re-exported and modified ExecStatusType */
+typedef enum
+{
+	/* We received a single CopyBoth result */
+	WP_EXEC_SUCCESS_COPYBOTH,
+	/* Any success result other than a single CopyBoth was received. The specifics of the result
+	 * were already logged, but it may be useful to provide an error message indicating which
+	 * safekeeper messed up.
+	 *
+	 * Do not expect PQerrorMessage to be appropriately set. */
+	WP_EXEC_UNEXPECTED_SUCCESS,
+	/* No result available at this time. Wait until read-ready, then call again. Internally, this is
+	 * returned when PQisBusy indicates that PQgetResult would block. */
+	WP_EXEC_NEEDS_INPUT,
+	/* Catch-all failure. Check PQerrorMessage. */
+	WP_EXEC_FAILED,
+} WalProposerExecStatusType;
+
+/* Re-exported ConnStatusType */
+typedef enum
+{
+	WP_CONNECTION_OK,
+	WP_CONNECTION_BAD,
+
+	/*
+	 * The original ConnStatusType has many more tags, but requests that
+	 * they not be relied upon (except for displaying to the user). We
+	 * don't need that extra functionality, so we collect them into a
+	 * single tag here.
+	 */
+	WP_CONNECTION_IN_PROGRESS,
+} WalProposerConnStatusType;
+
+/* Re-exported PQerrorMessage */
+typedef char* (*walprop_error_message_fn) (WalProposerConn* conn);
+
+/* Re-exported PQstatus */
+typedef WalProposerConnStatusType (*walprop_status_fn) (WalProposerConn* conn);
+
+/* Re-exported PQconnectStart */
+typedef WalProposerConn* (*walprop_connect_start_fn) (char* conninfo);
+
+/* Re-exported PQconectPoll */
+typedef WalProposerConnectPollStatusType (*walprop_connect_poll_fn) (WalProposerConn* conn);
+
+/* Blocking wrapper around PQsendQuery */
+typedef bool (*walprop_send_query_fn) (WalProposerConn* conn, char* query);
+
+/* Wrapper around PQconsumeInput + PQisBusy + PQgetResult */
+typedef WalProposerExecStatusType (*walprop_get_query_result_fn) (WalProposerConn* conn);
+
+/* Re-exported PQsocket */
+typedef pgsocket (*walprop_socket_fn) (WalProposerConn* conn);
+
+/* Wrapper around PQconsumeInput (if socket's read-ready) + PQflush */
+typedef int (*walprop_flush_fn) (WalProposerConn* conn);
+
+/* Re-exported PQfinish */
+typedef void (*walprop_finish_fn) (WalProposerConn* conn);
+
+/*
+ * Ergonomic wrapper around PGgetCopyData
+ *
+ * Reads a CopyData block from a safekeeper, setting *amount to the number
+ * of bytes returned.
+ *
+ * This function is allowed to assume certain properties specific to the
+ * protocol with the safekeepers, so it should not be used as-is for any
+ * other purpose.
+ *
+ * Note: If possible, using <AsyncRead> is generally preferred, because it
+ * performs a bit of extra checking work that's always required and is normally
+ * somewhat verbose.
+ */
+typedef PGAsyncReadResult (*walprop_async_read_fn) (WalProposerConn* conn,
+													char** buf,
+													int* amount);
+
+/*
+ * Ergonomic wrapper around PQputCopyData + PQflush
+ *
+ * Starts to write a CopyData block to a safekeeper.
+ *
+ * For information on the meaning of return codes, refer to PGAsyncWriteResult.
+ */
+typedef PGAsyncWriteResult (*walprop_async_write_fn) (WalProposerConn* conn,
+													  void const* buf,
+													  size_t size);
+
+/*
+ * Blocking equivalent to walprop_async_write_fn
+ *
+ * Returns 'true' if successful, 'false' on failure.
+ */
+typedef bool (*walprop_blocking_write_fn) (WalProposerConn* conn, void const* buf, size_t size);
+
+/* All libpqwalproposer exported functions collected together. */
+typedef struct WalProposerFunctionsType
+{
+	walprop_error_message_fn	walprop_error_message;
+	walprop_status_fn			walprop_status;
+	walprop_connect_start_fn	walprop_connect_start;
+	walprop_connect_poll_fn		walprop_connect_poll;
+	walprop_send_query_fn		walprop_send_query;
+	walprop_get_query_result_fn	walprop_get_query_result;
+	walprop_socket_fn			walprop_socket;
+	walprop_flush_fn			walprop_flush;
+	walprop_finish_fn			walprop_finish;
+	walprop_async_read_fn		walprop_async_read;
+	walprop_async_write_fn		walprop_async_write;
+	walprop_blocking_write_fn   walprop_blocking_write;
+} WalProposerFunctionsType;
+
+/* Allow the above functions to be "called" with normal syntax */
+#define walprop_error_message(conn) \
+	WalProposerFunctions->walprop_error_message(conn)
+#define walprop_status(conn) \
+	WalProposerFunctions->walprop_status(conn)
+#define walprop_connect_start(conninfo) \
+	WalProposerFunctions->walprop_connect_start(conninfo)
+#define walprop_connect_poll(conn) \
+	WalProposerFunctions->walprop_connect_poll(conn)
+#define walprop_send_query(conn, query) \
+	WalProposerFunctions->walprop_send_query(conn, query)
+#define walprop_get_query_result(conn) \
+	WalProposerFunctions->walprop_get_query_result(conn)
+#define walprop_set_nonblocking(conn, arg) \
+	WalProposerFunctions->walprop_set_nonblocking(conn, arg)
+#define walprop_socket(conn) \
+	WalProposerFunctions->walprop_socket(conn)
+#define walprop_flush(conn) \
+	WalProposerFunctions->walprop_flush(conn)
+#define walprop_finish(conn) \
+	WalProposerFunctions->walprop_finish(conn)
+#define walprop_async_read(conn, buf, amount) \
+	WalProposerFunctions->walprop_async_read(conn, buf, amount)
+#define walprop_async_write(conn, buf, size) \
+	WalProposerFunctions->walprop_async_write(conn, buf, size)
+#define walprop_blocking_write(conn, buf, size) \
+	WalProposerFunctions->walprop_blocking_write(conn, buf, size)
+
+/*
+ * The runtime location of the libpqwalproposer functions.
+ *
+ * This pointer is set by the initializer in libpqwalproposer, so that we
+ * can use it later.
+ */
+extern PGDLLIMPORT WalProposerFunctionsType *WalProposerFunctions;
+
+#endif /* __NEON_WALPROPOSER_H__ */

pgxn/neon/walproposer_utils.h

@@ -0,0 +1,19 @@
+#ifndef __NEON_WALPROPOSER_UTILS_H__
+#define __NEON_WALPROPOSER_UTILS_H__
+
+#include "walproposer.h"
+
+int        CompareLsn(const void *a, const void *b);
+char*      FormatSafekeeperState(SafekeeperState state);
+void       AssertEventsOkForState(uint32 events, Safekeeper* sk);
+uint32     SafekeeperStateDesiredEvents(SafekeeperState state);
+char*      FormatEvents(uint32 events);
+bool       HexDecodeString(uint8 *result, char *input, int nbytes);
+uint32     pq_getmsgint32_le(StringInfo msg);
+uint64     pq_getmsgint64_le(StringInfo msg);
+void       pq_sendint32_le(StringInfo buf, uint32 i);
+void       pq_sendint64_le(StringInfo buf, uint64 i);
+void       XLogWalPropWrite(char *buf, Size nbytes, XLogRecPtr recptr);
+void       XLogWalPropClose(XLogRecPtr recptr);
+
+#endif /* __NEON_WALPROPOSER_UTILS_H__ */

pgxn/neon_test_utils/Makefile

@@ -0,0 +1,15 @@
+# pgxs/neon_test_utils/Makefile
+
+
+MODULE_big = neon_test_utils
+OBJS = \
+	$(WIN32RES) \
+	neontest.o
+
+EXTENSION = neon_test_utils
+DATA = neon_test_utils--1.0.sql
+PGFILEDESC = "neon_test_utils - helpers for neon testing and debugging"
+
+PG_CONFIG = pg_config
+PGXS := $(shell $(PG_CONFIG) --pgxs)
+include $(PGXS)

pgxn/neon_test_utils/neon_test_utils--1.0.sql

@@ -0,0 +1,29 @@
+-- complain if script is sourced in psql, rather than via CREATE EXTENSION
+\echo Use "CREATE EXTENSION neon_test_utils" to load this file. \quit
+
+CREATE FUNCTION test_consume_xids(nxids int)
+RETURNS VOID
+AS 'MODULE_PATHNAME', 'test_consume_xids'
+LANGUAGE C STRICT
+PARALLEL UNSAFE;
+
+CREATE FUNCTION clear_buffer_cache()
+RETURNS VOID
+AS 'MODULE_PATHNAME', 'clear_buffer_cache'
+LANGUAGE C STRICT
+PARALLEL UNSAFE;
+
+CREATE FUNCTION get_raw_page_at_lsn(relname text, forkname text, blocknum int8, lsn pg_lsn)
+RETURNS bytea
+AS 'MODULE_PATHNAME', 'get_raw_page_at_lsn'
+LANGUAGE C PARALLEL UNSAFE;
+
+CREATE FUNCTION get_raw_page_at_lsn(tbspc oid, db oid, relfilenode oid, forknum int8, blocknum int8, lsn pg_lsn)
+RETURNS bytea
+AS 'MODULE_PATHNAME', 'get_raw_page_at_lsn_ex'
+LANGUAGE C PARALLEL UNSAFE;
+
+CREATE FUNCTION neon_xlogflush(lsn pg_lsn)
+RETURNS VOID
+AS 'MODULE_PATHNAME', 'neon_xlogflush'
+LANGUAGE C PARALLEL UNSAFE;

pgxn/neon_test_utils/neon_test_utils.control

@@ -0,0 +1,5 @@
+# neon_test_utils extension
+comment = 'helpers for neon testing and debugging'
+default_version = '1.0'
+module_pathname = '$libdir/neon_test_utils'
+relocatable = true

pgxn/neon_test_utils/neontest.c

@@ -0,0 +1,304 @@
+/*-------------------------------------------------------------------------
+ *
+ * neontest.c
+ *	  Helpers for neon testing and debugging
+ *
+ * IDENTIFICATION
+ *	 contrib/neon_test_utils/neontest.c
+ *
+ *-------------------------------------------------------------------------
+ */
+#include "postgres.h"
+
+#include "access/relation.h"
+#include "access/xact.h"
+#include "access/xlog.h"
+#include "catalog/namespace.h"
+#include "fmgr.h"
+#include "funcapi.h"
+#include "miscadmin.h"
+#include "storage/buf_internals.h"
+#include "storage/bufmgr.h"
+#include "utils/builtins.h"
+#include "utils/pg_lsn.h"
+#include "utils/rel.h"
+#include "utils/varlena.h"
+#include "../neon/pagestore_client.h"
+
+PG_MODULE_MAGIC;
+
+extern void _PG_init(void);
+
+PG_FUNCTION_INFO_V1(test_consume_xids);
+PG_FUNCTION_INFO_V1(clear_buffer_cache);
+PG_FUNCTION_INFO_V1(get_raw_page_at_lsn);
+PG_FUNCTION_INFO_V1(get_raw_page_at_lsn_ex);
+PG_FUNCTION_INFO_V1(neon_xlogflush);
+
+/*
+ * Linkage to functions in zenith module.
+ * The signature here would need to be updated whenever function parameters change in pagestore_smgr.c
+ */
+typedef void (*zenith_read_at_lsn_type)(RelFileNode rnode, ForkNumber forkNum, BlockNumber blkno,
+			XLogRecPtr request_lsn, bool request_latest, char *buffer);
+
+static zenith_read_at_lsn_type zenith_read_at_lsn_ptr;
+
+/*
+ * Module initialize function: fetch function pointers for cross-module calls.
+ */
+void
+_PG_init(void)
+{
+	/* Asserts verify that typedefs above match original declarations */
+	AssertVariableIsOfType(&zenith_read_at_lsn, zenith_read_at_lsn_type);
+	zenith_read_at_lsn_ptr = (zenith_read_at_lsn_type)
+		load_external_function("$libdir/neon", "zenith_read_at_lsn",
+							   true, NULL);
+}
+
+#define zenith_read_at_lsn zenith_read_at_lsn_ptr
+
+/*
+ * test_consume_xids(int4), for rapidly consuming XIDs, to test wraparound.
+ */
+Datum
+test_consume_xids(PG_FUNCTION_ARGS)
+{
+	int32		nxids = PG_GETARG_INT32(0);
+	TransactionId topxid;
+	FullTransactionId fullxid;
+	TransactionId xid;
+	TransactionId targetxid;
+
+	/* make sure we have a top-XID first */
+	topxid = GetTopTransactionId();
+
+	xid = ReadNextTransactionId();
+
+	targetxid = xid + nxids;
+	while (targetxid < FirstNormalTransactionId)
+		targetxid++;
+
+	while (TransactionIdPrecedes(xid, targetxid))
+	{
+		fullxid = GetNewTransactionId(true);
+		xid = XidFromFullTransactionId(fullxid);
+		elog(DEBUG1, "topxid: %u xid: %u", topxid, xid);
+	}
+
+	PG_RETURN_VOID();
+}
+
+/*
+ * Flush the buffer cache, evicting all pages that are not currently pinned.
+ */
+Datum
+clear_buffer_cache(PG_FUNCTION_ARGS)
+{
+	bool		save_zenith_test_evict;
+
+	/*
+	 * Temporarily set the zenith_test_evict GUC, so that when we pin and
+	 * unpin a buffer, the buffer is evicted. We use that hack to evict all
+	 * buffers, as there is no explicit "evict this buffer" function in the
+	 * buffer manager.
+	 */
+	save_zenith_test_evict = zenith_test_evict;
+	zenith_test_evict = true;
+	PG_TRY();
+	{
+		/* Scan through all the buffers */
+		for (int i = 0; i < NBuffers; i++)
+		{
+			BufferDesc *bufHdr;
+			uint32		buf_state;
+			Buffer		bufferid;
+			bool		isvalid;
+			RelFileNode rnode;
+			ForkNumber	forknum;
+			BlockNumber blocknum;
+
+			/* Peek into the buffer header to see what page it holds. */
+			bufHdr = GetBufferDescriptor(i);
+			buf_state = LockBufHdr(bufHdr);
+
+			if ((buf_state & BM_VALID) && (buf_state & BM_TAG_VALID))
+				isvalid = true;
+			else
+				isvalid = false;
+			bufferid = BufferDescriptorGetBuffer(bufHdr);
+			rnode = bufHdr->tag.rnode;
+			forknum = bufHdr->tag.forkNum;
+			blocknum = bufHdr->tag.blockNum;
+
+			UnlockBufHdr(bufHdr, buf_state);
+
+			/*
+			 * Pin the buffer, and release it again. Because we have
+			 * zenith_test_evict==true, this will evict the page from
+			 * the buffer cache if no one else is holding a pin on it.
+			 */
+			if (isvalid)
+			{
+				if (ReadRecentBuffer(rnode, forknum, blocknum, bufferid))
+					ReleaseBuffer(bufferid);
+			}
+		}
+	}
+	PG_FINALLY();
+	{
+		/* restore the GUC */
+		zenith_test_evict = save_zenith_test_evict;
+	}
+	PG_END_TRY();
+
+	PG_RETURN_VOID();
+}
+
+
+/*
+ * Reads the page from page server without buffer cache
+ * usage mimics get_raw_page() in pageinspect, but offers reading versions at specific LSN
+ * NULL read lsn will result in reading the latest version.
+ *
+ * Note: reading latest version will result in waiting for latest changes to reach the page server,
+ *       if this is undesirable, use pageinspect' get_raw_page that uses buffered access to the latest page
+ */
+Datum
+get_raw_page_at_lsn(PG_FUNCTION_ARGS)
+{
+	bytea	   *raw_page;
+	ForkNumber	forknum;
+	RangeVar   *relrv;
+	Relation	rel;
+	char	   *raw_page_data;
+	text	   *relname;
+	text	   *forkname;
+	uint32		blkno;
+
+	bool request_latest = PG_ARGISNULL(3);
+	uint64 read_lsn = request_latest ? GetXLogInsertRecPtr() : PG_GETARG_INT64(3);
+
+	if (PG_ARGISNULL(0) || PG_ARGISNULL(1) || PG_ARGISNULL(2))
+		PG_RETURN_NULL();
+
+	relname = PG_GETARG_TEXT_PP(0);
+	forkname = PG_GETARG_TEXT_PP(1);
+	blkno = PG_GETARG_UINT32(2);
+
+	if (!superuser())
+		ereport(ERROR,
+				(errcode(ERRCODE_INSUFFICIENT_PRIVILEGE),
+				 errmsg("must be superuser to use raw page functions")));
+
+	relrv = makeRangeVarFromNameList(textToQualifiedNameList(relname));
+	rel = relation_openrv(relrv, AccessShareLock);
+
+	/* Check that this relation has storage */
+	if (rel->rd_rel->relkind == RELKIND_VIEW)
+		ereport(ERROR,
+				(errcode(ERRCODE_WRONG_OBJECT_TYPE),
+				 errmsg("cannot get raw page from view \"%s\"",
+						RelationGetRelationName(rel))));
+	if (rel->rd_rel->relkind == RELKIND_COMPOSITE_TYPE)
+		ereport(ERROR,
+				(errcode(ERRCODE_WRONG_OBJECT_TYPE),
+				 errmsg("cannot get raw page from composite type \"%s\"",
+						RelationGetRelationName(rel))));
+	if (rel->rd_rel->relkind == RELKIND_FOREIGN_TABLE)
+		ereport(ERROR,
+				(errcode(ERRCODE_WRONG_OBJECT_TYPE),
+				 errmsg("cannot get raw page from foreign table \"%s\"",
+						RelationGetRelationName(rel))));
+	if (rel->rd_rel->relkind == RELKIND_PARTITIONED_TABLE)
+		ereport(ERROR,
+				(errcode(ERRCODE_WRONG_OBJECT_TYPE),
+				 errmsg("cannot get raw page from partitioned table \"%s\"",
+						RelationGetRelationName(rel))));
+	if (rel->rd_rel->relkind == RELKIND_PARTITIONED_INDEX)
+		ereport(ERROR,
+				(errcode(ERRCODE_WRONG_OBJECT_TYPE),
+				 errmsg("cannot get raw page from partitioned index \"%s\"",
+						RelationGetRelationName(rel))));
+
+	/*
+	 * Reject attempts to read non-local temporary relations; we would be
+	 * likely to get wrong data since we have no visibility into the owning
+	 * session's local buffers.
+	 */
+	if (RELATION_IS_OTHER_TEMP(rel))
+		ereport(ERROR,
+				(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+				 errmsg("cannot access temporary tables of other sessions")));
+
+
+	forknum = forkname_to_number(text_to_cstring(forkname));
+
+	/* Initialize buffer to copy to */
+	raw_page = (bytea *) palloc(BLCKSZ + VARHDRSZ);
+	SET_VARSIZE(raw_page, BLCKSZ + VARHDRSZ);
+	raw_page_data = VARDATA(raw_page);
+
+	zenith_read_at_lsn(rel->rd_node, forknum, blkno, read_lsn, request_latest, raw_page_data);
+
+	relation_close(rel, AccessShareLock);
+
+	PG_RETURN_BYTEA_P(raw_page);
+}
+
+/*
+ * Another option to read a relation page from page server without cache
+ * this version doesn't validate input and allows reading blocks of dropped relations
+ *
+ * Note: reading latest version will result in waiting for latest changes to reach the page server,
+ *  if this is undesirable, use pageinspect' get_raw_page that uses buffered access to the latest page
+ */
+Datum
+get_raw_page_at_lsn_ex(PG_FUNCTION_ARGS)
+{
+	char	   *raw_page_data;
+
+	if (!superuser())
+		ereport(ERROR,
+				(errcode(ERRCODE_INSUFFICIENT_PRIVILEGE),
+				errmsg("must be superuser to use raw page functions")));
+
+	if (PG_ARGISNULL(0) || PG_ARGISNULL(1) || PG_ARGISNULL(2) ||
+		PG_ARGISNULL(3) || PG_ARGISNULL(4))
+		PG_RETURN_NULL();
+
+	{
+		RelFileNode rnode = {
+			.spcNode = PG_GETARG_OID(0),
+			.dbNode  = PG_GETARG_OID(1),
+			.relNode = PG_GETARG_OID(2)
+		};
+
+		ForkNumber forknum = PG_GETARG_UINT32(3);
+
+		uint32 blkno = PG_GETARG_UINT32(4);
+		bool request_latest = PG_ARGISNULL(5);
+		uint64 read_lsn = request_latest ? GetXLogInsertRecPtr() : PG_GETARG_INT64(5);
+
+
+		/* Initialize buffer to copy to */
+		bytea *raw_page = (bytea *) palloc(BLCKSZ + VARHDRSZ);
+		SET_VARSIZE(raw_page, BLCKSZ + VARHDRSZ);
+		raw_page_data = VARDATA(raw_page);
+
+		zenith_read_at_lsn(rnode, forknum, blkno, read_lsn, request_latest, raw_page_data);
+		PG_RETURN_BYTEA_P(raw_page);
+	}
+}
+
+/*
+ * Directly calls XLogFlush(lsn) to flush WAL buffers.
+ */
+Datum
+neon_xlogflush(PG_FUNCTION_ARGS)
+{
+	XLogRecPtr lsn = PG_GETARG_LSN(0);
+	XLogFlush(lsn);
+	PG_RETURN_VOID();
+}

pre-commit.py

@@ -1,11 +1,10 @@
 #!/usr/bin/env python3
 
-from typing import List
+import argparse
+import enum
 import subprocess
 import sys
-import enum
-import argparse
-import os
+from typing import List
 
 
 @enum.unique
@@ -37,15 +36,24 @@ def rustfmt(fix_inplace: bool = False, no_color: bool = False) -> str:
     return cmd
 
 
-def yapf(fix_inplace: bool) -> str:
-    cmd = "poetry run yapf --recursive"
-    if fix_inplace:
-        cmd += " --in-place"
-    else:
-        cmd += " --diff"
+def black(fix_inplace: bool) -> str:
+    cmd = "poetry run black"
+    if not fix_inplace:
+        cmd += " --diff --check"
     return cmd
 
 
+def isort(fix_inplace: bool) -> str:
+    cmd = "poetry run isort"
+    if not fix_inplace:
+        cmd += " --diff --check"
+    return cmd
+
+
+def flake8() -> str:
+    return "poetry run flake8"
+
+
 def mypy() -> str:
     return "poetry run mypy"
 
@@ -71,11 +79,13 @@ def check(name: str, suffix: str, cmd: str, changed_files: List[str], no_color:
         else:
             print("Please inspect the output below and run make fmt to fix automatically.")
         if suffix == ".py":
-            print("If the output is empty, ensure that you've installed Python tooling by\n"
-                  "running './scripts/pysync' in the current directory (no root needed)")
+            print(
+                "If the output is empty, ensure that you've installed Python tooling by\n"
+                "running './scripts/pysync' in the current directory (no root needed)"
+            )
         print()
         print(res.stdout.decode())
-        exit(1)
+        sys.exit(1)
 
     print(colorify("[OK]", Color.GREEN, no_color))
 
@@ -83,10 +93,12 @@ def check(name: str, suffix: str, cmd: str, changed_files: List[str], no_color:
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument("--fix-inplace", action="store_true", help="apply fixes inplace")
-    parser.add_argument("--no-color",
-                        action="store_true",
-                        help="disable colored output",
-                        default=not sys.stdout.isatty())
+    parser.add_argument(
+        "--no-color",
+        action="store_true",
+        help="disable colored output",
+        default=not sys.stdout.isatty(),
+    )
     args = parser.parse_args()
 
     files = get_commit_files()
@@ -101,9 +113,23 @@ if __name__ == "__main__":
         no_color=args.no_color,
     )
     check(
-        name="yapf",
+        name="isort",
         suffix=".py",
-        cmd=yapf(fix_inplace=args.fix_inplace),
+        cmd=isort(fix_inplace=args.fix_inplace),
+        changed_files=files,
+        no_color=args.no_color,
+    )
+    check(
+        name="black",
+        suffix=".py",
+        cmd=black(fix_inplace=args.fix_inplace),
+        changed_files=files,
+        no_color=args.no_color,
+    )
+    check(
+        name="flake8",
+        suffix=".py",
+        cmd=flake8(),
         changed_files=files,
         no_color=args.no_color,
     )

proxy/Cargo.toml

@@ -11,7 +11,7 @@ bstr = "0.2.17"
 bytes = { version = "1.0.1", features = ['serde'] }
 clap = "3.0"
 futures = "0.3.13"
-hashbrown = "0.11.2"
+hashbrown = "0.12"
 hex = "0.4.3"
 hmac = "0.12.1"
 hyper = "0.14"
@@ -23,7 +23,7 @@ rand = "0.8.3"
 reqwest = { version = "0.11", default-features = false, features = ["blocking", "json", "rustls-tls"] }
 routerify = "3"
 rustls = "0.20.0"
-rustls-pemfile = "0.2.1"
+rustls-pemfile = "1"
 scopeguard = "1.1.0"
 serde = "1"
 serde_json = "1"

pyproject.toml

@@ -27,12 +27,54 @@ prometheus-client = "^0.14.1"
 pytest-timeout = "^2.1.0"
 Werkzeug = "2.1.2"
 pytest-order = "^1.0.1"
+allure-pytest = "^2.9.45"
+pytest-asyncio = "^0.19.0"
 
 [tool.poetry.dev-dependencies]
-yapf = "==0.31.0"
-flake8 = "^3.9.2"
+flake8 = "^5.0.4"
 mypy = "==0.971"
+black = "^22.6.0"
+isort = "^5.10.1"
 
 [build-system]
 requires = ["poetry-core>=1.0.0"]
 build-backend = "poetry.core.masonry.api"
+
+[tool.black]
+line-length = 100
+extend-exclude = '''
+/(
+    vendor
+)/
+'''
+
+[tool.isort]
+profile = "black"
+line_length = 100
+skip_gitignore = true
+skip = [
+    "vendor",
+]
+
+[tool.mypy]
+# mypy uses regex
+exclude = "^vendor/"
+# some tests don't typecheck when this flag is set
+check_untyped_defs = false
+# Help mypy find imports when running against list of individual files.
+# Without this line it would behave differently when executed on the entire project.
+mypy_path = "$MYPY_CONFIG_FILE_DIR:$MYPY_CONFIG_FILE_DIR/test_runner"
+
+disallow_incomplete_defs = false
+disallow_untyped_calls = false
+disallow_untyped_decorators = false
+disallow_untyped_defs = false
+strict = true
+
+[[tool.mypy.overrides]]
+module = [
+    "asyncpg.*",
+    "cached_property.*",
+    "pg8000.*",
+]
+ignore_missing_imports = true

run_clippy.sh

@@ -13,10 +13,10 @@
 # avoid running regular linting script that checks every feature.
 if [[ "$OSTYPE" == "darwin"* ]]; then
     # no extra features to test currently, add more here when needed
-    cargo clippy --all --all-targets -- -A unknown_lints -D warnings
+    cargo clippy --locked --all --all-targets -- -A unknown_lints -D warnings
 else
     # * `-A unknown_lints` – do not warn about unknown lint suppressions
     #                        that people with newer toolchains might use
     # * `-D warnings`      - fail on any warnings (`cargo` returns non-zero exit status)
-    cargo clippy --all --all-targets --all-features -- -A unknown_lints -D warnings
+    cargo clippy --locked --all --all-targets --all-features -- -A unknown_lints -D warnings
 fi

safekeeper/src/safekeeper.rs

@@ -332,7 +332,7 @@ pub struct AppendRequestHeader {
 }
 
 /// Report safekeeper state to proposer
-#[derive(Debug, PartialEq, Serialize, Deserialize)]
+#[derive(Debug, Serialize, Deserialize)]
 pub struct AppendResponse {
     // Current term of the safekeeper; if it is higher than proposer's, the
     // compute is out of date.

safekeeper/src/send_wal.rs

@@ -36,7 +36,7 @@ const NEON_STATUS_UPDATE_TAG_BYTE: u8 = b'z';
 type FullTransactionId = u64;
 
 /// Hot standby feedback received from replica
-#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
+#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
 pub struct HotStandbyFeedback {
     pub ts: TimestampTz,
     pub xmin: FullTransactionId,
@@ -54,7 +54,7 @@ impl HotStandbyFeedback {
 }
 
 /// Standby status update
-#[derive(Debug, Clone, Serialize, Deserialize)]
+#[derive(Debug, Clone, Deserialize)]
 pub struct StandbyReply {
     pub write_lsn: Lsn, // last lsn received by pageserver
     pub flush_lsn: Lsn, // pageserver's disk consistent lSN

scripts/coverage

@@ -9,13 +9,6 @@
 #  * https://github.com/taiki-e/cargo-llvm-cov
 #  * https://github.com/llvm/llvm-project/tree/main/llvm/test/tools/llvm-cov
 
-from abc import ABC, abstractmethod
-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
-
 import argparse
 import hashlib
 import json
@@ -24,6 +17,12 @@ import shutil
 import socket
 import subprocess
 import sys
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from textwrap import dedent
+from typing import Any, Dict, Iterable, Iterator, List, Optional
 
 
 def file_mtime_or_zero(path: Path) -> int:

scripts/export_import_between_pageservers.py

@@ -20,20 +20,21 @@
 # For more context on how to use this, see:
 # https://github.com/neondatabase/cloud/wiki/Storage-format-migration
 
-import os
-from os import path
-import shutil
-from pathlib import Path
-import tempfile
-from contextlib import closing
-import psycopg2
-import subprocess
 import argparse
+import os
+import shutil
+import subprocess
+import tempfile
 import time
-import requests
 import uuid
+from contextlib import closing
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple, cast
+
+import psycopg2
+import requests
 from psycopg2.extensions import connection as PgConnection
-from typing import Any, Callable, Dict, Iterator, List, Optional, TypeVar, cast, Union, Tuple
+from psycopg2.extensions import parse_dsn
 
 ###############################################
 ### client-side utils copied from test fixtures
@@ -45,7 +46,7 @@ _global_counter = 0
 
 
 def global_counter() -> int:
-    """ A really dumb global counter.
+    """A really dumb global counter.
     This is useful for giving output files a unique number, so if we run the
     same command multiple times we can keep their output separate.
     """
@@ -55,7 +56,7 @@ def global_counter() -> int:
 
 
 def subprocess_capture(capture_dir: str, cmd: List[str], **kwargs: Any) -> str:
-    """ Run a process and capture its output
+    """Run a process and capture its output
     Output will go to files named "cmd_NNN.stdout" and "cmd_NNN.stderr"
     where "cmd" is the name of the program and NNN is an incrementing
     counter.
@@ -63,13 +64,13 @@ def subprocess_capture(capture_dir: str, cmd: List[str], **kwargs: Any) -> str:
     Returns basepath for files with captured output.
     """
     assert type(cmd) is list
-    base = os.path.basename(cmd[0]) + '_{}'.format(global_counter())
+    base = os.path.basename(cmd[0]) + "_{}".format(global_counter())
     basepath = os.path.join(capture_dir, base)
-    stdout_filename = basepath + '.stdout'
-    stderr_filename = basepath + '.stderr'
+    stdout_filename = basepath + ".stdout"
+    stderr_filename = basepath + ".stderr"
 
-    with open(stdout_filename, 'w') as stdout_f:
-        with open(stderr_filename, 'w') as stderr_f:
+    with open(stdout_filename, "w") as stdout_f:
+        with open(stderr_filename, "w") as stderr_f:
             print('(capturing output to "{}.stdout")'.format(base))
             subprocess.run(cmd, **kwargs, stdout=stdout_f, stderr=stderr_f)
 
@@ -77,15 +78,16 @@ def subprocess_capture(capture_dir: str, cmd: List[str], **kwargs: Any) -> str:
 
 
 class PgBin:
-    """ A helper class for executing postgres binaries """
+    """A helper class for executing postgres binaries"""
+
     def __init__(self, log_dir: Path, pg_distrib_dir):
         self.log_dir = log_dir
-        self.pg_bin_path = os.path.join(str(pg_distrib_dir), 'bin')
+        self.pg_bin_path = os.path.join(str(pg_distrib_dir), "bin")
         self.env = os.environ.copy()
-        self.env['LD_LIBRARY_PATH'] = os.path.join(str(pg_distrib_dir), 'lib')
+        self.env["LD_LIBRARY_PATH"] = os.path.join(str(pg_distrib_dir), "lib")
 
     def _fixpath(self, command: List[str]):
-        if '/' not in command[0]:
+        if "/" not in command[0]:
             command[0] = os.path.join(self.pg_bin_path, command[0])
 
     def _build_env(self, env_add: Optional[Env]) -> Env:
@@ -106,15 +108,17 @@ class PgBin:
         """
 
         self._fixpath(command)
-        print('Running command "{}"'.format(' '.join(command)))
+        print('Running command "{}"'.format(" ".join(command)))
         env = self._build_env(env)
         subprocess.run(command, env=env, cwd=cwd, check=True)
 
-    def run_capture(self,
-                    command: List[str],
-                    env: Optional[Env] = None,
-                    cwd: Optional[str] = None,
-                    **kwargs: Any) -> str:
+    def run_capture(
+        self,
+        command: List[str],
+        env: Optional[Env] = None,
+        cwd: Optional[str] = None,
+        **kwargs: Any,
+    ) -> str:
         """
         Run one of the postgres binaries, with stderr and stdout redirected to a file.
         This is just like `run`, but for chatty programs. Returns basepath for files
@@ -122,35 +126,31 @@ class PgBin:
         """
 
         self._fixpath(command)
-        print('Running command "{}"'.format(' '.join(command)))
+        print('Running command "{}"'.format(" ".join(command)))
         env = self._build_env(env)
-        return subprocess_capture(str(self.log_dir),
-                                  command,
-                                  env=env,
-                                  cwd=cwd,
-                                  check=True,
-                                  **kwargs)
+        return subprocess_capture(
+            str(self.log_dir), command, env=env, cwd=cwd, check=True, **kwargs
+        )
 
 
 class PgProtocol:
-    """ Reusable connection logic """
+    """Reusable connection logic"""
+
     def __init__(self, **kwargs):
         self.default_options = kwargs
 
     def conn_options(self, **kwargs):
         conn_options = self.default_options.copy()
-        if 'dsn' in kwargs:
-            conn_options.update(parse_dsn(kwargs['dsn']))
+        if "dsn" in kwargs:
+            conn_options.update(parse_dsn(kwargs["dsn"]))
         conn_options.update(kwargs)
 
         # Individual statement timeout in seconds. 2 minutes should be
         # enough for our tests, but if you need a longer, you can
         # change it by calling "SET statement_timeout" after
         # connecting.
-        if 'options' in conn_options:
-            conn_options['options'] = f"-cstatement_timeout=120s " + conn_options['options']
-        else:
-            conn_options['options'] = "-cstatement_timeout=120s"
+        conn_options["options"] = f"-cstatement_timeout=120s {conn_options.get('options', '')}"
+
         return conn_options
 
     # autocommit=True here by default because that's what we need most of the time
@@ -194,18 +194,18 @@ class PgProtocol:
 
 class VanillaPostgres(PgProtocol):
     def __init__(self, pgdatadir: Path, pg_bin: PgBin, port: int, init=True):
-        super().__init__(host='localhost', port=port, dbname='postgres')
+        super().__init__(host="localhost", port=port, dbname="postgres")
         self.pgdatadir = pgdatadir
         self.pg_bin = pg_bin
         self.running = False
         if init:
-            self.pg_bin.run_capture(['initdb', '-D', str(pgdatadir)])
+            self.pg_bin.run_capture(["initdb", "-D", str(pgdatadir)])
         self.configure([f"port = {port}\n"])
 
     def configure(self, options: List[str]):
         """Append lines into postgresql.conf file."""
         assert not self.running
-        with open(os.path.join(self.pgdatadir, 'postgresql.conf'), 'a') as conf_file:
+        with open(os.path.join(self.pgdatadir, "postgresql.conf"), "a") as conf_file:
             conf_file.write("\n".join(options))
 
     def start(self, log_path: Optional[str] = None):
@@ -216,12 +216,13 @@ class VanillaPostgres(PgProtocol):
             log_path = os.path.join(self.pgdatadir, "pg.log")
 
         self.pg_bin.run_capture(
-            ['pg_ctl', '-w', '-D', str(self.pgdatadir), '-l', log_path, 'start'])
+            ["pg_ctl", "-w", "-D", str(self.pgdatadir), "-l", log_path, "start"]
+        )
 
     def stop(self):
         assert self.running
         self.running = False
-        self.pg_bin.run_capture(['pg_ctl', '-w', '-D', str(self.pgdatadir), 'stop'])
+        self.pg_bin.run_capture(["pg_ctl", "-w", "-D", str(self.pgdatadir), "stop"])
 
     def __enter__(self):
         return self
@@ -246,9 +247,9 @@ class NeonPageserverHttpClient(requests.Session):
             res.raise_for_status()
         except requests.RequestException as e:
             try:
-                msg = res.json()['msg']
-            except:
-                msg = ''
+                msg = res.json()["msg"]
+            except:  # noqa: E722
+                msg = ""
             raise NeonPageserverApiException(msg) from e
 
     def check_status(self):
@@ -265,17 +266,17 @@ class NeonPageserverHttpClient(requests.Session):
         res = self.post(
             f"http://{self.host}:{self.port}/v1/tenant",
             json={
-                'new_tenant_id': new_tenant_id.hex,
+                "new_tenant_id": new_tenant_id.hex,
             },
         )
 
         if res.status_code == 409:
             if ok_if_exists:
-                print(f'could not create tenant: already exists for id {new_tenant_id}')
+                print(f"could not create tenant: already exists for id {new_tenant_id}")
             else:
                 res.raise_for_status()
         elif res.status_code == 201:
-            print(f'created tenant {new_tenant_id}')
+            print(f"created tenant {new_tenant_id}")
         else:
             self.verbose_error(res)
 
@@ -299,47 +300,55 @@ class NeonPageserverHttpClient(requests.Session):
 
 
 def lsn_to_hex(num: int) -> str:
-    """ Convert lsn from int to standard hex notation. """
-    return "{:X}/{:X}".format(num >> 32, num & 0xffffffff)
+    """Convert lsn from int to standard hex notation."""
+    return "{:X}/{:X}".format(num >> 32, num & 0xFFFFFFFF)
 
 
 def lsn_from_hex(lsn_hex: str) -> int:
-    """ Convert lsn from hex notation to int. """
-    l, r = lsn_hex.split('/')
+    """Convert lsn from hex notation to int."""
+    l, r = lsn_hex.split("/")
     return (int(l, 16) << 32) + int(r, 16)
 
 
-def remote_consistent_lsn(pageserver_http_client: NeonPageserverHttpClient,
-                          tenant: uuid.UUID,
-                          timeline: uuid.UUID) -> int:
+def remote_consistent_lsn(
+    pageserver_http_client: NeonPageserverHttpClient, tenant: uuid.UUID, timeline: uuid.UUID
+) -> int:
     detail = pageserver_http_client.timeline_detail(tenant, timeline)
 
-    if detail['remote'] is None:
+    if detail["remote"] is None:
         # No remote information at all. This happens right after creating
         # a timeline, before any part of it has been uploaded to remote
         # storage yet.
         return 0
     else:
-        lsn_str = detail['remote']['remote_consistent_lsn']
+        lsn_str = detail["remote"]["remote_consistent_lsn"]
         assert isinstance(lsn_str, str)
         return lsn_from_hex(lsn_str)
 
 
-def wait_for_upload(pageserver_http_client: NeonPageserverHttpClient,
-                    tenant: uuid.UUID,
-                    timeline: uuid.UUID,
-                    lsn: int):
+def wait_for_upload(
+    pageserver_http_client: NeonPageserverHttpClient,
+    tenant: uuid.UUID,
+    timeline: uuid.UUID,
+    lsn: int,
+):
     """waits for local timeline upload up to specified lsn"""
     for i in range(10):
         current_lsn = remote_consistent_lsn(pageserver_http_client, tenant, timeline)
         if current_lsn >= lsn:
             return
-        print("waiting for remote_consistent_lsn to reach {}, now {}, iteration {}".format(
-            lsn_to_hex(lsn), lsn_to_hex(current_lsn), i + 1))
+        print(
+            "waiting for remote_consistent_lsn to reach {}, now {}, iteration {}".format(
+                lsn_to_hex(lsn), lsn_to_hex(current_lsn), i + 1
+            )
+        )
         time.sleep(1)
 
-    raise Exception("timed out while waiting for remote_consistent_lsn to reach {}, was {}".format(
-        lsn_to_hex(lsn), lsn_to_hex(current_lsn)))
+    raise Exception(
+        "timed out while waiting for remote_consistent_lsn to reach {}, was {}".format(
+            lsn_to_hex(lsn), lsn_to_hex(current_lsn)
+        )
+    )
 
 
 ##############
@@ -399,7 +408,7 @@ def reconstruct_paths(log_dir, pg_bin, base_tar):
                             # Add all template0copy paths to template0
                             prefix = f"base/{oid}/"
                             if filepath.startswith(prefix):
-                                suffix = filepath[len(prefix):]
+                                suffix = filepath[len(prefix) :]
                                 yield f"base/{template0_oid}/{suffix}"
                             elif filepath.startswith("global"):
                                 print(f"skipping {database} global file {filepath}")
@@ -451,55 +460,53 @@ def get_rlsn(pageserver_connstr, tenant_id, timeline_id):
     return last_lsn, prev_lsn
 
 
-def import_timeline(args,
-                    psql_path,
-                    pageserver_connstr,
-                    pageserver_http,
-                    tenant_id,
-                    timeline_id,
-                    last_lsn,
-                    prev_lsn,
-                    tar_filename):
+def import_timeline(
+    args,
+    psql_path,
+    pageserver_connstr,
+    pageserver_http,
+    tenant_id,
+    timeline_id,
+    last_lsn,
+    prev_lsn,
+    tar_filename,
+):
     # Import timelines to new pageserver
     import_cmd = f"import basebackup {tenant_id} {timeline_id} {last_lsn} {last_lsn}"
     full_cmd = rf"""cat {tar_filename} | {psql_path} {pageserver_connstr} -c '{import_cmd}' """
 
-    stderr_filename2 = path.join(args.work_dir, f"import_{tenant_id}_{timeline_id}.stderr")
-    stdout_filename = path.join(args.work_dir, f"import_{tenant_id}_{timeline_id}.stdout")
+    stderr_filename2 = os.path.join(args.work_dir, f"import_{tenant_id}_{timeline_id}.stderr")
+    stdout_filename = os.path.join(args.work_dir, f"import_{tenant_id}_{timeline_id}.stdout")
 
     print(f"Running: {full_cmd}")
 
-    with open(stdout_filename, 'w') as stdout_f:
-        with open(stderr_filename2, 'w') as stderr_f:
+    with open(stdout_filename, "w") as stdout_f:
+        with open(stderr_filename2, "w") as stderr_f:
             print(f"(capturing output to {stdout_filename})")
             pg_bin = PgBin(args.work_dir, args.pg_distrib_dir)
-            subprocess.run(full_cmd,
-                           stdout=stdout_f,
-                           stderr=stderr_f,
-                           env=pg_bin._build_env(None),
-                           shell=True,
-                           check=True)
+            subprocess.run(
+                full_cmd,
+                stdout=stdout_f,
+                stderr=stderr_f,
+                env=pg_bin._build_env(None),
+                shell=True,
+                check=True,
+            )
 
-            print(f"Done import")
+            print("Done import")
 
     # Wait until pageserver persists the files
-    wait_for_upload(pageserver_http,
-                    uuid.UUID(tenant_id),
-                    uuid.UUID(timeline_id),
-                    lsn_from_hex(last_lsn))
+    wait_for_upload(
+        pageserver_http, uuid.UUID(tenant_id), uuid.UUID(timeline_id), lsn_from_hex(last_lsn)
+    )
 
 
-def export_timeline(args,
-                    psql_path,
-                    pageserver_connstr,
-                    tenant_id,
-                    timeline_id,
-                    last_lsn,
-                    prev_lsn,
-                    tar_filename):
+def export_timeline(
+    args, psql_path, pageserver_connstr, tenant_id, timeline_id, last_lsn, prev_lsn, tar_filename
+):
     # Choose filenames
     incomplete_filename = tar_filename + ".incomplete"
-    stderr_filename = path.join(args.work_dir, f"{tenant_id}_{timeline_id}.stderr")
+    stderr_filename = os.path.join(args.work_dir, f"{tenant_id}_{timeline_id}.stderr")
 
     # Construct export command
     query = f"fullbackup {tenant_id} {timeline_id} {last_lsn} {prev_lsn}"
@@ -507,15 +514,13 @@ def export_timeline(args,
 
     # Run export command
     print(f"Running: {cmd}")
-    with open(incomplete_filename, 'w') as stdout_f:
-        with open(stderr_filename, 'w') as stderr_f:
+    with open(incomplete_filename, "w") as stdout_f:
+        with open(stderr_filename, "w") as stderr_f:
             print(f"(capturing output to {incomplete_filename})")
             pg_bin = PgBin(args.work_dir, args.pg_distrib_dir)
-            subprocess.run(cmd,
-                           stdout=stdout_f,
-                           stderr=stderr_f,
-                           env=pg_bin._build_env(None),
-                           check=True)
+            subprocess.run(
+                cmd, stdout=stdout_f, stderr=stderr_f, env=pg_bin._build_env(None), check=True
+            )
 
     # Add missing rels
     pg_bin = PgBin(args.work_dir, args.pg_distrib_dir)
@@ -551,27 +556,28 @@ def main(args: argparse.Namespace):
 
         for timeline in timelines:
             # Skip timelines we don't need to export
-            if args.timelines and timeline['timeline_id'] not in args.timelines:
+            if args.timelines and timeline["timeline_id"] not in args.timelines:
                 print(f"Skipping timeline {timeline['timeline_id']}")
                 continue
 
             # Choose filenames
-            tar_filename = path.join(args.work_dir,
-                                     f"{timeline['tenant_id']}_{timeline['timeline_id']}.tar")
+            tar_filename = os.path.join(
+                args.work_dir, f"{timeline['tenant_id']}_{timeline['timeline_id']}.tar"
+            )
 
             # Export timeline from old pageserver
             if args.only_import is False:
                 last_lsn, prev_lsn = get_rlsn(
                     old_pageserver_connstr,
-                    timeline['tenant_id'],
-                    timeline['timeline_id'],
+                    timeline["tenant_id"],
+                    timeline["timeline_id"],
                 )
                 export_timeline(
                     args,
                     psql_path,
                     old_pageserver_connstr,
-                    timeline['tenant_id'],
-                    timeline['timeline_id'],
+                    timeline["tenant_id"],
+                    timeline["timeline_id"],
                     last_lsn,
                     prev_lsn,
                     tar_filename,
@@ -583,8 +589,8 @@ def main(args: argparse.Namespace):
                 psql_path,
                 new_pageserver_connstr,
                 new_http_client,
-                timeline['tenant_id'],
-                timeline['timeline_id'],
+                timeline["tenant_id"],
+                timeline["timeline_id"],
                 last_lsn,
                 prev_lsn,
                 tar_filename,
@@ -592,117 +598,118 @@ def main(args: argparse.Namespace):
 
             # Re-export and compare
             re_export_filename = tar_filename + ".reexport"
-            export_timeline(args,
-                            psql_path,
-                            new_pageserver_connstr,
-                            timeline['tenant_id'],
-                            timeline['timeline_id'],
-                            last_lsn,
-                            prev_lsn,
-                            re_export_filename)
+            export_timeline(
+                args,
+                psql_path,
+                new_pageserver_connstr,
+                timeline["tenant_id"],
+                timeline["timeline_id"],
+                last_lsn,
+                prev_lsn,
+                re_export_filename,
+            )
 
             # Check the size is the same
-            old_size = os.path.getsize(tar_filename),
-            new_size = os.path.getsize(re_export_filename),
+            old_size = (os.path.getsize(tar_filename),)
+            new_size = (os.path.getsize(re_export_filename),)
             if old_size != new_size:
                 raise AssertionError(f"Sizes don't match old: {old_size} new: {new_size}")
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument(
-        '--tenant-id',
-        dest='tenants',
+        "--tenant-id",
+        dest="tenants",
         required=True,
-        nargs='+',
-        help='Id of the tenant to migrate. You can pass multiple arguments',
+        nargs="+",
+        help="Id of the tenant to migrate. You can pass multiple arguments",
     )
     parser.add_argument(
-        '--timeline-id',
-        dest='timelines',
+        "--timeline-id",
+        dest="timelines",
         required=False,
-        nargs='+',
-        help='Id of the timeline to migrate. You can pass multiple arguments',
+        nargs="+",
+        help="Id of the timeline to migrate. You can pass multiple arguments",
     )
     parser.add_argument(
-        '--from-host',
-        dest='old_pageserver_host',
+        "--from-host",
+        dest="old_pageserver_host",
         required=True,
-        help='Host of the pageserver to migrate data from',
+        help="Host of the pageserver to migrate data from",
     )
     parser.add_argument(
-        '--from-http-port',
-        dest='old_pageserver_http_port',
+        "--from-http-port",
+        dest="old_pageserver_http_port",
         required=False,
         type=int,
         default=9898,
-        help='HTTP port of the pageserver to migrate data from. Default: 9898',
+        help="HTTP port of the pageserver to migrate data from. Default: 9898",
     )
     parser.add_argument(
-        '--from-pg-port',
-        dest='old_pageserver_pg_port',
+        "--from-pg-port",
+        dest="old_pageserver_pg_port",
         required=False,
         type=int,
         default=6400,
-        help='pg port of the pageserver to migrate data from. Default: 6400',
+        help="pg port of the pageserver to migrate data from. Default: 6400",
     )
     parser.add_argument(
-        '--to-host',
-        dest='new_pageserver_host',
+        "--to-host",
+        dest="new_pageserver_host",
         required=True,
-        help='Host of the pageserver to migrate data to',
+        help="Host of the pageserver to migrate data to",
     )
     parser.add_argument(
-        '--to-http-port',
-        dest='new_pageserver_http_port',
+        "--to-http-port",
+        dest="new_pageserver_http_port",
         required=False,
         default=9898,
         type=int,
-        help='HTTP port of the pageserver to migrate data to. Default: 9898',
+        help="HTTP port of the pageserver to migrate data to. Default: 9898",
     )
     parser.add_argument(
-        '--to-pg-port',
-        dest='new_pageserver_pg_port',
+        "--to-pg-port",
+        dest="new_pageserver_pg_port",
         required=False,
         default=6400,
         type=int,
-        help='pg port of the pageserver to migrate data to. Default: 6400',
+        help="pg port of the pageserver to migrate data to. Default: 6400",
     )
     parser.add_argument(
-        '--ignore-tenant-exists',
-        dest='ok_if_exists',
+        "--ignore-tenant-exists",
+        dest="ok_if_exists",
         required=False,
-        help=
-        'Ignore error if we are trying to create the tenant that already exists. It can be dangerous if existing tenant already contains some data.',
+        help="Ignore error if we are trying to create the tenant that already exists. It can be dangerous if existing tenant already contains some data.",
     )
     parser.add_argument(
-        '--pg-distrib-dir',
-        dest='pg_distrib_dir',
+        "--pg-distrib-dir",
+        dest="pg_distrib_dir",
         required=False,
-        default='/usr/local/',
-        help='Path where postgres binaries are installed. Default: /usr/local/',
+        default="/usr/local/",
+        help="Path where postgres binaries are installed. Default: /usr/local/",
     )
     parser.add_argument(
-        '--psql-path',
-        dest='psql_path',
+        "--psql-path",
+        dest="psql_path",
         required=False,
-        default='/usr/local/bin/psql',
-        help='Path to the psql binary. Default: /usr/local/bin/psql',
+        default="/usr/local/bin/psql",
+        help="Path to the psql binary. Default: /usr/local/bin/psql",
     )
     parser.add_argument(
-        '--only-import',
-        dest='only_import',
+        "--only-import",
+        dest="only_import",
         required=False,
         default=False,
-        action='store_true',
-        help='Skip export and tenant creation part',
+        action="store_true",
+        help="Skip export and tenant creation part",
     )
     parser.add_argument(
-        '--work-dir',
-        dest='work_dir',
+        "--work-dir",
+        dest="work_dir",
         required=True,
         default=False,
-        help='directory where temporary tar files are stored',
+        help="directory where temporary tar files are stored",
     )
     args = parser.parse_args()
     main(args)

scripts/generate_perf_report_page.py

@@ -1,31 +1,36 @@
 #!/usr/bin/env python3
 import argparse
+import json
 from dataclasses import dataclass
 from pathlib import Path
-import json
 from typing import Any, Dict, List, Optional, Tuple, cast
+
 from jinja2 import Template
 
 # skip 'input' columns. They are included in the header and just blow the table
-EXCLUDE_COLUMNS = frozenset({
-    'scale',
-    'duration',
-    'number_of_clients',
-    'number_of_threads',
-    'init_start_timestamp',
-    'init_end_timestamp',
-    'run_start_timestamp',
-    'run_end_timestamp',
-})
+EXCLUDE_COLUMNS = frozenset(
+    {
+        "scale",
+        "duration",
+        "number_of_clients",
+        "number_of_threads",
+        "init_start_timestamp",
+        "init_end_timestamp",
+        "run_start_timestamp",
+        "run_end_timestamp",
+    }
+)
 
-KEY_EXCLUDE_FIELDS = frozenset({
-    'init_start_timestamp',
-    'init_end_timestamp',
-    'run_start_timestamp',
-    'run_end_timestamp',
-})
-NEGATIVE_COLOR = 'negative'
-POSITIVE_COLOR = 'positive'
+KEY_EXCLUDE_FIELDS = frozenset(
+    {
+        "init_start_timestamp",
+        "init_end_timestamp",
+        "run_start_timestamp",
+        "run_end_timestamp",
+    }
+)
+NEGATIVE_COLOR = "negative"
+POSITIVE_COLOR = "positive"
 EPS = 1e-6
 
 
@@ -55,75 +60,76 @@ def get_columns(values: List[Dict[Any, Any]]) -> Tuple[List[Tuple[str, str]], Li
     value_columns = []
     common_columns = []
     for item in values:
-        if item['name'] in KEY_EXCLUDE_FIELDS:
+        if item["name"] in KEY_EXCLUDE_FIELDS:
             continue
-        if item['report'] != 'test_param':
-            value_columns.append(cast(str, item['name']))
+        if item["report"] != "test_param":
+            value_columns.append(cast(str, item["name"]))
         else:
-            common_columns.append((cast(str, item['name']), cast(str, item['value'])))
+            common_columns.append((cast(str, item["name"]), cast(str, item["value"])))
     value_columns.sort()
     common_columns.sort(key=lambda x: x[0])  # sort by name
     return common_columns, value_columns
 
 
 def format_ratio(ratio: float, report: str) -> Tuple[str, str]:
-    color = ''
-    sign = '+' if ratio > 0 else ''
+    color = ""
+    sign = "+" if ratio > 0 else ""
     if abs(ratio) < 0.05:
-        return f'&nbsp({sign}{ratio:.2f})', color
+        return f"&nbsp({sign}{ratio:.2f})", color
 
-    if report not in {'test_param', 'higher_is_better', 'lower_is_better'}:
-        raise ValueError(f'Unknown report type: {report}')
+    if report not in {"test_param", "higher_is_better", "lower_is_better"}:
+        raise ValueError(f"Unknown report type: {report}")
 
-    if report == 'test_param':
-        return f'{ratio:.2f}', color
+    if report == "test_param":
+        return f"{ratio:.2f}", color
 
     if ratio > 0:
-        if report == 'higher_is_better':
+        if report == "higher_is_better":
             color = POSITIVE_COLOR
-        elif report == 'lower_is_better':
+        elif report == "lower_is_better":
             color = NEGATIVE_COLOR
     elif ratio < 0:
-        if report == 'higher_is_better':
+        if report == "higher_is_better":
             color = NEGATIVE_COLOR
-        elif report == 'lower_is_better':
+        elif report == "lower_is_better":
             color = POSITIVE_COLOR
 
-    return f'&nbsp({sign}{ratio:.2f})', color
+    return f"&nbsp({sign}{ratio:.2f})", color
 
 
 def extract_value(name: str, suit_run: SuitRun) -> Optional[Dict[str, Any]]:
-    for item in suit_run.values['data']:
-        if item['name'] == name:
+    for item in suit_run.values["data"]:
+        if item["name"] == name:
             return cast(Dict[str, Any], item)
     return None
 
 
-def get_row_values(columns: List[str], run_result: SuitRun,
-                   prev_result: Optional[SuitRun]) -> List[RowValue]:
+def get_row_values(
+    columns: List[str], run_result: SuitRun, prev_result: Optional[SuitRun]
+) -> List[RowValue]:
     row_values = []
     for column in columns:
         current_value = extract_value(column, run_result)
         if current_value is None:
             # should never happen
-            raise ValueError(f'{column} not found in {run_result.values}')
+            raise ValueError(f"{column} not found in {run_result.values}")
 
         value = current_value["value"]
         if isinstance(value, float):
-            value = f'{value:.2f}'
+            value = f"{value:.2f}"
 
         if prev_result is None:
-            row_values.append(RowValue(value, '', ''))
+            row_values.append(RowValue(value, "", ""))
             continue
 
         prev_value = extract_value(column, prev_result)
         if prev_value is None:
             # this might happen when new metric is added and there is no value for it in previous run
             # let this be here, TODO add proper handling when this actually happens
-            raise ValueError(f'{column} not found in previous result')
+            raise ValueError(f"{column} not found in previous result")
         # adding `EPS` to each term to avoid ZeroDivisionError when the denominator is zero
-        ratio = (float(value) + EPS) / (float(prev_value['value']) + EPS) - 1
-        ratio_display, color = format_ratio(ratio, current_value['report'])
+        ratio = (float(value) + EPS) / (float(prev_value["value"]) + EPS) - 1
+        ratio_display, color = format_ratio(ratio, current_value["report"])
         row_values.append(RowValue(value, color, ratio_display))
     return row_values
 
@@ -139,8 +145,10 @@ def prepare_rows_from_runs(value_columns: List[str], runs: List[SuitRun]) -> Lis
     prev_run = None
     for run in runs:
         rows.append(
-            SuiteRunTableRow(revision=run.revision,
-                             values=get_row_values(value_columns, run, prev_run)))
+            SuiteRunTableRow(
+                revision=run.revision, values=get_row_values(value_columns, run, prev_run)
+            )
+        )
         prev_run = run
 
     return rows
@@ -152,27 +160,29 @@ def main(args: argparse.Namespace) -> None:
     # we have files in form: <ctr>_<rev>.json
     # fill them in the hashmap so we have grouped items for the
     # same run configuration (scale, duration etc.) ordered by counter.
-    for item in sorted(input_dir.iterdir(), key=lambda x: int(x.name.split('_')[0])):
+    for item in sorted(input_dir.iterdir(), key=lambda x: int(x.name.split("_")[0])):
         run_data = json.loads(item.read_text())
-        revision = run_data['revision']
+        revision = run_data["revision"]
 
-        for suit_result in run_data['result']:
-            key = "{}{}".format(run_data['platform'], suit_result['suit'])
+        for suit_result in run_data["result"]:
+            key = "{}{}".format(run_data["platform"], suit_result["suit"])
             # pack total duration as a synthetic value
-            total_duration = suit_result['total_duration']
-            suit_result['data'].append({
-                'name': 'total_duration',
-                'value': total_duration,
-                'unit': 's',
-                'report': 'lower_is_better',
-            })
-            common_columns, value_columns = get_columns(suit_result['data'])
+            total_duration = suit_result["total_duration"]
+            suit_result["data"].append(
+                {
+                    "name": "total_duration",
+                    "value": total_duration,
+                    "unit": "s",
+                    "report": "lower_is_better",
+                }
+            )
+            common_columns, value_columns = get_columns(suit_result["data"])
 
             grouped_runs.setdefault(
                 key,
                 SuitRuns(
-                    platform=run_data['platform'],
-                    suit=suit_result['suit'],
+                    platform=run_data["platform"],
+                    suit=suit_result["suit"],
                     common_columns=common_columns,
                     value_columns=value_columns,
                     runs=[],
@@ -184,26 +194,26 @@ def main(args: argparse.Namespace) -> None:
     for result in grouped_runs.values():
         suit = result.suit
         context[suit] = {
-            'common_columns': result.common_columns,
-            'value_columns': result.value_columns,
-            'platform': result.platform,
+            "common_columns": result.common_columns,
+            "value_columns": result.value_columns,
+            "platform": result.platform,
             # reverse the order so newest results are on top of the table
-            'rows': reversed(prepare_rows_from_runs(result.value_columns, result.runs)),
+            "rows": reversed(prepare_rows_from_runs(result.value_columns, result.runs)),
         }
 
-    template = Template((Path(__file__).parent / 'perf_report_template.html').read_text())
+    template = Template((Path(__file__).parent / "perf_report_template.html").read_text())
 
     Path(args.out).write_text(template.render(context=context))
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument(
-        '--input-dir',
-        dest='input_dir',
+        "--input-dir",
+        dest="input_dir",
         required=True,
-        help='Directory with jsons generated by the test suite',
+        help="Directory with jsons generated by the test suite",
     )
-    parser.add_argument('--out', required=True, help='Output html file path')
+    parser.add_argument("--out", required=True, help="Output html file path")
     args = parser.parse_args()
     main(args)

scripts/get_local_binaries.sh

@@ -0,0 +1,36 @@
+#!/usr/bin/env bash
+
+# This script builds deployable binaries from local code.
+
+set -e
+mkdir -p neon_install/bin
+
+# Building the image from scratch can take minutes, since we're recompiling everything,
+# including dependencies. To save time we remember the container id so it can be reused
+# for incremental compilation.
+CONTAINER_ID_PATH="neon_install/container_id"
+
+# Get the CONTAINER_ID, or create a container if one doesn't exist.
+if [ -e $CONTAINER_ID_PATH ]
+then
+    CONTAINER_ID=$(cat $CONTAINER_ID_PATH)
+    echo "Using existing container with id: $CONTAINER_ID"
+else
+    echo "No existing container found. Creating ..."
+    docker build \
+        --tag "neon:local" \
+        --build-arg "REPOSITORY=neondatabase" \
+        # Don't use cachepot, it will only slow things down. Instead
+        # maybe reuse the same container to get incremental compilation.
+        --build-arg "RUSTC_WRAPPER=''" \
+        .
+	CONTAINER_ID=$(docker create neon:local)
+fi
+
+# TODO if files changed locally, cp the code into the image and incrementally re-compile.
+# Or maybe mount code as a volume?
+
+# Copy files out of the container
+# TODO reuse the second part of .github/ansible/get_binaries.sh to get all the binaries
+set -x
+docker cp ${CONTAINER_ID}:/usr/local/bin/pageserver neon_install/bin/

scripts/git-upload

@@ -1,17 +1,16 @@
 #!/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
 import os
+import shlex
 import shutil
 import subprocess
 import sys
 import textwrap
+from contextlib import contextmanager
+from distutils.dir_util import copy_tree
+from pathlib import Path
+from tempfile import TemporaryDirectory
 from typing import Optional
 
 

scripts/ingest_perf_test_result.py

@@ -1,12 +1,13 @@
 #!/usr/bin/env python3
 import argparse
-from contextlib import contextmanager
 import json
 import os
+from contextlib import contextmanager
+from datetime import datetime
+from pathlib import Path
+
 import psycopg2
 import psycopg2.extras
-from pathlib import Path
-from datetime import datetime
 
 CREATE_TABLE = """
 CREATE TABLE IF NOT EXISTS perf_test_results (
@@ -24,15 +25,15 @@ CREATE TABLE IF NOT EXISTS perf_test_results (
 
 
 def err(msg):
-    print(f'error: {msg}')
+    print(f"error: {msg}")
     exit(1)
 
 
 @contextmanager
 def get_connection_cursor():
-    connstr = os.getenv('DATABASE_URL')
+    connstr = os.getenv("DATABASE_URL")
     if not connstr:
-        err('DATABASE_URL environment variable is not set')
+        err("DATABASE_URL environment variable is not set")
     with psycopg2.connect(connstr) as conn:
         with conn.cursor() as cur:
             yield cur
@@ -44,33 +45,35 @@ def create_table(cur):
 
 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']
+    revision = run_data["revision"]
+    platform = run_data["platform"]
 
-    run_result = run_data['result']
+    run_result = run_data["result"]
     args_list = []
 
     for suit_result in run_result:
-        suit = suit_result['suit']
-        total_duration = suit_result['total_duration']
+        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',
-        })
+        suit_result["data"].append(
+            {
+                "name": "total_duration",
+                "value": total_duration,
+                "unit": "s",
+                "report": "lower_is_better",
+            }
+        )
 
-        for metric in suit_result['data']:
+        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),
+                "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)
 
@@ -104,13 +107,16 @@ def ingest_perf_test_result(cursor, data_dile: Path, recorded_at_timestamp: int)
 
 
 def main():
-    parser = argparse.ArgumentParser(description='Perf test result uploader. \
-            Database connection string should be provided via DATABASE_URL environment variable', )
+    parser = argparse.ArgumentParser(
+        description="Perf test result uploader. \
+            Database connection string should be provided via DATABASE_URL environment variable",
+    )
     parser.add_argument(
-        '--ingest',
+        "--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')
+        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:
@@ -118,19 +124,19 @@ def main():
             create_table(cur)
 
         if not args.ingest.exists():
-            err(f'ingest path {args.ingest} does not exist')
+            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])
+                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}')
+                    print(f"Ingested {ingested} metric values from {item}")
             else:
-                recorded_at_timestamp = int(args.ingest.name.split('_')[0])
+                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}')
+                print(f"Ingested {ingested} metric values from {args.ingest}")
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     main()

setup.cfg

@@ -1,43 +1,8 @@
-# Just trying to gather linter settings in one file.
-# I wonder if there's a way to de-duplicate them...
-
 [flake8]
-max-line-length = 100
-
-[pycodestyle]
-max-line-length = 100
-
-[yapf]
-based_on_style = pep8
-column_limit = 100
-split_all_top_level_comma_separated_values = true
-
-[mypy]
-# mypy uses regex
-exclude = ^vendor/
-# some tests don't typecheck when this flag is set
-check_untyped_defs = false
-
-# Help mypy find imports when running against list of individual files.
-# Without this line it would behave differently when executed on the entire project.
-mypy_path = $MYPY_CONFIG_FILE_DIR:$MYPY_CONFIG_FILE_DIR/test_runner
-
-disallow_incomplete_defs = false
-disallow_untyped_calls = false
-disallow_untyped_decorators = false
-disallow_untyped_defs = false
-strict = true
-
-[mypy-asyncpg.*]
-# There is some work in progress, though: https://github.com/MagicStack/asyncpg/pull/577
-ignore_missing_imports = true
-
-[mypy-pg8000.*]
-# Used only in testing clients
-ignore_missing_imports = true
-
-[mypy-cached_property.*]
-ignore_missing_imports = true
-
-[mypy-pytest.*]
-ignore_missing_imports = true
+# Move config to pyproject.toml as soon as flake8 supports it
+# https://github.com/PyCQA/flake8/issues/234
+extend-ignore =
+    E203, # Whitespace before ':' -- conflicts with black
+    E266, # Too many leading '#' for block comment -- we use it for formatting sometimes
+    E501  # Line too long -- black sorts it out
+extend-exclude = vendor/

test_runner/batch_others/test_ancestor_branch.py

@@ -13,83 +13,90 @@ def test_ancestor_branch(neon_env_builder: NeonEnvBuilder):
     # Extend compaction_period and gc_period to disable background compaction and gc.
     tenant, _ = env.neon_cli.create_tenant(
         conf={
-            'gc_period': '10 m',
-            'gc_horizon': '1048576',
-            'checkpoint_distance': '4194304',
-            'compaction_period': '10 m',
-            'compaction_threshold': '2',
-            'compaction_target_size': '4194304',
-        })
+            "gc_period": "10 m",
+            "gc_horizon": "1048576",
+            "checkpoint_distance": "4194304",
+            "compaction_period": "10 m",
+            "compaction_threshold": "2",
+            "compaction_target_size": "4194304",
+        }
+    )
 
     env.pageserver.safe_psql("failpoints flush-frozen-before-sync=sleep(10000)")
 
-    pg_branch0 = env.postgres.create_start('main', tenant_id=tenant)
+    pg_branch0 = env.postgres.create_start("main", tenant_id=tenant)
     branch0_cur = pg_branch0.connect().cursor()
     branch0_timeline = query_scalar(branch0_cur, "SHOW neon.timeline_id")
     log.info(f"b0 timeline {branch0_timeline}")
 
     # Create table, and insert 100k rows.
-    branch0_lsn = query_scalar(branch0_cur, 'SELECT pg_current_wal_insert_lsn()')
+    branch0_lsn = query_scalar(branch0_cur, "SELECT pg_current_wal_insert_lsn()")
     log.info(f"b0 at lsn {branch0_lsn}")
 
-    branch0_cur.execute('CREATE TABLE foo (t text) WITH (autovacuum_enabled = off)')
-    branch0_cur.execute('''
+    branch0_cur.execute("CREATE TABLE foo (t text) WITH (autovacuum_enabled = off)")
+    branch0_cur.execute(
+        """
         INSERT INTO foo
             SELECT '00112233445566778899AABBCCDDEEFF' || ':branch0:' || g
             FROM generate_series(1, 100000) g
-    ''')
-    lsn_100 = query_scalar(branch0_cur, 'SELECT pg_current_wal_insert_lsn()')
-    log.info(f'LSN after 100k rows: {lsn_100}')
+    """
+    )
+    lsn_100 = query_scalar(branch0_cur, "SELECT pg_current_wal_insert_lsn()")
+    log.info(f"LSN after 100k rows: {lsn_100}")
 
     # Create branch1.
-    env.neon_cli.create_branch('branch1', 'main', tenant_id=tenant, ancestor_start_lsn=lsn_100)
-    pg_branch1 = env.postgres.create_start('branch1', tenant_id=tenant)
+    env.neon_cli.create_branch("branch1", "main", tenant_id=tenant, ancestor_start_lsn=lsn_100)
+    pg_branch1 = env.postgres.create_start("branch1", tenant_id=tenant)
     log.info("postgres is running on 'branch1' branch")
 
     branch1_cur = pg_branch1.connect().cursor()
     branch1_timeline = query_scalar(branch1_cur, "SHOW neon.timeline_id")
     log.info(f"b1 timeline {branch1_timeline}")
 
-    branch1_lsn = query_scalar(branch1_cur, 'SELECT pg_current_wal_insert_lsn()')
+    branch1_lsn = query_scalar(branch1_cur, "SELECT pg_current_wal_insert_lsn()")
     log.info(f"b1 at lsn {branch1_lsn}")
 
     # Insert 100k rows.
-    branch1_cur.execute('''
+    branch1_cur.execute(
+        """
         INSERT INTO foo
             SELECT '00112233445566778899AABBCCDDEEFF' || ':branch1:' || g
             FROM generate_series(1, 100000) g
-    ''')
-    lsn_200 = query_scalar(branch1_cur, 'SELECT pg_current_wal_insert_lsn()')
-    log.info(f'LSN after 200k rows: {lsn_200}')
+    """
+    )
+    lsn_200 = query_scalar(branch1_cur, "SELECT pg_current_wal_insert_lsn()")
+    log.info(f"LSN after 200k rows: {lsn_200}")
 
     # Create branch2.
-    env.neon_cli.create_branch('branch2', 'branch1', tenant_id=tenant, ancestor_start_lsn=lsn_200)
-    pg_branch2 = env.postgres.create_start('branch2', tenant_id=tenant)
+    env.neon_cli.create_branch("branch2", "branch1", tenant_id=tenant, ancestor_start_lsn=lsn_200)
+    pg_branch2 = env.postgres.create_start("branch2", tenant_id=tenant)
     log.info("postgres is running on 'branch2' branch")
     branch2_cur = pg_branch2.connect().cursor()
 
     branch2_timeline = query_scalar(branch2_cur, "SHOW neon.timeline_id")
     log.info(f"b2 timeline {branch2_timeline}")
 
-    branch2_lsn = query_scalar(branch2_cur, 'SELECT pg_current_wal_insert_lsn()')
+    branch2_lsn = query_scalar(branch2_cur, "SELECT pg_current_wal_insert_lsn()")
     log.info(f"b2 at lsn {branch2_lsn}")
 
     # Insert 100k rows.
-    branch2_cur.execute('''
+    branch2_cur.execute(
+        """
         INSERT INTO foo
             SELECT '00112233445566778899AABBCCDDEEFF' || ':branch2:' || g
             FROM generate_series(1, 100000) g
-    ''')
-    lsn_300 = query_scalar(branch2_cur, 'SELECT pg_current_wal_insert_lsn()')
-    log.info(f'LSN after 300k rows: {lsn_300}')
+    """
+    )
+    lsn_300 = query_scalar(branch2_cur, "SELECT pg_current_wal_insert_lsn()")
+    log.info(f"LSN after 300k rows: {lsn_300}")
 
     # Run compaction on branch1.
-    compact = f'compact {tenant.hex} {branch1_timeline} {lsn_200}'
+    compact = f"compact {tenant.hex} {branch1_timeline} {lsn_200}"
     log.info(compact)
     env.pageserver.safe_psql(compact)
 
-    assert query_scalar(branch0_cur, 'SELECT count(*) FROM foo') == 100000
+    assert query_scalar(branch0_cur, "SELECT count(*) FROM foo") == 100000
 
-    assert query_scalar(branch1_cur, 'SELECT count(*) FROM foo') == 200000
+    assert query_scalar(branch1_cur, "SELECT count(*) FROM foo") == 200000
 
-    assert query_scalar(branch2_cur, 'SELECT count(*) FROM foo') == 300000
+    assert query_scalar(branch2_cur, "SELECT count(*) FROM foo") == 300000

test_runner/batch_others/test_auth.py

@@ -1,7 +1,8 @@
 from contextlib import closing
 from uuid import uuid4
-from fixtures.neon_fixtures import NeonEnvBuilder, NeonPageserverApiException
+
 import pytest
+from fixtures.neon_fixtures import NeonEnvBuilder, NeonPageserverApiException
 
 
 def test_pageserver_auth(neon_env_builder: NeonEnvBuilder):
@@ -23,41 +24,46 @@ def test_pageserver_auth(neon_env_builder: NeonEnvBuilder):
     ps.safe_psql("set FOO", password=tenant_token)
     ps.safe_psql("set FOO", password=management_token)
 
-    new_timeline_id = env.neon_cli.create_branch('test_pageserver_auth',
-                                                 tenant_id=env.initial_tenant)
+    new_timeline_id = env.neon_cli.create_branch(
+        "test_pageserver_auth", tenant_id=env.initial_tenant
+    )
 
     # tenant can create branches
-    tenant_http_client.timeline_create(tenant_id=env.initial_tenant,
-                                       ancestor_timeline_id=new_timeline_id)
+    tenant_http_client.timeline_create(
+        tenant_id=env.initial_tenant, ancestor_timeline_id=new_timeline_id
+    )
     # console can create branches for tenant
-    management_http_client.timeline_create(tenant_id=env.initial_tenant,
-                                           ancestor_timeline_id=new_timeline_id)
+    management_http_client.timeline_create(
+        tenant_id=env.initial_tenant, ancestor_timeline_id=new_timeline_id
+    )
 
     # fail to create branch using token with different tenant_id
-    with pytest.raises(NeonPageserverApiException,
-                       match='Forbidden: Tenant id mismatch. Permission denied'):
-        invalid_tenant_http_client.timeline_create(tenant_id=env.initial_tenant,
-                                                   ancestor_timeline_id=new_timeline_id)
+    with pytest.raises(
+        NeonPageserverApiException, match="Forbidden: Tenant id mismatch. Permission denied"
+    ):
+        invalid_tenant_http_client.timeline_create(
+            tenant_id=env.initial_tenant, ancestor_timeline_id=new_timeline_id
+        )
 
     # create tenant using management token
     management_http_client.tenant_create()
 
     # fail to create tenant using tenant token
     with pytest.raises(
-            NeonPageserverApiException,
-            match='Forbidden: Attempt to access management api with tenant scope. Permission denied'
+        NeonPageserverApiException,
+        match="Forbidden: Attempt to access management api with tenant scope. Permission denied",
     ):
         tenant_http_client.tenant_create()
 
 
-@pytest.mark.parametrize('with_safekeepers', [False, True])
+@pytest.mark.parametrize("with_safekeepers", [False, True])
 def test_compute_auth_to_pageserver(neon_env_builder: NeonEnvBuilder, with_safekeepers: bool):
     neon_env_builder.auth_enabled = True
     if with_safekeepers:
         neon_env_builder.num_safekeepers = 3
     env = neon_env_builder.init_start()
 
-    branch = f'test_compute_auth_to_pageserver{with_safekeepers}'
+    branch = f"test_compute_auth_to_pageserver{with_safekeepers}"
     env.neon_cli.create_branch(branch)
     pg = env.postgres.create_start(branch)
 
@@ -65,7 +71,7 @@ def test_compute_auth_to_pageserver(neon_env_builder: NeonEnvBuilder, with_safek
         with conn.cursor() as cur:
             # we rely upon autocommit after each statement
             # as waiting for acceptors happens there
-            cur.execute('CREATE TABLE t(key int primary key, value text)')
+            cur.execute("CREATE TABLE t(key int primary key, value text)")
             cur.execute("INSERT INTO t SELECT generate_series(1,100000), 'payload'")
-            cur.execute('SELECT sum(key) FROM t')
-            assert cur.fetchone() == (5000050000, )
+            cur.execute("SELECT sum(key) FROM t")
+            assert cur.fetchone() == (5000050000,)

test_runner/batch_others/test_backpressure.py

@@ -1,13 +1,13 @@
+import threading
+import time
 from contextlib import closing, contextmanager
+
 import psycopg2.extras
 import pytest
-from fixtures.neon_fixtures import NeonEnvBuilder
 from fixtures.log_helper import log
-import time
-from fixtures.neon_fixtures import Postgres
-import threading
+from fixtures.neon_fixtures import NeonEnvBuilder, Postgres
 
-pytest_plugins = ("fixtures.neon_fixtures")
+pytest_plugins = "fixtures.neon_fixtures"
 
 
 @contextmanager
@@ -44,7 +44,8 @@ def check_backpressure(pg: Postgres, stop_event: threading.Event, polling_interv
     with pg_cur(pg) as cur:
         while not stop_event.is_set():
             try:
-                cur.execute('''
+                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,
@@ -52,16 +53,19 @@ def check_backpressure(pg: Postgres, stop_event: threading.Event, polling_interv
                 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]})")
+                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
@@ -71,7 +75,9 @@ def check_backpressure(pg: Postgres, stop_event: threading.Event, polling_interv
                 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
+                    assert (
+                        remote_consistent_lsn_lag < max_replication_apply_lag_bytes + lag_overflow
+                    )
 
                 time.sleep(polling_interval)
 
@@ -79,7 +85,7 @@ def check_backpressure(pg: Postgres, stop_event: threading.Event, polling_interv
                 log.info(f"backpressure check query failed: {e}")
                 stop_event.set()
 
-    log.info('check thread stopped')
+    log.info("check thread stopped")
 
 
 # This test illustrates how to tune backpressure to control the lag
@@ -94,10 +100,11 @@ def check_backpressure(pg: Postgres, stop_event: threading.Event, polling_interv
 def test_backpressure_received_lsn_lag(neon_env_builder: NeonEnvBuilder):
     env = neon_env_builder.init_start()
     # Create a branch for us
-    env.neon_cli.create_branch('test_backpressure')
+    env.neon_cli.create_branch("test_backpressure")
 
-    pg = env.postgres.create_start('test_backpressure',
-                                   config_lines=['max_replication_write_lag=30MB'])
+    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
@@ -131,23 +138,29 @@ def test_backpressure_received_lsn_lag(neon_env_builder: NeonEnvBuilder):
                 rows_inserted += 100000
             except Exception as e:
                 if check_thread.is_alive():
-                    log.info('stopping check thread')
+                    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"
+                    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."
+                    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')
+        log.info("stopping check thread")
         check_stop_event.set()
         check_thread.join()
-        log.info('check thread stopped')
+        log.info("check thread stopped")
     else:
-        assert False, "WAL lag overflowed configured threshold. That means backpressure doesn't work."
+        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
+# 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_basebackup_error.py

@@ -1,5 +1,4 @@
 import pytest
-
 from fixtures.neon_fixtures import NeonEnv
 
 
@@ -12,7 +11,7 @@ def test_basebackup_error(neon_simple_env: NeonEnv):
     env.neon_cli.create_branch("test_basebackup_error", "empty")
 
     # Introduce failpoint
-    env.pageserver.safe_psql(f"failpoints basebackup-before-control-file=return")
+    env.pageserver.safe_psql("failpoints basebackup-before-control-file=return")
 
     with pytest.raises(Exception, match="basebackup-before-control-file"):
-        pg = env.postgres.create_start('test_basebackup_error')
+        env.postgres.create_start("test_basebackup_error")

test_runner/batch_others/test_branch_and_gc.py

@@ -1,6 +1,7 @@
 import threading
-import pytest
 import time
+
+import pytest
 from fixtures.log_helper import log
 from fixtures.neon_fixtures import NeonEnv
 from fixtures.utils import lsn_from_hex, query_scalar
@@ -49,55 +50,52 @@ def test_branch_and_gc(neon_simple_env: NeonEnv):
     tenant, _ = env.neon_cli.create_tenant(
         conf={
             # disable background GC
-            'gc_period': '10 m',
-            'gc_horizon': f'{10 * 1024 ** 3}',
-
+            "gc_period": "10 m",
+            "gc_horizon": f"{10 * 1024 ** 3}",
             # small checkpoint distance to create more delta layer files
-            'checkpoint_distance': f'{1024 ** 2}',
-
+            "checkpoint_distance": f"{1024 ** 2}",
             # set the target size to be large to allow the image layer to cover the whole key space
-            'compaction_target_size': f'{1024 ** 3}',
-
+            "compaction_target_size": f"{1024 ** 3}",
             # tweak the default settings to allow quickly create image layers and L1 layers
-            'compaction_period': '1 s',
-            'compaction_threshold': '2',
-            'image_creation_threshold': '1',
-
+            "compaction_period": "1 s",
+            "compaction_threshold": "2",
+            "image_creation_threshold": "1",
             # set PITR interval to be small, so we can do GC
-            'pitr_interval': '1 s'
-        })
+            "pitr_interval": "1 s",
+        }
+    )
 
-    timeline_main = env.neon_cli.create_timeline(f'test_main', tenant_id=tenant)
-    pg_main = env.postgres.create_start('test_main', tenant_id=tenant)
+    timeline_main = env.neon_cli.create_timeline("test_main", tenant_id=tenant)
+    pg_main = env.postgres.create_start("test_main", tenant_id=tenant)
 
     main_cur = pg_main.connect().cursor()
 
     main_cur.execute(
         "CREATE TABLE foo(key serial primary key, t text default 'foooooooooooooooooooooooooooooooooooooooooooooooooooo')"
     )
-    main_cur.execute('INSERT INTO foo SELECT FROM generate_series(1, 100000)')
-    lsn1 = query_scalar(main_cur, 'SELECT pg_current_wal_insert_lsn()')
-    log.info(f'LSN1: {lsn1}')
+    main_cur.execute("INSERT INTO foo SELECT FROM generate_series(1, 100000)")
+    lsn1 = query_scalar(main_cur, "SELECT pg_current_wal_insert_lsn()")
+    log.info(f"LSN1: {lsn1}")
 
-    main_cur.execute('INSERT INTO foo SELECT FROM generate_series(1, 100000)')
-    lsn2 = query_scalar(main_cur, 'SELECT pg_current_wal_insert_lsn()')
-    log.info(f'LSN2: {lsn2}')
+    main_cur.execute("INSERT INTO foo SELECT FROM generate_series(1, 100000)")
+    lsn2 = query_scalar(main_cur, "SELECT pg_current_wal_insert_lsn()")
+    log.info(f"LSN2: {lsn2}")
 
     # Set the GC horizon so that lsn1 is inside the horizon, which means
     # we can create a new branch starting from lsn1.
     env.pageserver.safe_psql(
-        f'do_gc {tenant.hex} {timeline_main.hex} {lsn_from_hex(lsn2) - lsn_from_hex(lsn1) + 1024}')
+        f"do_gc {tenant.hex} {timeline_main.hex} {lsn_from_hex(lsn2) - lsn_from_hex(lsn1) + 1024}"
+    )
 
-    env.neon_cli.create_branch('test_branch',
-                               'test_main',
-                               tenant_id=tenant,
-                               ancestor_start_lsn=lsn1)
-    pg_branch = env.postgres.create_start('test_branch', tenant_id=tenant)
+    env.neon_cli.create_branch(
+        "test_branch", "test_main", tenant_id=tenant, ancestor_start_lsn=lsn1
+    )
+    pg_branch = env.postgres.create_start("test_branch", tenant_id=tenant)
 
     branch_cur = pg_branch.connect().cursor()
-    branch_cur.execute('INSERT INTO foo SELECT FROM generate_series(1, 100000)')
+    branch_cur.execute("INSERT INTO foo SELECT FROM generate_series(1, 100000)")
 
-    assert query_scalar(branch_cur, 'SELECT count(*) FROM foo') == 200000
+    assert query_scalar(branch_cur, "SELECT count(*) FROM foo") == 200000
 
 
 # This test simulates a race condition happening when branch creation and GC are performed concurrently.
@@ -120,38 +118,37 @@ def test_branch_creation_before_gc(neon_simple_env: NeonEnv):
     tenant, _ = env.neon_cli.create_tenant(
         conf={
             # disable background GC
-            'gc_period': '10 m',
-            'gc_horizon': f'{10 * 1024 ** 3}',
-
+            "gc_period": "10 m",
+            "gc_horizon": f"{10 * 1024 ** 3}",
             # small checkpoint distance to create more delta layer files
-            'checkpoint_distance': f'{1024 ** 2}',
-
+            "checkpoint_distance": f"{1024 ** 2}",
             # set the target size to be large to allow the image layer to cover the whole key space
-            'compaction_target_size': f'{1024 ** 3}',
-
+            "compaction_target_size": f"{1024 ** 3}",
             # tweak the default settings to allow quickly create image layers and L1 layers
-            'compaction_period': '1 s',
-            'compaction_threshold': '2',
-            'image_creation_threshold': '1',
-
+            "compaction_period": "1 s",
+            "compaction_threshold": "2",
+            "image_creation_threshold": "1",
             # set PITR interval to be small, so we can do GC
-            'pitr_interval': '0 s'
-        })
+            "pitr_interval": "0 s",
+        }
+    )
 
-    b0 = env.neon_cli.create_branch('b0', tenant_id=tenant)
-    pg0 = env.postgres.create_start('b0', tenant_id=tenant)
-    res = pg0.safe_psql_many(queries=[
-        "CREATE TABLE t(key serial primary key)",
-        "INSERT INTO t SELECT FROM generate_series(1, 100000)",
-        "SELECT pg_current_wal_insert_lsn()",
-        "INSERT INTO t SELECT FROM generate_series(1, 100000)",
-    ])
+    b0 = env.neon_cli.create_branch("b0", tenant_id=tenant)
+    pg0 = env.postgres.create_start("b0", tenant_id=tenant)
+    res = pg0.safe_psql_many(
+        queries=[
+            "CREATE TABLE t(key serial primary key)",
+            "INSERT INTO t SELECT FROM generate_series(1, 100000)",
+            "SELECT pg_current_wal_insert_lsn()",
+            "INSERT INTO t SELECT FROM generate_series(1, 100000)",
+        ]
+    )
     lsn = res[2][0][0]
 
     # Use `failpoint=sleep` and `threading` to make the GC iteration triggers *before* the
     # branch creation task but the individual timeline GC iteration happens *after*
     # the branch creation task.
-    env.pageserver.safe_psql(f"failpoints before-timeline-gc=sleep(2000)")
+    env.pageserver.safe_psql("failpoints before-timeline-gc=sleep(2000)")
 
     def do_gc():
         env.pageserver.safe_psql(f"do_gc {tenant.hex} {b0.hex} 0")
@@ -166,6 +163,6 @@ def test_branch_creation_before_gc(neon_simple_env: NeonEnv):
 
     # The starting LSN is invalid as the corresponding record is scheduled to be removed by in-queue GC.
     with pytest.raises(Exception, match="invalid branch start lsn"):
-        env.neon_cli.create_branch('b1', 'b0', tenant_id=tenant, ancestor_start_lsn=lsn)
+        env.neon_cli.create_branch("b1", "b0", tenant_id=tenant, ancestor_start_lsn=lsn)
 
     thread.join()

test_runner/batch_others/test_branch_behind.py

@@ -1,8 +1,8 @@
 import psycopg2.extras
 import pytest
 from fixtures.log_helper import log
-from fixtures.utils import print_gc_result, query_scalar
 from fixtures.neon_fixtures import NeonEnvBuilder
+from fixtures.utils import print_gc_result, query_scalar
 
 
 #
@@ -21,8 +21,8 @@ def test_branch_behind(neon_env_builder: NeonEnvBuilder):
     env = neon_env_builder.init_start()
 
     # Branch at the point where only 100 rows were inserted
-    env.neon_cli.create_branch('test_branch_behind')
-    pgmain = env.postgres.create_start('test_branch_behind')
+    env.neon_cli.create_branch("test_branch_behind")
+    pgmain = env.postgres.create_start("test_branch_behind")
     log.info("postgres is running on 'test_branch_behind' branch")
 
     main_cur = pgmain.connect().cursor()
@@ -30,80 +30,86 @@ def test_branch_behind(neon_env_builder: NeonEnvBuilder):
     timeline = query_scalar(main_cur, "SHOW neon.timeline_id")
 
     # Create table, and insert the first 100 rows
-    main_cur.execute('CREATE TABLE foo (t text)')
+    main_cur.execute("CREATE TABLE foo (t text)")
 
     # keep some early lsn to test branch creation on out of date lsn
-    gced_lsn = query_scalar(main_cur, 'SELECT pg_current_wal_insert_lsn()')
+    gced_lsn = query_scalar(main_cur, "SELECT pg_current_wal_insert_lsn()")
 
-    main_cur.execute('''
+    main_cur.execute(
+        """
         INSERT INTO foo
             SELECT 'long string to consume some space' || g
             FROM generate_series(1, 100) g
-    ''')
-    lsn_a = query_scalar(main_cur, 'SELECT pg_current_wal_insert_lsn()')
-    log.info(f'LSN after 100 rows: {lsn_a}')
+    """
+    )
+    lsn_a = query_scalar(main_cur, "SELECT pg_current_wal_insert_lsn()")
+    log.info(f"LSN after 100 rows: {lsn_a}")
 
     # Insert some more rows. (This generates enough WAL to fill a few segments.)
-    main_cur.execute('''
+    main_cur.execute(
+        """
         INSERT INTO foo
             SELECT 'long string to consume some space' || g
             FROM generate_series(1, 200000) g
-    ''')
-    lsn_b = query_scalar(main_cur, 'SELECT pg_current_wal_insert_lsn()')
-    log.info(f'LSN after 200100 rows: {lsn_b}')
+    """
+    )
+    lsn_b = query_scalar(main_cur, "SELECT pg_current_wal_insert_lsn()")
+    log.info(f"LSN after 200100 rows: {lsn_b}")
 
     # Branch at the point where only 100 rows were inserted
-    env.neon_cli.create_branch('test_branch_behind_hundred',
-                               'test_branch_behind',
-                               ancestor_start_lsn=lsn_a)
+    env.neon_cli.create_branch(
+        "test_branch_behind_hundred", "test_branch_behind", ancestor_start_lsn=lsn_a
+    )
 
     # Insert many more rows. This generates enough WAL to fill a few segments.
-    main_cur.execute('''
+    main_cur.execute(
+        """
         INSERT INTO foo
             SELECT 'long string to consume some space' || g
             FROM generate_series(1, 200000) g
-    ''')
-    lsn_c = query_scalar(main_cur, 'SELECT pg_current_wal_insert_lsn()')
+    """
+    )
+    lsn_c = query_scalar(main_cur, "SELECT pg_current_wal_insert_lsn()")
 
-    log.info(f'LSN after 400100 rows: {lsn_c}')
+    log.info(f"LSN after 400100 rows: {lsn_c}")
 
     # Branch at the point where only 200100 rows were inserted
-    env.neon_cli.create_branch('test_branch_behind_more',
-                               'test_branch_behind',
-                               ancestor_start_lsn=lsn_b)
+    env.neon_cli.create_branch(
+        "test_branch_behind_more", "test_branch_behind", ancestor_start_lsn=lsn_b
+    )
 
-    pg_hundred = env.postgres.create_start('test_branch_behind_hundred')
-    pg_more = env.postgres.create_start('test_branch_behind_more')
+    pg_hundred = env.postgres.create_start("test_branch_behind_hundred")
+    pg_more = env.postgres.create_start("test_branch_behind_more")
 
     # On the 'hundred' branch, we should see only 100 rows
     hundred_cur = pg_hundred.connect().cursor()
-    assert query_scalar(hundred_cur, 'SELECT count(*) FROM foo') == 100
+    assert query_scalar(hundred_cur, "SELECT count(*) FROM foo") == 100
 
     # On the 'more' branch, we should see 100200 rows
     more_cur = pg_more.connect().cursor()
-    assert query_scalar(more_cur, 'SELECT count(*) FROM foo') == 200100
+    assert query_scalar(more_cur, "SELECT count(*) FROM foo") == 200100
 
     # All the rows are visible on the main branch
-    assert query_scalar(main_cur, 'SELECT count(*) FROM foo') == 400100
+    assert query_scalar(main_cur, "SELECT count(*) FROM foo") == 400100
 
     # Check bad lsn's for branching
 
     # branch at segment boundary
-    env.neon_cli.create_branch('test_branch_segment_boundary',
-                               'test_branch_behind',
-                               ancestor_start_lsn="0/3000000")
-    pg = env.postgres.create_start('test_branch_segment_boundary')
-    assert pg.safe_psql('SELECT 1')[0][0] == 1
+    env.neon_cli.create_branch(
+        "test_branch_segment_boundary", "test_branch_behind", ancestor_start_lsn="0/3000000"
+    )
+    pg = env.postgres.create_start("test_branch_segment_boundary")
+    assert pg.safe_psql("SELECT 1")[0][0] == 1
 
     # branch at pre-initdb lsn
     with pytest.raises(Exception, match="invalid branch start lsn"):
-        env.neon_cli.create_branch('test_branch_preinitdb', ancestor_start_lsn="0/42")
+        env.neon_cli.create_branch("test_branch_preinitdb", ancestor_start_lsn="0/42")
 
     # branch at pre-ancestor lsn
     with pytest.raises(Exception, match="less than timeline ancestor lsn"):
-        env.neon_cli.create_branch('test_branch_preinitdb',
-                                   'test_branch_behind',
-                                   ancestor_start_lsn="0/42")
+        env.neon_cli.create_branch(
+            "test_branch_preinitdb", "test_branch_behind", ancestor_start_lsn="0/42"
+        )
 
     # check that we cannot create branch based on garbage collected data
     with env.pageserver.cursor(cursor_factory=psycopg2.extras.DictCursor) as pscur:
@@ -114,13 +120,13 @@ def test_branch_behind(neon_env_builder: NeonEnvBuilder):
 
     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.neon_cli.create_branch('test_branch_create_fail',
-                                   'test_branch_behind',
-                                   ancestor_start_lsn=gced_lsn)
+        env.neon_cli.create_branch(
+            "test_branch_create_fail", "test_branch_behind", ancestor_start_lsn=gced_lsn
+        )
 
     # check that after gc everything is still there
-    assert query_scalar(hundred_cur, 'SELECT count(*) FROM foo') == 100
+    assert query_scalar(hundred_cur, "SELECT count(*) FROM foo") == 100
 
-    assert query_scalar(more_cur, 'SELECT count(*) FROM foo') == 200100
+    assert query_scalar(more_cur, "SELECT count(*) FROM foo") == 200100
 
-    assert query_scalar(main_cur, 'SELECT count(*) FROM foo') == 400100
+    assert query_scalar(main_cur, "SELECT count(*) FROM foo") == 400100

test_runner/batch_others/test_branching.py

@@ -1,10 +1,11 @@
-from typing import List
-import threading
-import pytest
-from fixtures.neon_fixtures import NeonEnv, PgBin, Postgres
-import time
 import random
+import threading
+import time
+from typing import List
+
+import pytest
 from fixtures.log_helper import log
+from fixtures.neon_fixtures import NeonEnv, PgBin, Postgres
 from performance.test_perf_pgbench import get_scales_matrix
 
 
@@ -20,38 +21,37 @@ from performance.test_perf_pgbench import get_scales_matrix
 @pytest.mark.parametrize("n_branches", [10])
 @pytest.mark.parametrize("scale", get_scales_matrix(1))
 @pytest.mark.parametrize("ty", ["cascade", "flat"])
-def test_branching_with_pgbench(neon_simple_env: NeonEnv,
-                                pg_bin: PgBin,
-                                n_branches: int,
-                                scale: int,
-                                ty: str):
+def test_branching_with_pgbench(
+    neon_simple_env: NeonEnv, pg_bin: PgBin, n_branches: int, scale: int, ty: str
+):
     env = neon_simple_env
 
     # Use aggressive GC and checkpoint settings, so that we also exercise GC during the test
     tenant, _ = env.neon_cli.create_tenant(
-         conf={
-             'gc_period': '5 s',
-             'gc_horizon': f'{1024 ** 2}',
-             'checkpoint_distance': f'{1024 ** 2}',
-             'compaction_target_size': f'{1024 ** 2}',
-             # set PITR interval to be small, so we can do GC
-             'pitr_interval': '5 s'
-         })
+        conf={
+            "gc_period": "5 s",
+            "gc_horizon": f"{1024 ** 2}",
+            "checkpoint_distance": f"{1024 ** 2}",
+            "compaction_target_size": f"{1024 ** 2}",
+            # set PITR interval to be small, so we can do GC
+            "pitr_interval": "5 s",
+        }
+    )
 
     def run_pgbench(pg: Postgres):
         connstr = pg.connstr()
 
         log.info(f"Start a pgbench workload on pg {connstr}")
 
-        pg_bin.run_capture(['pgbench', '-i', f'-s{scale}', connstr])
-        pg_bin.run_capture(['pgbench', '-T15', connstr])
+        pg_bin.run_capture(["pgbench", "-i", f"-s{scale}", connstr])
+        pg_bin.run_capture(["pgbench", "-T15", connstr])
 
-    env.neon_cli.create_branch('b0', tenant_id=tenant)
+    env.neon_cli.create_branch("b0", tenant_id=tenant)
     pgs: List[Postgres] = []
-    pgs.append(env.postgres.create_start('b0', tenant_id=tenant))
+    pgs.append(env.postgres.create_start("b0", tenant_id=tenant))
 
     threads: List[threading.Thread] = []
-    threads.append(threading.Thread(target=run_pgbench, args=(pgs[0], ), daemon=True))
+    threads.append(threading.Thread(target=run_pgbench, args=(pgs[0],), daemon=True))
     threads[-1].start()
 
     thread_limit = 4
@@ -72,18 +72,18 @@ def test_branching_with_pgbench(neon_simple_env: NeonEnv,
             threads = []
 
         if ty == "cascade":
-            env.neon_cli.create_branch('b{}'.format(i + 1), 'b{}'.format(i), tenant_id=tenant)
+            env.neon_cli.create_branch("b{}".format(i + 1), "b{}".format(i), tenant_id=tenant)
         else:
-            env.neon_cli.create_branch('b{}'.format(i + 1), 'b0', tenant_id=tenant)
+            env.neon_cli.create_branch("b{}".format(i + 1), "b0", tenant_id=tenant)
 
-        pgs.append(env.postgres.create_start('b{}'.format(i + 1), tenant_id=tenant))
+        pgs.append(env.postgres.create_start("b{}".format(i + 1), tenant_id=tenant))
 
-        threads.append(threading.Thread(target=run_pgbench, args=(pgs[-1], ), daemon=True))
+        threads.append(threading.Thread(target=run_pgbench, args=(pgs[-1],), daemon=True))
         threads[-1].start()
 
     for thread in threads:
         thread.join()
 
     for pg in pgs:
-        res = pg.safe_psql('SELECT count(*) from pgbench_accounts')
-        assert res[0] == (100000 * scale, )
+        res = pg.safe_psql("SELECT count(*) from pgbench_accounts")
+        assert res[0] == (100000 * scale,)

test_runner/batch_others/test_broken_timeline.py

@@ -1,12 +1,10 @@
-from typing import List, Tuple
-from uuid import UUID
-import pytest
 import concurrent.futures
-from contextlib import closing
-from fixtures.neon_fixtures import NeonEnvBuilder, NeonEnv, Postgres
-from fixtures.log_helper import log
 import os
+from typing import List, Tuple
 
+import pytest
+from fixtures.log_helper import log
+from fixtures.neon_fixtures import NeonEnv, NeonEnvBuilder, Postgres
 from fixtures.utils import query_scalar
 
 
@@ -24,7 +22,7 @@ def test_broken_timeline(neon_env_builder: NeonEnvBuilder):
         tenant_id = tenant_id_uuid.hex
         timeline_id = timeline_id_uuid.hex
 
-        pg = env.postgres.create_start(f'main', tenant_id=tenant_id_uuid)
+        pg = env.postgres.create_start("main", tenant_id=tenant_id_uuid)
         with pg.cursor() as cur:
             cur.execute("CREATE TABLE t(key int primary key, value text)")
             cur.execute("INSERT INTO t SELECT generate_series(1,100), 'payload'")
@@ -42,7 +40,7 @@ def test_broken_timeline(neon_env_builder: NeonEnvBuilder):
     # Corrupt metadata file on timeline 1
     (tenant1, timeline1, pg1) = tenant_timelines[1]
     metadata_path = "{}/tenants/{}/timelines/{}/metadata".format(env.repo_dir, tenant1, timeline1)
-    print(f'overwriting metadata file at {metadata_path}')
+    print(f"overwriting metadata file at {metadata_path}")
     f = open(metadata_path, "w")
     f.write("overwritten with garbage!")
     f.close()
@@ -52,17 +50,17 @@ def test_broken_timeline(neon_env_builder: NeonEnvBuilder):
     (tenant2, timeline2, pg2) = tenant_timelines[2]
     timeline_path = "{}/tenants/{}/timelines/{}/".format(env.repo_dir, tenant2, timeline2)
     for filename in os.listdir(timeline_path):
-        if filename.startswith('00000'):
+        if filename.startswith("00000"):
             # Looks like a layer file. Remove it
-            os.remove(f'{timeline_path}/{filename}')
+            os.remove(f"{timeline_path}/{filename}")
 
     # Corrupt layer files file on timeline 3
     (tenant3, timeline3, pg3) = tenant_timelines[3]
     timeline_path = "{}/tenants/{}/timelines/{}/".format(env.repo_dir, tenant3, timeline3)
     for filename in os.listdir(timeline_path):
-        if filename.startswith('00000'):
+        if filename.startswith("00000"):
             # Looks like a layer file. Corrupt it
-            f = open(f'{timeline_path}/{filename}', "w")
+            f = open(f"{timeline_path}/{filename}", "w")
             f.write("overwritten with garbage!")
             f.close()
 
@@ -77,7 +75,7 @@ def test_broken_timeline(neon_env_builder: NeonEnvBuilder):
         (tenant, timeline, pg) = tenant_timelines[n]
         with pytest.raises(Exception, match="Cannot load local timeline") as err:
             pg.start()
-        log.info(f'compute startup failed as expected: {err}')
+        log.info(f"compute startup failed as expected: {err}")
 
 
 def test_create_multiple_timelines_parallel(neon_simple_env: NeonEnv):
@@ -87,9 +85,10 @@ def test_create_multiple_timelines_parallel(neon_simple_env: NeonEnv):
 
     with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor:
         futures = [
-            executor.submit(env.neon_cli.create_timeline,
-                            f"test-create-multiple-timelines-{i}",
-                            tenant_id) for i in range(4)
+            executor.submit(
+                env.neon_cli.create_timeline, f"test-create-multiple-timelines-{i}", tenant_id
+            )
+            for i in range(4)
         ]
         for future in futures:
             future.result()
@@ -101,7 +100,7 @@ def test_fix_broken_timelines_on_startup(neon_simple_env: NeonEnv):
     tenant_id, _ = env.neon_cli.create_tenant()
 
     # Introduce failpoint when creating a new timeline
-    env.pageserver.safe_psql(f"failpoints before-checkpoint-new-timeline=return")
+    env.pageserver.safe_psql("failpoints before-checkpoint-new-timeline=return")
     with pytest.raises(Exception, match="before-checkpoint-new-timeline"):
         _ = env.neon_cli.create_timeline("test_fix_broken_timelines", tenant_id)
 

test_runner/batch_others/test_clog_truncate.py

@@ -1,10 +1,8 @@
-import time
 import os
+import time
 
-from contextlib import closing
-
-from fixtures.neon_fixtures import NeonEnv
 from fixtures.log_helper import log
+from fixtures.neon_fixtures import NeonEnv
 from fixtures.utils import query_scalar
 
 
@@ -13,60 +11,60 @@ from fixtures.utils import query_scalar
 #
 def test_clog_truncate(neon_simple_env: NeonEnv):
     env = neon_simple_env
-    env.neon_cli.create_branch('test_clog_truncate', 'empty')
+    env.neon_cli.create_branch("test_clog_truncate", "empty")
 
     # set aggressive autovacuum to make sure that truncation will happen
     config = [
-        'autovacuum_max_workers=10',
-        'autovacuum_vacuum_threshold=0',
-        'autovacuum_vacuum_insert_threshold=0',
-        'autovacuum_vacuum_cost_delay=0',
-        'autovacuum_vacuum_cost_limit=10000',
-        'autovacuum_naptime =1s',
-        'autovacuum_freeze_max_age=100000'
+        "autovacuum_max_workers=10",
+        "autovacuum_vacuum_threshold=0",
+        "autovacuum_vacuum_insert_threshold=0",
+        "autovacuum_vacuum_cost_delay=0",
+        "autovacuum_vacuum_cost_limit=10000",
+        "autovacuum_naptime =1s",
+        "autovacuum_freeze_max_age=100000",
     ]
 
-    pg = env.postgres.create_start('test_clog_truncate', config_lines=config)
-    log.info('postgres is running on test_clog_truncate branch')
+    pg = env.postgres.create_start("test_clog_truncate", config_lines=config)
+    log.info("postgres is running on test_clog_truncate branch")
 
     # Install extension containing function needed for test
-    pg.safe_psql('CREATE EXTENSION neon_test_utils')
+    pg.safe_psql("CREATE EXTENSION neon_test_utils")
 
     # Consume many xids to advance clog
     with pg.cursor() as cur:
-        cur.execute('select test_consume_xids(1000*1000*10);')
-        log.info('xids consumed')
+        cur.execute("select test_consume_xids(1000*1000*10);")
+        log.info("xids consumed")
 
         # call a checkpoint to trigger TruncateSubtrans
-        cur.execute('CHECKPOINT;')
+        cur.execute("CHECKPOINT;")
 
         # ensure WAL flush
-        cur.execute('select txid_current()')
+        cur.execute("select txid_current()")
         log.info(cur.fetchone())
 
     # wait for autovacuum to truncate the pg_xact
     # XXX Is it worth to add a timeout here?
-    pg_xact_0000_path = os.path.join(pg.pg_xact_dir_path(), '0000')
+    pg_xact_0000_path = os.path.join(pg.pg_xact_dir_path(), "0000")
     log.info(f"pg_xact_0000_path = {pg_xact_0000_path}")
 
     while os.path.isfile(pg_xact_0000_path):
-        log.info(f"file exists. wait for truncation. " "pg_xact_0000_path = {pg_xact_0000_path}")
+        log.info(f"file exists. wait for truncation: {pg_xact_0000_path=}")
         time.sleep(5)
 
     # checkpoint to advance latest lsn
     with pg.cursor() as cur:
-        cur.execute('CHECKPOINT;')
-        lsn_after_truncation = query_scalar(cur, 'select pg_current_wal_insert_lsn()')
+        cur.execute("CHECKPOINT;")
+        lsn_after_truncation = query_scalar(cur, "select pg_current_wal_insert_lsn()")
 
     # 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.neon_cli.create_branch('test_clog_truncate_new',
-                               'test_clog_truncate',
-                               ancestor_start_lsn=lsn_after_truncation)
-    pg2 = env.postgres.create_start('test_clog_truncate_new')
-    log.info('postgres is running on test_clog_truncate_new branch')
+    log.info(f"create branch at lsn_after_truncation {lsn_after_truncation}")
+    env.neon_cli.create_branch(
+        "test_clog_truncate_new", "test_clog_truncate", ancestor_start_lsn=lsn_after_truncation
+    )
+    pg2 = env.postgres.create_start("test_clog_truncate_new")
+    log.info("postgres is running on test_clog_truncate_new branch")
 
     # check that new node doesn't contain truncated segment
-    pg_xact_0000_path_new = os.path.join(pg2.pg_xact_dir_path(), '0000')
+    pg_xact_0000_path_new = os.path.join(pg2.pg_xact_dir_path(), "0000")
     log.info(f"pg_xact_0000_path_new = {pg_xact_0000_path_new}")
     assert os.path.isfile(pg_xact_0000_path_new) is False

test_runner/batch_others/test_close_fds.py

@@ -1,18 +1,18 @@
-from contextlib import closing
-import shutil
-import time
-import subprocess
 import os.path
+import shutil
+import subprocess
+import time
+from contextlib import closing
 
 from cached_property import threading
-from fixtures.neon_fixtures import NeonEnv
 from fixtures.log_helper import log
+from fixtures.neon_fixtures import NeonEnv
 
 
 def lsof_path() -> str:
     path_output = shutil.which("lsof")
     if path_output is None:
-        raise RuntimeError('lsof not found in PATH')
+        raise RuntimeError("lsof not found in PATH")
     else:
         return path_output
 
@@ -36,16 +36,18 @@ def test_lsof_pageserver_pid(neon_simple_env: NeonEnv):
     path = os.path.join(env.repo_dir, "pageserver.pid")
     lsof = lsof_path()
     while workload_thread.is_alive():
-        res = subprocess.run([lsof, path],
-                             check=False,
-                             universal_newlines=True,
-                             stdout=subprocess.PIPE,
-                             stderr=subprocess.PIPE)
+        res = subprocess.run(
+            [lsof, path],
+            check=False,
+            universal_newlines=True,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+        )
 
         # parse the `lsof` command's output to get only the list of commands
-        commands = [line.split(' ')[0] for line in res.stdout.strip().split('\n')[1:]]
+        commands = [line.split(" ")[0] for line in res.stdout.strip().split("\n")[1:]]
         if len(commands) > 0:
             log.info(f"lsof commands: {commands}")
-            assert commands == ['pageserve']
+            assert commands == ["pageserve"]
 
         time.sleep(1.0)

test_runner/batch_others/test_config.py

@@ -1,7 +1,7 @@
 from contextlib import closing
 
-from fixtures.neon_fixtures import NeonEnv
 from fixtures.log_helper import log
+from fixtures.neon_fixtures import NeonEnv
 
 
 #
@@ -12,19 +12,21 @@ def test_config(neon_simple_env: NeonEnv):
     env.neon_cli.create_branch("test_config", "empty")
 
     # change config
-    pg = env.postgres.create_start('test_config', config_lines=['log_min_messages=debug1'])
-    log.info('postgres is running on test_config branch')
+    pg = env.postgres.create_start("test_config", config_lines=["log_min_messages=debug1"])
+    log.info("postgres is running on test_config branch")
 
     with closing(pg.connect()) as conn:
         with conn.cursor() as cur:
-            cur.execute('''
+            cur.execute(
+                """
                 SELECT setting
                 FROM pg_settings
                 WHERE
                     source != 'default'
                     AND source != 'override'
                     AND name = 'log_min_messages'
-            ''')
+            """
+            )
 
             # check that config change was applied
-            assert cur.fetchone() == ('debug1', )
+            assert cur.fetchone() == ("debug1",)

test_runner/batch_others/test_crafted_wal_end.py

@@ -1,34 +1,38 @@
-from fixtures.neon_fixtures import NeonEnvBuilder, WalCraft
-from fixtures.log_helper import log
 import pytest
+from fixtures.log_helper import log
+from fixtures.neon_fixtures import NeonEnvBuilder, WalCraft
 
 # Restart nodes with WAL end having specially crafted shape, like last record
 # crossing segment boundary, to test decoding issues.
 
 
-@pytest.mark.parametrize('wal_type',
-                         [
-                             'simple',
-                             'last_wal_record_xlog_switch',
-                             'last_wal_record_xlog_switch_ends_on_page_boundary',
-                             'last_wal_record_crossing_segment',
-                             'wal_record_crossing_segment_followed_by_small_one',
-                         ])
+@pytest.mark.parametrize(
+    "wal_type",
+    [
+        "simple",
+        "last_wal_record_xlog_switch",
+        "last_wal_record_xlog_switch_ends_on_page_boundary",
+        "last_wal_record_crossing_segment",
+        "wal_record_crossing_segment_followed_by_small_one",
+    ],
+)
 def test_crafted_wal_end(neon_env_builder: NeonEnvBuilder, wal_type: str):
     neon_env_builder.num_safekeepers = 1
     env = neon_env_builder.init_start()
-    env.neon_cli.create_branch('test_crafted_wal_end')
+    env.neon_cli.create_branch("test_crafted_wal_end")
 
-    pg = env.postgres.create('test_crafted_wal_end')
+    pg = env.postgres.create("test_crafted_wal_end")
     wal_craft = WalCraft(env)
     pg.config(wal_craft.postgres_config())
     pg.start()
-    res = pg.safe_psql_many(queries=[
-        'CREATE TABLE keys(key int primary key)',
-        'INSERT INTO keys SELECT generate_series(1, 100)',
-        'SELECT SUM(key) FROM keys'
-    ])
-    assert res[-1][0] == (5050, )
+    res = pg.safe_psql_many(
+        queries=[
+            "CREATE TABLE keys(key int primary key)",
+            "INSERT INTO keys SELECT generate_series(1, 100)",
+            "SELECT SUM(key) FROM keys",
+        ]
+    )
+    assert res[-1][0] == (5050,)
 
     wal_craft.in_existing(wal_type, pg.connstr())
 
@@ -39,13 +43,15 @@ def test_crafted_wal_end(neon_env_builder: NeonEnvBuilder, wal_type: str):
     env.pageserver.start()
 
     log.info("Trying more queries")
-    res = pg.safe_psql_many(queries=[
-        'SELECT SUM(key) FROM keys',
-        'INSERT INTO keys SELECT generate_series(101, 200)',
-        'SELECT SUM(key) FROM keys',
-    ])
-    assert res[0][0] == (5050, )
-    assert res[-1][0] == (20100, )
+    res = pg.safe_psql_many(
+        queries=[
+            "SELECT SUM(key) FROM keys",
+            "INSERT INTO keys SELECT generate_series(101, 200)",
+            "SELECT SUM(key) FROM keys",
+        ]
+    )
+    assert res[0][0] == (5050,)
+    assert res[-1][0] == (20100,)
 
     log.info("Restarting all safekeepers and pageservers (again)")
     env.pageserver.stop()
@@ -54,10 +60,12 @@ def test_crafted_wal_end(neon_env_builder: NeonEnvBuilder, wal_type: str):
     env.pageserver.start()
 
     log.info("Trying more queries (again)")
-    res = pg.safe_psql_many(queries=[
-        'SELECT SUM(key) FROM keys',
-        'INSERT INTO keys SELECT generate_series(201, 300)',
-        'SELECT SUM(key) FROM keys',
-    ])
-    assert res[0][0] == (20100, )
-    assert res[-1][0] == (45150, )
+    res = pg.safe_psql_many(
+        queries=[
+            "SELECT SUM(key) FROM keys",
+            "INSERT INTO keys SELECT generate_series(201, 300)",
+            "SELECT SUM(key) FROM keys",
+        ]
+    )
+    assert res[0][0] == (20100,)
+    assert res[-1][0] == (45150,)