Separate Postgres build dir from installation dir.

Previously, Postgres was built in 'tmp_install/build', and installed
into 'tmp_install'. In the CI, the 'build' directory was included in
the final neon.tar.zst artifact that includes all the necessary
binaries. That was unnecessary, the intermediate build results are not
needed, only the final binaries. Separate the build directory so that
the Postgres build happens in 'build', and it is installed into
'tmp_install'. That makes the final neon.tar.zst artifact smaller.

The changes to the python tests are needed to find the 'pg_regress'
binary in the installation directory. Previously, they would use the
'pg_regress' binary from the build directory, not the final
installation location.

.circleci/config.yml

@@ -0,0 +1,369 @@
+version: 2.1
+
+executors:
+  neon-xlarge-executor:
+    resource_class: xlarge
+    docker:
+      # NB: when changed, do not forget to update rust image tag in all Dockerfiles
+      - image: neondatabase/rust:1.58
+  neon-executor:
+    docker:
+      - image: neondatabase/rust:1.58
+
+jobs:
+  # A job to build postgres
+  build-postgres:
+    executor: neon-xlarge-executor
+    parameters:
+      build_type:
+        type: enum
+        enum: ["debug", "release"]
+    environment:
+      BUILD_TYPE: << parameters.build_type >>
+    steps:
+        # Checkout the git repo (circleci doesn't have a flag to enable submodules here)
+      - checkout
+
+        # Grab the postgres git revision to build a cache key.
+        # Append makefile as it could change the way postgres is built.
+        # Note this works even though the submodule hasn't been checkout out yet.
+      - run:
+          name: Get postgres cache key
+          command: |
+              git rev-parse HEAD:vendor/postgres > /tmp/cache-key-postgres
+              cat Makefile >> /tmp/cache-key-postgres
+
+      - restore_cache:
+          name: Restore postgres cache
+          keys:
+            # Restore ONLY if the rev key matches exactly
+            - v05-postgres-cache-<< parameters.build_type >>-{{ checksum "/tmp/cache-key-postgres" }}
+
+        # Build postgres if the restore_cache didn't find a build.
+        # `make` can't figure out whether the cache is valid, since
+        # it only compares file timestamps.
+      - run:
+          name: build postgres
+          command: |
+            if [ ! -e tmp_install/bin/postgres ]; then
+              # "depth 1" saves some time by not cloning the whole repo
+              git submodule update --init --depth 1
+              # bail out on any warnings
+              COPT='-Werror' mold -run make postgres -j$(nproc)
+            fi
+
+      - save_cache:
+          name: Save postgres cache
+          key: v05-postgres-cache-<< parameters.build_type >>-{{ checksum "/tmp/cache-key-postgres" }}
+          paths:
+            - tmp_install
+
+  # A job to build Neon rust code
+  build-neon:
+    executor: neon-xlarge-executor
+    parameters:
+      build_type:
+        type: enum
+        enum: ["debug", "release"]
+    environment:
+      BUILD_TYPE: << parameters.build_type >>
+    steps:
+        # Checkout the git repo (without submodules)
+      - checkout
+
+        # Grab the postgres git revision to build a cache key.
+        # Append makefile as it could change the way postgres is built.
+        # Note this works even though the submodule hasn't been checkout out yet.
+      - run:
+          name: Get postgres cache key
+          command: |
+            git rev-parse HEAD:vendor/postgres > /tmp/cache-key-postgres
+            cat Makefile >> /tmp/cache-key-postgres
+
+
+      - restore_cache:
+          name: Restore postgres cache
+          keys:
+            # Restore ONLY if the rev key matches exactly
+            - v05-postgres-cache-<< parameters.build_type >>-{{ checksum "/tmp/cache-key-postgres" }}
+
+      - restore_cache:
+          name: Restore rust cache
+          keys:
+            # Require an exact match. While an out of date cache might speed up the build,
+            # there's no way to clean out old packages, so the cache grows every time something
+            # changes.
+            - v05-rust-cache-deps-<< parameters.build_type >>-{{ checksum "Cargo.lock" }}
+
+        # Build the rust code, including test binaries
+      - run:
+          name: Rust build << parameters.build_type >>
+          command: |
+            if [[ $BUILD_TYPE == "debug" ]]; then
+              CARGO_FLAGS=
+            elif [[ $BUILD_TYPE == "release" ]]; then
+              CARGO_FLAGS="--release --features profiling"
+            fi
+
+            export CARGO_INCREMENTAL=0
+            export CACHEPOT_BUCKET=zenith-rust-cachepot
+            export RUSTC_WRAPPER=""
+            export AWS_ACCESS_KEY_ID="${CACHEPOT_AWS_ACCESS_KEY_ID}"
+            export AWS_SECRET_ACCESS_KEY="${CACHEPOT_AWS_SECRET_ACCESS_KEY}"
+            mold -run cargo build $CARGO_FLAGS --features failpoints --bins --tests
+            cachepot -s
+
+      - save_cache:
+          name: Save rust cache
+          key: v05-rust-cache-deps-<< parameters.build_type >>-{{ checksum "Cargo.lock" }}
+          paths:
+            - ~/.cargo/registry
+            - ~/.cargo/git
+            - target
+
+        # Run rust unit tests
+      - run:
+          name: cargo test
+          command: |
+            if [[ $BUILD_TYPE == "debug" ]]; then
+              CARGO_FLAGS=
+            elif [[ $BUILD_TYPE == "release" ]]; then
+              CARGO_FLAGS=--release
+            fi
+
+            cargo test $CARGO_FLAGS
+
+        # Install the rust binaries, for use by test jobs
+      - run:
+          name: Install rust binaries
+          command: |
+            binaries=$(
+              cargo metadata --format-version=1 --no-deps |
+              jq -r '.packages[].targets[] | select(.kind | index("bin")) | .name'
+            )
+
+            mkdir -p /tmp/zenith/bin
+            mkdir -p /tmp/zenith/test_bin
+            mkdir -p /tmp/zenith/etc
+
+            # Install target binaries
+            for bin in $binaries; do
+              SRC=target/$BUILD_TYPE/$bin
+              DST=/tmp/zenith/bin/$bin
+              cp $SRC $DST
+            done
+
+        # Install the postgres binaries, for use by test jobs
+      - run:
+          name: Install postgres binaries
+          command: |
+            cp -a tmp_install /tmp/zenith/pg_install
+
+      # Save rust binaries for other jobs in the workflow
+      - persist_to_workspace:
+          root: /tmp/zenith
+          paths:
+            - "*"
+
+  check-codestyle-python:
+    executor: neon-executor
+    steps:
+      - checkout
+      - restore_cache:
+          keys:
+            - v2-python-deps-{{ checksum "poetry.lock" }}
+      - run:
+          name: Install deps
+          command: ./scripts/pysync
+      - save_cache:
+          key: v2-python-deps-{{ checksum "poetry.lock" }}
+          paths:
+            - /home/circleci/.cache/pypoetry/virtualenvs
+      - run:
+          name: Print versions
+          when: always
+          command: |
+              poetry run python --version
+              poetry show
+      - run:
+          name: Run yapf to ensure code format
+          when: always
+          command: poetry run yapf --recursive --diff .
+      - run:
+          name: Run mypy to check types
+          when: always
+          command: poetry run mypy .
+
+  run-pytest:
+    executor: neon-executor
+    parameters:
+      # pytest args to specify the tests to run.
+      #
+      # This can be a test file name, e.g. 'test_pgbench.py, or a subdirectory,
+      # or '-k foobar' to run tests containing string 'foobar'. See pytest man page
+      # section SPECIFYING TESTS / SELECTING TESTS for details.
+      #
+      # Select the type of Rust build. Must be "release" or "debug".
+      build_type:
+        type: string
+        default: "debug"
+      # This parameter is required, to prevent the mistake of running all tests in one job.
+      test_selection:
+        type: string
+        default: ""
+      # Arbitrary parameters to pytest. For example "-s" to prevent capturing stdout/stderr
+      extra_params:
+        type: string
+        default: ""
+      needs_postgres_source:
+        type: boolean
+        default: false
+      run_in_parallel:
+        type: boolean
+        default: true
+      save_perf_report:
+        type: boolean
+        default: false
+    environment:
+      BUILD_TYPE: << parameters.build_type >>
+    steps:
+      - attach_workspace:
+          at: /tmp/zenith
+      - checkout
+      - when:
+          condition: << parameters.needs_postgres_source >>
+          steps:
+            - run: git submodule update --init --depth 1
+      - restore_cache:
+          keys:
+            - v2-python-deps-{{ checksum "poetry.lock" }}
+      - run:
+          name: Install deps
+          command: ./scripts/pysync
+      - save_cache:
+          key: v2-python-deps-{{ checksum "poetry.lock" }}
+          paths:
+            - /home/circleci/.cache/pypoetry/virtualenvs
+      - run:
+          name: Run pytest
+          # pytest doesn't output test logs in real time, so CI job may fail with
+          # `Too long with no output` error, if a test is running for a long time.
+          # In that case, tests should have internal timeouts that are less than
+          # no_output_timeout, specified here.
+          no_output_timeout: 10m
+          environment:
+            - NEON_BIN: /tmp/zenith/bin
+            - POSTGRES_DISTRIB_DIR: /tmp/zenith/pg_install
+            - TEST_OUTPUT: /tmp/test_output
+            # this variable will be embedded in perf test report
+            # and is needed to distinguish different environments
+            - PLATFORM: zenith-local-ci
+          command: |
+            PERF_REPORT_DIR="$(realpath test_runner/perf-report-local)"
+            rm -rf $PERF_REPORT_DIR
+
+            TEST_SELECTION="test_runner/<< parameters.test_selection >>"
+            EXTRA_PARAMS="<< parameters.extra_params >>"
+            if [ -z "$TEST_SELECTION" ]; then
+              echo "test_selection must be set"
+              exit 1
+            fi
+            if << parameters.run_in_parallel >>; then
+              EXTRA_PARAMS="-n4 $EXTRA_PARAMS"
+            fi
+            if << parameters.save_perf_report >>; then
+              if [[ $CIRCLE_BRANCH == "main" ]]; then
+                mkdir -p "$PERF_REPORT_DIR"
+                EXTRA_PARAMS="--out-dir $PERF_REPORT_DIR $EXTRA_PARAMS"
+              fi
+            fi
+
+            export GITHUB_SHA=$CIRCLE_SHA1
+
+            # Run the tests.
+            #
+            # The junit.xml file allows CircleCI to display more fine-grained test information
+            # in its "Tests" tab in the results page.
+            # --verbose prints name of each test (helpful when there are
+            # multiple tests in one file)
+            # -rA prints summary in the end
+            # -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
+            ./scripts/pytest \
+              --junitxml=$TEST_OUTPUT/junit.xml \
+              --tb=short \
+              --verbose \
+              -m "not remote_cluster" \
+              -rA $TEST_SELECTION $EXTRA_PARAMS
+
+            if << parameters.save_perf_report >>; then
+              if [[ $CIRCLE_BRANCH == "main" ]]; then
+                export REPORT_FROM="$PERF_REPORT_DIR"
+                export REPORT_TO=local
+                scripts/generate_and_push_perf_report.sh
+              fi
+            fi
+      - run:
+          # CircleCI artifacts are preserved one file at a time, so skipping
+          # this step isn't a good idea. If you want to extract the
+          # pageserver state, perhaps a tarball would be a better idea.
+          name: Delete all data but logs
+          when: always
+          command: |
+            du -sh /tmp/test_output/*
+            find /tmp/test_output -type f ! -name "*.log" ! -name "regression.diffs" ! -name "junit.xml" ! -name "*.filediff" ! -name "*.stdout" ! -name "*.stderr" ! -name "flamegraph.svg" ! -name "*.metrics" -delete
+            du -sh /tmp/test_output/*
+      - store_artifacts:
+          path: /tmp/test_output
+      # The store_test_results step tells CircleCI where to find the junit.xml file.
+      - store_test_results:
+          path: /tmp/test_output
+      # Save data (if any)
+      - persist_to_workspace:
+          root: /tmp/zenith
+          paths:
+            - "*"
+
+workflows:
+  build_and_test:
+    jobs:
+      - check-codestyle-python
+      - build-postgres:
+          name: build-postgres-<< matrix.build_type >>
+          matrix:
+            parameters:
+              build_type: ["debug", "release"]
+      - build-neon:
+          name: build-neon-<< matrix.build_type >>
+          matrix:
+            parameters:
+              build_type: ["debug", "release"]
+          requires:
+            - build-postgres-<< matrix.build_type >>
+      - run-pytest:
+          name: pg_regress-tests-<< matrix.build_type >>
+          matrix:
+            parameters:
+              build_type: ["debug", "release"]
+          test_selection: batch_pg_regress
+          needs_postgres_source: true
+          requires:
+            - build-neon-<< matrix.build_type >>
+      - run-pytest:
+          name: other-tests-<< matrix.build_type >>
+          matrix:
+            parameters:
+              build_type: ["debug", "release"]
+          test_selection: batch_others
+          requires:
+            - build-neon-<< matrix.build_type >>
+      - run-pytest:
+          name: benchmarks
+          context: PERF_TEST_RESULT_CONNSTR
+          build_type: release
+          test_selection: performance
+          run_in_parallel: false
+          save_perf_report: true
+          requires:
+            - build-neon-release

.github/actions/download/action.yml

@@ -1,56 +0,0 @@
-name: "Download an artifact"
-description: "Custom download action"
-inputs:
-  name:
-    description: "Artifact name"
-    required: true
-  path:
-    description: "A directory to put artifact into"
-    default: "."
-    required: false
-  skip-if-does-not-exist:
-    description: "Allow to skip if file doesn't exist, fail otherwise"
-    default: false
-    required: false
-
-runs:
-  using: "composite"
-  steps:
-    - name: Download artifact
-      id: download-artifact
-      shell: bash -euxo pipefail {0}
-      env:
-        TARGET: ${{ inputs.path }}
-        ARCHIVE: /tmp/downloads/${{ inputs.name }}.tar.zst
-        SKIP_IF_DOES_NOT_EXIST: ${{ inputs.skip-if-does-not-exist }}
-      run: |
-        BUCKET=neon-github-public-dev
-        PREFIX=artifacts/${GITHUB_RUN_ID}
-        FILENAME=$(basename $ARCHIVE)
-
-        S3_KEY=$(aws s3api list-objects-v2 --bucket ${BUCKET} --prefix ${PREFIX} | jq -r '.Contents[].Key' | grep ${FILENAME} | sort --version-sort | tail -1 || true)
-        if [ -z "${S3_KEY}" ]; then
-          if [ "${SKIP_IF_DOES_NOT_EXIST}" = "true" ]; then
-            echo '::set-output name=SKIPPED::true'
-            exit 0
-          else
-            echo 2>&1 "Neither s3://${BUCKET}/${PREFIX}/${GITHUB_RUN_ATTEMPT}/${FILENAME} nor its version from previous attempts exist"
-            exit 1
-          fi
-        fi
-
-        echo '::set-output name=SKIPPED::false'
-
-        mkdir -p $(dirname $ARCHIVE)
-        time aws s3 cp --only-show-errors s3://${BUCKET}/${S3_KEY} ${ARCHIVE}
-
-    - name: Extract artifact
-      if: ${{ steps.download-artifact.outputs.SKIPPED == 'false' }}
-      shell: bash -euxo pipefail {0}
-      env:
-        TARGET: ${{ inputs.path }}
-        ARCHIVE: /tmp/downloads/${{ inputs.name }}.tar.zst
-      run: |
-        mkdir -p ${TARGET}
-        time tar -xf ${ARCHIVE} -C ${TARGET}
-        rm -f ${ARCHIVE}

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

@@ -27,36 +27,10 @@ inputs:
     description: 'Whether to upload the performance report'
     required: false
     default: 'false'
-  run_with_real_s3:
-    description: 'Whether to pass real s3 credentials to the test suite'
-    required: false
-    default: 'false'
-  real_s3_bucket:
-    description: 'Bucket name for real s3 tests'
-    required: false
-    default: ''
-  real_s3_region:
-    description: 'Region name for real s3 tests'
-    required: false
-    default: ''
-  real_s3_access_key_id:
-    description: 'Access key id'
-    required: false
-    default: ''
-  real_s3_secret_access_key:
-    description: 'Secret access key'
-    required: false
-    default: ''
 
 runs:
   using: "composite"
   steps:
-    - name: Get Neon artifact
-      uses: ./.github/actions/download
-      with:
-        name: neon-${{ runner.os }}-${{ inputs.build_type }}-${{ inputs.rust_toolchain }}-artifact
-        path: /tmp/neon
-
     - name: Checkout
       if: inputs.needs_postgres_source == 'true'
       uses: actions/checkout@v3
@@ -64,6 +38,25 @@ runs:
         submodules: true
         fetch-depth: 1
 
+    - name: Get Neon artifact for restoration
+      uses: actions/download-artifact@v3
+      with:
+        name: neon-${{ runner.os }}-${{ inputs.build_type }}-${{ inputs.rust_toolchain }}-artifact
+        path: ./neon-artifact/
+
+    - name: Extract Neon artifact
+      shell: bash -ex {0}
+      run: |
+        mkdir -p /tmp/neon/
+        tar -xf ./neon-artifact/neon.tar.zst -C /tmp/neon/
+        rm -rf ./neon-artifact/
+
+        # Restore the parts of the 'build' directory that were included in the
+        # tarball. This includes the regression test modules in
+        # src/test/regress/*.so.
+        mkdir -p build/
+        cp -a /tmp/neon/pg_build/* build/
+
     - name: Cache poetry deps
       id: cache_poetry
       uses: actions/cache@v3
@@ -72,7 +65,7 @@ runs:
         key: v1-${{ runner.os }}-python-deps-${{ hashFiles('poetry.lock') }}
 
     - name: Install Python deps
-      shell: bash -euxo pipefail {0}
+      shell: bash -ex {0}
       run: ./scripts/pysync
 
     - name: Run pytest
@@ -83,10 +76,7 @@ runs:
         # this variable will be embedded in perf test report
         # and is needed to distinguish different environments
         PLATFORM: github-actions-selfhosted
-        BUILD_TYPE: ${{ inputs.build_type }}
-        AWS_ACCESS_KEY_ID: ${{ inputs.real_s3_access_key_id }}
-        AWS_SECRET_ACCESS_KEY: ${{ inputs.real_s3_secret_access_key }}
-      shell: bash -euxo pipefail {0}
+      shell: bash -ex {0}
       run: |
         PERF_REPORT_DIR="$(realpath test_runner/perf-report-local)"
         rm -rf $PERF_REPORT_DIR
@@ -100,14 +90,6 @@ runs:
         if [[ "${{ inputs.run_in_parallel }}" == "true" ]]; then
           EXTRA_PARAMS="-n4 $EXTRA_PARAMS"
         fi
-
-        if [[ "${{ inputs.run_with_real_s3 }}" == "true" ]]; then
-          echo "REAL S3 ENABLED"
-          export ENABLE_REAL_S3_REMOTE_STORAGE=nonempty
-          export REMOTE_STORAGE_S3_BUCKET=${{ inputs.real_s3_bucket }}
-          export REMOTE_STORAGE_S3_REGION=${{ inputs.real_s3_region }}
-        fi
-
         if [[ "${{ inputs.save_perf_report }}" == "true" ]]; then
           if [[ "$GITHUB_REF" == "refs/heads/main" ]]; then
             mkdir -p "$PERF_REPORT_DIR"
@@ -123,7 +105,7 @@ runs:
 
         # Run the tests.
         #
-        # The junit.xml file allows CI tools to display more fine-grained test information
+        # The junit.xml file allows CircleCI to display more fine-grained test information
         # in its "Tests" tab in the results page.
         # --verbose prints name of each test (helpful when there are
         # multiple tests in one file)
@@ -147,7 +129,7 @@ runs:
         fi
 
     - name: Delete all data but logs
-      shell: bash -euxo pipefail {0}
+      shell: bash -ex {0}
       if: always()
       run: |
         du -sh /tmp/test_output/*
@@ -156,7 +138,9 @@ runs:
 
     - name: Upload python test logs
       if: always()
-      uses: ./.github/actions/upload
+      uses: actions/upload-artifact@v3
       with:
+        retention-days: 7
+        if-no-files-found: error
         name: python-test-${{ inputs.test_selection }}-${{ runner.os }}-${{ inputs.build_type }}-${{ inputs.rust_toolchain }}-logs
         path: /tmp/test_output/

.github/actions/save-coverage-data/action.yml

@@ -5,18 +5,13 @@ runs:
   using: "composite"
   steps:
     - name: Merge coverage data
-      shell: bash -euxo pipefail {0}
+      shell: bash -ex {0}
       run: scripts/coverage "--profraw-prefix=$GITHUB_JOB" --dir=/tmp/coverage merge
 
-    - name: Download previous coverage data into the same directory
-      uses: ./.github/actions/download
-      with:
-        name: coverage-data-artifact
-        path: /tmp/coverage
-        skip-if-does-not-exist: true # skip if there's no previous coverage to download
-
     - name: Upload coverage data
-      uses: ./.github/actions/upload
+      uses: actions/upload-artifact@v3
       with:
+        retention-days: 7
+        if-no-files-found: error
         name: coverage-data-artifact
-        path: /tmp/coverage
+        path: /tmp/coverage/

.github/actions/upload/action.yml

@@ -1,55 +0,0 @@
-name: "Upload an artifact"
-description: "Custom upload action"
-inputs:
-  name:
-    description: "Artifact name"
-    required: true
-  path:
-    description: "A directory or file to upload"
-    required: true
-
-runs:
-  using: "composite"
-  steps:
-    - name: Prepare artifact
-      shell: bash -euxo pipefail {0}
-      env:
-        SOURCE: ${{ inputs.path }}
-        ARCHIVE: /tmp/uploads/${{ inputs.name }}.tar.zst
-      run: |
-        mkdir -p $(dirname $ARCHIVE)
-
-        if [ -f ${ARCHIVE} ]; then
-          echo 2>&1 "File ${ARCHIVE} already exist. Something went wrong before"
-          exit 1
-        fi
-
-        ZSTD_NBTHREADS=0
-        if [ -d  ${SOURCE} ]; then
-          time tar -C ${SOURCE} -cf ${ARCHIVE} --zstd .
-        elif [ -f ${SOURCE} ]; then
-          time tar -cf ${ARCHIVE} --zstd ${SOURCE}
-        elif ! ls ${SOURCE} > /dev/null 2>&1; then
-          echo 2>&1 "${SOURCE} does not exist"
-          exit 2
-        else
-          echo 2>&1 "${SOURCE} is neither a directory nor a file, do not know how to handle it"
-          exit 3
-        fi
-
-    - name: Upload artifact
-      shell: bash -euxo pipefail {0}
-      env:
-        SOURCE: ${{ inputs.path }}
-        ARCHIVE: /tmp/uploads/${{ inputs.name }}.tar.zst
-      run: |
-        BUCKET=neon-github-public-dev
-        PREFIX=artifacts/${GITHUB_RUN_ID}
-        FILENAME=$(basename $ARCHIVE)
-
-        FILESIZE=$(du -sh ${ARCHIVE} | cut -f1)
-
-        time aws s3 mv --only-show-errors ${ARCHIVE} s3://${BUCKET}/${PREFIX}/${GITHUB_RUN_ATTEMPT}/${FILENAME}
-
-        # Ref https://docs.github.com/en/actions/using-workflows/workflow-commands-for-github-actions#adding-a-job-summary
-        echo "[${FILENAME}](https://${BUCKET}.s3.amazonaws.com/${PREFIX}/${GITHUB_RUN_ATTEMPT}/${FILENAME}) ${FILESIZE}" >> ${GITHUB_STEP_SUMMARY}

.github/ansible/get_binaries.sh

@@ -2,14 +2,30 @@
 
 set -e
 
-if [ -n "${DOCKER_TAG}" ]; then
-  # Verson is DOCKER_TAG but without prefix
-  VERSION=$(echo $DOCKER_TAG | sed 's/^.*-//g')
+RELEASE=${RELEASE:-false}
+
+# look at docker hub for latest tag for neon docker image
+if [ "${RELEASE}" = "true" ]; then
+    echo "search latest release tag"
+    VERSION=$(curl -s https://registry.hub.docker.com/v1/repositories/neondatabase/neon/tags |jq -r -S '.[].name' | grep release | sed 's/release-//g' | grep -E '^[0-9]+$' | sort -n | tail -1)
+    if [ -z "${VERSION}" ]; then
+        echo "no any docker tags found, exiting..."
+        exit 1
+    else
+        TAG="release-${VERSION}"
+    fi
 else
-  echo "Please set DOCKER_TAG environment variable"
-  exit 1
+    echo "search latest dev tag"
+    VERSION=$(curl -s https://registry.hub.docker.com/v1/repositories/neondatabase/neon/tags |jq -r -S '.[].name' | grep -E '^[0-9]+$' | sort -n | tail -1)
+    if [ -z "${VERSION}" ]; then
+        echo "no any docker tags found, exiting..."
+        exit 1
+    else
+        TAG="${VERSION}"
+    fi
 fi
 
+echo "found ${VERSION}"
 
 # do initial cleanup
 rm -rf neon_install postgres_install.tar.gz neon_install.tar.gz .neon_current_version
@@ -17,8 +33,8 @@ mkdir neon_install
 
 # retrieve binaries from docker image
 echo "getting binaries from docker image"
-docker pull --quiet neondatabase/neon:${DOCKER_TAG}
-ID=$(docker create neondatabase/neon:${DOCKER_TAG})
+docker pull --quiet neondatabase/neon:${TAG}
+ID=$(docker create neondatabase/neon:${TAG})
 docker cp ${ID}:/data/postgres_install.tar.gz .
 tar -xzf postgres_install.tar.gz -C neon_install
 docker cp ${ID}:/usr/local/bin/pageserver neon_install/bin/

.github/ansible/production.hosts

@@ -17,4 +17,4 @@ env_name = prod-1
 console_mgmt_base_url = http://console-release.local
 bucket_name           = zenith-storage-oregon
 bucket_region         = us-west-2
-etcd_endpoints        = zenith-1-etcd.local:2379
+etcd_endpoints        = etcd-release.local:2379

.github/ansible/scripts/init_safekeeper.sh

@@ -12,9 +12,10 @@ cat <<EOF | tee /tmp/payload
   "version": 1,
   "host": "${HOST}",
   "port": 6500,
-  "http_port": 7676,
   "region_id": {{ console_region_id }},
-  "instance_id": "${INSTANCE_ID}"
+  "instance_id": "${INSTANCE_ID}",
+  "http_host": "${HOST}",
+  "http_port": 7676
 }
 EOF
 

.github/ansible/staging.hosts

@@ -17,4 +17,4 @@ env_name = us-stage
 console_mgmt_base_url = http://console-staging.local
 bucket_name           = zenith-staging-storage-us-east-1
 bucket_region         = us-east-1
-etcd_endpoints        = zenith-us-stage-etcd.local:2379
+etcd_endpoints        = etcd-staging.local:2379

.github/workflows/benchmarking.yml

@@ -1,4 +1,4 @@
-name: Benchmarking
+name: benchmarking
 
 on:
   # uncomment to run on push for debugging your PR
@@ -15,15 +15,6 @@ on:
 
   workflow_dispatch: # adds ability to run this manually
 
-defaults:
-  run:
-    shell: bash -euxo pipefail {0}
-
-concurrency:
-  # Allow only one workflow per any non-`main` branch.
-  group: ${{ github.workflow }}-${{ github.ref }}-${{ github.ref == 'refs/heads/main' && github.sha || 'anysha' }}
-  cancel-in-progress: true
-
 jobs:
   bench:
     # this workflow runs on self hosteed runner
@@ -69,6 +60,7 @@ jobs:
     - name: Setup cluster
       env:
         BENCHMARK_CONNSTR: "${{ secrets.BENCHMARK_STAGING_CONNSTR }}"
+      shell: bash
       run: |
         set -e
 
@@ -104,9 +96,7 @@ jobs:
         # since it might generate duplicates when calling ingest_perf_test_result.py
         rm -rf perf-report-staging
         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 3600
+        ./scripts/pytest test_runner/performance/ -v -m "remote_cluster" --skip-interfering-proc-check --out-dir perf-report-staging --timeout 3600
 
     - name: Submit result
       env:
@@ -123,106 +113,3 @@ jobs:
         slack-message: "Periodic perf testing: ${{ job.status }}\n${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}"
       env:
         SLACK_BOT_TOKEN: ${{ secrets.SLACK_BOT_TOKEN }}
-
-  pgbench-compare:
-    env:
-      TEST_PG_BENCH_DURATIONS_MATRIX: "60m"
-      TEST_PG_BENCH_SCALES_MATRIX: "10gb"
-      REMOTE_ENV: "1"
-      POSTGRES_DISTRIB_DIR: /usr
-      TEST_OUTPUT: /tmp/test_output
-
-    strategy:
-      fail-fast: false
-      matrix:
-        connstr: [ BENCHMARK_CAPTEST_CONNSTR, BENCHMARK_RDS_CONNSTR ]
-
-    runs-on: dev
-    container: 369495373322.dkr.ecr.eu-central-1.amazonaws.com/rustlegacy:2817580636
-
-    timeout-minutes: 360 # 6h
-
-    steps:
-    - uses: actions/checkout@v3
-
-    - name: Cache poetry deps
-      id: cache_poetry
-      uses: actions/cache@v3
-      with:
-        path: ~/.cache/pypoetry/virtualenvs
-        key: v2-${{ runner.os }}-python-deps-${{ hashFiles('poetry.lock') }}
-
-    - name: Install Python deps
-      run: ./scripts/pysync
-
-    - name: Calculate platform
-      id: calculate-platform
-      env:
-        CONNSTR: ${{ matrix.connstr }}
-      run: |
-        if [ "${CONNSTR}" = "BENCHMARK_CAPTEST_CONNSTR" ]; then
-          PLATFORM=neon-captest
-        elif [ "${CONNSTR}" = "BENCHMARK_RDS_CONNSTR" ]; then
-          PLATFORM=rds-aurora
-        else
-          echo 2>&1 "Unknown CONNSTR=${CONNSTR}. Allowed are BENCHMARK_CAPTEST_CONNSTR, and BENCHMARK_RDS_CONNSTR only"
-          exit 1
-        fi
-
-        echo "::set-output name=PLATFORM::${PLATFORM}"
-
-    - name: Install Deps
-      run: |
-        echo "deb http://apt.postgresql.org/pub/repos/apt focal-pgdg main" | sudo tee /etc/apt/sources.list.d/pgdg.list
-        wget --quiet -O - https://www.postgresql.org/media/keys/ACCC4CF8.asc | sudo apt-key add -
-        sudo apt -y update
-        sudo apt install -y postgresql-14 postgresql-client-14
-
-    - name: Benchmark init
-      env:
-        PLATFORM: ${{ steps.calculate-platform.outputs.PLATFORM }}
-        BENCHMARK_CONNSTR: ${{ secrets[matrix.connstr] }}
-      run: |
-        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
-
-    - name: Benchmark simple-update
-      env:
-        PLATFORM: ${{ steps.calculate-platform.outputs.PLATFORM }}
-        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
-
-    - name: Benchmark select-only
-      env:
-        PLATFORM: ${{ steps.calculate-platform.outputs.PLATFORM }}
-        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
-
-    - name: Submit result
-      env:
-        VIP_VAP_ACCESS_TOKEN: "${{ secrets.VIP_VAP_ACCESS_TOKEN }}"
-        PERF_TEST_RESULT_CONNSTR: "${{ secrets.PERF_TEST_RESULT_CONNSTR }}"
-      run: |
-        REPORT_FROM=$(realpath perf-report-captest) REPORT_TO=staging scripts/generate_and_push_perf_report.sh
-
-    - name: Upload logs
-      if: always()
-      uses: ./.github/actions/upload
-      with:
-        name: bench-captest-${{ steps.calculate-platform.outputs.PLATFORM }}
-        path: /tmp/test_output/
-
-    - name: Post to a Slack channel
-      if: ${{ github.event.schedule && failure() }}
-      uses: slackapi/slack-github-action@v1
-      with:
-        channel-id: "C033QLM5P7D" # dev-staging-stream
-        slack-message: "Periodic perf testing: ${{ job.status }}\n${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}"
-      env:
-        SLACK_BOT_TOKEN: ${{ secrets.SLACK_BOT_TOKEN }}

.github/workflows/build_and_test.yml

@@ -3,13 +3,13 @@ name: Test and Deploy
 on:
   push:
     branches:
-      - main
-      - release
+    - main
+    - release
   pull_request:
 
 defaults:
   run:
-    shell: bash -euxo pipefail {0}
+    shell: bash -ex {0}
 
 concurrency:
   # Allow only one workflow per any non-`main` branch.
@@ -22,8 +22,7 @@ env:
 
 jobs:
   build-neon:
-    runs-on: dev
-    container: 369495373322.dkr.ecr.eu-central-1.amazonaws.com/rustlegacy:2746987948
+    runs-on: [ self-hosted, Linux, k8s-runner ]
     strategy:
       fail-fast: false
       matrix:
@@ -32,19 +31,8 @@ jobs:
 
     env:
       BUILD_TYPE: ${{ matrix.build_type }}
-      GIT_VERSION: ${{ github.sha }}
 
     steps:
-      - name: Fix git ownerwhip
-        run: |
-          # Workaround for `fatal: detected dubious ownership in repository at ...`
-          #
-          # Use both ${{ github.workspace }} and ${GITHUB_WORKSPACE} because they're different on host and in containers
-          #   Ref https://github.com/actions/checkout/issues/785
-          #
-          git config --global --add safe.directory ${{ github.workspace }}
-          git config --global --add safe.directory ${GITHUB_WORKSPACE}
-
       - name: Checkout
         uses: actions/checkout@v3
         with:
@@ -100,7 +88,9 @@ jobs:
         id: cache_pg
         uses: actions/cache@v3
         with:
-          path: tmp_install/
+          path: |
+            tmp_install/
+            build/src/test/regress/*.so
           key: v1-${{ runner.os }}-${{ matrix.build_type }}-pg-${{ steps.pg_ver.outputs.pg_rev }}-${{ hashFiles('Makefile') }}
 
       - name: Build postgres
@@ -135,7 +125,6 @@ jobs:
             mkdir -p /tmp/coverage/
 
             mkdir -p /tmp/neon/test_bin/
-
             test_exe_paths=$(
               ${cov_prefix} cargo test $CARGO_FLAGS --message-format=json --no-run |
               jq -r '.executable | select(. != null)'
@@ -156,22 +145,32 @@ jobs:
           fi
 
       - name: Install postgres binaries
-        run: cp -a tmp_install /tmp/neon/pg_install
+        run: |
+          cp -a tmp_install /tmp/neon/pg_install
 
-      - name: Upload Neon artifact
-        uses: ./.github/actions/upload
+          # Include modules needed by the Postgres regression tests
+          mkdir -p /tmp/neon/pg_build/src/test/regress
+          cp -a build/src/test/regress/*.so /tmp/neon/pg_build/src/test/regress
+
+      - name: Prepare neon artifact
+        run: ZSTD_NBTHREADS=0 tar -C /tmp/neon/ -cf ./neon.tar.zst --zstd .
+
+      - name: Upload neon binaries
+        uses: actions/upload-artifact@v3
         with:
+          retention-days: 7
+          if-no-files-found: error
           name: neon-${{ runner.os }}-${{ matrix.build_type }}-${{ matrix.rust_toolchain }}-artifact
-          path: /tmp/neon
+          path: ./neon.tar.zst
 
       # XXX: keep this after the binaries.list is formed, so the coverage can properly work later
       - name: Merge and upload coverage data
         if: matrix.build_type == 'debug'
         uses: ./.github/actions/save-coverage-data
 
+
   pg_regress-tests:
-    runs-on: dev
-    container: 369495373322.dkr.ecr.eu-central-1.amazonaws.com/rustlegacy:2746987948
+    runs-on: [ self-hosted, Linux, k8s-runner ]
     needs: [ build-neon ]
     strategy:
       fail-fast: false
@@ -198,8 +197,7 @@ jobs:
         uses: ./.github/actions/save-coverage-data
 
   other-tests:
-    runs-on: dev
-    container: 369495373322.dkr.ecr.eu-central-1.amazonaws.com/rustlegacy:2746987948
+    runs-on: [ self-hosted, Linux, k8s-runner ]
     needs: [ build-neon ]
     strategy:
       fail-fast: false
@@ -219,20 +217,14 @@ jobs:
           build_type: ${{ matrix.build_type }}
           rust_toolchain: ${{ matrix.rust_toolchain }}
           test_selection: batch_others
-          run_with_real_s3: true
-          real_s3_bucket: ci-tests-s3
-          real_s3_region: us-west-2
-          real_s3_access_key_id: "${{ secrets.AWS_ACCESS_KEY_ID_CI_TESTS_S3 }}"
-          real_s3_secret_access_key: "${{ secrets.AWS_SECRET_ACCESS_KEY_CI_TESTS_S3 }}"
+
       - name: Merge and upload coverage data
         if: matrix.build_type == 'debug'
         uses: ./.github/actions/save-coverage-data
 
   benchmarks:
-    runs-on: dev
-    container: 369495373322.dkr.ecr.eu-central-1.amazonaws.com/rustlegacy:2746987948
+    runs-on: [ self-hosted, Linux, k8s-runner ]
     needs: [ build-neon ]
-    if: github.ref_name == 'main' || contains(github.event.pull_request.labels.*.name, 'run-benchmarks')
     strategy:
       fail-fast: false
       matrix:
@@ -260,8 +252,7 @@ jobs:
       # while coverage is currently collected for the debug ones
 
   coverage-report:
-    runs-on: dev
-    container: 369495373322.dkr.ecr.eu-central-1.amazonaws.com/rustlegacy:2746987948
+    runs-on: [ self-hosted, Linux, k8s-runner ]
     needs: [ other-tests, pg_regress-tests ]
     strategy:
       fail-fast: false
@@ -286,17 +277,23 @@ jobs:
             target/
           key: v3-${{ runner.os }}-${{ matrix.build_type }}-cargo-${{ matrix.rust_toolchain }}-${{ hashFiles('Cargo.lock') }}
 
-      - name: Get Neon artifact
-        uses: ./.github/actions/download
+      - name: Get Neon artifact for restoration
+        uses: actions/download-artifact@v3
         with:
           name: neon-${{ runner.os }}-${{ matrix.build_type }}-${{ matrix.rust_toolchain }}-artifact
-          path: /tmp/neon
+          path: ./neon-artifact/
 
-      - name: Get coverage artifact
-        uses: ./.github/actions/download
+      - name: Extract Neon artifact
+        run: |
+          mkdir -p /tmp/neon/
+          tar -xf ./neon-artifact/neon.tar.zst -C /tmp/neon/
+          rm -rf ./neon-artifact/
+
+      - name: Restore coverage data
+        uses: actions/download-artifact@v3
         with:
           name: coverage-data-artifact
-          path: /tmp/coverage
+          path: /tmp/coverage/
 
       - name: Merge coverage data
         run: scripts/coverage "--profraw-prefix=$GITHUB_JOB" --dir=/tmp/coverage merge
@@ -334,40 +331,40 @@ jobs:
             }"
 
   trigger-e2e-tests:
-    runs-on: [ self-hosted, Linux, k8s-runner ]
-    needs: [ build-neon ]
-    steps:
-      - name: Set PR's status to pending and request a remote CI test
-        run: |
-          COMMIT_SHA=${{ github.event.pull_request.head.sha }}
-          COMMIT_SHA=${COMMIT_SHA:-${{ github.sha }}}
+   runs-on: [ self-hosted, Linux, k8s-runner ]
+   needs: [ build-neon ]
+   steps:
+     - name: Set PR's status to pending and request a remote CI test
+       run: |
+         COMMIT_SHA=${{ github.event.pull_request.head.sha }}
+         COMMIT_SHA=${COMMIT_SHA:-${{ github.sha }}}
 
-          REMOTE_REPO="${{ github.repository_owner }}/cloud"
+         REMOTE_REPO="${{ github.repository_owner }}/cloud"
 
-          curl -f -X POST \
-          https://api.github.com/repos/${{ github.repository }}/statuses/$COMMIT_SHA \
-          -H "Accept: application/vnd.github.v3+json" \
-          --user "${{ secrets.CI_ACCESS_TOKEN }}" \
-          --data \
-            "{
-              \"state\": \"pending\",
-              \"context\": \"neon-cloud-e2e\",
-              \"description\": \"[$REMOTE_REPO] Remote CI job is about to start\"
-            }"
+         curl -f -X POST \
+         https://api.github.com/repos/${{ github.repository }}/statuses/$COMMIT_SHA \
+         -H "Accept: application/vnd.github.v3+json" \
+         --user "${{ secrets.CI_ACCESS_TOKEN }}" \
+         --data \
+           "{
+             \"state\": \"pending\",
+             \"context\": \"neon-cloud-e2e\",
+             \"description\": \"[$REMOTE_REPO] Remote CI job is about to start\"
+           }"
 
-          curl -f -X POST \
-          https://api.github.com/repos/$REMOTE_REPO/actions/workflows/testing.yml/dispatches \
-          -H "Accept: application/vnd.github.v3+json" \
-          --user "${{ secrets.CI_ACCESS_TOKEN }}" \
-          --data \
-            "{
-              \"ref\": \"main\",
-              \"inputs\": {
-                \"ci_job_name\": \"neon-cloud-e2e\",
-                \"commit_hash\": \"$COMMIT_SHA\",
-                \"remote_repo\": \"${{ github.repository }}\"
-              }
-            }"
+         curl -f -X POST \
+         https://api.github.com/repos/$REMOTE_REPO/actions/workflows/testing.yml/dispatches \
+         -H "Accept: application/vnd.github.v3+json" \
+         --user "${{ secrets.CI_ACCESS_TOKEN }}" \
+         --data \
+           "{
+             \"ref\": \"main\",
+             \"inputs\": {
+               \"ci_job_name\": \"neon-cloud-e2e\",
+               \"commit_hash\": \"$COMMIT_SHA\",
+               \"remote_repo\": \"${{ github.repository }}\"
+             }
+           }"
 
   docker-image:
     runs-on: [ self-hosted, Linux, k8s-runner ]
@@ -527,7 +524,7 @@ jobs:
           if [[ "$GITHUB_REF_NAME" == "main" ]]; then
             STAGING='{"env_name": "staging", "proxy_job": "neon-proxy", "proxy_config": "staging.proxy", "kubeconfig_secret": "STAGING_KUBECONFIG_DATA"}'
             NEON_STRESS='{"env_name": "neon-stress", "proxy_job": "neon-stress-proxy", "proxy_config": "neon-stress.proxy", "kubeconfig_secret": "NEON_STRESS_KUBECONFIG_DATA"}'
-            echo "::set-output name=include::[$STAGING, $NEON_STRESS]"
+            echo "::set-output name=include::[$STAGING]"
           elif [[ "$GITHUB_REF_NAME" == "release" ]]; then
             PRODUCTION='{"env_name": "production", "proxy_job": "neon-proxy", "proxy_config": "production.proxy", "kubeconfig_secret": "PRODUCTION_KUBECONFIG_DATA"}'
             echo "::set-output name=include::[$PRODUCTION]"
@@ -562,7 +559,6 @@ jobs:
 
       - name: Redeploy
         run: |
-          export DOCKER_TAG=${{needs.docker-image.outputs.build-tag}}
           cd "$(pwd)/.github/ansible"
 
           if [[ "$GITHUB_REF_NAME" == "main" ]]; then

.github/workflows/codestyle.yml

@@ -8,7 +8,7 @@ on:
 
 defaults:
   run:
-    shell: bash -euxo pipefail {0}
+    shell: bash -ex {0}
 
 concurrency:
   # Allow only one workflow per any non-`main` branch.
@@ -27,7 +27,7 @@ jobs:
         # Rust toolchains (e.g. nightly or 1.37.0), add them here.
         rust_toolchain: [1.58]
         os: [ubuntu-latest, macos-latest]
-    timeout-minutes: 60
+    timeout-minutes: 50
     name: run regression test suite
     runs-on: ${{ matrix.os }}
 

.github/workflows/pg_clients.yml

@@ -19,12 +19,8 @@ concurrency:
 
 jobs:
   test-postgres-client-libs:
-    # TODO: switch to gen2 runner, requires docker
     runs-on: [ ubuntu-latest ]
 
-    env:
-      TEST_OUTPUT: /tmp/test_output
-
     steps:
     - name: Checkout
       uses: actions/checkout@v3
@@ -44,16 +40,16 @@ jobs:
         key: v1-${{ runner.os }}-python-deps-${{ hashFiles('poetry.lock') }}
 
     - name: Install Python deps
-      shell: bash -euxo pipefail {0}
+      shell: bash -ex {0}
       run: ./scripts/pysync
 
     - name: Run pytest
       env:
         REMOTE_ENV: 1
         BENCHMARK_CONNSTR: "${{ secrets.BENCHMARK_STAGING_CONNSTR }}"
-
+        TEST_OUTPUT: /tmp/test_output
         POSTGRES_DISTRIB_DIR: /tmp/neon/pg_install
-      shell: bash -euxo pipefail {0}
+      shell: bash -ex {0}
       run: |
         # Test framework expects we have psql binary;
         # but since we don't really need it in this test, let's mock it
@@ -65,18 +61,9 @@ jobs:
           -m "remote_cluster" \
           -rA "test_runner/pg_clients"
 
-    # We use GitHub's action upload-artifact because `ubuntu-latest` doesn't have configured AWS CLI.
-    # It will be fixed after switching to gen2 runner
-    - name: Upload python test logs
-      if: always()
-      uses: actions/upload-artifact@v3
-      with:
-        retention-days: 7
-        name: python-test-pg_clients-${{ runner.os }}-stage-logs
-        path: ${{ env.TEST_OUTPUT }}
-
     - name: Post to a Slack channel
-      if: ${{ github.event.schedule && failure() }}
+      if: failure()
+      id: slack
       uses: slackapi/slack-github-action@v1
       with:
         channel-id: "C033QLM5P7D" # dev-staging-stream

CONTRIBUTING.md

@@ -11,15 +11,17 @@ than it was before.
 
 ## Submitting changes
 
-1. Get at least one +1 on your PR before you push.
+1. Make a PR for every change.
+
+   Even seemingly trivial patches can break things in surprising ways.
+Use of common sense is OK. If you're only fixing a typo in a comment,
+it's probably fine to just push it. But if in doubt, open a PR.
+
+2. Get at least one +1 on your PR before you push.
 
    For simple patches, it will only take a minute for someone to review
 it.
 
-2. Don't force push small changes after making the PR ready for review.
-Doing so will force readers to re-read your entire PR, which will delay
-the review process.
-
 3. Always keep the CI green.
 
    Do not push, if the CI failed on your PR. Even if you think it's not

Cargo.lock

@@ -154,9 +154,9 @@ checksum = "d468802bab17cbc0cc575e9b053f41e72aa36bfa6b7f55e3529ffa43161b97fa"
 
 [[package]]
 name = "axum"
-version = "0.5.13"
+version = "0.5.12"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "6b9496f0c1d1afb7a2af4338bbe1d969cddfead41d87a9fb3aaa6d0bbc7af648"
+checksum = "d16705af05732b7d3258ec0f7b73c03a658a28925e050d8852d5b568ee8bcf4e"
 dependencies = [
  "async-trait",
  "axum-core",
@@ -317,6 +317,15 @@ dependencies = [
  "serde",
 ]
 
+[[package]]
+name = "cast"
+version = "0.2.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4c24dab4283a142afa2fdca129b80ad2c6284e073930f964c3a1293c225ee39a"
+dependencies = [
+ "rustc_version",
+]
+
 [[package]]
 name = "cast"
 version = "0.3.0"
@@ -458,6 +467,7 @@ dependencies = [
  "clap 3.2.12",
  "env_logger",
  "hyper",
+ "libc",
  "log",
  "postgres",
  "regex",
@@ -495,8 +505,8 @@ name = "control_plane"
 version = "0.1.0"
 dependencies = [
  "anyhow",
+ "lazy_static",
  "nix",
- "once_cell",
  "pageserver",
  "postgres",
  "regex",
@@ -507,6 +517,7 @@ dependencies = [
  "tar",
  "thiserror",
  "toml",
+ "url",
  "utils",
  "workspace_hack",
 ]
@@ -570,7 +581,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "b01d6de93b2b6c65e17c634a26653a29d107b3c98c607c765bf38d041531cd8f"
 dependencies = [
  "atty",
- "cast",
+ "cast 0.3.0",
  "clap 2.34.0",
  "criterion-plot",
  "csv",
@@ -591,11 +602,11 @@ dependencies = [
 
 [[package]]
 name = "criterion-plot"
-version = "0.4.5"
+version = "0.4.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "2673cc8207403546f45f5fd319a974b1e6983ad1a3ee7e6041650013be041876"
+checksum = "d00996de9f2f7559f7f4dc286073197f83e92256a59ed395f9aac01fe717da57"
 dependencies = [
- "cast",
+ "cast 0.2.7",
  "itertools",
 ]
 
@@ -671,9 +682,9 @@ dependencies = [
 
 [[package]]
 name = "crypto-common"
-version = "0.1.6"
+version = "0.1.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "1bfb12502f3fc46cca1bb51ac28df9d618d813cdc3d2f25b9fe775a34af26bb3"
+checksum = "2ccfd8c0ee4cce11e45b3fd6f9d5e69e0cc62912aa6a0cb1bf4617b0eba5a12f"
 dependencies = [
  "generic-array",
  "typenum",
@@ -1107,9 +1118,9 @@ dependencies = [
 
 [[package]]
 name = "gimli"
-version = "0.26.2"
+version = "0.26.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "22030e2c5a68ec659fde1e949a745124b48e6fa8b045b7ed5bd1fe4ccc5c4e5d"
+checksum = "78cc372d058dcf6d5ecd98510e7fbc9e5aec4d21de70f65fea8fecebcd881bd4"
 
 [[package]]
 name = "git-version"
@@ -1175,9 +1186,9 @@ dependencies = [
 
 [[package]]
 name = "hashbrown"
-version = "0.12.3"
+version = "0.12.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "8a9ee70c43aaf417c914396645a0fa852624801b24ebb7ae78fe8272889ac888"
+checksum = "607c8a29735385251a339424dd462993c0fed8fa09d378f259377df08c126022"
 
 [[package]]
 name = "heck"
@@ -1379,7 +1390,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "10a35a97730320ffe8e2d410b5d3b69279b98d2c14bdb8b70ea89ecf7888d41e"
 dependencies = [
  "autocfg",
- "hashbrown 0.12.3",
+ "hashbrown 0.12.2",
 ]
 
 [[package]]
@@ -1591,6 +1602,7 @@ dependencies = [
 name = "metrics"
 version = "0.1.0"
 dependencies = [
+ "lazy_static",
  "libc",
  "once_cell",
  "prometheus",
@@ -1665,6 +1677,7 @@ dependencies = [
  "git-version",
  "pageserver",
  "postgres",
+ "postgres_ffi",
  "safekeeper",
  "serde_json",
  "utils",
@@ -1842,9 +1855,9 @@ dependencies = [
 
 [[package]]
 name = "os_str_bytes"
-version = "6.2.0"
+version = "6.1.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "648001efe5d5c0102d8cea768e348da85d90af8ba91f0bea908f157951493cd4"
+checksum = "21326818e99cfe6ce1e524c2a805c189a99b5ae555a35d19f9a284b427d86afa"
 
 [[package]]
 name = "pageserver"
@@ -1870,6 +1883,7 @@ dependencies = [
  "humantime-serde",
  "hyper",
  "itertools",
+ "lazy_static",
  "metrics",
  "nix",
  "once_cell",
@@ -1891,6 +1905,7 @@ dependencies = [
  "thiserror",
  "tokio",
  "tokio-postgres",
+ "tokio-stream",
  "toml_edit",
  "tracing",
  "url",
@@ -2115,9 +2130,9 @@ dependencies = [
  "crc32c",
  "env_logger",
  "hex",
+ "lazy_static",
  "log",
  "memoffset",
- "once_cell",
  "postgres",
  "rand",
  "regex",
@@ -2269,7 +2284,6 @@ dependencies = [
  "anyhow",
  "async-trait",
  "base64",
- "bstr",
  "bytes",
  "clap 3.2.12",
  "futures",
@@ -2278,9 +2292,9 @@ dependencies = [
  "hex",
  "hmac 0.12.1",
  "hyper",
+ "lazy_static",
  "md5",
  "metrics",
- "once_cell",
  "parking_lot 0.12.1",
  "pin-project-lite",
  "rand",
@@ -2726,9 +2740,9 @@ dependencies = [
 
 [[package]]
 name = "rustversion"
-version = "1.0.8"
+version = "1.0.7"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "24c8ad4f0c00e1eb5bc7614d236a7f1300e3dbd76b68cac8e06fb00b015ad8d8"
+checksum = "a0a5f7c728f5d284929a1cccb5bc19884422bfe6ef4d6c409da2c41838983fcf"
 
 [[package]]
 name = "ryu"
@@ -2750,10 +2764,12 @@ dependencies = [
  "daemonize",
  "etcd_broker",
  "fs2",
+ "futures",
  "git-version",
  "hex",
  "humantime",
  "hyper",
+ "lazy_static",
  "metrics",
  "once_cell",
  "postgres",
@@ -2768,10 +2784,12 @@ dependencies = [
  "tempfile",
  "tokio",
  "tokio-postgres",
+ "tokio-util",
  "toml_edit",
  "tracing",
  "url",
  "utils",
+ "walkdir",
  "workspace_hack",
 ]
 
@@ -3607,9 +3625,9 @@ checksum = "099b7128301d285f79ddd55b9a83d5e6b9e97c92e0ea0daebee7263e932de992"
 
 [[package]]
 name = "unicode-ident"
-version = "1.0.2"
+version = "1.0.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "15c61ba63f9235225a22310255a29b806b907c9b8c964bcbd0a2c70f3f2deea7"
+checksum = "5bd2fe26506023ed7b5e1e315add59d6f584c621d037f9368fea9cfb988f368c"
 
 [[package]]
 name = "unicode-normalization"
@@ -3670,9 +3688,9 @@ dependencies = [
  "hex-literal",
  "hyper",
  "jsonwebtoken",
+ "lazy_static",
  "metrics",
  "nix",
- "once_cell",
  "pin-project-lite",
  "postgres",
  "postgres-protocol",

Dockerfile

@@ -17,10 +17,6 @@ RUN set -e \
 FROM neondatabase/rust:1.58 AS build
 ARG GIT_VERSION=local
 
-# Enable https://github.com/paritytech/cachepot to cache Rust crates' compilation results in Docker builds.
-# Set up cachepot to use an AWS S3 bucket for cache results, to reuse it between `docker build` invocations.
-# cachepot falls back to local filesystem if S3 is misconfigured, not failing the build.
-ARG RUSTC_WRAPPER=cachepot
 ARG CACHEPOT_BUCKET=zenith-rust-cachepot
 ARG AWS_ACCESS_KEY_ID
 ARG AWS_SECRET_ACCESS_KEY

Dockerfile.compute-tools

@@ -1,11 +1,7 @@
 # First transient image to build compute_tools binaries
-# NB: keep in sync with rust image version in .github/workflows/build_and_test.yml
+# NB: keep in sync with rust image version in .circle/config.yml
 FROM neondatabase/rust:1.58 AS rust-build
 
-# Enable https://github.com/paritytech/cachepot to cache Rust crates' compilation results in Docker builds.
-# Set up cachepot to use an AWS S3 bucket for cache results, to reuse it between `docker build` invocations.
-# cachepot falls back to local filesystem if S3 is misconfigured, not failing the build.
-ARG RUSTC_WRAPPER=cachepot
 ARG CACHEPOT_BUCKET=zenith-rust-cachepot
 ARG AWS_ACCESS_KEY_ID
 ARG AWS_SECRET_ACCESS_KEY

Makefile

@@ -29,11 +29,9 @@ else
 endif
 
 # macOS with brew-installed openssl requires explicit paths
-# It can be configured with OPENSSL_PREFIX variable
 UNAME_S := $(shell uname -s)
 ifeq ($(UNAME_S),Darwin)
-    OPENSSL_PREFIX ?= $(shell brew --prefix openssl@3)
-    PG_CONFIGURE_OPTS += --with-includes=$(OPENSSL_PREFIX)/include --with-libraries=$(OPENSSL_PREFIX)/lib
+    PG_CONFIGURE_OPTS += --with-includes=$(HOMEBREW_PREFIX)/opt/openssl/include --with-libraries=$(HOMEBREW_PREFIX)/opt/openssl/lib
 endif
 
 # Choose whether we should be silent or verbose
@@ -48,24 +46,28 @@ CARGO_CMD_PREFIX += $(if $(filter n,$(MAKEFLAGS)),,+)
 CARGO_CMD_PREFIX += CARGO_TERM_PROGRESS_WHEN=never CI=1
 
 #
-# Top level Makefile to build Zenith and PostgreSQL
+# Top level Makefile to build Neon and PostgreSQL
 #
 .PHONY: all
-all: zenith postgres
+all: neon postgres
 
-### Zenith Rust bits
+### Neon Rust bits
 #
 # The 'postgres_ffi' depends on the Postgres headers.
-.PHONY: zenith
-zenith: postgres-headers
-	+@echo "Compiling Zenith"
+.PHONY: neon
+neon: postgres-headers
+	+@echo "Compiling Neon"
 	$(CARGO_CMD_PREFIX) cargo build $(CARGO_BUILD_FLAGS)
 
 ### PostgreSQL parts
-$(POSTGRES_INSTALL_DIR)/build/config.status:
+#
+# Postgres is built in the 'build' directory, and installed into
+# $(POSTGRES_INSTALL_DIR), which defaults to 'tmp_install'
+#
+build/config.status:
 	+@echo "Configuring postgres build"
-	mkdir -p $(POSTGRES_INSTALL_DIR)/build
-	(cd $(POSTGRES_INSTALL_DIR)/build && \
+	mkdir -p build
+	(cd build && \
 	$(ROOT_PROJECT_DIR)/vendor/postgres/configure CFLAGS='$(PG_CFLAGS)' \
 		$(PG_CONFIGURE_OPTS) \
 		$(SECCOMP) \
@@ -73,44 +75,44 @@ $(POSTGRES_INSTALL_DIR)/build/config.status:
 
 # nicer alias for running 'configure'
 .PHONY: postgres-configure
-postgres-configure: $(POSTGRES_INSTALL_DIR)/build/config.status
+postgres-configure: build/config.status
 
 # Install the PostgreSQL header files into $(POSTGRES_INSTALL_DIR)/include
 .PHONY: postgres-headers
 postgres-headers: postgres-configure
 	+@echo "Installing PostgreSQL headers"
-	$(MAKE) -C $(POSTGRES_INSTALL_DIR)/build/src/include MAKELEVEL=0 install
+	$(MAKE) -C build/src/include MAKELEVEL=0 install
 
 # Compile and install PostgreSQL and contrib/neon
 .PHONY: postgres
 postgres: postgres-configure \
-		  postgres-headers # to prevent `make install` conflicts with zenith's `postgres-headers`
+		  postgres-headers # to prevent `make install` conflicts with neon's `postgres-headers`
 	+@echo "Compiling PostgreSQL"
-	$(MAKE) -C $(POSTGRES_INSTALL_DIR)/build MAKELEVEL=0 install
+	$(MAKE) -C build MAKELEVEL=0 install
 	+@echo "Compiling contrib/neon"
-	$(MAKE) -C $(POSTGRES_INSTALL_DIR)/build/contrib/neon install
+	$(MAKE) -C build/contrib/neon install
 	+@echo "Compiling contrib/neon_test_utils"
-	$(MAKE) -C $(POSTGRES_INSTALL_DIR)/build/contrib/neon_test_utils install
+	$(MAKE) -C build/contrib/neon_test_utils install
 	+@echo "Compiling pg_buffercache"
-	$(MAKE) -C $(POSTGRES_INSTALL_DIR)/build/contrib/pg_buffercache install
+	$(MAKE) -C build/contrib/pg_buffercache install
 	+@echo "Compiling pageinspect"
-	$(MAKE) -C $(POSTGRES_INSTALL_DIR)/build/contrib/pageinspect install
+	$(MAKE) -C build/contrib/pageinspect install
 
 
 .PHONY: postgres-clean
 postgres-clean:
-	$(MAKE) -C $(POSTGRES_INSTALL_DIR)/build MAKELEVEL=0 clean
+	$(MAKE) -C build MAKELEVEL=0 clean
 
 # This doesn't remove the effects of 'configure'.
 .PHONY: clean
 clean:
-	cd $(POSTGRES_INSTALL_DIR)/build && $(MAKE) clean
+	cd build && $(MAKE) clean
 	$(CARGO_CMD_PREFIX) cargo clean
 
 # This removes everything
 .PHONY: distclean
 distclean:
-	rm -rf $(POSTGRES_INSTALL_DIR)
+	rm -rf build $(POSTGRES_INSTALL_DIR)
 	$(CARGO_CMD_PREFIX) cargo clean
 
 .PHONY: fmt

README.md

@@ -1,6 +1,6 @@
 # Neon
 
-Neon is a serverless open-source alternative to AWS Aurora Postgres. It separates storage and compute and substitutes the PostgreSQL storage layer by redistributing data across a cluster of nodes.
+Neon is a serverless open source alternative to AWS Aurora Postgres. It separates storage and compute and substitutes PostgreSQL storage layer by redistributing data across a cluster of nodes.
 
 The project used to be called "Zenith". Many of the commands and code comments
 still refer to "zenith", but we are in the process of renaming things.
@@ -12,31 +12,32 @@ Alternatively, compile and run the project [locally](#running-local-installation
 
 ## Architecture overview
 
-A Neon installation consists of compute nodes and a Neon storage engine.
+A Neon installation consists of compute nodes and Neon storage engine.
 
-Compute nodes are stateless PostgreSQL nodes backed by the Neon storage engine.
+Compute nodes are stateless PostgreSQL nodes, backed by Neon storage engine.
 
-The Neon storage engine consists of two major components:
-- Pageserver. Scalable storage backend for the compute nodes.
-- WAL service. The service receives WAL from the compute node and ensures that it is stored durably.
+Neon storage engine consists of two major components:
+- Pageserver. Scalable storage backend for compute nodes.
+- WAL service. The service that receives WAL from compute node and ensures that it is stored durably.
 
 Pageserver consists of:
 - Repository - Neon storage implementation.
 - WAL receiver - service that receives WAL from WAL service and stores it in the repository.
 - Page service - service that communicates with compute nodes and responds with pages from the repository.
-- WAL redo - service that builds pages from base images and WAL records on Page service request
+- WAL redo - service that builds pages from base images and WAL records on Page service request.
+
 ## Running local installation
 
 
 #### Installing dependencies on Linux
-1. Install build dependencies and other applicable packages
+1. Install build dependencies and other useful packages
 
-* On Ubuntu or Debian, this set of packages should be sufficient to build the code:
+* On Ubuntu or Debian this set of packages should be sufficient to build the code:
 ```bash
 apt install build-essential libtool libreadline-dev zlib1g-dev flex bison libseccomp-dev \
 libssl-dev clang pkg-config libpq-dev etcd cmake postgresql-client
 ```
-* On Fedora, these packages are needed:
+* On Fedora these packages are needed:
 ```bash
 dnf install flex bison readline-devel zlib-devel openssl-devel \
   libseccomp-devel perl clang cmake etcd postgresql postgresql-contrib
@@ -68,7 +69,7 @@ brew install libpq
 brew link --force libpq
 ```
 
-#### Building on Linux
+#### Building on Linux and OSX
 
 1. Build neon and patched postgres
 ```
@@ -79,35 +80,19 @@ cd neon
 
 # The preferred and default is to make a debug build. This will create a 
 # demonstrably slower build than a release build. If you want to use a release
-# build, utilize "BUILD_TYPE=release make -j`nproc`" 
+# build, utilize "`BUILD_TYPE=release make -j`nproc``" 
 
 make -j`nproc`
 ```
 
-#### Building on OSX
-
-1. Build neon and patched postgres
-```
-# Note: The path to the neon sources can not contain a space.
-
-git clone --recursive https://github.com/neondatabase/neon.git
-cd neon
-
-# The preferred and default is to make a debug build. This will create a 
-# demonstrably slower build than a release build. If you want to use a release
-# build, utilize "BUILD_TYPE=release make -j`sysctl -n hw.logicalcpu`" 
-
-make -j`sysctl -n hw.logicalcpu`
-```
-
-#### Dependency installation notes
+#### dependency installation notes
 To run the `psql` client, install the `postgresql-client` package or modify `PATH` and `LD_LIBRARY_PATH` to include `tmp_install/bin` and `tmp_install/lib`, respectively.
 
 To run the integration tests or Python scripts (not required to use the code), install
-Python (3.9 or higher), and install python3 packages using `./scripts/pysync` (requires [poetry](https://python-poetry.org/)) in the project directory.
+Python (3.9 or higher), and install python3 packages using `./scripts/pysync` (requires poetry) in the project directory.
 
 
-#### Running neon database
+#### running neon database
 1. Start pageserver and postgres on top of it (should be called from repo root):
 ```sh
 # Create repository in .neon with proper paths to binaries and data
@@ -138,7 +123,7 @@ Starting postgres node at 'host=127.0.0.1 port=55432 user=cloud_admin dbname=pos
  main  127.0.0.1:55432  de200bd42b49cc1814412c7e592dd6e9  main         0/16B5BA8  running
 ```
 
-2. Now, it is possible to connect to postgres and run some queries:
+2. Now it is possible to connect to postgres and run some queries:
 ```text
 > psql -p55432 -h 127.0.0.1 -U cloud_admin postgres
 postgres=# CREATE TABLE t(key int primary key, value text);
@@ -196,16 +181,14 @@ postgres=# select * from t;
 (1 row)
 ```
 
-4. If you want to run tests afterward (see below), you must stop all the running of the pageserver, safekeeper, and postgres instances
-   you have just started. You can terminate them all with one command:
+4. If you want to run tests afterwards (see below), you have to stop all the running the pageserver, safekeeper and postgres instances
+   you have just started. You can stop them all with one command:
 ```sh
 > ./target/debug/neon_local stop
 ```
 
 ## Running tests
 
-Ensure your dependencies are installed as described [here](https://github.com/neondatabase/neon#dependency-installation-notes).
-
 ```sh
 git clone --recursive https://github.com/neondatabase/neon.git
 make # builds also postgres and installs it to ./tmp_install
@@ -222,8 +205,8 @@ To view your `rustdoc` documentation in a browser, try running `cargo doc --no-d
 
 ### Postgres-specific terms
 
-Due to Neon's very close relation with PostgreSQL internals, numerous specific terms are used.
-The same applies to certain spelling: i.e. we use MB to denote 1024 * 1024 bytes, while MiB would be technically more correct, it's inconsistent with what PostgreSQL code and its documentation use.
+Due to Neon's very close relation with PostgreSQL internals, there are numerous specific terms used.
+Same applies to certain spelling: i.e. we use MB to denote 1024 * 1024 bytes, while MiB would be technically more correct, it's inconsistent with what PostgreSQL code and its documentation use.
 
 To get more familiar with this aspect, refer to:
 

compute_tools/Cargo.toml

@@ -4,6 +4,7 @@ version = "0.1.0"
 edition = "2021"
 
 [dependencies]
+libc = "0.2"
 anyhow = "1.0"
 chrono = "0.4"
 clap = "3.0"

compute_tools/src/bin/compute_ctl.rs

@@ -157,7 +157,7 @@ fn main() -> Result<()> {
             exit(code)
         }
         Err(error) => {
-            error!("could not start the compute node: {:?}", error);
+            error!("could not start the compute node: {}", error);
 
             let mut state = compute.state.write().unwrap();
             state.error = Some(format!("{:?}", error));

control_plane/Cargo.toml

@@ -9,11 +9,12 @@ postgres = { git = "https://github.com/zenithdb/rust-postgres.git", rev="d052ee8
 serde = { version = "1.0", features = ["derive"] }
 serde_with = "1.12.0"
 toml = "0.5"
-once_cell = "1.13.0"
+lazy_static = "1.4"
 regex = "1"
 anyhow = "1.0"
 thiserror = "1"
 nix = "0.23"
+url = "2.2.2"
 reqwest = { version = "0.11", default-features = false, features = ["blocking", "json", "rustls-tls"] }
 
 pageserver = { path = "../pageserver" }

control_plane/src/etcd.rs

@@ -30,14 +30,14 @@ pub fn start_etcd_process(env: &local_env::LocalEnv) -> anyhow::Result<()> {
     let etcd_stdout_file =
         fs::File::create(etcd_data_dir.join("etcd.stdout.log")).with_context(|| {
             format!(
-                "Failed to create etcd stout file in directory {}",
+                "Failed to create ectd stout file in directory {}",
                 etcd_data_dir.display()
             )
         })?;
     let etcd_stderr_file =
         fs::File::create(etcd_data_dir.join("etcd.stderr.log")).with_context(|| {
             format!(
-                "Failed to create etcd stderr file in directory {}",
+                "Failed to create ectd stderr file in directory {}",
                 etcd_data_dir.display()
             )
         })?;

control_plane/src/lib.rs

@@ -51,11 +51,7 @@ fn fill_rust_env_vars(cmd: &mut Command) -> &mut Command {
 }
 
 fn fill_aws_secrets_vars(mut cmd: &mut Command) -> &mut Command {
-    for env_key in [
-        "AWS_ACCESS_KEY_ID",
-        "AWS_SECRET_ACCESS_KEY",
-        "AWS_SESSION_TOKEN",
-    ] {
+    for env_key in ["AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY"] {
         if let Ok(value) = std::env::var(env_key) {
             cmd = cmd.env(env_key, value);
         }

control_plane/src/postgresql_conf.rs

@@ -5,7 +5,7 @@
 /// enough to extract a few settings we need in Zenith, assuming you don't do
 /// funny stuff like include-directives or funny escaping.
 use anyhow::{bail, Context, Result};
-use once_cell::sync::Lazy;
+use lazy_static::lazy_static;
 use regex::Regex;
 use std::collections::HashMap;
 use std::fmt;
@@ -19,7 +19,9 @@ pub struct PostgresConf {
     hash: HashMap<String, String>,
 }
 
-static CONF_LINE_RE: Lazy<Regex> = Lazy::new(|| Regex::new(r"^((?:\w|\.)+)\s*=\s*(\S+)$").unwrap());
+lazy_static! {
+    static ref CONF_LINE_RE: Regex = Regex::new(r"^((?:\w|\.)+)\s*=\s*(\S+)$").unwrap();
+}
 
 impl PostgresConf {
     pub fn new() -> PostgresConf {
@@ -137,10 +139,10 @@ fn escape_str(s: &str) -> String {
     //
     // This regex is a bit more conservative than the rules in guc-file.l, so we quote some
     // strings that PostgreSQL would accept without quoting, but that's OK.
-
-    static UNQUOTED_RE: Lazy<Regex> =
-        Lazy::new(|| Regex::new(r"(^[-+]?[0-9]+[a-zA-Z]*$)|(^[a-zA-Z][a-zA-Z0-9]*$)").unwrap());
-
+    lazy_static! {
+        static ref UNQUOTED_RE: Regex =
+            Regex::new(r"(^[-+]?[0-9]+[a-zA-Z]*$)|(^[a-zA-Z][a-zA-Z0-9]*$)").unwrap();
+    }
     if UNQUOTED_RE.is_match(s) {
         s.to_string()
     } else {

control_plane/src/safekeeper.rs

@@ -247,7 +247,7 @@ impl SafekeeperNode {
         // Shutting down may take a long time,
         // if safekeeper flushes a lot of data
         let mut tcp_stopped = false;
-        for i in 0..600 {
+        for _ in 0..100 {
             if !tcp_stopped {
                 if let Err(err) = TcpStream::connect(&address) {
                     tcp_stopped = true;
@@ -272,11 +272,9 @@ impl SafekeeperNode {
                     }
                 }
             }
-            if i % 10 == 0 {
-                print!(".");
-                io::stdout().flush().unwrap();
-            }
-            thread::sleep(Duration::from_millis(100));
+            print!(".");
+            io::stdout().flush().unwrap();
+            thread::sleep(Duration::from_secs(1));
         }
 
         bail!("Failed to stop safekeeper with pid {}", pid);
@@ -306,9 +304,10 @@ impl SafekeeperNode {
         Ok(self
             .http_request(
                 Method::POST,
-                format!("{}/tenant/{}/timeline", self.http_base_url, tenant_id),
+                format!("{}/{}", self.http_base_url, "timeline"),
             )
             .json(&TimelineCreateRequest {
+                tenant_id,
                 timeline_id,
                 peer_ids,
             })

control_plane/src/storage.rs

@@ -12,9 +12,9 @@ use anyhow::{bail, Context};
 use nix::errno::Errno;
 use nix::sys::signal::{kill, Signal};
 use nix::unistd::Pid;
-use pageserver::http::models::{
-    TenantConfigRequest, TenantCreateRequest, TenantInfo, TimelineCreateRequest, TimelineInfo,
-};
+use pageserver::http::models::{TenantConfigRequest, TenantCreateRequest, TimelineCreateRequest};
+use pageserver::tenant_mgr::TenantInfo;
+use pageserver::timelines::TimelineInfo;
 use postgres::{Config, NoTls};
 use reqwest::blocking::{Client, RequestBuilder, Response};
 use reqwest::{IntoUrl, Method};
@@ -318,7 +318,7 @@ impl PageServerNode {
         // Shutting down may take a long time,
         // if pageserver checkpoints a lot of data
         let mut tcp_stopped = false;
-        for i in 0..600 {
+        for _ in 0..100 {
             if !tcp_stopped {
                 if let Err(err) = TcpStream::connect(&address) {
                     tcp_stopped = true;
@@ -344,11 +344,9 @@ impl PageServerNode {
                     }
                 }
             }
-            if i % 10 == 0 {
-                print!(".");
-                io::stdout().flush().unwrap();
-            }
-            thread::sleep(Duration::from_millis(100));
+            print!(".");
+            io::stdout().flush().unwrap();
+            thread::sleep(Duration::from_secs(1));
         }
 
         bail!("Failed to stop pageserver with pid {}", pid);
@@ -401,7 +399,6 @@ impl PageServerNode {
                     .get("checkpoint_distance")
                     .map(|x| x.parse::<u64>())
                     .transpose()?,
-                checkpoint_timeout: settings.get("checkpoint_timeout").map(|x| x.to_string()),
                 compaction_target_size: settings
                     .get("compaction_target_size")
                     .map(|x| x.parse::<u64>())
@@ -456,7 +453,6 @@ impl PageServerNode {
                     .map(|x| x.parse::<u64>())
                     .transpose()
                     .context("Failed to parse 'checkpoint_distance' as an integer")?,
-                checkpoint_timeout: settings.get("checkpoint_timeout").map(|x| x.to_string()),
                 compaction_target_size: settings
                     .get("compaction_target_size")
                     .map(|x| x.parse::<u64>())

docker-entrypoint.sh

@@ -1,8 +1,6 @@
 #!/bin/sh
 set -eux
 
-pageserver_id_param="${NODE_ID:-10}"
-
 broker_endpoints_param="${BROKER_ENDPOINT:-absent}"
 if [ "$broker_endpoints_param" != "absent" ]; then
     broker_endpoints_param="-c broker_endpoints=['$broker_endpoints_param']"
@@ -10,12 +8,10 @@ else
     broker_endpoints_param=''
 fi
 
-remote_storage_param="${REMOTE_STORAGE:-}"
-
 if [ "$1" = 'pageserver' ]; then
     if [ ! -d "/data/tenants" ]; then
         echo "Initializing pageserver data directory"
-        pageserver --init -D /data -c "pg_distrib_dir='/usr/local'" -c "id=${pageserver_id_param}" $broker_endpoints_param $remote_storage_param
+        pageserver --init -D /data -c "pg_distrib_dir='/usr/local'" -c "id=10" $broker_endpoints_param
     fi
     echo "Staring pageserver at 0.0.0.0:6400"
     pageserver -c "listen_pg_addr='0.0.0.0:6400'" -c "listen_http_addr='0.0.0.0:9898'" $broker_endpoints_param -D /data

docs/SUMMARY.md

@@ -52,8 +52,10 @@
 - [multitenancy.md](./multitenancy.md) — how multitenancy is organized in the pageserver and Zenith CLI.
 - [settings.md](./settings.md)
 #FIXME: move these under sourcetree.md
+#- [pageserver/README.md](/pageserver/README.md)
 #- [postgres_ffi/README.md](/libs/postgres_ffi/README.md)
 #- [test_runner/README.md](/test_runner/README.md)
+#- [safekeeper/README.md](/safekeeper/README.md)
 
 
 # RFCs

docs/glossary.md

@@ -75,7 +75,7 @@ layer's Segment and range of LSNs.
 There are two kinds of layers, in-memory and on-disk layers. In-memory
 layers are used to ingest incoming WAL, and provide fast access
 to the recent page versions. On-disk layers are stored as files on disk, and
-are immutable. See [pageserver-storage.md](./pageserver-storage.md) for more.
+are immutable. See pageserver/src/layered_repository/README.md for more.
 
 ### Layer file (on-disk layer)
 
@@ -111,7 +111,7 @@ PostgreSQL LSNs and functions to monitor them:
 * `pg_last_wal_replay_lsn ()` - Returns the last write-ahead log location that has been replayed during recovery. If recovery is still in progress this will increase monotonically.
 [source PostgreSQL documentation](https://www.postgresql.org/docs/devel/functions-admin.html):
 
-Neon safekeeper LSNs. See [safekeeper protocol section](safekeeper-protocol.md) for more information.
+Neon safekeeper LSNs. For more check [safekeeper/README_PROTO.md](/safekeeper/README_PROTO.md)
 * `CommitLSN`: position in WAL confirmed by quorum safekeepers.
 * `RestartLSN`: position in WAL confirmed by all safekeepers.
 * `FlushLSN`: part of WAL persisted to the disk by safekeeper.

docs/pageserver-services.md

@@ -68,6 +68,8 @@ There are the following implementations present:
 * local filesystem — to use in tests mainly
 * AWS S3           - to use in production
 
+Implementation details are covered in the [backup readme](./src/remote_storage/README.md) and corresponding Rust file docs, parameters documentation can be found at [settings docs](../docs/settings.md).
+
 The backup service is disabled by default and can be enabled to interact with a single remote storage.
 
 CLI examples:
@@ -116,7 +118,7 @@ implemented by the LayeredRepository object in
 `layered_repository.rs`. There is only that one implementation of the
 Repository trait, but it's still a useful abstraction that keeps the
 interface for the low-level storage functionality clean. The layered
-storage format is described in [pageserver-storage.md](./pageserver-storage.md).
+storage format is described in layered_repository/README.md.
 
 Each repository consists of multiple Timelines. Timeline is a
 workhorse that accepts page changes from the WAL, and serves

docs/settings.md

@@ -15,7 +15,7 @@ listen_pg_addr = '127.0.0.1:64000'
 listen_http_addr = '127.0.0.1:9898'
 
 checkpoint_distance = '268435456' # in bytes
-checkpoint_timeout = '10m'
+checkpoint_period = '1 s'
 
 gc_period = '100 s'
 gc_horizon = '67108864'
@@ -46,7 +46,7 @@ Note the `[remote_storage]` section: it's a [table](https://toml.io/en/v1.0.0#ta
 
 All values can be passed as an argument to the pageserver binary, using the `-c` parameter and specified as a valid TOML string. All tables should be passed in the inline form.
 
-Example: `${PAGESERVER_BIN} -c "checkpoint_timeout = '10 m'" -c "remote_storage={local_path='/some/local/path/'}"`
+Example: `${PAGESERVER_BIN} -c "checkpoint_period = '100 s'" -c "remote_storage={local_path='/some/local/path/'}"`
 
 Note that TOML distinguishes between strings and integers, the former require single or double quotes around them.
 
@@ -82,14 +82,6 @@ S3.
 
 The unit is # of bytes.
 
-#### checkpoint_timeout
-
-Apart from `checkpoint_distance`, open layer flushing is also triggered
-`checkpoint_timeout` after the last flush. This makes WAL eventually uploaded to
-s3 when activity is stopped.
-
-The default is 10m.
-
 #### compaction_period
 
 Every `compaction_period` seconds, the page server checks if

docs/sourcetree.md

@@ -28,7 +28,7 @@ The pageserver has a few different duties:
 - Receive WAL from the WAL service and decode it.
 - Replay WAL that's applicable to the chunks that the Page Server maintains
 
-For more detailed info, see [pageserver-services.md](./pageserver-services.md)
+For more detailed info, see [/pageserver/README](/pageserver/README.md)
 
 `/proxy`:
 
@@ -57,7 +57,7 @@ PostgreSQL extension that contains functions needed for testing and debugging.
 The zenith WAL service that receives WAL from a primary compute nodes and streams it to the pageserver.
 It acts as a holding area and redistribution center for recently generated WAL.
 
-For more detailed info, see [walservice.md](./walservice.md)
+For more detailed info, see [/safekeeper/README](/safekeeper/README.md)
 
 `/workspace_hack`:
 The workspace_hack crate exists only to pin down some dependencies.

docs/walservice.md

@@ -75,8 +75,8 @@ safekeepers. The Paxos and crash recovery algorithm ensures that only
 one primary node can be actively streaming WAL to the quorum of
 safekeepers.
 
-See [this section](safekeeper-protocol.md) for a more detailed description of
-the consensus protocol. spec/ contains TLA+ specification of it.
+See README_PROTO.md for a more detailed description of the consensus
+protocol. spec/ contains TLA+ specification of it.
 
 # Q&A
 

libs/etcd_broker/Cargo.toml

@@ -9,7 +9,7 @@
  serde = { version = "1.0", features = ["derive"] }
  serde_json = "1"
  serde_with = "1.12.0"
- once_cell = "1.13.0"
+ once_cell = "1.8.0"
 
  utils = { path = "../utils" }
  workspace_hack = { version = "0.1", path = "../../workspace_hack" }

libs/metrics/Cargo.toml

@@ -6,5 +6,6 @@ edition = "2021"
 [dependencies]
 prometheus = {version = "0.13", default_features=false, features = ["process"]} # removes protobuf dependency
 libc = "0.2"
-once_cell = "1.13.0"
+lazy_static = "1.4"
+once_cell = "1.8.0"
 workspace_hack = { version = "0.1", path = "../../workspace_hack" }

libs/metrics/src/lib.rs

@@ -2,10 +2,7 @@
 //! make sure that we use the same dep version everywhere.
 //! Otherwise, we might not see all metrics registered via
 //! a default registry.
-use once_cell::sync::Lazy;
-use prometheus::core::{AtomicU64, GenericGauge, GenericGaugeVec};
-pub use prometheus::opts;
-pub use prometheus::register;
+use lazy_static::lazy_static;
 pub use prometheus::{core, default_registry, proto};
 pub use prometheus::{exponential_buckets, linear_buckets};
 pub use prometheus::{register_gauge, Gauge};
@@ -21,17 +18,6 @@ pub use prometheus::{Encoder, TextEncoder};
 mod wrappers;
 pub use wrappers::{CountedReader, CountedWriter};
 
-pub type UIntGauge = GenericGauge<AtomicU64>;
-pub type UIntGaugeVec = GenericGaugeVec<AtomicU64>;
-
-#[macro_export]
-macro_rules! register_uint_gauge_vec {
-    ($NAME:expr, $HELP:expr, $LABELS_NAMES:expr $(,)?) => {{
-        let gauge_vec = UIntGaugeVec::new($crate::opts!($NAME, $HELP), $LABELS_NAMES).unwrap();
-        $crate::register(Box::new(gauge_vec.clone())).map(|_| gauge_vec)
-    }};
-}
-
 /// Gathers all Prometheus metrics and records the I/O stats just before that.
 ///
 /// Metrics gathering is a relatively simple and standalone operation, so
@@ -41,22 +27,19 @@ pub fn gather() -> Vec<prometheus::proto::MetricFamily> {
     prometheus::gather()
 }
 
-static DISK_IO_BYTES: Lazy<IntGaugeVec> = Lazy::new(|| {
-    register_int_gauge_vec!(
+lazy_static! {
+    static ref DISK_IO_BYTES: IntGaugeVec = register_int_gauge_vec!(
         "libmetrics_disk_io_bytes_total",
         "Bytes written and read from disk, grouped by the operation (read|write)",
         &["io_operation"]
     )
-    .expect("Failed to register disk i/o bytes int gauge vec")
-});
-
-static MAXRSS_KB: Lazy<IntGauge> = Lazy::new(|| {
-    register_int_gauge!(
+    .expect("Failed to register disk i/o bytes int gauge vec");
+    static ref MAXRSS_KB: IntGauge = register_int_gauge!(
         "libmetrics_maxrss_kb",
         "Memory usage (Maximum Resident Set Size)"
     )
-    .expect("Failed to register maxrss_kb int gauge")
-});
+    .expect("Failed to register maxrss_kb int gauge");
+}
 
 pub const DISK_WRITE_SECONDS_BUCKETS: &[f64] = &[
     0.000_050, 0.000_100, 0.000_500, 0.001, 0.003, 0.005, 0.01, 0.05, 0.1, 0.3, 0.5,

libs/metrics/src/wrappers.rs

@@ -10,13 +10,13 @@ use std::io::{Read, Result, Write};
 /// # use std::io::{Result, Read};
 /// # use metrics::{register_int_counter, IntCounter};
 /// # use metrics::CountedReader;
-/// # use once_cell::sync::Lazy;
 /// #
-/// # static INT_COUNTER: Lazy<IntCounter> = Lazy::new( || { register_int_counter!(
+/// # lazy_static::lazy_static! {
+/// #     static ref INT_COUNTER: IntCounter = register_int_counter!(
 /// #         "int_counter",
 /// #         "let's count something!"
-/// #     ).unwrap()
-/// # });
+/// #     ).unwrap();
+/// # }
 /// #
 /// fn do_some_reads(stream: impl Read, count: usize) -> Result<Vec<u8>> {
 ///     let mut reader = CountedReader::new(stream, |cnt| {
@@ -85,13 +85,13 @@ impl<T: Read> Read for CountedReader<'_, T> {
 /// # use std::io::{Result, Write};
 /// # use metrics::{register_int_counter, IntCounter};
 /// # use metrics::CountedWriter;
-/// # use once_cell::sync::Lazy;
 /// #
-/// # static INT_COUNTER: Lazy<IntCounter> = Lazy::new( || { register_int_counter!(
+/// # lazy_static::lazy_static! {
+/// #     static ref INT_COUNTER: IntCounter = register_int_counter!(
 /// #         "int_counter",
 /// #         "let's count something!"
-/// #     ).unwrap()
-/// # });
+/// #     ).unwrap();
+/// # }
 /// #
 /// fn do_some_writes(stream: impl Write, payload: &[u8]) -> Result<()> {
 ///     let mut writer = CountedWriter::new(stream, |cnt| {

libs/postgres_ffi/Cargo.toml

@@ -12,7 +12,7 @@ byteorder = "1.4.3"
 anyhow = "1.0"
 crc32c = "0.6.0"
 hex = "0.4.3"
-once_cell = "1.13.0"
+lazy_static = "1.4"
 log = "0.4.14"
 memoffset = "0.6.2"
 thiserror = "1.0"

libs/postgres_ffi/src/relfile_utils.rs

@@ -2,7 +2,7 @@
 //! Common utilities for dealing with PostgreSQL relation files.
 //!
 use crate::pg_constants;
-use once_cell::sync::OnceCell;
+use lazy_static::lazy_static;
 use regex::Regex;
 
 #[derive(Debug, Clone, thiserror::Error, PartialEq)]
@@ -54,14 +54,11 @@ pub fn forknumber_to_name(forknum: u8) -> Option<&'static str> {
 /// See functions relpath() and _mdfd_segpath() in PostgreSQL sources.
 ///
 pub fn parse_relfilename(fname: &str) -> Result<(u32, u8, u32), FilePathError> {
-    static RELFILE_RE: OnceCell<Regex> = OnceCell::new();
-    RELFILE_RE.get_or_init(|| {
-        Regex::new(r"^(?P<relnode>\d+)(_(?P<forkname>[a-z]+))?(\.(?P<segno>\d+))?$").unwrap()
-    });
-
+    lazy_static! {
+        static ref RELFILE_RE: Regex =
+            Regex::new(r"^(?P<relnode>\d+)(_(?P<forkname>[a-z]+))?(\.(?P<segno>\d+))?$").unwrap();
+    }
     let caps = RELFILE_RE
-        .get()
-        .unwrap()
         .captures(fname)
         .ok_or(FilePathError::InvalidFileName)?;
 

libs/postgres_ffi/src/waldecoder.rs

@@ -13,30 +13,24 @@ use super::xlog_utils::*;
 use super::XLogLongPageHeaderData;
 use super::XLogPageHeaderData;
 use super::XLogRecord;
-use super::XLOG_PAGE_MAGIC;
 use bytes::{Buf, BufMut, Bytes, BytesMut};
 use crc32c::*;
 use log::*;
 use std::cmp::min;
-use std::num::NonZeroU32;
 use thiserror::Error;
 use utils::lsn::Lsn;
 
-enum State {
-    WaitingForRecord,
-    ReassemblingRecord {
-        recordbuf: BytesMut,
-        contlen: NonZeroU32,
-    },
-    SkippingEverything {
-        skip_until_lsn: Lsn,
-    },
-}
-
 pub struct WalStreamDecoder {
     lsn: Lsn,
+
+    startlsn: Lsn, // LSN where this record starts
+    contlen: u32,
+    padlen: u32,
+
     inputbuf: BytesMut,
-    state: State,
+
+    /// buffer used to reassemble records that cross page boundaries.
+    recordbuf: BytesMut,
 }
 
 #[derive(Error, Debug, Clone)]
@@ -54,8 +48,13 @@ impl WalStreamDecoder {
     pub fn new(lsn: Lsn) -> WalStreamDecoder {
         WalStreamDecoder {
             lsn,
+
+            startlsn: Lsn(0),
+            contlen: 0,
+            padlen: 0,
+
             inputbuf: BytesMut::new(),
-            state: State::WaitingForRecord,
+            recordbuf: BytesMut::new(),
         }
     }
 
@@ -68,58 +67,6 @@ impl WalStreamDecoder {
         self.inputbuf.extend_from_slice(buf);
     }
 
-    fn validate_page_header(&self, hdr: &XLogPageHeaderData) -> Result<(), WalDecodeError> {
-        let validate_impl = || {
-            if hdr.xlp_magic != XLOG_PAGE_MAGIC as u16 {
-                return Err(format!(
-                    "invalid xlog page header: xlp_magic={}, expected {}",
-                    hdr.xlp_magic, XLOG_PAGE_MAGIC
-                ));
-            }
-            if hdr.xlp_pageaddr != self.lsn.0 {
-                return Err(format!(
-                    "invalid xlog page header: xlp_pageaddr={}, expected {}",
-                    hdr.xlp_pageaddr, self.lsn
-                ));
-            }
-            match self.state {
-                State::WaitingForRecord => {
-                    // if hdr.xlp_info & XLP_FIRST_IS_CONTRECORD != 0 {
-                    //     return Err(
-                    //         "invalid xlog page header: unexpected XLP_FIRST_IS_CONTRECORD".into(),
-                    //     );
-                    // }
-                    if hdr.xlp_rem_len != 0 {
-                        return Err(format!(
-                            "invalid xlog page header: xlp_rem_len={}, but it's not a contrecord",
-                            hdr.xlp_rem_len
-                        ));
-                    }
-                }
-                State::ReassemblingRecord { contlen, .. } => {
-                    if hdr.xlp_info & XLP_FIRST_IS_CONTRECORD == 0 {
-                        return Err(
-                            "invalid xlog page header: XLP_FIRST_IS_CONTRECORD expected, not found"
-                                .into(),
-                        );
-                    }
-                    if hdr.xlp_rem_len != contlen.get() {
-                        return Err(format!(
-                            "invalid xlog page header: xlp_rem_len={}, expected {}",
-                            hdr.xlp_rem_len,
-                            contlen.get()
-                        ));
-                    }
-                }
-                State::SkippingEverything { .. } => {
-                    panic!("Should not be validating page header in the SkippingEverything state");
-                }
-            };
-            Ok(())
-        };
-        validate_impl().map_err(|msg| WalDecodeError { msg, lsn: self.lsn })
-    }
-
     /// Attempt to decode another WAL record from the input that has been fed to the
     /// decoder so far.
     ///
@@ -129,121 +76,128 @@ impl WalStreamDecoder {
     ///     Err(WalDecodeError): an error occurred while decoding, meaning the input was invalid.
     ///
     pub fn poll_decode(&mut self) -> Result<Option<(Lsn, Bytes)>, WalDecodeError> {
+        let recordbuf;
+
         // Run state machine that validates page headers, and reassembles records
         // that cross page boundaries.
         loop {
             // parse and verify page boundaries as we go
-            // However, we may have to skip some page headers if we're processing the XLOG_SWITCH record or skipping padding for whatever reason.
-            match self.state {
-                State::WaitingForRecord | State::ReassemblingRecord { .. } => {
-                    if self.lsn.segment_offset(pg_constants::WAL_SEGMENT_SIZE) == 0 {
-                        // parse long header
-
-                        if self.inputbuf.remaining() < XLOG_SIZE_OF_XLOG_LONG_PHD {
-                            return Ok(None);
-                        }
-
-                        let hdr = XLogLongPageHeaderData::from_bytes(&mut self.inputbuf).map_err(
-                            |e| WalDecodeError {
-                                msg: format!("long header deserialization failed {}", e),
-                                lsn: self.lsn,
-                            },
-                        )?;
-
-                        self.validate_page_header(&hdr.std)?;
-
-                        self.lsn += XLOG_SIZE_OF_XLOG_LONG_PHD as u64;
-                    } else if self.lsn.block_offset() == 0 {
-                        if self.inputbuf.remaining() < XLOG_SIZE_OF_XLOG_SHORT_PHD {
-                            return Ok(None);
-                        }
-
-                        let hdr =
-                            XLogPageHeaderData::from_bytes(&mut self.inputbuf).map_err(|e| {
-                                WalDecodeError {
-                                    msg: format!("header deserialization failed {}", e),
-                                    lsn: self.lsn,
-                                }
-                            })?;
-
-                        self.validate_page_header(&hdr)?;
-
-                        self.lsn += XLOG_SIZE_OF_XLOG_SHORT_PHD as u64;
-                    }
+            if self.padlen > 0 {
+                // We should first skip padding, as we may have to skip some page headers if we're processing the XLOG_SWITCH record.
+                if self.inputbuf.remaining() < self.padlen as usize {
+                    return Ok(None);
                 }
-                State::SkippingEverything { .. } => {}
-            }
-            match &mut self.state {
-                State::WaitingForRecord => {
-                    // need to have at least the xl_tot_len field
-                    if self.inputbuf.remaining() < 4 {
-                        return Ok(None);
-                    }
 
-                    // peek xl_tot_len at the beginning of the record.
-                    // FIXME: assumes little-endian
-                    let xl_tot_len = (&self.inputbuf[0..4]).get_u32_le();
-                    if (xl_tot_len as usize) < XLOG_SIZE_OF_XLOG_RECORD {
-                        return Err(WalDecodeError {
-                            msg: format!("invalid xl_tot_len {}", xl_tot_len),
-                            lsn: self.lsn,
-                        });
-                    }
-                    // Fast path for the common case that the whole record fits on the page.
-                    let pageleft = self.lsn.remaining_in_block() as u32;
-                    if self.inputbuf.remaining() >= xl_tot_len as usize && xl_tot_len <= pageleft {
-                        self.lsn += xl_tot_len as u64;
-                        let recordbuf = self.inputbuf.copy_to_bytes(xl_tot_len as usize);
-                        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
-                        // to 'recordbuf'.  Subsequent iterations will skip page headers, and
-                        // append the continuations from the next pages to 'recordbuf'.
-                        self.state = State::ReassemblingRecord {
-                            recordbuf: BytesMut::with_capacity(xl_tot_len as usize),
-                            contlen: NonZeroU32::new(xl_tot_len).unwrap(),
-                        }
-                    }
+                // skip padding
+                self.inputbuf.advance(self.padlen as usize);
+                self.lsn += self.padlen as u64;
+                self.padlen = 0;
+            } else if self.lsn.segment_offset(pg_constants::WAL_SEGMENT_SIZE) == 0 {
+                // parse long header
+
+                if self.inputbuf.remaining() < XLOG_SIZE_OF_XLOG_LONG_PHD {
+                    return Ok(None);
                 }
-                State::ReassemblingRecord { recordbuf, contlen } => {
-                    // we're continuing a record, possibly from previous page.
-                    let pageleft = self.lsn.remaining_in_block() as u32;
 
-                    // read the rest of the record, or as much as fits on this page.
-                    let n = min(contlen.get(), pageleft) as usize;
-
-                    if self.inputbuf.remaining() < n {
-                        return Ok(None);
+                let hdr = XLogLongPageHeaderData::from_bytes(&mut self.inputbuf).map_err(|e| {
+                    WalDecodeError {
+                        msg: format!("long header deserialization failed {}", e),
+                        lsn: self.lsn,
                     }
+                })?;
 
-                    recordbuf.put(self.inputbuf.split_to(n));
-                    self.lsn += n as u64;
-                    *contlen = match NonZeroU32::new(contlen.get() - n as u32) {
-                        Some(x) => x,
-                        None => {
-                            // The record is now complete.
-                            let recordbuf = std::mem::replace(recordbuf, BytesMut::new()).freeze();
-                            return Ok(Some(self.complete_record(recordbuf)?));
-                        }
-                    }
+                if hdr.std.xlp_pageaddr != self.lsn.0 {
+                    return Err(WalDecodeError {
+                        msg: "invalid xlog segment header".into(),
+                        lsn: self.lsn,
+                    });
                 }
-                State::SkippingEverything { skip_until_lsn } => {
-                    assert!(*skip_until_lsn >= self.lsn);
-                    let n = skip_until_lsn.0 - self.lsn.0;
-                    if self.inputbuf.remaining() < n as usize {
-                        return Ok(None);
-                    }
-                    self.inputbuf.advance(n as usize);
-                    self.lsn += n;
-                    self.state = State::WaitingForRecord;
+                // TODO: verify the remaining fields in the header
+
+                self.lsn += XLOG_SIZE_OF_XLOG_LONG_PHD as u64;
+                continue;
+            } else if self.lsn.block_offset() == 0 {
+                if self.inputbuf.remaining() < XLOG_SIZE_OF_XLOG_SHORT_PHD {
+                    return Ok(None);
                 }
+
+                let hdr = XLogPageHeaderData::from_bytes(&mut self.inputbuf).map_err(|e| {
+                    WalDecodeError {
+                        msg: format!("header deserialization failed {}", e),
+                        lsn: self.lsn,
+                    }
+                })?;
+
+                if hdr.xlp_pageaddr != self.lsn.0 {
+                    return Err(WalDecodeError {
+                        msg: "invalid xlog page header".into(),
+                        lsn: self.lsn,
+                    });
+                }
+                // TODO: verify the remaining fields in the header
+
+                self.lsn += XLOG_SIZE_OF_XLOG_SHORT_PHD as u64;
+                continue;
+            } else if self.contlen == 0 {
+                assert!(self.recordbuf.is_empty());
+
+                // need to have at least the xl_tot_len field
+                if self.inputbuf.remaining() < 4 {
+                    return Ok(None);
+                }
+
+                // peek xl_tot_len at the beginning of the record.
+                // FIXME: assumes little-endian
+                self.startlsn = self.lsn;
+                let xl_tot_len = (&self.inputbuf[0..4]).get_u32_le();
+                if (xl_tot_len as usize) < XLOG_SIZE_OF_XLOG_RECORD {
+                    return Err(WalDecodeError {
+                        msg: format!("invalid xl_tot_len {}", xl_tot_len),
+                        lsn: self.lsn,
+                    });
+                }
+
+                // Fast path for the common case that the whole record fits on the page.
+                let pageleft = self.lsn.remaining_in_block() as u32;
+                if self.inputbuf.remaining() >= xl_tot_len as usize && xl_tot_len <= pageleft {
+                    // Take the record from the 'inputbuf', and validate it.
+                    recordbuf = self.inputbuf.copy_to_bytes(xl_tot_len as usize);
+                    self.lsn += xl_tot_len as u64;
+                    break;
+                } 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
+                    // to 'recordbuf'.  Subsequent iterations will skip page headers, and
+                    // append the continuations from the next pages to 'recordbuf'.
+                    self.recordbuf.reserve(xl_tot_len as usize);
+                    self.contlen = xl_tot_len;
+                    continue;
+                }
+            } else {
+                // we're continuing a record, possibly from previous page.
+                let pageleft = self.lsn.remaining_in_block() as u32;
+
+                // read the rest of the record, or as much as fits on this page.
+                let n = min(self.contlen, pageleft) as usize;
+
+                if self.inputbuf.remaining() < n {
+                    return Ok(None);
+                }
+
+                self.recordbuf.put(self.inputbuf.split_to(n));
+                self.lsn += n as u64;
+                self.contlen -= n as u32;
+
+                if self.contlen == 0 {
+                    // The record is now complete.
+                    recordbuf = std::mem::replace(&mut self.recordbuf, BytesMut::new()).freeze();
+                    break;
+                }
+                continue;
             }
         }
-    }
 
-    fn complete_record(&mut self, recordbuf: Bytes) -> Result<(Lsn, Bytes), WalDecodeError> {
         // We now have a record in the 'recordbuf' local variable.
         let xlogrec =
             XLogRecord::from_slice(&recordbuf[0..XLOG_SIZE_OF_XLOG_RECORD]).map_err(|e| {
@@ -265,20 +219,18 @@ impl WalStreamDecoder {
 
         // XLOG_SWITCH records are special. If we see one, we need to skip
         // to the next WAL segment.
-        let next_lsn = if xlogrec.is_xlog_switch_record() {
+        if xlogrec.is_xlog_switch_record() {
             trace!("saw xlog switch record at {}", self.lsn);
-            self.lsn + self.lsn.calc_padding(pg_constants::WAL_SEGMENT_SIZE as u64)
+            self.padlen = self.lsn.calc_padding(pg_constants::WAL_SEGMENT_SIZE as u64) as u32;
         } else {
             // Pad to an 8-byte boundary
-            self.lsn.align()
-        };
-        self.state = State::SkippingEverything {
-            skip_until_lsn: next_lsn,
-        };
+            self.padlen = self.lsn.calc_padding(8u32) as u32;
+        }
 
         // We should return LSN of the next record, not the last byte of this record or
         // the byte immediately after. Note that this handles both XLOG_SWITCH and usual
         // records, the former "spans" until the next WAL segment (see test_xlog_switch).
-        Ok((next_lsn, recordbuf))
+        let result = (self.lsn + self.padlen as u64, recordbuf);
+        Ok(Some(result))
     }
 }

libs/postgres_ffi/src/xlog_utils.rs

@@ -16,22 +16,22 @@ use crate::XLogRecord;
 use crate::XLOG_PAGE_MAGIC;
 
 use crate::pg_constants::WAL_SEGMENT_SIZE;
-use crate::waldecoder::WalStreamDecoder;
-
+use anyhow::{bail, ensure};
+use byteorder::{ByteOrder, LittleEndian};
 use bytes::BytesMut;
 use bytes::{Buf, Bytes};
-
+use crc32c::*;
 use log::*;
-
-use std::fs::File;
+use std::cmp::max;
+use std::cmp::min;
+use std::fs::{self, File};
 use std::io::prelude::*;
-use std::io::ErrorKind;
 use std::io::SeekFrom;
 use std::path::{Path, PathBuf};
 use std::time::SystemTime;
 use utils::bin_ser::DeserializeError;
 use utils::bin_ser::SerializeError;
-
+use utils::const_assert;
 use utils::lsn::Lsn;
 
 pub const XLOG_FNAME_LEN: usize = 24;
@@ -140,93 +140,335 @@ pub fn to_pg_timestamp(time: SystemTime) -> TimestampTz {
     }
 }
 
-// Returns (aligned) end_lsn of the last record in data_dir with WAL segments.
-// start_lsn must point to some previously known record boundary (beginning of
-// the next record). If no valid record after is found, start_lsn is returned
-// back.
-pub fn find_end_of_wal(
+/// Return offset of the last valid record in the segment segno, starting
+/// looking at start_offset. Returns start_offset if no records found.
+fn find_end_of_wal_segment(
     data_dir: &Path,
+    segno: XLogSegNo,
+    tli: TimeLineID,
     wal_seg_size: usize,
-    start_lsn: Lsn, // start reading WAL at this point; must point at record start_lsn.
-) -> anyhow::Result<Lsn> {
-    let mut result = start_lsn;
-    let mut curr_lsn = start_lsn;
+    start_offset: usize, // start reading at this point
+) -> anyhow::Result<u32> {
+    // step back to the beginning of the page to read it in...
+    let mut offs: usize = start_offset - start_offset % XLOG_BLCKSZ;
+    let mut skipping_first_contrecord: bool = false;
+    let mut contlen: usize = 0;
+    let mut xl_crc: u32 = 0;
+    let mut crc: u32 = 0;
+    let mut rec_offs: usize = 0;
     let mut buf = [0u8; XLOG_BLCKSZ];
-    let mut decoder = WalStreamDecoder::new(start_lsn);
+    let file_name = XLogFileName(tli, segno, wal_seg_size);
+    let mut last_valid_rec_pos: usize = start_offset; // assume at given start_offset begins new record
+    let mut file = File::open(data_dir.join(file_name.clone() + ".partial")).unwrap();
+    file.seek(SeekFrom::Start(offs as u64))?;
+    // xl_crc is the last field in XLogRecord, will not be read into rec_hdr
+    const_assert!(XLOG_RECORD_CRC_OFFS + 4 == XLOG_SIZE_OF_XLOG_RECORD);
+    let mut rec_hdr = [0u8; XLOG_RECORD_CRC_OFFS];
 
-    // loop over segments
-    loop {
-        let segno = curr_lsn.segment_number(wal_seg_size);
-        let seg_file_name = XLogFileName(PG_TLI, segno, wal_seg_size);
-        let seg_file_path = data_dir.join(seg_file_name);
-        match open_wal_segment(&seg_file_path)? {
-            None => {
-                // no more segments
-                info!(
-                    "find_end_of_wal reached end at {:?}, segment {:?} doesn't exist",
-                    result, seg_file_path
+    trace!("find_end_of_wal_segment(data_dir={}, segno={}, tli={}, wal_seg_size={}, start_offset=0x{:x})", data_dir.display(), segno, tli, wal_seg_size, start_offset);
+    while offs < wal_seg_size {
+        // we are at the beginning of the page; read it in
+        if offs % XLOG_BLCKSZ == 0 {
+            trace!("offs=0x{:x}: new page", offs);
+            let bytes_read = file.read(&mut buf)?;
+            if bytes_read != buf.len() {
+                bail!(
+                    "failed to read {} bytes from {} at {}",
+                    XLOG_BLCKSZ,
+                    file_name,
+                    offs
                 );
-                return Ok(result);
             }
-            Some(mut segment) => {
-                let seg_offs = curr_lsn.segment_offset(wal_seg_size);
-                segment.seek(SeekFrom::Start(seg_offs as u64))?;
-                // loop inside segment
-                loop {
-                    let bytes_read = segment.read(&mut buf)?;
-                    if bytes_read == 0 {
-                        break; // EOF
-                    }
-                    curr_lsn += bytes_read as u64;
-                    decoder.feed_bytes(&buf[0..bytes_read]);
 
-                    // advance result past all completely read records
-                    loop {
-                        match decoder.poll_decode() {
-                            Ok(Some(record)) => result = record.0,
-                            Err(e) => {
-                                info!(
-                                    "find_end_of_wal reached end at {:?}, decode error: {:?}",
-                                    result, e
-                                );
-                                return Ok(result);
-                            }
-                            Ok(None) => break, // need more data
-                        }
+            let xlp_magic = LittleEndian::read_u16(&buf[0..2]);
+            let xlp_info = LittleEndian::read_u16(&buf[2..4]);
+            let xlp_rem_len = LittleEndian::read_u32(&buf[XLP_REM_LEN_OFFS..XLP_REM_LEN_OFFS + 4]);
+            trace!(
+                "  xlp_magic=0x{:x}, xlp_info=0x{:x}, xlp_rem_len={}",
+                xlp_magic,
+                xlp_info,
+                xlp_rem_len
+            );
+            // this is expected in current usage when valid WAL starts after page header
+            if xlp_magic != XLOG_PAGE_MAGIC as u16 {
+                trace!(
+                    "  invalid WAL file {}.partial magic {} at {:?}",
+                    file_name,
+                    xlp_magic,
+                    Lsn(XLogSegNoOffsetToRecPtr(segno, offs as u32, wal_seg_size)),
+                );
+            }
+            if offs == 0 {
+                offs += XLOG_SIZE_OF_XLOG_LONG_PHD;
+                if (xlp_info & XLP_FIRST_IS_CONTRECORD) != 0 {
+                    trace!("  first record is contrecord");
+                    skipping_first_contrecord = true;
+                    contlen = xlp_rem_len as usize;
+                    if offs < start_offset {
+                        // Pre-condition failed: the beginning of the segment is unexpectedly corrupted.
+                        ensure!(start_offset - offs >= contlen,
+                            "start_offset is in the middle of the first record (which happens to be a contrecord), \
+                             expected to be on a record boundary. Is beginning of the segment corrupted?");
+                        contlen = 0;
+                        // keep skipping_first_contrecord to avoid counting the contrecord as valid, we did not check it.
                     }
+                } else {
+                    trace!("  first record is not contrecord");
+                }
+            } else {
+                offs += XLOG_SIZE_OF_XLOG_SHORT_PHD;
+            }
+            // ... and step forward again if asked
+            trace!("  skipped header to 0x{:x}", offs);
+            offs = max(offs, start_offset);
+        // beginning of the next record
+        } else if contlen == 0 {
+            let page_offs = offs % XLOG_BLCKSZ;
+            let xl_tot_len = LittleEndian::read_u32(&buf[page_offs..page_offs + 4]) as usize;
+            trace!("offs=0x{:x}: new record, xl_tot_len={}", offs, xl_tot_len);
+            if xl_tot_len == 0 {
+                info!(
+                    "find_end_of_wal_segment reached zeros at {:?}, last records ends at {:?}",
+                    Lsn(XLogSegNoOffsetToRecPtr(segno, offs as u32, wal_seg_size)),
+                    Lsn(XLogSegNoOffsetToRecPtr(
+                        segno,
+                        last_valid_rec_pos as u32,
+                        wal_seg_size
+                    ))
+                );
+                break; // zeros, reached the end
+            }
+            if skipping_first_contrecord {
+                skipping_first_contrecord = false;
+                trace!("  first contrecord has been just completed");
+            } else {
+                trace!(
+                    "  updating last_valid_rec_pos: 0x{:x} --> 0x{:x}",
+                    last_valid_rec_pos,
+                    offs
+                );
+                last_valid_rec_pos = offs;
+            }
+            offs += 4;
+            rec_offs = 4;
+            contlen = xl_tot_len - 4;
+            trace!(
+                "  reading rec_hdr[0..4] <-- [0x{:x}; 0x{:x})",
+                page_offs,
+                page_offs + 4
+            );
+            rec_hdr[0..4].copy_from_slice(&buf[page_offs..page_offs + 4]);
+        } else {
+            // we're continuing a record, possibly from previous page.
+            let page_offs = offs % XLOG_BLCKSZ;
+            let pageleft = XLOG_BLCKSZ - page_offs;
+
+            // read the rest of the record, or as much as fits on this page.
+            let n = min(contlen, pageleft);
+            trace!(
+                "offs=0x{:x}, record continuation, pageleft={}, contlen={}",
+                offs,
+                pageleft,
+                contlen
+            );
+            // fill rec_hdr header up to (but not including) xl_crc field
+            trace!(
+                "  rec_offs={}, XLOG_RECORD_CRC_OFFS={}, XLOG_SIZE_OF_XLOG_RECORD={}",
+                rec_offs,
+                XLOG_RECORD_CRC_OFFS,
+                XLOG_SIZE_OF_XLOG_RECORD
+            );
+            if rec_offs < XLOG_RECORD_CRC_OFFS {
+                let len = min(XLOG_RECORD_CRC_OFFS - rec_offs, n);
+                trace!(
+                    "  reading rec_hdr[{}..{}] <-- [0x{:x}; 0x{:x})",
+                    rec_offs,
+                    rec_offs + len,
+                    page_offs,
+                    page_offs + len
+                );
+                rec_hdr[rec_offs..rec_offs + len].copy_from_slice(&buf[page_offs..page_offs + len]);
+            }
+            if rec_offs <= XLOG_RECORD_CRC_OFFS && rec_offs + n >= XLOG_SIZE_OF_XLOG_RECORD {
+                let crc_offs = page_offs - rec_offs + XLOG_RECORD_CRC_OFFS;
+                // All records are aligned on 8-byte boundary, so their 8-byte frames
+                // cannot be split between pages. As xl_crc is the last field,
+                // its content is always on the same page.
+                const_assert!(XLOG_RECORD_CRC_OFFS % 8 == 4);
+                // We should always start reading aligned records even in incorrect WALs so if
+                // the condition is false it is likely a bug. However, it is localized somewhere
+                // in this function, hence we do not crash and just report failure instead.
+                ensure!(crc_offs % 8 == 4, "Record is not aligned properly (bug?)");
+                xl_crc = LittleEndian::read_u32(&buf[crc_offs..crc_offs + 4]);
+                trace!(
+                    "  reading xl_crc: [0x{:x}; 0x{:x}) = 0x{:x}",
+                    crc_offs,
+                    crc_offs + 4,
+                    xl_crc
+                );
+                crc = crc32c_append(0, &buf[crc_offs + 4..page_offs + n]);
+                trace!(
+                    "  initializing crc: [0x{:x}; 0x{:x}); crc = 0x{:x}",
+                    crc_offs + 4,
+                    page_offs + n,
+                    crc
+                );
+            } else if rec_offs > XLOG_RECORD_CRC_OFFS {
+                // As all records are 8-byte aligned, the header is already fully read and `crc` is initialized in the branch above.
+                ensure!(rec_offs >= XLOG_SIZE_OF_XLOG_RECORD);
+                let old_crc = crc;
+                crc = crc32c_append(crc, &buf[page_offs..page_offs + n]);
+                trace!(
+                    "  appending to crc: [0x{:x}; 0x{:x}); 0x{:x} --> 0x{:x}",
+                    page_offs,
+                    page_offs + n,
+                    old_crc,
+                    crc
+                );
+            } else {
+                // Correct because of the way conditions are written above.
+                assert!(rec_offs + n < XLOG_SIZE_OF_XLOG_RECORD);
+                // If `skipping_first_contrecord == true`, we may be reading from a middle of a record
+                // which started in the previous segment. Hence there is no point in validating the header.
+                if !skipping_first_contrecord && rec_offs + n > XLOG_RECORD_CRC_OFFS {
+                    info!(
+                        "Curiously corrupted WAL: a record stops inside the header; \
+                             offs=0x{:x}, record continuation, pageleft={}, contlen={}",
+                        offs, pageleft, contlen
+                    );
+                    break;
+                }
+                // Do nothing: we are still reading the header. It's accounted in CRC in the end of the record.
+            }
+            rec_offs += n;
+            offs += n;
+            contlen -= n;
+
+            if contlen == 0 {
+                trace!("  record completed at 0x{:x}", offs);
+                crc = crc32c_append(crc, &rec_hdr);
+                offs = (offs + 7) & !7; // pad on 8 bytes boundary */
+                trace!(
+                    "  padded offs to 0x{:x}, crc is {:x}, expected crc is {:x}",
+                    offs,
+                    crc,
+                    xl_crc
+                );
+                if skipping_first_contrecord {
+                    // do nothing, the flag will go down on next iteration when we're reading new record
+                    trace!("  first conrecord has been just completed");
+                } else if crc == xl_crc {
+                    // record is valid, advance the result to its end (with
+                    // alignment to the next record taken into account)
+                    trace!(
+                        "  updating last_valid_rec_pos: 0x{:x} --> 0x{:x}",
+                        last_valid_rec_pos,
+                        offs
+                    );
+                    last_valid_rec_pos = offs;
+                } else {
+                    info!(
+                        "CRC mismatch {} vs {} at {}",
+                        crc, xl_crc, last_valid_rec_pos
+                    );
+                    break;
                 }
             }
         }
     }
+    trace!("last_valid_rec_pos=0x{:x}", last_valid_rec_pos);
+    Ok(last_valid_rec_pos as u32)
 }
 
-// Open .partial or full WAL segment file, if present.
-fn open_wal_segment(seg_file_path: &Path) -> anyhow::Result<Option<File>> {
-    let mut partial_path = seg_file_path.to_owned();
-    partial_path.set_extension("partial");
-    match File::open(partial_path) {
-        Ok(file) => Ok(Some(file)),
-        Err(e) => match e.kind() {
-            ErrorKind::NotFound => {
-                // .partial not found, try full
-                match File::open(seg_file_path) {
-                    Ok(file) => Ok(Some(file)),
-                    Err(e) => match e.kind() {
-                        ErrorKind::NotFound => Ok(None),
-                        _ => Err(e.into()),
-                    },
-                }
-            }
-            _ => Err(e.into()),
-        },
+///
+/// Scan a directory that contains PostgreSQL WAL files, for the end of WAL.
+/// If precise, returns end LSN (next insertion point, basically);
+/// otherwise, start of the last segment.
+/// Returns (0, 0) if there is no WAL.
+///
+pub fn find_end_of_wal(
+    data_dir: &Path,
+    wal_seg_size: usize,
+    precise: bool,
+    start_lsn: Lsn, // start reading WAL at this point or later
+) -> anyhow::Result<(XLogRecPtr, TimeLineID)> {
+    let mut high_segno: XLogSegNo = 0;
+    let mut high_tli: TimeLineID = 0;
+    let mut high_ispartial = false;
+
+    for entry in fs::read_dir(data_dir).unwrap().flatten() {
+        let ispartial: bool;
+        let entry_name = entry.file_name();
+        let fname = entry_name.to_str().unwrap();
+        /*
+         * Check if the filename looks like an xlog file, or a .partial file.
+         */
+        if IsXLogFileName(fname) {
+            ispartial = false;
+        } else if IsPartialXLogFileName(fname) {
+            ispartial = true;
+        } else {
+            continue;
+        }
+        let (segno, tli) = XLogFromFileName(fname, wal_seg_size);
+        if !ispartial && entry.metadata().unwrap().len() != wal_seg_size as u64 {
+            continue;
+        }
+        if segno > high_segno
+            || (segno == high_segno && tli > high_tli)
+            || (segno == high_segno && tli == high_tli && high_ispartial && !ispartial)
+        {
+            high_segno = segno;
+            high_tli = tli;
+            high_ispartial = ispartial;
+        }
     }
+    if high_segno > 0 {
+        let mut high_offs = 0;
+        /*
+         * Move the starting pointer to the start of the next segment, if the
+         * highest one we saw was completed.
+         */
+        if !high_ispartial {
+            high_segno += 1;
+        } else if precise {
+            /* otherwise locate last record in last partial segment */
+            if start_lsn.segment_number(wal_seg_size) > high_segno {
+                bail!(
+                    "provided start_lsn {:?} is beyond highest segno {:?} available",
+                    start_lsn,
+                    high_segno,
+                );
+            }
+            let start_offset = if start_lsn.segment_number(wal_seg_size) == high_segno {
+                start_lsn.segment_offset(wal_seg_size)
+            } else {
+                0
+            };
+            high_offs = find_end_of_wal_segment(
+                data_dir,
+                high_segno,
+                high_tli,
+                wal_seg_size,
+                start_offset,
+            )?;
+        }
+        let high_ptr = XLogSegNoOffsetToRecPtr(high_segno, high_offs, wal_seg_size);
+        return Ok((high_ptr, high_tli));
+    }
+    Ok((0, 0))
 }
 
 pub fn main() {
     let mut data_dir = PathBuf::new();
     data_dir.push(".");
-    let wal_end = find_end_of_wal(&data_dir, WAL_SEGMENT_SIZE, Lsn(0)).unwrap();
-    println!("wal_end={:?}", wal_end);
+    let (wal_end, tli) = find_end_of_wal(&data_dir, WAL_SEGMENT_SIZE, true, Lsn(0)).unwrap();
+    println!(
+        "wal_end={:>08X}{:>08X}, tli={}",
+        (wal_end >> 32) as u32,
+        wal_end as u32,
+        tli
+    );
 }
 
 impl XLogRecord {
@@ -350,10 +592,7 @@ pub fn generate_wal_segment(segno: u64, system_id: u64) -> Result<Bytes, Seriali
 mod tests {
     use super::*;
     use regex::Regex;
-    use std::cmp::min;
-    use std::fs;
     use std::{env, str::FromStr};
-    use utils::const_assert;
 
     fn init_logging() {
         let _ = env_logger::Builder::from_env(
@@ -364,7 +603,10 @@ mod tests {
         .try_init();
     }
 
-    fn test_end_of_wal<C: wal_craft::Crafter>(test_name: &str) {
+    fn test_end_of_wal<C: wal_craft::Crafter>(
+        test_name: &str,
+        expected_end_of_wal_non_partial: Lsn,
+    ) {
         use wal_craft::*;
         // Craft some WAL
         let top_path = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
@@ -385,7 +627,7 @@ mod tests {
             .iter()
             .map(|&lsn| u64::from(lsn).into())
             .collect();
-        let expected_end_of_wal: Lsn = u64::from(expected_end_of_wal_partial).into();
+        let expected_end_of_wal_partial: Lsn = u64::from(expected_end_of_wal_partial).into();
         srv.kill();
 
         // Check find_end_of_wal on the initial WAL
@@ -397,10 +639,10 @@ mod tests {
             .filter(|fname| IsXLogFileName(fname))
             .max()
             .unwrap();
-        check_pg_waldump_end_of_wal(&cfg, &last_segment, expected_end_of_wal);
-        for start_lsn in intermediate_lsns
-            .iter()
-            .chain(std::iter::once(&expected_end_of_wal))
+        check_pg_waldump_end_of_wal(&cfg, &last_segment, expected_end_of_wal_partial);
+        for start_lsn in std::iter::once(Lsn(0))
+            .chain(intermediate_lsns)
+            .chain(std::iter::once(expected_end_of_wal_partial))
         {
             // Erase all WAL before `start_lsn` to ensure it's not used by `find_end_of_wal`.
             // We assume that `start_lsn` is non-decreasing.
@@ -415,7 +657,7 @@ mod tests {
                 }
                 let (segno, _) = XLogFromFileName(&fname, WAL_SEGMENT_SIZE);
                 let seg_start_lsn = XLogSegNoOffsetToRecPtr(segno, 0, WAL_SEGMENT_SIZE);
-                if seg_start_lsn > u64::from(*start_lsn) {
+                if seg_start_lsn > u64::from(start_lsn) {
                     continue;
                 }
                 let mut f = File::options().write(true).open(file.path()).unwrap();
@@ -423,12 +665,18 @@ mod tests {
                 f.write_all(
                     &ZEROS[0..min(
                         WAL_SEGMENT_SIZE,
-                        (u64::from(*start_lsn) - seg_start_lsn) as usize,
+                        (u64::from(start_lsn) - seg_start_lsn) as usize,
                     )],
                 )
                 .unwrap();
             }
-            check_end_of_wal(&cfg, &last_segment, *start_lsn, expected_end_of_wal);
+            check_end_of_wal(
+                &cfg,
+                &last_segment,
+                start_lsn,
+                expected_end_of_wal_non_partial,
+                expected_end_of_wal_partial,
+            );
         }
     }
 
@@ -465,15 +713,18 @@ mod tests {
         cfg: &wal_craft::Conf,
         last_segment: &str,
         start_lsn: Lsn,
-        expected_end_of_wal: Lsn,
+        expected_end_of_wal_non_partial: Lsn,
+        expected_end_of_wal_partial: Lsn,
     ) {
         // Check end_of_wal on non-partial WAL segment (we treat it as fully populated)
-        // let wal_end = find_end_of_wal(&cfg.wal_dir(), WAL_SEGMENT_SIZE, start_lsn).unwrap();
-        // info!(
-        //     "find_end_of_wal returned wal_end={} with non-partial WAL segment",
-        //     wal_end
-        // );
-        // assert_eq!(wal_end, expected_end_of_wal_non_partial);
+        let (wal_end, tli) =
+            find_end_of_wal(&cfg.wal_dir(), WAL_SEGMENT_SIZE, true, start_lsn).unwrap();
+        let wal_end = Lsn(wal_end);
+        info!(
+            "find_end_of_wal returned (wal_end={}, tli={}) with non-partial WAL segment",
+            wal_end, tli
+        );
+        assert_eq!(wal_end, expected_end_of_wal_non_partial);
 
         // Rename file to partial to actually find last valid lsn, then rename it back.
         fs::rename(
@@ -481,12 +732,14 @@ mod tests {
             cfg.wal_dir().join(format!("{}.partial", last_segment)),
         )
         .unwrap();
-        let wal_end = find_end_of_wal(&cfg.wal_dir(), WAL_SEGMENT_SIZE, start_lsn).unwrap();
+        let (wal_end, tli) =
+            find_end_of_wal(&cfg.wal_dir(), WAL_SEGMENT_SIZE, true, start_lsn).unwrap();
+        let wal_end = Lsn(wal_end);
         info!(
-            "find_end_of_wal returned wal_end={} with partial WAL segment",
-            wal_end
+            "find_end_of_wal returned (wal_end={}, tli={}) with partial WAL segment",
+            wal_end, tli
         );
-        assert_eq!(wal_end, expected_end_of_wal);
+        assert_eq!(wal_end, expected_end_of_wal_partial);
         fs::rename(
             cfg.wal_dir().join(format!("{}.partial", last_segment)),
             cfg.wal_dir().join(last_segment),
@@ -499,7 +752,10 @@ mod tests {
     #[test]
     pub fn test_find_end_of_wal_simple() {
         init_logging();
-        test_end_of_wal::<wal_craft::Simple>("test_find_end_of_wal_simple");
+        test_end_of_wal::<wal_craft::Simple>(
+            "test_find_end_of_wal_simple",
+            "0/2000000".parse::<Lsn>().unwrap(),
+        );
     }
 
     #[test]
@@ -507,14 +763,17 @@ mod tests {
         init_logging();
         test_end_of_wal::<wal_craft::WalRecordCrossingSegmentFollowedBySmallOne>(
             "test_find_end_of_wal_crossing_segment_followed_by_small_one",
+            "0/3000000".parse::<Lsn>().unwrap(),
         );
     }
 
     #[test]
+    #[ignore = "not yet fixed, needs correct parsing of pre-last segments"] // TODO
     pub fn test_find_end_of_wal_last_crossing_segment() {
         init_logging();
         test_end_of_wal::<wal_craft::LastWalRecordCrossingSegment>(
             "test_find_end_of_wal_last_crossing_segment",
+            "0/3000000".parse::<Lsn>().unwrap(),
         );
     }
 

libs/postgres_ffi/wal_craft/Cargo.toml

@@ -10,7 +10,7 @@ anyhow = "1.0"
 clap = "3.0"
 env_logger = "0.9"
 log = "0.4"
-once_cell = "1.13.0"
+once_cell = "1.8.0"
 postgres = { git = "https://github.com/zenithdb/rust-postgres.git", rev="d052ee8b86fff9897c77b0fe89ea9daba0e1fa38" }
 postgres_ffi = { path = "../" }
 tempfile = "3.2"

libs/remote_storage/Cargo.toml

@@ -7,7 +7,7 @@ edition = "2021"
 anyhow = { version = "1.0", features = ["backtrace"] }
 async-trait = "0.1"
 metrics = { version = "0.1", path = "../metrics" }
-once_cell = "1.13.0"
+once_cell = "1.8.0"
 rusoto_core = "0.48"
 rusoto_s3 = "0.48"
 serde = { version = "1.0", features = ["derive"] }

libs/remote_storage/src/lib.rs

@@ -66,9 +66,6 @@ pub trait RemoteStorage: Send + Sync {
     async fn list(&self) -> anyhow::Result<Vec<Self::RemoteObjectId>>;
 
     /// Lists all top level subdirectories for a given prefix
-    /// Note: here we assume that if the prefix is passed it was obtained via remote_object_id
-    /// which already takes into account any kind of global prefix (prefix_in_bucket for S3 or storage_root for LocalFS)
-    /// so this method doesnt need to.
     async fn list_prefixes(
         &self,
         prefix: Option<Self::RemoteObjectId>,

libs/remote_storage/src/local_fs.rs

@@ -116,7 +116,7 @@ impl RemoteStorage for LocalFs {
         prefix: Option<Self::RemoteObjectId>,
     ) -> anyhow::Result<Vec<Self::RemoteObjectId>> {
         let path = match prefix {
-            Some(prefix) => Cow::Owned(prefix),
+            Some(prefix) => Cow::Owned(self.storage_root.join(prefix)),
             None => Cow::Borrowed(&self.storage_root),
         };
         get_all_files(path.as_ref(), false).await

libs/remote_storage/src/s3_bucket.rs

@@ -171,25 +171,17 @@ impl S3Bucket {
 
         let access_key_id = std::env::var("AWS_ACCESS_KEY_ID").ok();
         let secret_access_key = std::env::var("AWS_SECRET_ACCESS_KEY").ok();
-        // session token is used when authorizing through sso
-        // which is typically the case when testing locally on developer machine
-        let session_token = std::env::var("AWS_SESSION_TOKEN").ok();
 
         let client = if access_key_id.is_none() && secret_access_key.is_none() {
             debug!("Using IAM-based AWS access");
             S3Client::new_with(request_dispatcher, InstanceMetadataProvider::new(), region)
         } else {
-            debug!(
-                "Using credentials-based AWS access. Session token is set: {}",
-                session_token.is_some()
-            );
+            debug!("Using credentials-based AWS access");
             S3Client::new_with(
                 request_dispatcher,
-                StaticProvider::new(
+                StaticProvider::new_minimal(
                     access_key_id.unwrap_or_default(),
                     secret_access_key.unwrap_or_default(),
-                    session_token,
-                    None,
                 ),
                 region,
             )
@@ -312,24 +304,32 @@ impl RemoteStorage for S3Bucket {
         Ok(document_keys)
     }
 
-    /// See the doc for `RemoteStorage::list_prefixes`
     /// Note: it wont include empty "directories"
     async fn list_prefixes(
         &self,
         prefix: Option<Self::RemoteObjectId>,
     ) -> anyhow::Result<Vec<Self::RemoteObjectId>> {
-        // get the passed prefix or if it is not set use prefix_in_bucket value
-        let list_prefix = prefix
-            .map(|p| p.0)
-            .or_else(|| self.prefix_in_bucket.clone())
-            .map(|mut p| {
+        let list_prefix = match prefix {
+            Some(prefix) => {
+                let mut prefix_in_bucket = self.prefix_in_bucket.clone().unwrap_or_default();
+                // if there is no trailing / in default prefix and
+                // supplied prefix does not start with "/" insert it
+                if !(prefix_in_bucket.ends_with(S3_PREFIX_SEPARATOR)
+                    || prefix.0.starts_with(S3_PREFIX_SEPARATOR))
+                {
+                    prefix_in_bucket.push(S3_PREFIX_SEPARATOR);
+                }
+
+                prefix_in_bucket.push_str(&prefix.0);
                 // required to end with a separator
                 // otherwise request will return only the entry of a prefix
-                if !p.ends_with(S3_PREFIX_SEPARATOR) {
-                    p.push(S3_PREFIX_SEPARATOR);
+                if !prefix_in_bucket.ends_with(S3_PREFIX_SEPARATOR) {
+                    prefix_in_bucket.push(S3_PREFIX_SEPARATOR);
                 }
-                p
-            });
+                Some(prefix_in_bucket)
+            }
+            None => self.prefix_in_bucket.clone(),
+        };
 
         let mut document_keys = Vec::new();
 

libs/utils/Cargo.toml

@@ -8,6 +8,7 @@ anyhow = "1.0"
 bincode = "1.3"
 bytes = "1.0.1"
 hyper = { version = "0.14.7", features = ["full"] }
+lazy_static = "1.4.0"
 pin-project-lite = "0.2.7"
 postgres = { git = "https://github.com/zenithdb/rust-postgres.git", rev="d052ee8b86fff9897c77b0fe89ea9daba0e1fa38" }
 postgres-protocol = { git = "https://github.com/zenithdb/rust-postgres.git", rev="d052ee8b86fff9897c77b0fe89ea9daba0e1fa38" }
@@ -27,8 +28,6 @@ rustls = "0.20.2"
 rustls-split = "0.3.0"
 git-version = "0.3.5"
 serde_with = "1.12.0"
-once_cell = "1.13.0"
-
 
 metrics = { path = "../metrics" }
 workspace_hack = { version = "0.1", path = "../../workspace_hack" }

libs/utils/src/http/endpoint.rs

@@ -4,8 +4,8 @@ use crate::zid::ZTenantId;
 use anyhow::anyhow;
 use hyper::header::AUTHORIZATION;
 use hyper::{header::CONTENT_TYPE, Body, Request, Response, Server};
+use lazy_static::lazy_static;
 use metrics::{register_int_counter, Encoder, IntCounter, TextEncoder};
-use once_cell::sync::Lazy;
 use routerify::ext::RequestExt;
 use routerify::RequestInfo;
 use routerify::{Middleware, Router, RouterBuilder, RouterService};
@@ -16,13 +16,13 @@ use std::net::TcpListener;
 
 use super::error::ApiError;
 
-static SERVE_METRICS_COUNT: Lazy<IntCounter> = Lazy::new(|| {
-    register_int_counter!(
+lazy_static! {
+    static ref SERVE_METRICS_COUNT: IntCounter = register_int_counter!(
         "libmetrics_metric_handler_requests_total",
         "Number of metric requests made"
     )
-    .expect("failed to define a metric")
-});
+    .expect("failed to define a metric");
+}
 
 async fn logger(res: Response<Body>, info: RequestInfo) -> Result<Response<Body>, ApiError> {
     info!("{} {} {}", info.method(), info.uri().path(), res.status(),);

libs/utils/src/pq_proto.rs

@@ -47,12 +47,10 @@ pub enum FeStartupPacket {
     StartupMessage {
         major_version: u32,
         minor_version: u32,
-        params: StartupMessageParams,
+        params: HashMap<String, String>,
     },
 }
 
-pub type StartupMessageParams = HashMap<String, String>;
-
 #[derive(Debug, Hash, PartialEq, Eq, Clone, Copy)]
 pub struct CancelKeyData {
     pub backend_pid: i32,

libs/utils/tests/ssl_test.rs

@@ -7,7 +7,7 @@ use std::{
 
 use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt};
 use bytes::{Buf, BufMut, Bytes, BytesMut};
-use once_cell::sync::Lazy;
+use lazy_static::lazy_static;
 
 use utils::postgres_backend::{AuthType, Handler, PostgresBackend};
 
@@ -19,15 +19,16 @@ fn make_tcp_pair() -> (TcpStream, TcpStream) {
     (server_stream, client_stream)
 }
 
-static KEY: Lazy<rustls::PrivateKey> = Lazy::new(|| {
-    let mut cursor = Cursor::new(include_bytes!("key.pem"));
-    rustls::PrivateKey(rustls_pemfile::rsa_private_keys(&mut cursor).unwrap()[0].clone())
-});
-
-static CERT: Lazy<rustls::Certificate> = Lazy::new(|| {
-    let mut cursor = Cursor::new(include_bytes!("cert.pem"));
-    rustls::Certificate(rustls_pemfile::certs(&mut cursor).unwrap()[0].clone())
-});
+lazy_static! {
+    static ref KEY: rustls::PrivateKey = {
+        let mut cursor = Cursor::new(include_bytes!("key.pem"));
+        rustls::PrivateKey(rustls_pemfile::rsa_private_keys(&mut cursor).unwrap()[0].clone())
+    };
+    static ref CERT: rustls::Certificate = {
+        let mut cursor = Cursor::new(include_bytes!("cert.pem"));
+        rustls::Certificate(rustls_pemfile::certs(&mut cursor).unwrap()[0].clone())
+    };
+}
 
 #[test]
 fn ssl() {

neon_local/Cargo.toml

@@ -15,5 +15,6 @@ git-version = "0.3.5"
 pageserver = { path = "../pageserver" }
 control_plane = { path = "../control_plane" }
 safekeeper = { path = "../safekeeper" }
+postgres_ffi = { path = "../libs/postgres_ffi" }
 utils = { path = "../libs/utils" }
 workspace_hack = { version = "0.1", path = "../workspace_hack" }

neon_local/src/main.rs

@@ -9,7 +9,6 @@ use pageserver::config::defaults::{
     DEFAULT_HTTP_LISTEN_ADDR as DEFAULT_PAGESERVER_HTTP_ADDR,
     DEFAULT_PG_LISTEN_ADDR as DEFAULT_PAGESERVER_PG_ADDR,
 };
-use pageserver::http::models::TimelineInfo;
 use safekeeper::defaults::{
     DEFAULT_HTTP_LISTEN_PORT as DEFAULT_SAFEKEEPER_HTTP_PORT,
     DEFAULT_PG_LISTEN_PORT as DEFAULT_SAFEKEEPER_PG_PORT,
@@ -26,6 +25,8 @@ use utils::{
     zid::{NodeId, ZTenantId, ZTenantTimelineId, ZTimelineId},
 };
 
+use pageserver::timelines::TimelineInfo;
+
 // Default id of a safekeeper node, if not specified on the command line.
 const DEFAULT_SAFEKEEPER_ID: NodeId = NodeId(1);
 const DEFAULT_PAGESERVER_ID: NodeId = NodeId(1);
@@ -884,7 +885,7 @@ fn handle_pageserver(sub_match: &ArgMatches, env: &local_env::LocalEnv) -> Resul
     match sub_match.subcommand() {
         Some(("start", start_match)) => {
             if let Err(e) = pageserver.start(&pageserver_config_overrides(start_match)) {
-                eprintln!("pageserver start failed: {e}");
+                eprintln!("pageserver start failed: {}", e);
                 exit(1);
             }
         }
@@ -906,19 +907,10 @@ fn handle_pageserver(sub_match: &ArgMatches, env: &local_env::LocalEnv) -> Resul
             }
 
             if let Err(e) = pageserver.start(&pageserver_config_overrides(restart_match)) {
-                eprintln!("pageserver start failed: {e}");
+                eprintln!("pageserver start failed: {}", e);
                 exit(1);
             }
         }
-
-        Some(("status", _)) => match PageServerNode::from_env(env).check_status() {
-            Ok(_) => println!("Page server is up and running"),
-            Err(err) => {
-                eprintln!("Page server is not available: {}", err);
-                exit(1);
-            }
-        },
-
         Some((sub_name, _)) => bail!("Unexpected pageserver subcommand '{}'", sub_name),
         None => bail!("no pageserver subcommand provided"),
     }

pageserver/Cargo.toml

@@ -21,6 +21,7 @@ futures = "0.3.13"
 hex = "0.4.3"
 hyper = "0.14"
 itertools = "0.10.3"
+lazy_static = "1.4.0"
 clap = "3.0"
 daemonize = "0.4.1"
 tokio = { version = "1.17", features = ["process", "sync", "macros", "fs", "rt", "io-util", "time"] }
@@ -28,6 +29,7 @@ postgres-types = { git = "https://github.com/zenithdb/rust-postgres.git", rev="d
 postgres-protocol = { git = "https://github.com/zenithdb/rust-postgres.git", rev="d052ee8b86fff9897c77b0fe89ea9daba0e1fa38" }
 postgres = { git = "https://github.com/zenithdb/rust-postgres.git", rev="d052ee8b86fff9897c77b0fe89ea9daba0e1fa38" }
 tokio-postgres = { git = "https://github.com/zenithdb/rust-postgres.git", rev="d052ee8b86fff9897c77b0fe89ea9daba0e1fa38" }
+tokio-stream = "0.1.8"
 anyhow = { version = "1.0", features = ["backtrace"] }
 crc32c = "0.6.0"
 thiserror = "1.0"
@@ -47,7 +49,7 @@ tracing = "0.1.27"
 signal-hook = "0.3.10"
 url = "2"
 nix = "0.23"
-once_cell = "1.13.0"
+once_cell = "1.8.0"
 crossbeam-utils = "0.8.5"
 fail = "0.5.0"
 git-version = "0.3.5"

pageserver/src/basebackup.rs

@@ -23,7 +23,8 @@ use tar::{Builder, EntryType, Header};
 use tracing::*;
 
 use crate::reltag::{RelTag, SlruKind};
-use crate::DatadirTimeline;
+use crate::repository::Timeline;
+use crate::DatadirTimelineImpl;
 use postgres_ffi::xlog_utils::*;
 use postgres_ffi::*;
 use utils::lsn::Lsn;
@@ -31,13 +32,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<DatadirTimelineImpl>,
     pub lsn: Lsn,
     prev_record_lsn: Lsn,
     full_backup: bool,
@@ -52,18 +52,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<DatadirTimelineImpl>,
         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
@@ -80,13 +79,13 @@ where
         let (backup_prev, backup_lsn) = if let Some(req_lsn) = req_lsn {
             // Backup was requested at a particular LSN. Wait for it to arrive.
             info!("waiting for {}", req_lsn);
-            timeline.wait_lsn(req_lsn)?;
+            timeline.tline.wait_lsn(req_lsn)?;
 
             // If the requested point is the end of the timeline, we can
             // provide prev_lsn. (get_last_record_rlsn() might return it as
             // zero, though, if no WAL has been generated on this timeline
             // yet.)
-            let end_of_timeline = timeline.get_last_record_rlsn();
+            let end_of_timeline = timeline.tline.get_last_record_rlsn();
             if req_lsn == end_of_timeline.last {
                 (end_of_timeline.prev, req_lsn)
             } else {
@@ -94,7 +93,7 @@ where
             }
         } else {
             // Backup was requested at end of the timeline.
-            let end_of_timeline = timeline.get_last_record_rlsn();
+            let end_of_timeline = timeline.tline.get_last_record_rlsn();
             (end_of_timeline.prev, end_of_timeline.last)
         };
 
@@ -372,7 +371,7 @@ where
         // add zenith.signal file
         let mut zenith_signal = String::new();
         if self.prev_record_lsn == Lsn(0) {
-            if self.lsn == self.timeline.get_ancestor_lsn() {
+            if self.lsn == self.timeline.tline.get_ancestor_lsn() {
                 write!(zenith_signal, "PREV LSN: none")?;
             } else {
                 write!(zenith_signal, "PREV LSN: invalid")?;
@@ -403,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/config.rs

@@ -59,7 +59,6 @@ pub mod defaults {
 
 # [tenant_config]
 #checkpoint_distance = {DEFAULT_CHECKPOINT_DISTANCE} # in bytes
-#checkpoint_timeout = {DEFAULT_CHECKPOINT_TIMEOUT}
 #compaction_target_size = {DEFAULT_COMPACTION_TARGET_SIZE} # in bytes
 #compaction_period = '{DEFAULT_COMPACTION_PERIOD}'
 #compaction_threshold = '{DEFAULT_COMPACTION_THRESHOLD}'
@@ -453,13 +452,6 @@ impl PageServerConf {
                 Some(parse_toml_u64("checkpoint_distance", checkpoint_distance)?);
         }
 
-        if let Some(checkpoint_timeout) = item.get("checkpoint_timeout") {
-            t_conf.checkpoint_timeout = Some(parse_toml_duration(
-                "checkpoint_timeout",
-                checkpoint_timeout,
-            )?);
-        }
-
         if let Some(compaction_target_size) = item.get("compaction_target_size") {
             t_conf.compaction_target_size = Some(parse_toml_u64(
                 "compaction_target_size",

pageserver/src/http/models.rs

@@ -7,10 +7,6 @@ use utils::{
     zid::{NodeId, ZTenantId, ZTimelineId},
 };
 
-// These enums are used in the API response fields.
-use crate::repository::LocalTimelineState;
-use crate::tenant_mgr::TenantState;
-
 #[serde_as]
 #[derive(Serialize, Deserialize)]
 pub struct TimelineCreateRequest {
@@ -32,7 +28,6 @@ pub struct TenantCreateRequest {
     #[serde_as(as = "Option<DisplayFromStr>")]
     pub new_tenant_id: Option<ZTenantId>,
     pub checkpoint_distance: Option<u64>,
-    pub checkpoint_timeout: Option<String>,
     pub compaction_target_size: Option<u64>,
     pub compaction_period: Option<String>,
     pub compaction_threshold: Option<usize>,
@@ -71,7 +66,6 @@ pub struct TenantConfigRequest {
     #[serde(default)]
     #[serde_as(as = "Option<DisplayFromStr>")]
     pub checkpoint_distance: Option<u64>,
-    pub checkpoint_timeout: Option<String>,
     pub compaction_target_size: Option<u64>,
     pub compaction_period: Option<String>,
     pub compaction_threshold: Option<usize>,
@@ -89,7 +83,6 @@ impl TenantConfigRequest {
         TenantConfigRequest {
             tenant_id,
             checkpoint_distance: None,
-            checkpoint_timeout: None,
             compaction_target_size: None,
             compaction_period: None,
             compaction_threshold: None,
@@ -104,59 +97,14 @@ impl TenantConfigRequest {
     }
 }
 
-#[serde_as]
-#[derive(Serialize, Deserialize, Clone)]
-pub struct TenantInfo {
-    #[serde_as(as = "DisplayFromStr")]
-    pub id: ZTenantId,
-    pub state: Option<TenantState>,
-    pub current_physical_size: Option<u64>, // physical size is only included in `tenant_status` endpoint
-    pub has_in_progress_downloads: Option<bool>,
-}
-
+/// A WAL receiver's data stored inside the global `WAL_RECEIVERS`.
+/// We keep one WAL receiver active per timeline.
 #[serde_as]
 #[derive(Debug, Serialize, Deserialize, Clone)]
-pub struct LocalTimelineInfo {
-    #[serde_as(as = "Option<DisplayFromStr>")]
-    pub ancestor_timeline_id: Option<ZTimelineId>,
-    #[serde_as(as = "Option<DisplayFromStr>")]
-    pub ancestor_lsn: Option<Lsn>,
-    #[serde_as(as = "DisplayFromStr")]
-    pub last_record_lsn: Lsn,
-    #[serde_as(as = "Option<DisplayFromStr>")]
-    pub prev_record_lsn: Option<Lsn>,
-    #[serde_as(as = "DisplayFromStr")]
-    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_physical_size_non_incremental: Option<u64>,
-    pub timeline_state: LocalTimelineState,
-
-    pub wal_source_connstr: Option<String>,
+pub struct WalReceiverEntry {
+    pub wal_producer_connstr: Option<String>,
     #[serde_as(as = "Option<DisplayFromStr>")]
     pub last_received_msg_lsn: Option<Lsn>,
     /// the timestamp (in microseconds) of the last received message
     pub last_received_msg_ts: Option<u128>,
 }
-
-#[serde_as]
-#[derive(Debug, Serialize, Deserialize, Clone)]
-pub struct RemoteTimelineInfo {
-    #[serde_as(as = "DisplayFromStr")]
-    pub remote_consistent_lsn: Lsn,
-    pub awaits_download: bool,
-}
-
-#[serde_as]
-#[derive(Debug, Serialize, Deserialize, Clone)]
-pub struct TimelineInfo {
-    #[serde_as(as = "DisplayFromStr")]
-    pub tenant_id: ZTenantId,
-    #[serde_as(as = "DisplayFromStr")]
-    pub timeline_id: ZTimelineId,
-    pub local: Option<LocalTimelineInfo>,
-    pub remote: Option<RemoteTimelineInfo>,
-}

pageserver/src/http/openapi_spec.yml

@@ -78,11 +78,6 @@ paths:
         schema:
           type: string
           description: Controls calculation of current_logical_size_non_incremental
-      - name: include-non-incremental-physical-size
-        in: query
-        schema:
-          type: string
-          description: Controls calculation of current_physical_size_non_incremental
     get:
       description: Get timelines for tenant
       responses:
@@ -141,11 +136,6 @@ paths:
           schema:
             type: string
           description: Controls calculation of current_logical_size_non_incremental
-        - name: include-non-incremental-physical-size
-          in: query
-          schema:
-            type: string
-            description: Controls calculation of current_physical_size_non_incremental
       responses:
         "200":
           description: TimelineInfo
@@ -207,6 +197,54 @@ paths:
               schema:
                 $ref: "#/components/schemas/Error"
 
+  /v1/tenant/{tenant_id}/timeline/{timeline_id}/wal_receiver:
+    parameters:
+      - name: tenant_id
+        in: path
+        required: true
+        schema:
+          type: string
+          format: hex
+      - name: timeline_id
+        in: path
+        required: true
+        schema:
+          type: string
+          format: hex
+    get:
+      description: Get wal receiver's data attached to the timeline
+      responses:
+        "200":
+          description: WalReceiverEntry
+          content:
+            application/json:
+              schema:
+                $ref: "#/components/schemas/WalReceiverEntry"
+        "401":
+          description: Unauthorized Error
+          content:
+            application/json:
+              schema:
+                $ref: "#/components/schemas/UnauthorizedError"
+        "403":
+          description: Forbidden Error
+          content:
+            application/json:
+              schema:
+                $ref: "#/components/schemas/ForbiddenError"
+        "404":
+          description: Error when no wal receiver is running or found
+          content:
+            application/json:
+              schema:
+                $ref: "#/components/schemas/NotFoundError"
+        "500":
+          description: Generic operation error
+          content:
+            application/json:
+              schema:
+                $ref: "#/components/schemas/Error"
+
   /v1/tenant/{tenant_id}/attach:
     parameters:
       - name: tenant_id
@@ -539,8 +577,6 @@ components:
           type: string
         state:
           type: string
-        current_physical_size:
-          type: integer
         has_in_progress_downloads:
           type: boolean
     TenantCreateInfo:
@@ -560,8 +596,6 @@ components:
           type: string
         checkpoint_distance:
           type: integer
-        checkpoint_timeout:
-          type: string
         compaction_period:
           type: string
         compaction_threshold:
@@ -580,8 +614,6 @@ components:
           type: string
         checkpoint_distance:
           type: integer
-        checkpoint_timeout:
-          type: string
         compaction_period:
           type: string
         compaction_threshold:
@@ -639,13 +671,18 @@ components:
           format: hex
         current_logical_size:
           type: integer
-        current_physical_size:
-          type: integer
         current_logical_size_non_incremental:
           type: integer
-        current_physical_size_non_incremental:
+
+    WalReceiverEntry:
+      type: object
+      required:
+        - thread_id
+        - wal_producer_connstr
+      properties:
+        thread_id:
           type: integer
-        wal_source_connstr:
+        wal_producer_connstr:
           type: string
         last_received_msg_lsn:
           type: string

pageserver/src/http/routes.rs

@@ -6,19 +6,16 @@ use hyper::{Body, Request, Response, Uri};
 use remote_storage::GenericRemoteStorage;
 use tracing::*;
 
-use super::models::{LocalTimelineInfo, RemoteTimelineInfo, TimelineInfo};
 use super::models::{
-    StatusResponse, TenantConfigRequest, TenantCreateRequest, TenantCreateResponse, TenantInfo,
+    StatusResponse, TenantConfigRequest, TenantCreateRequest, TenantCreateResponse,
     TimelineCreateRequest,
 };
-use crate::layered_repository::metadata::TimelineMetadata;
-use crate::pgdatadir_mapping::DatadirTimeline;
-use crate::repository::{LocalTimelineState, RepositoryTimeline};
-use crate::repository::{Repository, Timeline};
+use crate::repository::Repository;
 use crate::storage_sync;
 use crate::storage_sync::index::{RemoteIndex, RemoteTimeline};
 use crate::tenant_config::TenantConfOpt;
-use crate::TimelineImpl;
+use crate::tenant_mgr::TenantInfo;
+use crate::timelines::{LocalTimelineInfo, RemoteTimelineInfo, TimelineInfo};
 use crate::{config::PageServerConf, tenant_mgr, timelines};
 use utils::{
     auth::JwtAuth,
@@ -29,7 +26,6 @@ use utils::{
         request::parse_request_param,
         RequestExt, RouterBuilder,
     },
-    lsn::Lsn,
     zid::{ZTenantId, ZTenantTimelineId, ZTimelineId},
 };
 
@@ -83,123 +79,6 @@ fn get_config(request: &Request<Body>) -> &'static PageServerConf {
     get_state(request).conf
 }
 
-// Helper functions to construct a LocalTimelineInfo struct for a timeline
-
-fn local_timeline_info_from_loaded_timeline(
-    timeline: &TimelineImpl,
-    include_non_incremental_logical_size: bool,
-    include_non_incremental_physical_size: bool,
-) -> anyhow::Result<LocalTimelineInfo> {
-    let last_record_lsn = timeline.get_last_record_lsn();
-    let (wal_source_connstr, last_received_msg_lsn, last_received_msg_ts) = {
-        let guard = timeline.last_received_wal.lock().unwrap();
-        if let Some(info) = guard.as_ref() {
-            (
-                Some(info.wal_source_connstr.clone()),
-                Some(info.last_received_msg_lsn),
-                Some(info.last_received_msg_ts),
-            )
-        } else {
-            (None, None, None)
-        }
-    };
-
-    let info = LocalTimelineInfo {
-        ancestor_timeline_id: timeline.get_ancestor_timeline_id(),
-        ancestor_lsn: {
-            match timeline.get_ancestor_lsn() {
-                Lsn(0) => None,
-                lsn @ Lsn(_) => Some(lsn),
-            }
-        },
-        disk_consistent_lsn: timeline.get_disk_consistent_lsn(),
-        last_record_lsn,
-        prev_record_lsn: Some(timeline.get_prev_record_lsn()),
-        latest_gc_cutoff_lsn: *timeline.get_latest_gc_cutoff_lsn(),
-        timeline_state: LocalTimelineState::Loaded,
-        current_logical_size: Some(timeline.get_current_logical_size()),
-        current_physical_size: Some(timeline.get_physical_size()),
-        current_logical_size_non_incremental: if include_non_incremental_logical_size {
-            Some(timeline.get_current_logical_size_non_incremental(last_record_lsn)?)
-        } else {
-            None
-        },
-        current_physical_size_non_incremental: if include_non_incremental_physical_size {
-            Some(timeline.get_physical_size_non_incremental()?)
-        } else {
-            None
-        },
-        wal_source_connstr,
-        last_received_msg_lsn,
-        last_received_msg_ts,
-    };
-    Ok(info)
-}
-
-fn local_timeline_info_from_unloaded_timeline(metadata: &TimelineMetadata) -> LocalTimelineInfo {
-    LocalTimelineInfo {
-        ancestor_timeline_id: metadata.ancestor_timeline(),
-        ancestor_lsn: {
-            match metadata.ancestor_lsn() {
-                Lsn(0) => None,
-                lsn @ Lsn(_) => Some(lsn),
-            }
-        },
-        disk_consistent_lsn: metadata.disk_consistent_lsn(),
-        last_record_lsn: metadata.disk_consistent_lsn(),
-        prev_record_lsn: metadata.prev_record_lsn(),
-        latest_gc_cutoff_lsn: metadata.latest_gc_cutoff_lsn(),
-        timeline_state: LocalTimelineState::Unloaded,
-        current_logical_size: None,
-        current_physical_size: None,
-        current_logical_size_non_incremental: None,
-        current_physical_size_non_incremental: None,
-        wal_source_connstr: None,
-        last_received_msg_lsn: None,
-        last_received_msg_ts: None,
-    }
-}
-
-fn local_timeline_info_from_repo_timeline(
-    repo_timeline: &RepositoryTimeline<TimelineImpl>,
-    include_non_incremental_logical_size: bool,
-    include_non_incremental_physical_size: bool,
-) -> anyhow::Result<LocalTimelineInfo> {
-    match repo_timeline {
-        RepositoryTimeline::Loaded(timeline) => local_timeline_info_from_loaded_timeline(
-            &*timeline,
-            include_non_incremental_logical_size,
-            include_non_incremental_physical_size,
-        ),
-        RepositoryTimeline::Unloaded { metadata } => {
-            Ok(local_timeline_info_from_unloaded_timeline(metadata))
-        }
-    }
-}
-
-fn list_local_timelines(
-    tenant_id: ZTenantId,
-    include_non_incremental_logical_size: bool,
-    include_non_incremental_physical_size: bool,
-) -> Result<Vec<(ZTimelineId, LocalTimelineInfo)>> {
-    let repo = tenant_mgr::get_repository_for_tenant(tenant_id)
-        .with_context(|| format!("Failed to get repo for tenant {}", tenant_id))?;
-    let repo_timelines = repo.list_timelines();
-
-    let mut local_timeline_info = Vec::with_capacity(repo_timelines.len());
-    for (timeline_id, repository_timeline) in repo_timelines {
-        local_timeline_info.push((
-            timeline_id,
-            local_timeline_info_from_repo_timeline(
-                &repository_timeline,
-                include_non_incremental_logical_size,
-                include_non_incremental_physical_size,
-            )?,
-        ))
-    }
-    Ok(local_timeline_info)
-}
-
 // healthcheck handler
 async fn status_handler(request: Request<Body>) -> Result<Response<Body>, ApiError> {
     let config = get_config(&request);
@@ -214,30 +93,16 @@ async fn timeline_create_handler(mut request: Request<Body>) -> Result<Response<
 
     let new_timeline_info = tokio::task::spawn_blocking(move || {
         let _enter = info_span!("/timeline_create", tenant = %tenant_id, new_timeline = ?request_data.new_timeline_id, lsn=?request_data.ancestor_start_lsn).entered();
-
-        match timelines::create_timeline(
+        timelines::create_timeline(
             get_config(&request),
             tenant_id,
             request_data.new_timeline_id.map(ZTimelineId::from),
             request_data.ancestor_timeline_id.map(ZTimelineId::from),
             request_data.ancestor_start_lsn,
-        ) {
-            Ok(Some((new_timeline_id, new_timeline))) => {
-                // Created. Construct a TimelineInfo for it.
-                let local_info = local_timeline_info_from_loaded_timeline(new_timeline.as_ref(), false, false)?;
-                Ok(Some(TimelineInfo {
-                    tenant_id,
-                    timeline_id: new_timeline_id,
-                    local: Some(local_info),
-                    remote: None,
-                }))
-            }
-            Ok(None) => Ok(None), // timeline already exists
-            Err(err) => Err(err),
-        }
+        )
     })
     .await
-        .map_err(ApiError::from_err)??;
+    .map_err(ApiError::from_err)??;
 
     Ok(match new_timeline_info {
         Some(info) => json_response(StatusCode::CREATED, info)?,
@@ -248,17 +113,10 @@ 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");
+    let include_non_incremental_logical_size = get_include_non_incremental_logical_size(&request);
     let local_timeline_infos = tokio::task::spawn_blocking(move || {
         let _enter = info_span!("timeline_list", tenant = %tenant_id).entered();
-        list_local_timelines(
-            tenant_id,
-            include_non_incremental_logical_size,
-            include_non_incremental_physical_size,
-        )
+        crate::timelines::get_local_timelines(tenant_id, include_non_incremental_logical_size)
     })
     .await
     .map_err(ApiError::from_err)??;
@@ -287,15 +145,17 @@ async fn timeline_list_handler(request: Request<Body>) -> Result<Response<Body>,
     json_response(StatusCode::OK, response_data)
 }
 
-/// Checks if a query param is present in the request's URL
-fn query_param_present(request: &Request<Body>, param: &str) -> bool {
+// Gate non incremental logical size calculation behind a flag
+// after pgbench -i -s100 calculation took 28ms so if multiplied by the number of timelines
+// and tenants it can take noticeable amount of time. Also the value currently used only in tests
+fn get_include_non_incremental_logical_size(request: &Request<Body>) -> bool {
     request
         .uri()
         .query()
         .map(|v| {
             url::form_urlencoded::parse(v.as_bytes())
                 .into_owned()
-                .any(|(p, _)| p == param)
+                .any(|(param, _)| param == "include-non-incremental-logical-size")
         })
         .unwrap_or(false)
 }
@@ -305,10 +165,7 @@ async fn timeline_detail_handler(request: Request<Body>) -> Result<Response<Body
     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");
+    let include_non_incremental_logical_size = get_include_non_incremental_logical_size(&request);
 
     let (local_timeline_info, remote_timeline_info) = async {
         // any error here will render local timeline as None
@@ -319,10 +176,11 @@ async fn timeline_detail_handler(request: Request<Body>) -> Result<Response<Body
                 repo.get_timeline(timeline_id)
                     .as_ref()
                     .map(|timeline| {
-                        local_timeline_info_from_repo_timeline(
+                        LocalTimelineInfo::from_repo_timeline(
+                            tenant_id,
+                            timeline_id,
                             timeline,
                             include_non_incremental_logical_size,
-                            include_non_incremental_physical_size,
                         )
                     })
                     .transpose()?
@@ -367,6 +225,23 @@ async fn timeline_detail_handler(request: Request<Body>) -> Result<Response<Body
     json_response(StatusCode::OK, timeline_info)
 }
 
+async fn wal_receiver_get_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 wal_receiver_entry = crate::walreceiver::get_wal_receiver_entry(tenant_id, timeline_id)
+        .instrument(info_span!("wal_receiver_get", tenant = %tenant_id, timeline = %timeline_id))
+        .await
+        .ok_or_else(|| {
+            ApiError::NotFound(format!(
+                "WAL receiver data not found for tenant {tenant_id} and timeline {timeline_id}"
+            ))
+        })?;
+
+    json_response(StatusCode::OK, &wal_receiver_entry)
+}
+
 // TODO makes sense to provide tenant config right away the same way as it handled in tenant_create
 async fn tenant_attach_handler(request: Request<Body>) -> Result<Response<Body>, ApiError> {
     let tenant_id: ZTenantId = parse_request_param(&request, "tenant_id")?;
@@ -554,36 +429,14 @@ async fn tenant_status(request: Request<Body>) -> Result<Response<Body>, ApiErro
     let index_accessor = remote_index.read().await;
     let has_in_progress_downloads = index_accessor
         .tenant_entry(&tenant_id)
-        .map(|t| t.has_in_progress_downloads())
-        .unwrap_or_else(|| {
-            info!("Tenant {tenant_id} not found in remote index");
-            false
-        });
-
-    let current_physical_size =
-        match tokio::task::spawn_blocking(move || list_local_timelines(tenant_id, false, false))
-            .await
-            .map_err(ApiError::from_err)?
-        {
-            Err(err) => {
-                // Getting local timelines can fail when no local repo is on disk (e.g, when tenant data is being downloaded).
-                // In that case, put a warning message into log and operate normally.
-                warn!("Failed to get local timelines for tenant {tenant_id}: {err}");
-                None
-            }
-            Ok(local_timeline_infos) => Some(
-                local_timeline_infos
-                    .into_iter()
-                    .fold(0, |acc, x| acc + x.1.current_physical_size.unwrap()),
-            ),
-        };
+        .ok_or_else(|| ApiError::NotFound("Tenant not found in remote index".to_string()))?
+        .has_in_progress_downloads();
 
     json_response(
         StatusCode::OK,
         TenantInfo {
             id: tenant_id,
             state: tenant_state,
-            current_physical_size,
             has_in_progress_downloads: Some(has_in_progress_downloads),
         },
     )
@@ -623,11 +476,6 @@ async fn tenant_create_handler(mut request: Request<Body>) -> Result<Response<Bo
     }
 
     tenant_conf.checkpoint_distance = request_data.checkpoint_distance;
-    if let Some(checkpoint_timeout) = request_data.checkpoint_timeout {
-        tenant_conf.checkpoint_timeout =
-            Some(humantime::parse_duration(&checkpoint_timeout).map_err(ApiError::from_err)?);
-    }
-
     tenant_conf.compaction_target_size = request_data.compaction_target_size;
     tenant_conf.compaction_threshold = request_data.compaction_threshold;
 
@@ -688,10 +536,6 @@ async fn tenant_config_handler(mut request: Request<Body>) -> Result<Response<Bo
     }
 
     tenant_conf.checkpoint_distance = request_data.checkpoint_distance;
-    if let Some(checkpoint_timeout) = request_data.checkpoint_timeout {
-        tenant_conf.checkpoint_timeout =
-            Some(humantime::parse_duration(&checkpoint_timeout).map_err(ApiError::from_err)?);
-    }
     tenant_conf.compaction_target_size = request_data.compaction_target_size;
     tenant_conf.compaction_threshold = request_data.compaction_threshold;
 
@@ -762,5 +606,9 @@ pub fn make_router(
             "/v1/tenant/:tenant_id/timeline/:timeline_id/detach",
             timeline_delete_handler,
         )
+        .get(
+            "/v1/tenant/:tenant_id/timeline/:timeline_id/wal_receiver",
+            wal_receiver_get_handler,
+        )
         .any(handler_404))
 }

pageserver/src/import_datadir.rs

@@ -13,8 +13,9 @@ use walkdir::WalkDir;
 
 use crate::pgdatadir_mapping::*;
 use crate::reltag::{RelTag, SlruKind};
+use crate::repository::Repository;
+use crate::repository::Timeline;
 use crate::walingest::WalIngest;
-use crate::walrecord::DecodedWALRecord;
 use postgres_ffi::relfile_utils::*;
 use postgres_ffi::waldecoder::*;
 use postgres_ffi::xlog_utils::*;
@@ -28,9 +29,9 @@ use utils::lsn::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<R: Repository>(
     path: &Path,
-    tline: &T,
+    tline: &mut DatadirTimeline<R>,
     lsn: Lsn,
 ) -> Result<()> {
     let mut pg_control: Option<ControlFileData> = None;
@@ -88,8 +89,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<R: Repository, Reader: Read>(
+    modification: &mut DatadirModification<R>,
     path: &Path,
     spcoid: Oid,
     dboid: Oid,
@@ -168,8 +169,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<R: Repository, Reader: Read>(
+    modification: &mut DatadirModification<R>,
     slru: SlruKind,
     path: &Path,
     mut reader: Reader,
@@ -224,9 +225,9 @@ 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>(
+fn import_wal<R: Repository>(
     walpath: &Path,
-    tline: &T,
+    tline: &mut DatadirTimeline<R>,
     startpoint: Lsn,
     endpoint: Lsn,
 ) -> Result<()> {
@@ -267,11 +268,9 @@ fn import_wal<T: DatadirTimeline>(
         waldecoder.feed_bytes(&buf);
 
         let mut nrecords = 0;
-        let mut modification = tline.begin_modification(endpoint);
-        let mut decoded = DecodedWALRecord::default();
         while last_lsn <= endpoint {
             if let Some((lsn, recdata)) = waldecoder.poll_decode()? {
-                walingest.ingest_record(recdata, lsn, &mut modification, &mut decoded)?;
+                walingest.ingest_record(tline, recdata, lsn)?;
                 last_lsn = lsn;
 
                 nrecords += 1;
@@ -295,8 +294,8 @@ fn import_wal<T: DatadirTimeline>(
     Ok(())
 }
 
-pub fn import_basebackup_from_tar<T: DatadirTimeline, Reader: Read>(
-    tline: &T,
+pub fn import_basebackup_from_tar<R: Repository, Reader: Read>(
+    tline: &mut DatadirTimeline<R>,
     reader: Reader,
     base_lsn: Lsn,
 ) -> Result<()> {
@@ -337,8 +336,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<R: Repository, Reader: Read>(
+    tline: &mut DatadirTimeline<R>,
     reader: Reader,
     start_lsn: Lsn,
     end_lsn: Lsn,
@@ -385,11 +384,9 @@ pub fn import_wal_from_tar<T: DatadirTimeline, Reader: Read>(
 
         waldecoder.feed_bytes(&bytes[offset..]);
 
-        let mut modification = tline.begin_modification(end_lsn);
-        let mut decoded = DecodedWALRecord::default();
         while last_lsn <= end_lsn {
             if let Some((lsn, recdata)) = waldecoder.poll_decode()? {
-                walingest.ingest_record(recdata, lsn, &mut modification, &mut decoded)?;
+                walingest.ingest_record(tline, recdata, lsn)?;
                 last_lsn = lsn;
 
                 debug!("imported record at {} (end {})", lsn, end_lsn);
@@ -418,8 +415,8 @@ pub fn import_wal_from_tar<T: DatadirTimeline, Reader: Read>(
     Ok(())
 }
 
-pub fn import_file<T: DatadirTimeline, Reader: Read>(
-    modification: &mut DatadirModification<T>,
+pub fn import_file<R: Repository, Reader: Read>(
+    modification: &mut DatadirModification<R>,
     file_path: &Path,
     reader: Reader,
     len: usize,
@@ -538,7 +535,7 @@ pub fn import_file<T: DatadirTimeline, Reader: Read>(
         // zenith.signal is not necessarily the last file, that we handle
         // but it is ok to call `finish_write()`, because final `modification.commit()`
         // will update lsn once more to the final one.
-        let writer = modification.tline.writer();
+        let writer = modification.tline.tline.writer();
         writer.finish_write(prev_lsn);
 
         debug!("imported zenith signal {}", prev_lsn);

pageserver/src/layered_repository/block_io.rs

@@ -5,7 +5,7 @@
 use crate::page_cache;
 use crate::page_cache::{ReadBufResult, PAGE_SZ};
 use bytes::Bytes;
-use once_cell::sync::Lazy;
+use lazy_static::lazy_static;
 use std::ops::{Deref, DerefMut};
 use std::os::unix::fs::FileExt;
 use std::sync::atomic::AtomicU64;
@@ -117,7 +117,9 @@ where
     }
 }
 
-static NEXT_ID: Lazy<AtomicU64> = Lazy::new(|| AtomicU64::new(1));
+lazy_static! {
+    static ref NEXT_ID: AtomicU64 = AtomicU64::new(1);
+}
 
 /// An adapter for reading a (virtual) file using the page cache.
 ///

pageserver/src/layered_repository/delta_layer.rs

@@ -316,18 +316,6 @@ impl Layer for DeltaLayer {
         }
     }
 
-    fn key_iter<'a>(&'a self) -> Box<dyn Iterator<Item = (Key, Lsn, u64)> + 'a> {
-        let inner = match self.load() {
-            Ok(inner) => inner,
-            Err(e) => panic!("Failed to load a delta layer: {e:?}"),
-        };
-
-        match DeltaKeyIter::new(inner) {
-            Ok(iter) => Box::new(iter),
-            Err(e) => panic!("Layer index is corrupted: {e:?}"),
-        }
-    }
-
     fn delete(&self) -> Result<()> {
         // delete underlying file
         fs::remove_file(self.path())?;
@@ -672,21 +660,11 @@ impl DeltaLayerWriter {
     /// The values must be appended in key, lsn order.
     ///
     pub fn put_value(&mut self, key: Key, lsn: Lsn, val: Value) -> Result<()> {
-        self.put_value_bytes(key, lsn, &Value::ser(&val)?, val.will_init())
-    }
-
-    pub fn put_value_bytes(
-        &mut self,
-        key: Key,
-        lsn: Lsn,
-        val: &[u8],
-        will_init: bool,
-    ) -> Result<()> {
         assert!(self.lsn_range.start <= lsn);
 
-        let off = self.blob_writer.write_blob(val)?;
+        let off = self.blob_writer.write_blob(&Value::ser(&val)?)?;
 
-        let blob_ref = BlobRef::new(off, will_init);
+        let blob_ref = BlobRef::new(off, val.will_init());
 
         let delta_key = DeltaKey::from_key_lsn(&key, lsn);
         self.tree.append(&delta_key.0, blob_ref.0)?;
@@ -844,75 +822,3 @@ impl<'a> DeltaValueIter<'a> {
         }
     }
 }
-///
-/// Iterator over all keys stored in a delta layer
-///
-/// FIXME: This creates a Vector to hold all keys.
-/// That takes up quite a lot of memory. Should do this in a more streaming
-/// fashion.
-///
-struct DeltaKeyIter {
-    all_keys: Vec<(DeltaKey, u64)>,
-    next_idx: usize,
-}
-
-impl Iterator for DeltaKeyIter {
-    type Item = (Key, Lsn, u64);
-
-    fn next(&mut self) -> Option<Self::Item> {
-        if self.next_idx < self.all_keys.len() {
-            let (delta_key, size) = &self.all_keys[self.next_idx];
-
-            let key = delta_key.key();
-            let lsn = delta_key.lsn();
-
-            self.next_idx += 1;
-            Some((key, lsn, *size))
-        } else {
-            None
-        }
-    }
-}
-
-impl<'a> DeltaKeyIter {
-    fn new(inner: RwLockReadGuard<'a, DeltaLayerInner>) -> Result<Self> {
-        let file = inner.file.as_ref().unwrap();
-        let tree_reader = DiskBtreeReader::<_, DELTA_KEY_SIZE>::new(
-            inner.index_start_blk,
-            inner.index_root_blk,
-            file,
-        );
-
-        let mut all_keys: Vec<(DeltaKey, u64)> = Vec::new();
-        tree_reader.visit(
-            &[0u8; DELTA_KEY_SIZE],
-            VisitDirection::Forwards,
-            |key, value| {
-                let delta_key = DeltaKey::from_slice(key);
-                let pos = BlobRef(value).pos();
-                if let Some(last) = all_keys.last_mut() {
-                    if last.0.key() == delta_key.key() {
-                        return true;
-                    } else {
-                        // subtract offset of new key BLOB and first blob of this key
-                        // to get total size if values associated with this key
-                        let first_pos = last.1;
-                        last.1 = pos - first_pos;
-                    }
-                }
-                all_keys.push((delta_key, pos));
-                true
-            },
-        )?;
-        if let Some(last) = all_keys.last_mut() {
-            // Last key occupies all space till end of layer
-            last.1 = std::fs::metadata(&file.file.path)?.len() - last.1;
-        }
-        let iter = DeltaKeyIter {
-            all_keys,
-            next_idx: 0,
-        };
-
-        Ok(iter)
-    }
-}

pageserver/src/layered_repository/ephemeral_file.rs

@@ -8,7 +8,7 @@ use crate::page_cache;
 use crate::page_cache::PAGE_SZ;
 use crate::page_cache::{ReadBufResult, WriteBufResult};
 use crate::virtual_file::VirtualFile;
-use once_cell::sync::Lazy;
+use lazy_static::lazy_static;
 use std::cmp::min;
 use std::collections::HashMap;
 use std::fs::OpenOptions;
@@ -21,15 +21,15 @@ use utils::zid::{ZTenantId, ZTimelineId};
 
 use std::os::unix::fs::FileExt;
 
-///
-/// This is the global cache of file descriptors (File objects).
-///
-static EPHEMERAL_FILES: Lazy<RwLock<EphemeralFiles>> = Lazy::new(|| {
-    RwLock::new(EphemeralFiles {
+lazy_static! {
+    ///
+    /// This is the global cache of file descriptors (File objects).
+    ///
+    static ref EPHEMERAL_FILES: RwLock<EphemeralFiles> = RwLock::new(EphemeralFiles {
         next_file_id: 1,
         files: HashMap::new(),
-    })
-});
+    });
+}
 
 pub struct EphemeralFiles {
     next_file_id: u64,
@@ -43,7 +43,7 @@ pub struct EphemeralFile {
     _timelineid: ZTimelineId,
     file: Arc<VirtualFile>,
 
-    pub size: u64,
+    size: u64,
 }
 
 impl EphemeralFile {

pageserver/src/layered_repository/inmemory_layer.rs

@@ -15,7 +15,6 @@ use crate::layered_repository::storage_layer::{
 use crate::repository::{Key, Value};
 use crate::walrecord;
 use anyhow::{bail, ensure, Result};
-use std::cell::RefCell;
 use std::collections::HashMap;
 use tracing::*;
 use utils::{
@@ -31,12 +30,6 @@ use std::ops::Range;
 use std::path::PathBuf;
 use std::sync::RwLock;
 
-thread_local! {
-    /// A buffer for serializing object during [`InMemoryLayer::put_value`].
-    /// This buffer is reused for each serialization to avoid additional malloc calls.
-    static SER_BUFFER: RefCell<Vec<u8>> = RefCell::new(Vec::new());
-}
-
 pub struct InMemoryLayer {
     conf: &'static PageServerConf,
     tenantid: ZTenantId,
@@ -240,14 +233,6 @@ impl Layer for InMemoryLayer {
 }
 
 impl InMemoryLayer {
-    ///
-    /// Get layer size on the disk
-    ///
-    pub fn size(&self) -> Result<u64> {
-        let inner = self.inner.read().unwrap();
-        Ok(inner.file.size)
-    }
-
     ///
     /// Create a new, empty, in-memory layer
     ///
@@ -285,17 +270,10 @@ impl InMemoryLayer {
     pub fn put_value(&self, key: Key, lsn: Lsn, val: &Value) -> Result<()> {
         trace!("put_value key {} at {}/{}", key, self.timelineid, lsn);
         let mut inner = self.inner.write().unwrap();
+
         inner.assert_writeable();
 
-        let off = {
-            SER_BUFFER.with(|x| -> Result<_> {
-                let mut buf = x.borrow_mut();
-                buf.clear();
-                val.ser_into(&mut (*buf))?;
-                let off = inner.file.write_blob(&buf)?;
-                Ok(off)
-            })?
-        };
+        let off = inner.file.write_blob(&Value::ser(val)?)?;
 
         let vec_map = inner.index.entry(key).or_default();
         let old = vec_map.append_or_update_last(lsn, off).unwrap().0;
@@ -364,8 +342,8 @@ impl InMemoryLayer {
             // Write all page versions
             for (lsn, pos) in vec_map.as_slice() {
                 cursor.read_blob_into_buf(*pos, &mut buf)?;
-                let will_init = Value::des(&buf)?.will_init();
-                delta_layer_writer.put_value_bytes(key, *lsn, &buf, will_init)?;
+                let val = Value::des(&buf)?;
+                delta_layer_writer.put_value(key, *lsn, val)?;
             }
         }
 

pageserver/src/layered_repository/layer_map.rs

@@ -10,23 +10,24 @@
 //! corresponding files are written to disk.
 //!
 
-use crate::layered_repository::inmemory_layer::InMemoryLayer;
 use crate::layered_repository::storage_layer::Layer;
 use crate::layered_repository::storage_layer::{range_eq, range_overlaps};
+use crate::layered_repository::InMemoryLayer;
 use crate::repository::Key;
 use anyhow::Result;
+use lazy_static::lazy_static;
 use metrics::{register_int_gauge, IntGauge};
-use once_cell::sync::Lazy;
 use std::collections::VecDeque;
 use std::ops::Range;
 use std::sync::Arc;
 use tracing::*;
 use utils::lsn::Lsn;
 
-static NUM_ONDISK_LAYERS: Lazy<IntGauge> = Lazy::new(|| {
-    register_int_gauge!("pageserver_ondisk_layers", "Number of layers on-disk")
-        .expect("failed to define a metric")
-});
+lazy_static! {
+    static ref NUM_ONDISK_LAYERS: IntGauge =
+        register_int_gauge!("pageserver_ondisk_layers", "Number of layers on-disk")
+            .expect("failed to define a metric");
+}
 
 ///
 /// LayerMap tracks what layers exist on a timeline.

pageserver/src/layered_repository/storage_layer.rs

@@ -139,12 +139,6 @@ pub trait Layer: Send + Sync {
     /// Iterate through all keys and values stored in the layer
     fn iter(&self) -> Box<dyn Iterator<Item = Result<(Key, Lsn, Value)>> + '_>;
 
-    /// Iterate through all keys stored in the layer. Returns key, lsn and value size
-    /// It is used only for compaction and so is currently implemented only for DeltaLayer
-    fn key_iter(&self) -> Box<dyn Iterator<Item = (Key, Lsn, u64)> + '_> {
-        panic!("Not implemented")
-    }
-
     /// Permanently remove this layer from disk.
     fn delete(&self) -> Result<()>;
 

pageserver/src/lib.rs

@@ -22,7 +22,7 @@ pub mod walreceiver;
 pub mod walrecord;
 pub mod walredo;
 
-use once_cell::sync::Lazy;
+use lazy_static::lazy_static;
 use tracing::info;
 
 use crate::thread_mgr::ThreadKind;
@@ -42,14 +42,14 @@ pub const STORAGE_FORMAT_VERSION: u16 = 3;
 pub const IMAGE_FILE_MAGIC: u16 = 0x5A60;
 pub const DELTA_FILE_MAGIC: u16 = 0x5A61;
 
-static LIVE_CONNECTIONS_COUNT: Lazy<IntGaugeVec> = Lazy::new(|| {
-    register_int_gauge_vec!(
+lazy_static! {
+    static ref LIVE_CONNECTIONS_COUNT: IntGaugeVec = register_int_gauge_vec!(
         "pageserver_live_connections",
         "Number of live network connections",
         &["pageserver_connection_kind"]
     )
-    .expect("failed to define a metric")
-});
+    .expect("failed to define a metric");
+}
 
 pub const LOG_FILE_NAME: &str = "pageserver.log";
 
@@ -63,7 +63,8 @@ pub enum CheckpointConfig {
 }
 
 pub type RepositoryImpl = LayeredRepository;
-pub type TimelineImpl = <LayeredRepository as repository::Repository>::Timeline;
+
+pub type DatadirTimelineImpl = DatadirTimeline<RepositoryImpl>;
 
 pub fn shutdown_pageserver(exit_code: i32) {
     // Shut down the libpq endpoint thread. This prevents new connections from
@@ -93,56 +94,3 @@ pub fn shutdown_pageserver(exit_code: i32) {
     info!("Shut down successfully completed");
     std::process::exit(exit_code);
 }
-
-const DEFAULT_BASE_BACKOFF_SECONDS: f64 = 0.1;
-const DEFAULT_MAX_BACKOFF_SECONDS: f64 = 3.0;
-
-async fn exponential_backoff(n: u32, base_increment: f64, max_seconds: f64) {
-    let backoff_duration_seconds =
-        exponential_backoff_duration_seconds(n, base_increment, max_seconds);
-    if backoff_duration_seconds > 0.0 {
-        info!(
-            "Backoff: waiting {backoff_duration_seconds} seconds before processing with the task",
-        );
-        tokio::time::sleep(std::time::Duration::from_secs_f64(backoff_duration_seconds)).await;
-    }
-}
-
-fn exponential_backoff_duration_seconds(n: u32, base_increment: f64, max_seconds: f64) -> f64 {
-    if n == 0 {
-        0.0
-    } else {
-        (1.0 + base_increment).powf(f64::from(n)).min(max_seconds)
-    }
-}
-
-#[cfg(test)]
-mod backoff_defaults_tests {
-    use super::*;
-
-    #[test]
-    fn backoff_defaults_produce_growing_backoff_sequence() {
-        let mut current_backoff_value = None;
-
-        for i in 0..10_000 {
-            let new_backoff_value = exponential_backoff_duration_seconds(
-                i,
-                DEFAULT_BASE_BACKOFF_SECONDS,
-                DEFAULT_MAX_BACKOFF_SECONDS,
-            );
-
-            if let Some(old_backoff_value) = current_backoff_value.replace(new_backoff_value) {
-                assert!(
-                    old_backoff_value <= new_backoff_value,
-                    "{i}th backoff value {new_backoff_value} is smaller than the previous one {old_backoff_value}"
-                )
-            }
-        }
-
-        assert_eq!(
-            current_backoff_value.expect("Should have produced backoff values to compare"),
-            DEFAULT_MAX_BACKOFF_SECONDS,
-            "Given big enough of retries, backoff should reach its allowed max value"
-        );
-    }
-}

pageserver/src/page_service.rs

@@ -11,7 +11,7 @@
 
 use anyhow::{bail, ensure, Context, Result};
 use bytes::{Buf, BufMut, Bytes, BytesMut};
-use once_cell::sync::Lazy;
+use lazy_static::lazy_static;
 use regex::Regex;
 use std::io::{self, Read};
 use std::net::TcpListener;
@@ -30,6 +30,7 @@ use utils::{
 use crate::basebackup;
 use crate::config::{PageServerConf, ProfilingConfig};
 use crate::import_datadir::{import_basebackup_from_tar, import_wal_from_tar};
+use crate::layered_repository::LayeredRepository;
 use crate::pgdatadir_mapping::{DatadirTimeline, LsnForTimestamp};
 use crate::profiling::profpoint_start;
 use crate::reltag::RelTag;
@@ -434,15 +435,15 @@ const TIME_BUCKETS: &[f64] = &[
     0.1,  // 1/10 s
 ];
 
-static SMGR_QUERY_TIME: Lazy<HistogramVec> = Lazy::new(|| {
-    register_histogram_vec!(
+lazy_static! {
+    static ref SMGR_QUERY_TIME: HistogramVec = register_histogram_vec!(
         "pageserver_smgr_query_seconds",
         "Time spent on smgr query handling",
         &["smgr_query_type", "tenant_id", "timeline_id"],
         TIME_BUCKETS.into()
     )
-    .expect("failed to define a metric")
-});
+    .expect("failed to define a metric");
+}
 
 impl PageServerHandler {
     pub fn new(conf: &'static PageServerConf, auth: Option<Arc<JwtAuth>>) -> Self {
@@ -554,6 +555,9 @@ impl PageServerHandler {
         info!("creating new timeline");
         let repo = tenant_mgr::get_repository_for_tenant(tenant_id)?;
         let timeline = repo.create_empty_timeline(timeline_id, base_lsn)?;
+        let repartition_distance = repo.get_checkpoint_distance();
+        let mut datadir_timeline =
+            DatadirTimeline::<LayeredRepository>::new(timeline, repartition_distance);
 
         // TODO mark timeline as not ready until it reaches end_lsn.
         // We might have some wal to import as well, and we should prevent compute
@@ -569,7 +573,7 @@ impl PageServerHandler {
         info!("importing basebackup");
         pgb.write_message(&BeMessage::CopyInResponse)?;
         let reader = CopyInReader::new(pgb);
-        import_basebackup_from_tar(&*timeline, reader, base_lsn)?;
+        import_basebackup_from_tar(&mut datadir_timeline, reader, base_lsn)?;
 
         // TODO check checksum
         // Meanwhile you can verify client-side by taking fullbackup
@@ -579,7 +583,7 @@ impl PageServerHandler {
 
         // Flush data to disk, then upload to s3
         info!("flushing layers");
-        timeline.checkpoint(CheckpointConfig::Flush)?;
+        datadir_timeline.tline.checkpoint(CheckpointConfig::Flush)?;
 
         info!("done");
         Ok(())
@@ -601,6 +605,10 @@ impl PageServerHandler {
         let timeline = repo.get_timeline_load(timeline_id)?;
         ensure!(timeline.get_last_record_lsn() == start_lsn);
 
+        let repartition_distance = repo.get_checkpoint_distance();
+        let mut datadir_timeline =
+            DatadirTimeline::<LayeredRepository>::new(timeline, repartition_distance);
+
         // TODO leave clean state on error. For now you can use detach to clean
         // up broken state from a failed import.
 
@@ -608,16 +616,16 @@ impl PageServerHandler {
         info!("importing wal");
         pgb.write_message(&BeMessage::CopyInResponse)?;
         let reader = CopyInReader::new(pgb);
-        import_wal_from_tar(&*timeline, reader, start_lsn, end_lsn)?;
+        import_wal_from_tar(&mut datadir_timeline, reader, start_lsn, end_lsn)?;
 
         // TODO Does it make sense to overshoot?
-        ensure!(timeline.get_last_record_lsn() >= end_lsn);
+        ensure!(datadir_timeline.tline.get_last_record_lsn() >= end_lsn);
 
         // Flush data to disk, then upload to s3. No need for a forced checkpoint.
         // We only want to persist the data, and it doesn't matter if it's in the
         // shape of deltas or images.
         info!("flushing layers");
-        timeline.checkpoint(CheckpointConfig::Flush)?;
+        datadir_timeline.tline.checkpoint(CheckpointConfig::Flush)?;
 
         info!("done");
         Ok(())
@@ -635,8 +643,8 @@ 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<R: Repository>(
+        timeline: &DatadirTimeline<R>,
         mut lsn: Lsn,
         latest: bool,
         latest_gc_cutoff_lsn: &RwLockReadGuard<Lsn>,
@@ -663,7 +671,7 @@ impl PageServerHandler {
             if lsn <= last_record_lsn {
                 lsn = last_record_lsn;
             } else {
-                timeline.wait_lsn(lsn)?;
+                timeline.tline.wait_lsn(lsn)?;
                 // Since we waited for 'lsn' to arrive, that is now the last
                 // record LSN. (Or close enough for our purposes; the
                 // last-record LSN can advance immediately after we return
@@ -673,7 +681,7 @@ impl PageServerHandler {
             if lsn == Lsn(0) {
                 bail!("invalid LSN(0) in request");
             }
-            timeline.wait_lsn(lsn)?;
+            timeline.tline.wait_lsn(lsn)?;
         }
         ensure!(
             lsn >= **latest_gc_cutoff_lsn,
@@ -683,14 +691,14 @@ impl PageServerHandler {
         Ok(lsn)
     }
 
-    fn handle_get_rel_exists_request<T: DatadirTimeline>(
+    fn handle_get_rel_exists_request<R: Repository>(
         &self,
-        timeline: &T,
+        timeline: &DatadirTimeline<R>,
         req: &PagestreamExistsRequest,
     ) -> Result<PagestreamBeMessage> {
         let _enter = info_span!("get_rel_exists", rel = %req.rel, req_lsn = %req.lsn).entered();
 
-        let latest_gc_cutoff_lsn = timeline.get_latest_gc_cutoff_lsn();
+        let latest_gc_cutoff_lsn = timeline.tline.get_latest_gc_cutoff_lsn();
         let lsn = Self::wait_or_get_last_lsn(timeline, req.lsn, req.latest, &latest_gc_cutoff_lsn)?;
 
         let exists = timeline.get_rel_exists(req.rel, lsn)?;
@@ -700,13 +708,13 @@ impl PageServerHandler {
         }))
     }
 
-    fn handle_get_nblocks_request<T: DatadirTimeline>(
+    fn handle_get_nblocks_request<R: Repository>(
         &self,
-        timeline: &T,
+        timeline: &DatadirTimeline<R>,
         req: &PagestreamNblocksRequest,
     ) -> Result<PagestreamBeMessage> {
         let _enter = info_span!("get_nblocks", rel = %req.rel, req_lsn = %req.lsn).entered();
-        let latest_gc_cutoff_lsn = timeline.get_latest_gc_cutoff_lsn();
+        let latest_gc_cutoff_lsn = timeline.tline.get_latest_gc_cutoff_lsn();
         let lsn = Self::wait_or_get_last_lsn(timeline, req.lsn, req.latest, &latest_gc_cutoff_lsn)?;
 
         let n_blocks = timeline.get_rel_size(req.rel, lsn)?;
@@ -716,13 +724,13 @@ impl PageServerHandler {
         }))
     }
 
-    fn handle_db_size_request<T: DatadirTimeline>(
+    fn handle_db_size_request<R: Repository>(
         &self,
-        timeline: &T,
+        timeline: &DatadirTimeline<R>,
         req: &PagestreamDbSizeRequest,
     ) -> Result<PagestreamBeMessage> {
         let _enter = info_span!("get_db_size", dbnode = %req.dbnode, req_lsn = %req.lsn).entered();
-        let latest_gc_cutoff_lsn = timeline.get_latest_gc_cutoff_lsn();
+        let latest_gc_cutoff_lsn = timeline.tline.get_latest_gc_cutoff_lsn();
         let lsn = Self::wait_or_get_last_lsn(timeline, req.lsn, req.latest, &latest_gc_cutoff_lsn)?;
 
         let total_blocks =
@@ -735,14 +743,14 @@ impl PageServerHandler {
         }))
     }
 
-    fn handle_get_page_at_lsn_request<T: DatadirTimeline>(
+    fn handle_get_page_at_lsn_request<R: Repository>(
         &self,
-        timeline: &T,
+        timeline: &DatadirTimeline<R>,
         req: &PagestreamGetPageRequest,
     ) -> Result<PagestreamBeMessage> {
         let _enter = info_span!("get_page", rel = %req.rel, blkno = &req.blkno, req_lsn = %req.lsn)
             .entered();
-        let latest_gc_cutoff_lsn = timeline.get_latest_gc_cutoff_lsn();
+        let latest_gc_cutoff_lsn = timeline.tline.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
@@ -775,7 +783,7 @@ impl PageServerHandler {
         // check that the timeline exists
         let timeline = tenant_mgr::get_local_timeline_with_load(tenantid, timelineid)
             .context("Cannot load local timeline")?;
-        let latest_gc_cutoff_lsn = timeline.get_latest_gc_cutoff_lsn();
+        let latest_gc_cutoff_lsn = timeline.tline.get_latest_gc_cutoff_lsn();
         if let Some(lsn) = lsn {
             timeline
                 .check_lsn_is_in_scope(lsn, &latest_gc_cutoff_lsn)
@@ -913,7 +921,7 @@ impl postgres_backend::Handler for PageServerHandler {
             let timeline = tenant_mgr::get_local_timeline_with_load(tenantid, timelineid)
                 .context("Cannot load local timeline")?;
 
-            let end_of_timeline = timeline.get_last_record_rlsn();
+            let end_of_timeline = timeline.tline.get_last_record_rlsn();
 
             pgb.write_message_noflush(&BeMessage::RowDescription(&[
                 RowDescriptor::text_col(b"prev_lsn"),
@@ -1044,7 +1052,6 @@ impl postgres_backend::Handler for PageServerHandler {
             let repo = tenant_mgr::get_repository_for_tenant(tenantid)?;
             pgb.write_message_noflush(&BeMessage::RowDescription(&[
                 RowDescriptor::int8_col(b"checkpoint_distance"),
-                RowDescriptor::int8_col(b"checkpoint_timeout"),
                 RowDescriptor::int8_col(b"compaction_target_size"),
                 RowDescriptor::int8_col(b"compaction_period"),
                 RowDescriptor::int8_col(b"compaction_threshold"),
@@ -1055,12 +1062,6 @@ impl postgres_backend::Handler for PageServerHandler {
             ]))?
             .write_message_noflush(&BeMessage::DataRow(&[
                 Some(repo.get_checkpoint_distance().to_string().as_bytes()),
-                Some(
-                    repo.get_checkpoint_timeout()
-                        .as_secs()
-                        .to_string()
-                        .as_bytes(),
-                ),
                 Some(repo.get_compaction_target_size().to_string().as_bytes()),
                 Some(
                     repo.get_compaction_period()
@@ -1138,7 +1139,7 @@ impl postgres_backend::Handler for PageServerHandler {
             let timelineid = ZTimelineId::from_str(caps.get(2).unwrap().as_str())?;
             let timeline = tenant_mgr::get_local_timeline_with_load(tenantid, timelineid)
                 .context("Couldn't load timeline")?;
-            timeline.compact()?;
+            timeline.tline.compact()?;
 
             pgb.write_message_noflush(&SINGLE_COL_ROWDESC)?
                 .write_message_noflush(&BeMessage::CommandComplete(b"SELECT 1"))?;
@@ -1158,8 +1159,13 @@ impl postgres_backend::Handler for PageServerHandler {
             let timeline = tenant_mgr::get_local_timeline_with_load(tenantid, timelineid)
                 .context("Cannot load local timeline")?;
 
-            // Checkpoint the timeline and also compact it (due to `CheckpointConfig::Forced`).
-            timeline.checkpoint(CheckpointConfig::Forced)?;
+            timeline.tline.checkpoint(CheckpointConfig::Forced)?;
+
+            // Also compact it.
+            //
+            // FIXME: This probably shouldn't be part of a "checkpoint" command, but a
+            // separate operation. Update the tests if you change this.
+            timeline.tline.compact()?;
 
             pgb.write_message_noflush(&SINGLE_COL_ROWDESC)?
                 .write_message_noflush(&BeMessage::CommandComplete(b"SELECT 1"))?;

pageserver/src/pgdatadir_mapping.rs

@@ -6,10 +6,10 @@
 //! walingest.rs handles a few things like implicit relation creation and extension.
 //! Clarify that)
 //!
-use crate::keyspace::{KeySpace, KeySpaceAccum};
+use crate::keyspace::{KeyPartitioning, KeySpace, KeySpaceAccum};
 use crate::reltag::{RelTag, SlruKind};
-use crate::repository::Timeline;
 use crate::repository::*;
+use crate::repository::{Repository, Timeline};
 use crate::walrecord::ZenithWalRecord;
 use anyhow::{bail, ensure, Result};
 use bytes::{Buf, Bytes};
@@ -18,12 +18,34 @@ use postgres_ffi::{pg_constants, Oid, TransactionId};
 use serde::{Deserialize, Serialize};
 use std::collections::{HashMap, HashSet};
 use std::ops::Range;
-use tracing::{debug, trace, warn};
+use std::sync::atomic::{AtomicIsize, Ordering};
+use std::sync::{Arc, Mutex, RwLockReadGuard};
+use tracing::{debug, error, trace, warn};
 use utils::{bin_ser::BeSer, lsn::Lsn};
 
 /// Block number within a relation or SLRU. This matches PostgreSQL's BlockNumber type.
 pub type BlockNumber = u32;
 
+pub struct DatadirTimeline<R>
+where
+    R: Repository,
+{
+    /// The underlying key-value store. Callers should not read or modify the
+    /// data in the underlying store directly. However, it is exposed to have
+    /// access to information like last-LSN, ancestor, and operations like
+    /// compaction.
+    pub tline: Arc<R::Timeline>,
+
+    /// When did we last calculate the partitioning?
+    partitioning: Mutex<(KeyPartitioning, Lsn)>,
+
+    /// Configuration: how often should the partitioning be recalculated.
+    repartition_threshold: u64,
+
+    /// Current logical size of the "datadir", at the last LSN.
+    current_logical_size: AtomicIsize,
+}
+
 #[derive(Debug)]
 pub enum LsnForTimestamp {
     Present(Lsn),
@@ -32,29 +54,33 @@ pub enum LsnForTimestamp {
     NoData(Lsn),
 }
 
-///
-/// This trait 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.
-///
-/// 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 {
+impl<R: Repository> DatadirTimeline<R> {
+    pub fn new(tline: Arc<R::Timeline>, repartition_threshold: u64) -> Self {
+        DatadirTimeline {
+            tline,
+            partitioning: Mutex::new((KeyPartitioning::new(), Lsn(0))),
+            current_logical_size: AtomicIsize::new(0),
+            repartition_threshold,
+        }
+    }
+
+    /// (Re-)calculate the logical size of the database at the latest LSN.
+    ///
+    /// This can be a slow operation.
+    pub fn init_logical_size(&self) -> Result<()> {
+        let last_lsn = self.tline.get_last_record_lsn();
+        self.current_logical_size.store(
+            self.get_current_logical_size_non_incremental(last_lsn)? as isize,
+            Ordering::SeqCst,
+        );
+        Ok(())
+    }
+
     /// Start ingesting a WAL record, or other atomic modification of
     /// the timeline.
     ///
     /// This provides a transaction-like interface to perform a bunch
-    /// of modifications atomically.
+    /// of modifications atomically, all stamped with one LSN.
     ///
     /// To ingest a WAL record, call begin_modification(lsn) to get a
     /// DatadirModification object. Use the functions in the object to
@@ -62,27 +88,18 @@ pub trait DatadirTimeline: Timeline {
     /// that the WAL record affects. When you're done, call commit() to
     /// commit the changes.
     ///
-    /// Lsn stored in modification is advanced by `ingest_record` and
-    /// is used by `commit()` to update `last_record_lsn`.
-    ///
-    /// Calling commit() will flush all the changes and reset the state,
-    /// so the `DatadirModification` struct can be reused to perform the next modification.
-    ///
     /// Note that any pending modifications you make through the
     /// modification object won't be visible to calls to the 'get' and list
     /// 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>
-    where
-        Self: Sized,
-    {
+    pub fn begin_modification(&self, lsn: Lsn) -> DatadirModification<R> {
         DatadirModification {
             tline: self,
+            lsn,
             pending_updates: HashMap::new(),
             pending_deletions: Vec::new(),
             pending_nblocks: 0,
-            lsn,
         }
     }
 
@@ -91,7 +108,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)?;
@@ -104,11 +121,11 @@ pub trait DatadirTimeline: Timeline {
         }
 
         let key = rel_block_to_key(tag, blknum);
-        self.get(key, lsn)
+        self.tline.get(key, lsn)
     }
 
     // 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)?;
@@ -121,13 +138,9 @@ 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) {
-            return Ok(nblocks);
-        }
-
         if (tag.forknum == pg_constants::FSM_FORKNUM
             || tag.forknum == pg_constants::VISIBILITYMAP_FORKNUM)
             && !self.get_rel_exists(tag, lsn)?
@@ -140,25 +153,17 @@ pub trait DatadirTimeline: Timeline {
         }
 
         let key = rel_size_to_key(tag);
-        let mut buf = self.get(key, lsn)?;
-        let nblocks = buf.get_u32_le();
-
-        // Update relation size cache
-        self.update_cached_rel_size(tag, lsn, nblocks);
-        Ok(nblocks)
+        let mut buf = self.tline.get(key, lsn)?;
+        Ok(buf.get_u32_le())
     }
 
     /// 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
-        if let Some(_nblocks) = self.get_cached_rel_size(&tag, lsn) {
-            return Ok(true);
-        }
         // fetch directory listing
         let key = rel_dir_to_key(tag.spcnode, tag.dbnode);
-        let buf = self.get(key, lsn)?;
+        let buf = self.tline.get(key, lsn)?;
         let dir = RelDirectory::des(&buf)?;
 
         let exists = dir.rels.get(&(tag.relnode, tag.forknum)).is_some();
@@ -167,10 +172,10 @@ 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)?;
+        let buf = self.tline.get(key, lsn)?;
         let dir = RelDirectory::des(&buf)?;
 
         let rels: HashSet<RelTag> =
@@ -185,7 +190,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,
@@ -193,21 +198,26 @@ pub trait DatadirTimeline: Timeline {
         lsn: Lsn,
     ) -> Result<Bytes> {
         let key = slru_block_to_key(kind, segno, blknum);
-        self.get(key, lsn)
+        self.tline.get(key, lsn)
     }
 
     /// 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)?;
+        let mut buf = self.tline.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)?;
+        let buf = self.tline.get(key, lsn)?;
         let dir = SlruSegmentDirectory::des(&buf)?;
 
         let exists = dir.segments.get(&segno).is_some();
@@ -221,10 +231,10 @@ 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> {
-        let gc_cutoff_lsn_guard = self.get_latest_gc_cutoff_lsn();
+    pub fn find_lsn_for_timestamp(&self, search_timestamp: TimestampTz) -> Result<LsnForTimestamp> {
+        let gc_cutoff_lsn_guard = self.tline.get_latest_gc_cutoff_lsn();
         let min_lsn = *gc_cutoff_lsn_guard;
-        let max_lsn = self.get_last_record_lsn();
+        let max_lsn = self.tline.get_last_record_lsn();
 
         // LSNs are always 8-byte aligned. low/mid/high represent the
         // LSN divided by 8.
@@ -315,51 +325,88 @@ 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);
 
-        let buf = self.get(key, lsn)?;
+        let buf = self.tline.get(key, lsn)?;
         let dir = SlruSegmentDirectory::des(&buf)?;
 
         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)?;
+        let buf = self.tline.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 buf = self.tline.get(DBDIR_KEY, lsn)?;
         let dir = DbDirectory::des(&buf)?;
 
         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)?;
+        let buf = self.tline.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 buf = self.tline.get(TWOPHASEDIR_KEY, lsn)?;
         let dir = TwoPhaseDirectory::des(&buf)?;
 
         Ok(dir.xids)
     }
 
-    fn get_control_file(&self, lsn: Lsn) -> Result<Bytes> {
-        self.get(CONTROLFILE_KEY, lsn)
+    pub fn get_control_file(&self, lsn: Lsn) -> Result<Bytes> {
+        self.tline.get(CONTROLFILE_KEY, lsn)
     }
 
-    fn get_checkpoint(&self, lsn: Lsn) -> Result<Bytes> {
-        self.get(CHECKPOINT_KEY, lsn)
+    pub fn get_checkpoint(&self, lsn: Lsn) -> Result<Bytes> {
+        self.tline.get(CHECKPOINT_KEY, lsn)
+    }
+
+    /// Get the LSN of the last ingested WAL record.
+    ///
+    /// This is just a convenience wrapper that calls through to the underlying
+    /// repository.
+    pub fn get_last_record_lsn(&self) -> Lsn {
+        self.tline.get_last_record_lsn()
+    }
+
+    /// Check that it is valid to request operations with that lsn.
+    ///
+    /// This is just a convenience wrapper that calls through to the underlying
+    /// repository.
+    pub fn check_lsn_is_in_scope(
+        &self,
+        lsn: Lsn,
+        latest_gc_cutoff_lsn: &RwLockReadGuard<Lsn>,
+    ) -> Result<()> {
+        self.tline.check_lsn_is_in_scope(lsn, latest_gc_cutoff_lsn)
+    }
+
+    /// Retrieve current logical size of the timeline
+    ///
+    /// NOTE: counted incrementally, includes ancestors,
+    pub fn get_current_logical_size(&self) -> usize {
+        let current_logical_size = self.current_logical_size.load(Ordering::Acquire);
+        match usize::try_from(current_logical_size) {
+            Ok(sz) => sz,
+            Err(_) => {
+                error!(
+                    "current_logical_size is out of range: {}",
+                    current_logical_size
+                );
+                0
+            }
+        }
     }
 
     /// Does the same as get_current_logical_size but counted on demand.
@@ -367,16 +414,16 @@ 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<usize> {
         // Fetch list of database dirs and iterate them
-        let buf = self.get(DBDIR_KEY, lsn)?;
+        let buf = self.tline.get(DBDIR_KEY, lsn)?;
         let dbdir = DbDirectory::des(&buf)?;
 
         let mut total_size: usize = 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 mut buf = self.tline.get(relsize_key, lsn)?;
                 let relsize = buf.get_u32_le();
 
                 total_size += relsize as usize;
@@ -397,7 +444,7 @@ pub trait DatadirTimeline: Timeline {
         result.add_key(DBDIR_KEY);
 
         // Fetch list of database dirs and iterate them
-        let buf = self.get(DBDIR_KEY, lsn)?;
+        let buf = self.tline.get(DBDIR_KEY, lsn)?;
         let dbdir = DbDirectory::des(&buf)?;
 
         let mut dbs: Vec<(Oid, Oid)> = dbdir.dbdirs.keys().cloned().collect();
@@ -414,7 +461,7 @@ pub trait DatadirTimeline: Timeline {
             rels.sort_unstable();
             for rel in rels {
                 let relsize_key = rel_size_to_key(rel);
-                let mut buf = self.get(relsize_key, lsn)?;
+                let mut buf = self.tline.get(relsize_key, lsn)?;
                 let relsize = buf.get_u32_le();
 
                 result.add_range(rel_block_to_key(rel, 0)..rel_block_to_key(rel, relsize));
@@ -430,13 +477,13 @@ pub trait DatadirTimeline: Timeline {
         ] {
             let slrudir_key = slru_dir_to_key(kind);
             result.add_key(slrudir_key);
-            let buf = self.get(slrudir_key, lsn)?;
+            let buf = self.tline.get(slrudir_key, lsn)?;
             let dir = SlruSegmentDirectory::des(&buf)?;
             let mut segments: Vec<u32> = dir.segments.iter().cloned().collect();
             segments.sort_unstable();
             for segno in segments {
                 let segsize_key = slru_segment_size_to_key(kind, segno);
-                let mut buf = self.get(segsize_key, lsn)?;
+                let mut buf = self.tline.get(segsize_key, lsn)?;
                 let segsize = buf.get_u32_le();
 
                 result.add_range(
@@ -448,7 +495,7 @@ pub trait DatadirTimeline: Timeline {
 
         // Then pg_twophase
         result.add_key(TWOPHASEDIR_KEY);
-        let buf = self.get(TWOPHASEDIR_KEY, lsn)?;
+        let buf = self.tline.get(TWOPHASEDIR_KEY, lsn)?;
         let twophase_dir = TwoPhaseDirectory::des(&buf)?;
         let mut xids: Vec<TransactionId> = twophase_dir.xids.iter().cloned().collect();
         xids.sort_unstable();
@@ -462,31 +509,31 @@ pub trait DatadirTimeline: Timeline {
         Ok(result.to_keyspace())
     }
 
-    /// Get cached size of relation if it not updated after specified LSN
-    fn get_cached_rel_size(&self, tag: &RelTag, lsn: Lsn) -> Option<BlockNumber>;
-
-    /// Update cached relation size if there is no more recent update
-    fn update_cached_rel_size(&self, tag: RelTag, lsn: Lsn, nblocks: BlockNumber);
-
-    /// Store cached relation size
-    fn set_cached_rel_size(&self, tag: RelTag, lsn: Lsn, nblocks: BlockNumber);
-
-    /// Remove cached relation size
-    fn remove_cached_rel_size(&self, tag: &RelTag);
+    pub fn repartition(&self, lsn: Lsn, partition_size: u64) -> Result<(KeyPartitioning, Lsn)> {
+        let mut partitioning_guard = self.partitioning.lock().unwrap();
+        if partitioning_guard.1 == Lsn(0)
+            || lsn.0 - partitioning_guard.1 .0 > self.repartition_threshold
+        {
+            let keyspace = self.collect_keyspace(lsn)?;
+            let partitioning = keyspace.partition(partition_size);
+            *partitioning_guard = (partitioning, lsn);
+            return Ok((partitioning_guard.0.clone(), lsn));
+        }
+        Ok((partitioning_guard.0.clone(), partitioning_guard.1))
+    }
 }
 
 /// 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, R: Repository> {
     /// 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 DatadirTimeline<R>,
 
-    /// Lsn assigned by begin_modification
-    pub lsn: Lsn,
+    lsn: Lsn,
 
     // The modifications are not applied directly to the underlying key-value store.
     // The put-functions add the modifications here, and they are flushed to the
@@ -496,7 +543,7 @@ pub struct DatadirModification<'a, T: DatadirTimeline> {
     pending_nblocks: isize,
 }
 
-impl<'a, T: DatadirTimeline> DatadirModification<'a, T> {
+impl<'a, R: Repository> DatadirModification<'a, R> {
     /// Initialize a completely new repository.
     ///
     /// This inserts the directory metadata entries that are assumed to
@@ -697,36 +744,26 @@ impl<'a, T: DatadirTimeline> DatadirModification<'a, T> {
 
         self.pending_nblocks += nblocks as isize;
 
-        // Update relation size cache
-        self.tline.set_cached_rel_size(rel, self.lsn, nblocks);
-
         // Even if nblocks > 0, we don't insert any actual blocks here. That's up to the
         // caller.
+
         Ok(())
     }
 
     /// Truncate relation
     pub fn put_rel_truncation(&mut self, rel: RelTag, nblocks: BlockNumber) -> Result<()> {
         ensure!(rel.relnode != 0, "invalid relnode");
-        let last_lsn = self.tline.get_last_record_lsn();
-        if self.tline.get_rel_exists(rel, last_lsn)? {
-            let size_key = rel_size_to_key(rel);
-            // Fetch the old size first
-            let old_size = self.get(size_key)?.get_u32_le();
+        let size_key = rel_size_to_key(rel);
 
-            // Update the entry with the new size.
-            let buf = nblocks.to_le_bytes();
-            self.put(size_key, Value::Image(Bytes::from(buf.to_vec())));
+        // Fetch the old size first
+        let old_size = self.get(size_key)?.get_u32_le();
 
-            // Update relation size cache
-            self.tline.set_cached_rel_size(rel, self.lsn, nblocks);
+        // Update the entry with the new size.
+        let buf = nblocks.to_le_bytes();
+        self.put(size_key, Value::Image(Bytes::from(buf.to_vec())));
 
-            // Update relation size cache
-            self.tline.set_cached_rel_size(rel, self.lsn, nblocks);
-
-            // Update logical database size.
-            self.pending_nblocks -= old_size as isize - nblocks as isize;
-        }
+        // Update logical database size.
+        self.pending_nblocks -= old_size as isize - nblocks as isize;
         Ok(())
     }
 
@@ -744,9 +781,6 @@ impl<'a, T: DatadirTimeline> DatadirModification<'a, T> {
             let buf = nblocks.to_le_bytes();
             self.put(size_key, Value::Image(Bytes::from(buf.to_vec())));
 
-            // Update relation size cache
-            self.tline.set_cached_rel_size(rel, self.lsn, nblocks);
-
             self.pending_nblocks += nblocks as isize - old_size as isize;
         }
         Ok(())
@@ -772,9 +806,6 @@ impl<'a, T: DatadirTimeline> DatadirModification<'a, T> {
         let old_size = self.get(size_key)?.get_u32_le();
         self.pending_nblocks -= old_size as isize;
 
-        // Remove enty from relation size cache
-        self.tline.remove_cached_rel_size(&rel);
-
         // Delete size entry, as well as all blocks
         self.delete(rel_key_range(rel));
 
@@ -897,7 +928,7 @@ impl<'a, T: DatadirTimeline> DatadirModification<'a, T> {
             return Ok(());
         }
 
-        let writer = self.tline.writer();
+        let writer = self.tline.tline.writer();
 
         // Flush relation and  SLRU data blocks, keep metadata.
         let mut result: Result<()> = Ok(());
@@ -912,7 +943,10 @@ impl<'a, T: DatadirTimeline> DatadirModification<'a, T> {
         result?;
 
         if pending_nblocks != 0 {
-            writer.update_current_logical_size(pending_nblocks * pg_constants::BLCKSZ as isize);
+            self.tline.current_logical_size.fetch_add(
+                pending_nblocks * pg_constants::BLCKSZ as isize,
+                Ordering::SeqCst,
+            );
             self.pending_nblocks = 0;
         }
 
@@ -922,25 +956,26 @@ impl<'a, T: DatadirTimeline> DatadirModification<'a, T> {
     ///
     /// Finish this atomic update, writing all the updated keys to the
     /// underlying timeline.
-    /// All the modifications in this atomic update are stamped by the specified LSN.
     ///
-    pub fn commit(&mut self) -> Result<()> {
-        let writer = self.tline.writer();
-        let lsn = self.lsn;
+    pub fn commit(self) -> Result<()> {
+        let writer = self.tline.tline.writer();
+
         let pending_nblocks = self.pending_nblocks;
-        self.pending_nblocks = 0;
 
-        for (key, value) in self.pending_updates.drain() {
-            writer.put(key, lsn, &value)?;
+        for (key, value) in self.pending_updates {
+            writer.put(key, self.lsn, &value)?;
         }
-        for key_range in self.pending_deletions.drain(..) {
-            writer.delete(key_range, lsn)?;
+        for key_range in self.pending_deletions {
+            writer.delete(key_range.clone(), self.lsn)?;
         }
 
-        writer.finish_write(lsn);
+        writer.finish_write(self.lsn);
 
         if pending_nblocks != 0 {
-            writer.update_current_logical_size(pending_nblocks * pg_constants::BLCKSZ as isize);
+            self.tline.current_logical_size.fetch_add(
+                pending_nblocks * pg_constants::BLCKSZ as isize,
+                Ordering::SeqCst,
+            );
         }
 
         Ok(())
@@ -966,8 +1001,8 @@ impl<'a, T: DatadirTimeline> DatadirModification<'a, T> {
                 bail!("unexpected pending WAL record");
             }
         } else {
-            let lsn = Lsn::max(self.tline.get_last_record_lsn(), self.lsn);
-            self.tline.get(key, lsn)
+            let last_lsn = self.tline.get_last_record_lsn();
+            self.tline.tline.get(key, last_lsn)
         }
     }
 
@@ -1369,12 +1404,13 @@ fn is_slru_block_key(key: Key) -> bool {
 pub fn create_test_timeline<R: Repository>(
     repo: R,
     timeline_id: utils::zid::ZTimelineId,
-) -> Result<std::sync::Arc<R::Timeline>> {
+) -> Result<Arc<crate::DatadirTimeline<R>>> {
     let tline = repo.create_empty_timeline(timeline_id, Lsn(8))?;
+    let tline = DatadirTimeline::new(tline, 256 * 1024);
     let mut m = tline.begin_modification(Lsn(8));
     m.init_empty()?;
     m.commit()?;
-    Ok(tline)
+    Ok(Arc::new(tline))
 }
 
 #[allow(clippy::bool_assert_comparison)]
@@ -1447,7 +1483,7 @@ mod tests {
             .contains(&TESTREL_A));
 
         // Run checkpoint and garbage collection and check that it's still not visible
-        newtline.checkpoint(CheckpointConfig::Forced)?;
+        newtline.tline.checkpoint(CheckpointConfig::Forced)?;
         repo.gc_iteration(Some(NEW_TIMELINE_ID), 0, true)?;
 
         assert!(!newtline

pageserver/src/repository.rs

@@ -185,7 +185,7 @@ 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;
+    type Timeline: Timeline;
 
     /// Updates timeline based on the `TimelineSyncStatusUpdate`, received from the remote storage synchronization.
     /// See [`crate::remote_storage`] for more details about the synchronization.
@@ -277,6 +277,15 @@ pub enum LocalTimelineState {
     Unloaded,
 }
 
+impl<'a, T> From<&'a RepositoryTimeline<T>> for LocalTimelineState {
+    fn from(local_timeline_entry: &'a RepositoryTimeline<T>) -> Self {
+        match local_timeline_entry {
+            RepositoryTimeline::Loaded(_) => LocalTimelineState::Loaded,
+            RepositoryTimeline::Unloaded { .. } => LocalTimelineState::Unloaded,
+        }
+    }
+}
+
 ///
 /// Result of performing GC
 ///
@@ -373,11 +382,6 @@ pub trait Timeline: Send + Sync {
         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.
@@ -401,14 +405,12 @@ pub trait TimelineWriter<'a> {
     /// 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 lazy_static::lazy_static;
     use std::sync::{Arc, RwLock, RwLockReadGuard, RwLockWriteGuard};
     use std::{fs, path::PathBuf};
 
@@ -439,13 +441,14 @@ pub mod repo_harness {
         buf.freeze()
     }
 
-    static LOCK: Lazy<RwLock<()>> = Lazy::new(|| RwLock::new(()));
+    lazy_static! {
+        static ref LOCK: RwLock<()> = 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),
@@ -588,10 +591,11 @@ mod tests {
     //use std::sync::Arc;
     use bytes::BytesMut;
     use hex_literal::hex;
-    use once_cell::sync::Lazy;
+    use lazy_static::lazy_static;
 
-    static TEST_KEY: Lazy<Key> =
-        Lazy::new(|| Key::from_slice(&hex!("112222222233333333444444445500000001")));
+    lazy_static! {
+        static ref TEST_KEY: Key = Key::from_slice(&hex!("112222222233333333444444445500000001"));
+    }
 
     #[test]
     fn test_basic() -> Result<()> {

pageserver/src/storage_sync.rs

@@ -155,7 +155,8 @@ use std::{
 
 use anyhow::{anyhow, bail, Context};
 use futures::stream::{FuturesUnordered, StreamExt};
-use once_cell::sync::{Lazy, OnceCell};
+use lazy_static::lazy_static;
+use once_cell::sync::OnceCell;
 use remote_storage::{GenericRemoteStorage, RemoteStorage};
 use tokio::{
     fs,
@@ -172,10 +173,10 @@ use self::{
 };
 use crate::{
     config::PageServerConf,
-    exponential_backoff,
     layered_repository::{
         ephemeral_file::is_ephemeral_file,
         metadata::{metadata_path, TimelineMetadata, METADATA_FILE_NAME},
+        LayeredRepository,
     },
     storage_sync::{self, index::RemoteIndex},
     tenant_mgr::attach_downloaded_tenants,
@@ -184,8 +185,8 @@ use crate::{
 };
 
 use metrics::{
-    register_histogram_vec, register_int_counter_vec, register_int_gauge, HistogramVec,
-    IntCounterVec, IntGauge,
+    register_histogram_vec, register_int_counter, register_int_counter_vec, register_int_gauge,
+    HistogramVec, IntCounter, IntCounterVec, IntGauge,
 };
 use utils::zid::{ZTenantId, ZTenantTimelineId, ZTimelineId};
 
@@ -193,33 +194,32 @@ use self::download::download_index_parts;
 pub use self::download::gather_tenant_timelines_index_parts;
 pub use self::download::TEMP_DOWNLOAD_EXTENSION;
 
-static REMAINING_SYNC_ITEMS: Lazy<IntGauge> = Lazy::new(|| {
-    register_int_gauge!(
+lazy_static! {
+    static ref REMAINING_SYNC_ITEMS: IntGauge = register_int_gauge!(
         "pageserver_remote_storage_remaining_sync_items",
         "Number of storage sync items left in the queue"
     )
-    .expect("failed to register pageserver remote storage remaining sync items int gauge")
-});
-
-static IMAGE_SYNC_TIME: Lazy<HistogramVec> = Lazy::new(|| {
-    register_histogram_vec!(
+    .expect("failed to register pageserver remote storage remaining sync items int gauge");
+    static ref FATAL_TASK_FAILURES: IntCounter = register_int_counter!(
+        "pageserver_remote_storage_fatal_task_failures_total",
+        "Number of critically failed tasks"
+    )
+    .expect("failed to register pageserver remote storage remaining sync items int gauge");
+    static ref IMAGE_SYNC_TIME: HistogramVec = register_histogram_vec!(
         "pageserver_remote_storage_image_sync_seconds",
         "Time took to synchronize (download or upload) a whole pageserver image. \
         Grouped by tenant and timeline ids, `operation_kind` (upload|download) and `status` (success|failure)",
         &["tenant_id", "timeline_id", "operation_kind", "status"],
         vec![0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 3.0, 10.0, 20.0]
     )
-    .expect("failed to register pageserver image sync time histogram vec")
-});
-
-static REMOTE_INDEX_UPLOAD: Lazy<IntCounterVec> = Lazy::new(|| {
-    register_int_counter_vec!(
+    .expect("failed to register pageserver image sync time histogram vec");
+    static ref REMOTE_INDEX_UPLOAD: IntCounterVec = register_int_counter_vec!(
         "pageserver_remote_storage_remote_index_uploads_total",
         "Number of remote index uploads",
         &["tenant_id", "timeline_id"],
     )
-    .expect("failed to register pageserver remote index upload vec")
-});
+    .expect("failed to register pageserver remote index upload vec");
+}
 
 static SYNC_QUEUE: OnceCell<SyncQueue> = OnceCell::new();
 
@@ -970,19 +970,14 @@ fn storage_sync_loop<P, S>(
     }
 }
 
+// needed to check whether the download happened
+// more informative than just a bool
 #[derive(Debug)]
-enum DownloadStatus {
+enum DownloadMarker {
     Downloaded,
     Nothing,
 }
 
-#[derive(Debug)]
-enum UploadStatus {
-    Uploaded,
-    Failed(anyhow::Error),
-    Nothing,
-}
-
 async fn process_batches<P, S>(
     conf: &'static PageServerConf,
     max_sync_errors: NonZeroU32,
@@ -1022,7 +1017,7 @@ where
             "Finished storage sync task for sync id {sync_id} download marker {:?}",
             download_marker
         );
-        if matches!(download_marker, DownloadStatus::Downloaded) {
+        if matches!(download_marker, DownloadMarker::Downloaded) {
             downloaded_timelines.insert(sync_id.tenant_id);
         }
     }
@@ -1036,7 +1031,7 @@ async fn process_sync_task_batch<P, S>(
     max_sync_errors: NonZeroU32,
     sync_id: ZTenantTimelineId,
     batch: SyncTaskBatch,
-) -> DownloadStatus
+) -> DownloadMarker
 where
     P: Debug + Send + Sync + 'static,
     S: RemoteStorage<RemoteObjectId = P> + Send + Sync + 'static,
@@ -1053,110 +1048,36 @@ where
     // When operating in a system without tasks failing over the error threshold,
     // current batching and task processing systems aim to update the layer set and metadata files (remote and local),
     // without "losing" such layer files.
-    let (upload_status, download_status) = tokio::join!(
+    let (upload_result, status_update) = tokio::join!(
         async {
             if let Some(upload_data) = upload_data {
-                let upload_retries = upload_data.retries;
-                match validate_task_retries(upload_retries, max_sync_errors)
+                match validate_task_retries(upload_data, max_sync_errors)
                     .instrument(info_span!("retries_validation"))
                     .await
                 {
-                    ControlFlow::Continue(()) => {
+                    ControlFlow::Continue(new_upload_data) => {
                         upload_timeline_data(
                             conf,
                             (storage.as_ref(), &index, sync_queue),
                             current_remote_timeline.as_ref(),
                             sync_id,
-                            upload_data,
+                            new_upload_data,
                             sync_start,
                             "upload",
                         )
-                        .await
-                    }
-                    ControlFlow::Break(()) => match update_remote_data(
-                        conf,
-                        storage.as_ref(),
-                        &index,
-                        sync_id,
-                        RemoteDataUpdate::Upload {
-                            uploaded_data: upload_data.data,
-                            upload_failed: true,
-                        },
-                    )
-                    .await
-                    {
-                        Ok(()) => UploadStatus::Failed(anyhow::anyhow!(
-                            "Aborted after retries validation, current retries: {upload_retries}, max retries allowed: {max_sync_errors}"
-                        )),
-                        Err(e) => {
-                            error!("Failed to update remote timeline {sync_id}: {e:?}");
-                            UploadStatus::Failed(e)
-                        }
-                    },
-                }
-            } else {
-                UploadStatus::Nothing
-            }
-        }
-        .instrument(info_span!("upload_timeline_data")),
-        async {
-            if let Some(download_data) = download_data {
-                match validate_task_retries(download_data.retries, max_sync_errors)
-                    .instrument(info_span!("retries_validation"))
-                    .await
-                {
-                    ControlFlow::Continue(()) => {
-                        return download_timeline_data(
-                            conf,
-                            (storage.as_ref(), &index, sync_queue),
-                            current_remote_timeline.as_ref(),
-                            sync_id,
-                            download_data,
-                            sync_start,
-                            "download",
-                        )
                         .await;
+                        return Some(());
                     }
-                    ControlFlow::Break(()) => {
-                        index
-                            .write()
-                            .await
-                            .set_awaits_download(&sync_id, false)
-                            .ok();
-                    }
-                }
-            }
-            DownloadStatus::Nothing
-        }
-        .instrument(info_span!("download_timeline_data")),
-    );
-
-    if let Some(delete_data) = batch.delete {
-        match upload_status {
-            UploadStatus::Uploaded | UploadStatus::Nothing => {
-                match validate_task_retries(delete_data.retries, max_sync_errors)
-                    .instrument(info_span!("retries_validation"))
-                    .await
-                {
-                    ControlFlow::Continue(()) => {
-                        delete_timeline_data(
-                            conf,
-                            (storage.as_ref(), &index, sync_queue),
-                            sync_id,
-                            delete_data,
-                            sync_start,
-                            "delete",
-                        )
-                        .instrument(info_span!("delete_timeline_data"))
-                        .await;
-                    }
-                    ControlFlow::Break(()) => {
+                    ControlFlow::Break(failed_upload_data) => {
                         if let Err(e) = update_remote_data(
                             conf,
                             storage.as_ref(),
                             &index,
                             sync_id,
-                            RemoteDataUpdate::Delete(&delete_data.data.deleted_layers),
+                            RemoteDataUpdate::Upload {
+                                uploaded_data: failed_upload_data.data,
+                                upload_failed: true,
+                            },
                         )
                         .await
                         {
@@ -1165,14 +1086,80 @@ where
                     }
                 }
             }
-            UploadStatus::Failed(e) => {
-                warn!("Skipping delete task due to failed upload tasks, reenqueuing. Upload data: {:?}, delete data: {delete_data:?}. Upload failure: {e:#}", batch.upload);
-                sync_queue.push(sync_id, SyncTask::Delete(delete_data));
+            None
+        }
+        .instrument(info_span!("upload_timeline_data")),
+        async {
+            if let Some(download_data) = download_data {
+                match validate_task_retries(download_data, max_sync_errors)
+                    .instrument(info_span!("retries_validation"))
+                    .await
+                {
+                    ControlFlow::Continue(new_download_data) => {
+                        return download_timeline_data(
+                            conf,
+                            (storage.as_ref(), &index, sync_queue),
+                            current_remote_timeline.as_ref(),
+                            sync_id,
+                            new_download_data,
+                            sync_start,
+                            "download",
+                        )
+                        .await;
+                    }
+                    ControlFlow::Break(_) => {
+                        index
+                            .write()
+                            .await
+                            .set_awaits_download(&sync_id, false)
+                            .ok();
+                    }
+                }
             }
+            DownloadMarker::Nothing
+        }
+        .instrument(info_span!("download_timeline_data")),
+    );
+
+    if let Some(delete_data) = batch.delete {
+        if upload_result.is_some() {
+            match validate_task_retries(delete_data, max_sync_errors)
+                .instrument(info_span!("retries_validation"))
+                .await
+            {
+                ControlFlow::Continue(new_delete_data) => {
+                    delete_timeline_data(
+                        conf,
+                        (storage.as_ref(), &index, sync_queue),
+                        sync_id,
+                        new_delete_data,
+                        sync_start,
+                        "delete",
+                    )
+                    .instrument(info_span!("delete_timeline_data"))
+                    .await;
+                }
+                ControlFlow::Break(failed_delete_data) => {
+                    if let Err(e) = update_remote_data(
+                        conf,
+                        storage.as_ref(),
+                        &index,
+                        sync_id,
+                        RemoteDataUpdate::Delete(&failed_delete_data.data.deleted_layers),
+                    )
+                    .await
+                    {
+                        error!("Failed to update remote timeline {sync_id}: {e:?}");
+                    }
+                }
+            }
+        } else {
+            sync_queue.push(sync_id, SyncTask::Delete(delete_data));
+            warn!("Skipping delete task due to failed upload tasks, reenqueuing");
         }
     }
 
-    download_status
+    status_update
 }
 
 async fn download_timeline_data<P, S>(
@@ -1183,7 +1170,7 @@ async fn download_timeline_data<P, S>(
     new_download_data: SyncData<LayersDownload>,
     sync_start: Instant,
     task_name: &str,
-) -> DownloadStatus
+) -> DownloadMarker
 where
     P: Debug + Send + Sync + 'static,
     S: RemoteStorage<RemoteObjectId = P> + Send + Sync + 'static,
@@ -1212,7 +1199,7 @@ where
                 Ok(()) => match index.write().await.set_awaits_download(&sync_id, false) {
                     Ok(()) => {
                         register_sync_status(sync_id, sync_start, task_name, Some(true));
-                        return DownloadStatus::Downloaded;
+                        return DownloadMarker::Downloaded;
                     }
                     Err(e) => {
                         error!("Timeline {sync_id} was expected to be in the remote index after a successful download, but it's absent: {e:?}");
@@ -1228,7 +1215,7 @@ where
         }
     }
 
-    DownloadStatus::Nothing
+    DownloadMarker::Nothing
 }
 
 async fn update_local_metadata(
@@ -1270,13 +1257,7 @@ async fn update_local_metadata(
             timeline_id,
         } = sync_id;
         tokio::task::spawn_blocking(move || {
-            crate::layered_repository::save_metadata(
-                conf,
-                timeline_id,
-                tenant_id,
-                &cloned_metadata,
-                true,
-            )
+            LayeredRepository::save_metadata(conf, timeline_id, tenant_id, &cloned_metadata, true)
         })
         .await
         .with_context(|| {
@@ -1351,8 +1332,7 @@ async fn upload_timeline_data<P, S>(
     new_upload_data: SyncData<LayersUpload>,
     sync_start: Instant,
     task_name: &str,
-) -> UploadStatus
-where
+) where
     P: Debug + Send + Sync + 'static,
     S: RemoteStorage<RemoteObjectId = P> + Send + Sync + 'static,
 {
@@ -1365,9 +1345,9 @@ where
     )
     .await
     {
-        UploadedTimeline::FailedAndRescheduled(e) => {
+        UploadedTimeline::FailedAndRescheduled => {
             register_sync_status(sync_id, sync_start, task_name, Some(false));
-            return UploadStatus::Failed(e);
+            return;
         }
         UploadedTimeline::Successful(upload_data) => upload_data,
     };
@@ -1386,14 +1366,12 @@ where
     {
         Ok(()) => {
             register_sync_status(sync_id, sync_start, task_name, Some(true));
-            UploadStatus::Uploaded
         }
         Err(e) => {
             error!("Failed to update remote timeline {sync_id}: {e:?}");
             uploaded_data.retries += 1;
             sync_queue.push(sync_id, SyncTask::Upload(uploaded_data));
             register_sync_status(sync_id, sync_start, task_name, Some(false));
-            UploadStatus::Failed(e)
         }
     }
 }
@@ -1496,17 +1474,25 @@ where
         .context("Failed to upload new index part")
 }
 
-async fn validate_task_retries(
-    current_attempt: u32,
+async fn validate_task_retries<T>(
+    sync_data: SyncData<T>,
     max_sync_errors: NonZeroU32,
-) -> ControlFlow<(), ()> {
+) -> ControlFlow<SyncData<T>, SyncData<T>> {
+    let current_attempt = sync_data.retries;
     let max_sync_errors = max_sync_errors.get();
     if current_attempt >= max_sync_errors {
-        return ControlFlow::Break(());
+        error!(
+            "Aborting task that failed {current_attempt} times, exceeding retries threshold of {max_sync_errors}",
+        );
+        return ControlFlow::Break(sync_data);
     }
 
-    exponential_backoff(current_attempt, 1.0, 30.0).await;
-    ControlFlow::Continue(())
+    if current_attempt > 0 {
+        let seconds_to_wait = 2.0_f64.powf(current_attempt as f64 - 1.0).min(30.0);
+        info!("Waiting {seconds_to_wait} seconds before starting the task");
+        tokio::time::sleep(Duration::from_secs_f64(seconds_to_wait)).await;
+    }
+    ControlFlow::Continue(sync_data)
 }
 
 fn schedule_first_sync_tasks(

pageserver/src/storage_sync/delete.rs

@@ -95,8 +95,6 @@ where
         debug!("Reenqueuing failed delete task for timeline {sync_id}");
         delete_data.retries += 1;
         sync_queue.push(sync_id, SyncTask::Delete(delete_data));
-    } else {
-        info!("Successfully deleted all layers");
     }
     errored
 }

pageserver/src/storage_sync/download.rs

@@ -130,7 +130,6 @@ where
             tenant_path.display()
         )
     })?;
-
     let timelines = storage
         .list_prefixes(Some(tenant_storage_path))
         .await
@@ -141,13 +140,6 @@ where
             )
         })?;
 
-    if timelines.is_empty() {
-        anyhow::bail!(
-            "no timelines found on the remote storage for tenant {}",
-            tenant_id
-        )
-    }
-
     let mut sync_ids = HashSet::new();
 
     for timeline_remote_storage_key in timelines {

pageserver/src/storage_sync/upload.rs

@@ -4,7 +4,7 @@ use std::{fmt::Debug, path::PathBuf};
 
 use anyhow::Context;
 use futures::stream::{FuturesUnordered, StreamExt};
-use once_cell::sync::Lazy;
+use lazy_static::lazy_static;
 use remote_storage::RemoteStorage;
 use tokio::fs;
 use tracing::{debug, error, info, warn};
@@ -20,14 +20,14 @@ use crate::{
 };
 use metrics::{register_int_counter_vec, IntCounterVec};
 
-static NO_LAYERS_UPLOAD: Lazy<IntCounterVec> = Lazy::new(|| {
-    register_int_counter_vec!(
+lazy_static! {
+    static ref NO_LAYERS_UPLOAD: IntCounterVec = register_int_counter_vec!(
         "pageserver_remote_storage_no_layers_uploads_total",
         "Number of skipped uploads due to no layers",
         &["tenant_id", "timeline_id"],
     )
-    .expect("failed to register pageserver no layers upload vec")
-});
+    .expect("failed to register pageserver no layers upload vec");
+}
 
 /// Serializes and uploads the given index part data to the remote storage.
 pub(super) async fn upload_index_part<P, S>(
@@ -75,7 +75,7 @@ where
 #[derive(Debug)]
 pub(super) enum UploadedTimeline {
     /// Upload failed due to some error, the upload task is rescheduled for another retry.
-    FailedAndRescheduled(anyhow::Error),
+    FailedAndRescheduled,
     /// No issues happened during the upload, all task files were put into the remote storage.
     Successful(SyncData<LayersUpload>),
 }
@@ -179,7 +179,7 @@ where
         })
         .collect::<FuturesUnordered<_>>();
 
-    let mut errors = Vec::new();
+    let mut errors_happened = false;
     while let Some(upload_result) = upload_tasks.next().await {
         match upload_result {
             Ok(uploaded_path) => {
@@ -188,13 +188,13 @@ where
             }
             Err(e) => match e {
                 UploadError::Other(e) => {
+                    errors_happened = true;
                     error!("Failed to upload a layer for timeline {sync_id}: {e:?}");
-                    errors.push(format!("{e:#}"));
                 }
                 UploadError::MissingLocalFile(source_path, e) => {
                     if source_path.exists() {
+                        errors_happened = true;
                         error!("Failed to upload a layer for timeline {sync_id}: {e:?}");
-                        errors.push(format!("{e:#}"));
                     } else {
                         // We have run the upload sync task, but the file we wanted to upload is gone.
                         // This is "fine" due the asynchronous nature of the sync loop: it only reacts to events and might need to
@@ -217,17 +217,14 @@ where
         }
     }
 
-    if errors.is_empty() {
-        info!("Successfully uploaded all layers");
-        UploadedTimeline::Successful(upload_data)
-    } else {
+    if errors_happened {
         debug!("Reenqueuing failed upload task for timeline {sync_id}");
         upload_data.retries += 1;
         sync_queue.push(sync_id, SyncTask::Upload(upload_data));
-        UploadedTimeline::FailedAndRescheduled(anyhow::anyhow!(
-            "Errors appeared during layer uploads: {:?}",
-            errors
-        ))
+        UploadedTimeline::FailedAndRescheduled
+    } else {
+        info!("Successfully uploaded all layers");
+        UploadedTimeline::Successful(upload_data)
     }
 }
 

pageserver/src/tenant_config.rs

@@ -23,7 +23,6 @@ pub mod defaults {
     // which is good for now to trigger bugs.
     // This parameter actually determines L0 layer file size.
     pub const DEFAULT_CHECKPOINT_DISTANCE: u64 = 256 * 1024 * 1024;
-    pub const DEFAULT_CHECKPOINT_TIMEOUT: &str = "10 m";
 
     // Target file size, when creating image and delta layers.
     // This parameter determines L1 layer file size.
@@ -37,7 +36,7 @@ pub mod defaults {
     pub const DEFAULT_IMAGE_CREATION_THRESHOLD: usize = 3;
     pub const DEFAULT_PITR_INTERVAL: &str = "30 days";
     pub const DEFAULT_WALRECEIVER_CONNECT_TIMEOUT: &str = "2 seconds";
-    pub const DEFAULT_WALRECEIVER_LAGGING_WAL_TIMEOUT: &str = "3 seconds";
+    pub const DEFAULT_WALRECEIVER_LAGGING_WAL_TIMEOUT: &str = "10 seconds";
     pub const DEFAULT_MAX_WALRECEIVER_LSN_WAL_LAG: u64 = 10 * 1024 * 1024;
 }
 
@@ -49,9 +48,6 @@ pub struct TenantConf {
     // page server crashes.
     // This parameter actually determines L0 layer file size.
     pub checkpoint_distance: u64,
-    // Inmemory layer is also flushed at least once in checkpoint_timeout to
-    // eventually upload WAL after activity is stopped.
-    pub checkpoint_timeout: Duration,
     // Target file size, when creating image and delta layers.
     // This parameter determines L1 layer file size.
     pub compaction_target_size: u64,
@@ -94,7 +90,6 @@ pub struct TenantConf {
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
 pub struct TenantConfOpt {
     pub checkpoint_distance: Option<u64>,
-    pub checkpoint_timeout: Option<Duration>,
     pub compaction_target_size: Option<u64>,
     #[serde(with = "humantime_serde")]
     pub compaction_period: Option<Duration>,
@@ -118,9 +113,6 @@ impl TenantConfOpt {
             checkpoint_distance: self
                 .checkpoint_distance
                 .unwrap_or(global_conf.checkpoint_distance),
-            checkpoint_timeout: self
-                .checkpoint_timeout
-                .unwrap_or(global_conf.checkpoint_timeout),
             compaction_target_size: self
                 .compaction_target_size
                 .unwrap_or(global_conf.compaction_target_size),
@@ -150,9 +142,6 @@ impl TenantConfOpt {
         if let Some(checkpoint_distance) = other.checkpoint_distance {
             self.checkpoint_distance = Some(checkpoint_distance);
         }
-        if let Some(checkpoint_timeout) = other.checkpoint_timeout {
-            self.checkpoint_timeout = Some(checkpoint_timeout);
-        }
         if let Some(compaction_target_size) = other.compaction_target_size {
             self.compaction_target_size = Some(compaction_target_size);
         }
@@ -192,8 +181,6 @@ impl TenantConf {
 
         TenantConf {
             checkpoint_distance: DEFAULT_CHECKPOINT_DISTANCE,
-            checkpoint_timeout: humantime::parse_duration(DEFAULT_CHECKPOINT_TIMEOUT)
-                .expect("cannot parse default checkpoint timeout"),
             compaction_target_size: DEFAULT_COMPACTION_TARGET_SIZE,
             compaction_period: humantime::parse_duration(DEFAULT_COMPACTION_PERIOD)
                 .expect("cannot parse default compaction period"),
@@ -225,7 +212,6 @@ impl TenantConf {
     pub fn dummy_conf() -> Self {
         TenantConf {
             checkpoint_distance: defaults::DEFAULT_CHECKPOINT_DISTANCE,
-            checkpoint_timeout: Duration::from_secs(600),
             compaction_target_size: 4 * 1024 * 1024,
             compaction_period: Duration::from_secs(10),
             compaction_threshold: defaults::DEFAULT_COMPACTION_THRESHOLD,

pageserver/src/tenant_mgr.rs

@@ -2,8 +2,8 @@
 //! page server.
 
 use crate::config::PageServerConf;
-use crate::http::models::TenantInfo;
 use crate::layered_repository::{load_metadata, LayeredRepository};
+use crate::pgdatadir_mapping::DatadirTimeline;
 use crate::repository::Repository;
 use crate::storage_sync::index::{RemoteIndex, RemoteTimelineIndex};
 use crate::storage_sync::{self, LocalTimelineInitStatus, SyncStartupData};
@@ -12,9 +12,10 @@ use crate::thread_mgr::ThreadKind;
 use crate::timelines::CreateRepo;
 use crate::walredo::PostgresRedoManager;
 use crate::{thread_mgr, timelines, walreceiver};
-use crate::{RepositoryImpl, TimelineImpl};
+use crate::{DatadirTimelineImpl, RepositoryImpl};
 use anyhow::Context;
 use serde::{Deserialize, Serialize};
+use serde_with::{serde_as, DisplayFromStr};
 use std::collections::hash_map::Entry;
 use std::collections::{HashMap, HashSet};
 use std::fmt;
@@ -27,25 +28,23 @@ use utils::zid::{ZTenantId, ZTenantTimelineId, ZTimelineId};
 
 mod tenants_state {
     use anyhow::ensure;
-    use once_cell::sync::Lazy;
     use std::{
         collections::HashMap,
         sync::{RwLock, RwLockReadGuard, RwLockWriteGuard},
     };
     use tokio::sync::mpsc;
     use tracing::{debug, error};
+
     use utils::zid::ZTenantId;
 
     use crate::tenant_mgr::{LocalTimelineUpdate, Tenant};
 
-    static TENANTS: Lazy<RwLock<HashMap<ZTenantId, Tenant>>> =
-        Lazy::new(|| RwLock::new(HashMap::new()));
-
-    /// Sends updates to the local timelines (creation and deletion) to the WAL receiver,
-    /// so that it can enable/disable corresponding processes.
-    static TIMELINE_UPDATE_SENDER: Lazy<
-        RwLock<Option<mpsc::UnboundedSender<LocalTimelineUpdate>>>,
-    > = Lazy::new(|| RwLock::new(None));
+    lazy_static::lazy_static! {
+        static ref TENANTS: RwLock<HashMap<ZTenantId, Tenant>> = RwLock::new(HashMap::new());
+        /// Sends updates to the local timelines (creation and deletion) to the WAL receiver,
+        /// so that it can enable/disable corresponding processes.
+        static ref TIMELINE_UPDATE_SENDER: RwLock<Option<mpsc::UnboundedSender<LocalTimelineUpdate>>> = RwLock::new(None);
+    }
 
     pub(super) fn read_tenants() -> RwLockReadGuard<'static, HashMap<ZTenantId, Tenant>> {
         TENANTS
@@ -102,7 +101,7 @@ struct Tenant {
     ///
     /// 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<<RepositoryImpl as Repository>::Timeline>>,
+    local_timelines: HashMap<ZTimelineId, Arc<DatadirTimelineImpl>>,
 }
 
 #[derive(Debug, Serialize, Deserialize, Clone, Copy, PartialEq, Eq)]
@@ -179,7 +178,7 @@ pub enum LocalTimelineUpdate {
     },
     Attach {
         id: ZTenantTimelineId,
-        datadir: Arc<<RepositoryImpl as Repository>::Timeline>,
+        datadir: Arc<DatadirTimelineImpl>,
     },
 }
 
@@ -383,7 +382,7 @@ pub fn get_repository_for_tenant(tenant_id: ZTenantId) -> anyhow::Result<Arc<Rep
 pub fn get_local_timeline_with_load(
     tenant_id: ZTenantId,
     timeline_id: ZTimelineId,
-) -> anyhow::Result<Arc<TimelineImpl>> {
+) -> anyhow::Result<Arc<DatadirTimelineImpl>> {
     let mut m = tenants_state::write_tenants();
     let tenant = m
         .get_mut(&tenant_id)
@@ -490,23 +489,34 @@ pub fn detach_tenant(conf: &'static PageServerConf, tenant_id: ZTenantId) -> any
 fn load_local_timeline(
     repo: &RepositoryImpl,
     timeline_id: ZTimelineId,
-) -> anyhow::Result<Arc<TimelineImpl>> {
+) -> anyhow::Result<Arc<DatadirTimeline<LayeredRepository>>> {
     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()?;
+    let repartition_distance = repo.get_checkpoint_distance() / 10;
+    let page_tline = Arc::new(DatadirTimelineImpl::new(
+        inmem_timeline,
+        repartition_distance,
+    ));
+    page_tline.init_logical_size()?;
 
     tenants_state::try_send_timeline_update(LocalTimelineUpdate::Attach {
         id: ZTenantTimelineId::new(repo.tenant_id(), timeline_id),
-        datadir: Arc::clone(&inmem_timeline),
+        datadir: Arc::clone(&page_tline),
     });
 
-    Ok(inmem_timeline)
+    Ok(page_tline)
+}
+
+#[serde_as]
+#[derive(Serialize, Deserialize, Clone)]
+pub struct TenantInfo {
+    #[serde_as(as = "DisplayFromStr")]
+    pub id: ZTenantId,
+    pub state: Option<TenantState>,
+    pub has_in_progress_downloads: Option<bool>,
 }
 
-///
-/// Get list of tenants, for the mgmt API
-///
 pub fn list_tenants(remote_index: &RemoteTimelineIndex) -> Vec<TenantInfo> {
     tenants_state::read_tenants()
         .iter()
@@ -522,7 +532,6 @@ pub fn list_tenants(remote_index: &RemoteTimelineIndex) -> Vec<TenantInfo> {
             TenantInfo {
                 id: *id,
                 state: Some(tenant.state),
-                current_physical_size: None,
                 has_in_progress_downloads,
             }
         })

pageserver/src/tenant_tasks.rs

@@ -120,10 +120,6 @@ pub fn init_tenant_task_pool() -> anyhow::Result<()> {
     let runtime = tokio::runtime::Builder::new_multi_thread()
         .thread_name("tenant-task-worker")
         .enable_all()
-        .on_thread_start(|| {
-            thread_mgr::register(ThreadKind::TenantTaskWorker, "tenant-task-worker")
-        })
-        .on_thread_stop(thread_mgr::deregister)
         .build()?;
 
     let (gc_send, mut gc_recv) = mpsc::channel::<ZTenantId>(100);

pageserver/src/thread_mgr.rs

@@ -45,20 +45,21 @@ use tokio::sync::watch;
 
 use tracing::{debug, error, info, warn};
 
-use once_cell::sync::Lazy;
+use lazy_static::lazy_static;
 
 use utils::zid::{ZTenantId, ZTimelineId};
 
 use crate::shutdown_pageserver;
 
-/// Each thread that we track is associated with a "thread ID". It's just
-/// an increasing number that we assign, not related to any system thread
-/// id.
-static NEXT_THREAD_ID: Lazy<AtomicU64> = Lazy::new(|| AtomicU64::new(1));
+lazy_static! {
+    /// Each thread that we track is associated with a "thread ID". It's just
+    /// an increasing number that we assign, not related to any system thread
+    /// id.
+    static ref NEXT_THREAD_ID: AtomicU64 = AtomicU64::new(1);
 
-/// Global registry of threads
-static THREADS: Lazy<Mutex<HashMap<u64, Arc<PageServerThread>>>> =
-    Lazy::new(|| Mutex::new(HashMap::new()));
+    /// Global registry of threads
+    static ref THREADS: Mutex<HashMap<u64, Arc<PageServerThread>>> = Mutex::new(HashMap::new());
+}
 
 // There is a Tokio watch channel for each thread, which can be used to signal the
 // thread that it needs to shut down. This thread local variable holds the receiving
@@ -96,9 +97,6 @@ pub enum ThreadKind {
     // Thread that schedules new compaction and gc jobs
     TenantTaskManager,
 
-    // Worker thread for tenant tasks thread pool
-    TenantTaskWorker,
-
     // Thread that flushes frozen in-memory layers to disk
     LayerFlushThread,
 
@@ -107,20 +105,18 @@ pub enum ThreadKind {
     StorageSync,
 }
 
-#[derive(Default)]
 struct MutableThreadState {
     /// Tenant and timeline that this thread is associated with.
     tenant_id: Option<ZTenantId>,
     timeline_id: Option<ZTimelineId>,
 
     /// Handle for waiting for the thread to exit. It can be None, if the
-    /// the thread has already exited. OR if this thread is managed externally
-    /// and was not spawned through thread_mgr.rs::spawn function.
+    /// the thread has already exited.
     join_handle: Option<JoinHandle<()>>,
 }
 
 struct PageServerThread {
-    thread_id: u64,
+    _thread_id: u64,
 
     kind: ThreadKind,
 
@@ -151,7 +147,7 @@ where
     let (shutdown_tx, shutdown_rx) = watch::channel(());
     let thread_id = NEXT_THREAD_ID.fetch_add(1, Ordering::Relaxed);
     let thread = Arc::new(PageServerThread {
-        thread_id,
+        _thread_id: thread_id,
         kind,
         name: name.to_string(),
         shutdown_requested: AtomicBool::new(false),
@@ -319,10 +315,8 @@ pub fn shutdown_threads(
             drop(thread_mut);
             let _ = join_handle.join();
         } else {
-            // Possibly one of:
-            //  * The thread had not even fully started yet.
-            //  * It was shut down concurrently and already exited
-            //  * Is managed through `register`/`deregister` fns without providing a join handle
+            // The thread had not even fully started yet. Or it was shut down
+            // concurrently and already exited
         }
     }
 }
@@ -354,56 +348,3 @@ pub fn is_shutdown_requested() -> bool {
         }
     })
 }
-
-/// Needed to register threads that were not spawned through spawn function.
-/// For example tokio blocking threads. This function is expected to be used
-/// in tandem with `deregister`.
-/// NOTE: threads registered through this function cannot be joined
-pub fn register(kind: ThreadKind, name: &str) {
-    CURRENT_THREAD.with(|ct| {
-        let mut borrowed = ct.borrow_mut();
-        if borrowed.is_some() {
-            panic!("thread already registered")
-        };
-        let (shutdown_tx, shutdown_rx) = watch::channel(());
-        let thread_id = NEXT_THREAD_ID.fetch_add(1, Ordering::Relaxed);
-
-        let thread = Arc::new(PageServerThread {
-            thread_id,
-            kind,
-            name: name.to_owned(),
-            shutdown_requested: AtomicBool::new(false),
-            shutdown_tx,
-            mutable: Mutex::new(MutableThreadState {
-                tenant_id: None,
-                timeline_id: None,
-                join_handle: None,
-            }),
-        });
-
-        *borrowed = Some(Arc::clone(&thread));
-
-        SHUTDOWN_RX.with(|rx| {
-            *rx.borrow_mut() = Some(shutdown_rx);
-        });
-
-        THREADS.lock().unwrap().insert(thread_id, thread);
-    });
-}
-
-// Expected to be used in tandem with `register`. See the doc for `register` for more details
-pub fn deregister() {
-    CURRENT_THREAD.with(|ct| {
-        let mut borrowed = ct.borrow_mut();
-        let thread = match borrowed.take() {
-            Some(thread) => thread,
-            None => panic!("calling deregister on unregistered thread"),
-        };
-
-        SHUTDOWN_RX.with(|rx| {
-            *rx.borrow_mut() = None;
-        });
-
-        THREADS.lock().unwrap().remove(&thread.thread_id)
-    });
-}

pageserver/src/timelines.rs

@@ -4,6 +4,8 @@
 
 use anyhow::{bail, ensure, Context, Result};
 use postgres_ffi::ControlFileData;
+use serde::{Deserialize, Serialize};
+use serde_with::{serde_as, DisplayFromStr};
 use std::{
     fs,
     path::Path,
@@ -18,15 +20,123 @@ use utils::{
     zid::{ZTenantId, ZTimelineId},
 };
 
-use crate::tenant_mgr;
 use crate::{
-    config::PageServerConf, repository::Repository, storage_sync::index::RemoteIndex,
-    tenant_config::TenantConfOpt, RepositoryImpl, TimelineImpl,
+    config::PageServerConf,
+    layered_repository::metadata::TimelineMetadata,
+    repository::{LocalTimelineState, Repository},
+    storage_sync::index::RemoteIndex,
+    tenant_config::TenantConfOpt,
+    DatadirTimeline, RepositoryImpl,
 };
 use crate::{import_datadir, LOG_FILE_NAME};
 use crate::{layered_repository::LayeredRepository, walredo::WalRedoManager};
+use crate::{repository::RepositoryTimeline, tenant_mgr};
 use crate::{repository::Timeline, CheckpointConfig};
 
+#[serde_as]
+#[derive(Debug, Serialize, Deserialize, Clone)]
+pub struct LocalTimelineInfo {
+    #[serde_as(as = "Option<DisplayFromStr>")]
+    pub ancestor_timeline_id: Option<ZTimelineId>,
+    #[serde_as(as = "Option<DisplayFromStr>")]
+    pub ancestor_lsn: Option<Lsn>,
+    #[serde_as(as = "DisplayFromStr")]
+    pub last_record_lsn: Lsn,
+    #[serde_as(as = "Option<DisplayFromStr>")]
+    pub prev_record_lsn: Option<Lsn>,
+    #[serde_as(as = "DisplayFromStr")]
+    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_logical_size_non_incremental: Option<usize>,
+    pub timeline_state: LocalTimelineState,
+}
+
+impl LocalTimelineInfo {
+    pub fn from_loaded_timeline<R: Repository>(
+        datadir_tline: &DatadirTimeline<R>,
+        include_non_incremental_logical_size: bool,
+    ) -> anyhow::Result<Self> {
+        let last_record_lsn = datadir_tline.tline.get_last_record_lsn();
+        let info = LocalTimelineInfo {
+            ancestor_timeline_id: datadir_tline.tline.get_ancestor_timeline_id(),
+            ancestor_lsn: {
+                match datadir_tline.tline.get_ancestor_lsn() {
+                    Lsn(0) => None,
+                    lsn @ Lsn(_) => Some(lsn),
+                }
+            },
+            disk_consistent_lsn: datadir_tline.tline.get_disk_consistent_lsn(),
+            last_record_lsn,
+            prev_record_lsn: Some(datadir_tline.tline.get_prev_record_lsn()),
+            latest_gc_cutoff_lsn: *datadir_tline.tline.get_latest_gc_cutoff_lsn(),
+            timeline_state: LocalTimelineState::Loaded,
+            current_logical_size: Some(datadir_tline.get_current_logical_size()),
+            current_logical_size_non_incremental: if include_non_incremental_logical_size {
+                Some(datadir_tline.get_current_logical_size_non_incremental(last_record_lsn)?)
+            } else {
+                None
+            },
+        };
+        Ok(info)
+    }
+
+    pub fn from_unloaded_timeline(metadata: &TimelineMetadata) -> Self {
+        LocalTimelineInfo {
+            ancestor_timeline_id: metadata.ancestor_timeline(),
+            ancestor_lsn: {
+                match metadata.ancestor_lsn() {
+                    Lsn(0) => None,
+                    lsn @ Lsn(_) => Some(lsn),
+                }
+            },
+            disk_consistent_lsn: metadata.disk_consistent_lsn(),
+            last_record_lsn: metadata.disk_consistent_lsn(),
+            prev_record_lsn: metadata.prev_record_lsn(),
+            latest_gc_cutoff_lsn: metadata.latest_gc_cutoff_lsn(),
+            timeline_state: LocalTimelineState::Unloaded,
+            current_logical_size: None,
+            current_logical_size_non_incremental: None,
+        }
+    }
+
+    pub fn from_repo_timeline<T>(
+        tenant_id: ZTenantId,
+        timeline_id: ZTimelineId,
+        repo_timeline: &RepositoryTimeline<T>,
+        include_non_incremental_logical_size: bool,
+    ) -> anyhow::Result<Self> {
+        match repo_timeline {
+            RepositoryTimeline::Loaded(_) => {
+                let datadir_tline =
+                    tenant_mgr::get_local_timeline_with_load(tenant_id, timeline_id)?;
+                Self::from_loaded_timeline(&datadir_tline, include_non_incremental_logical_size)
+            }
+            RepositoryTimeline::Unloaded { metadata } => Ok(Self::from_unloaded_timeline(metadata)),
+        }
+    }
+}
+
+#[serde_as]
+#[derive(Debug, Serialize, Deserialize, Clone)]
+pub struct RemoteTimelineInfo {
+    #[serde_as(as = "DisplayFromStr")]
+    pub remote_consistent_lsn: Lsn,
+    pub awaits_download: bool,
+}
+
+#[serde_as]
+#[derive(Debug, Serialize, Deserialize, Clone)]
+pub struct TimelineInfo {
+    #[serde_as(as = "DisplayFromStr")]
+    pub tenant_id: ZTenantId,
+    #[serde_as(as = "DisplayFromStr")]
+    pub timeline_id: ZTimelineId,
+    pub local: Option<LocalTimelineInfo>,
+    pub remote: Option<RemoteTimelineInfo>,
+}
+
 #[derive(Debug, Clone, Copy)]
 pub struct PointInTime {
     pub timeline_id: ZTimelineId,
@@ -188,18 +298,19 @@ fn bootstrap_timeline<R: Repository>(
     // Initdb lsn will be equal to last_record_lsn which will be set after import.
     // Because we know it upfront avoid having an option or dummy zero value by passing it to create_empty_timeline.
     let timeline = repo.create_empty_timeline(tli, lsn)?;
-    import_datadir::import_timeline_from_postgres_datadir(&pgdata_path, &*timeline, lsn)?;
+    let mut page_tline: DatadirTimeline<R> = DatadirTimeline::new(timeline, u64::MAX);
+    import_datadir::import_timeline_from_postgres_datadir(&pgdata_path, &mut page_tline, lsn)?;
 
     fail::fail_point!("before-checkpoint-new-timeline", |_| {
         bail!("failpoint before-checkpoint-new-timeline");
     });
 
-    timeline.checkpoint(CheckpointConfig::Forced)?;
+    page_tline.tline.checkpoint(CheckpointConfig::Forced)?;
 
     info!(
         "created root timeline {} timeline.lsn {}",
         tli,
-        timeline.get_last_record_lsn()
+        page_tline.tline.get_last_record_lsn()
     );
 
     // Remove temp dir. We don't need it anymore
@@ -208,22 +319,36 @@ fn bootstrap_timeline<R: Repository>(
     Ok(())
 }
 
-///
-/// Create a new timeline.
-///
-/// Returns the new timeline ID and reference to its Timeline object.
-///
-/// If the caller specified the timeline ID to use (`new_timeline_id`), and timeline with
-/// the same timeline ID already exists, returns None. If `new_timeline_id` is not given,
-/// a new unique ID is generated.
-///
+pub(crate) fn get_local_timelines(
+    tenant_id: ZTenantId,
+    include_non_incremental_logical_size: bool,
+) -> Result<Vec<(ZTimelineId, LocalTimelineInfo)>> {
+    let repo = tenant_mgr::get_repository_for_tenant(tenant_id)
+        .with_context(|| format!("Failed to get repo for tenant {}", tenant_id))?;
+    let repo_timelines = repo.list_timelines();
+
+    let mut local_timeline_info = Vec::with_capacity(repo_timelines.len());
+    for (timeline_id, repository_timeline) in repo_timelines {
+        local_timeline_info.push((
+            timeline_id,
+            LocalTimelineInfo::from_repo_timeline(
+                tenant_id,
+                timeline_id,
+                &repository_timeline,
+                include_non_incremental_logical_size,
+            )?,
+        ))
+    }
+    Ok(local_timeline_info)
+}
+
 pub(crate) fn create_timeline(
     conf: &'static PageServerConf,
     tenant_id: ZTenantId,
     new_timeline_id: Option<ZTimelineId>,
     ancestor_timeline_id: Option<ZTimelineId>,
     mut ancestor_start_lsn: Option<Lsn>,
-) -> Result<Option<(ZTimelineId, Arc<TimelineImpl>)>> {
+) -> Result<Option<TimelineInfo>> {
     let new_timeline_id = new_timeline_id.unwrap_or_else(ZTimelineId::generate);
     let repo = tenant_mgr::get_repository_for_tenant(tenant_id)?;
 
@@ -232,7 +357,7 @@ pub(crate) fn create_timeline(
         return Ok(None);
     }
 
-    match ancestor_timeline_id {
+    let new_timeline_info = match ancestor_timeline_id {
         Some(ancestor_timeline_id) => {
             let ancestor_timeline = repo
                 .get_timeline_load(ancestor_timeline_id)
@@ -260,13 +385,26 @@ pub(crate) fn create_timeline(
                 }
             }
 
-            repo.branch_timeline(ancestor_timeline_id, new_timeline_id, ancestor_start_lsn)?
+            repo.branch_timeline(ancestor_timeline_id, new_timeline_id, ancestor_start_lsn)?;
+            // load the timeline into memory
+            let loaded_timeline =
+                tenant_mgr::get_local_timeline_with_load(tenant_id, new_timeline_id)?;
+            LocalTimelineInfo::from_loaded_timeline(&loaded_timeline, false)
+                .context("cannot fill timeline info")?
+        }
+        None => {
+            bootstrap_timeline(conf, tenant_id, new_timeline_id, repo.as_ref())?;
+            // load the timeline into memory
+            let new_timeline =
+                tenant_mgr::get_local_timeline_with_load(tenant_id, new_timeline_id)?;
+            LocalTimelineInfo::from_loaded_timeline(&new_timeline, false)
+                .context("cannot fill timeline info")?
         }
-        None => bootstrap_timeline(conf, tenant_id, new_timeline_id, repo.as_ref())?,
     };
-
-    // load the timeline into memory
-    let loaded_timeline = tenant_mgr::get_local_timeline_with_load(tenant_id, new_timeline_id)?;
-
-    Ok(Some((new_timeline_id, loaded_timeline)))
+    Ok(Some(TimelineInfo {
+        tenant_id,
+        timeline_id: new_timeline_id,
+        local: Some(new_timeline_info),
+        remote: None,
+    }))
 }

pageserver/src/virtual_file.rs

@@ -10,7 +10,7 @@
 //! This is similar to PostgreSQL's virtual file descriptor facility in
 //! src/backend/storage/file/fd.c
 //!
-use once_cell::sync::Lazy;
+use lazy_static::lazy_static;
 use once_cell::sync::OnceCell;
 use std::fs::{File, OpenOptions};
 use std::io::{Error, ErrorKind, Read, Seek, SeekFrom, Write};
@@ -32,24 +32,23 @@ const STORAGE_IO_TIME_BUCKETS: &[f64] = &[
     1.0,      // 1 sec
 ];
 
-static STORAGE_IO_TIME: Lazy<HistogramVec> = Lazy::new(|| {
-    register_histogram_vec!(
+lazy_static! {
+    static ref STORAGE_IO_TIME: HistogramVec = register_histogram_vec!(
         "pageserver_io_operations_seconds",
         "Time spent in IO operations",
         &["operation", "tenant_id", "timeline_id"],
         STORAGE_IO_TIME_BUCKETS.into()
     )
-    .expect("failed to define a metric")
-});
-
-static STORAGE_IO_SIZE: Lazy<IntGaugeVec> = Lazy::new(|| {
-    register_int_gauge_vec!(
+    .expect("failed to define a metric");
+}
+lazy_static! {
+    static ref STORAGE_IO_SIZE: IntGaugeVec = register_int_gauge_vec!(
         "pageserver_io_operations_bytes_total",
         "Total amount of bytes read/written in IO operations",
         &["operation", "tenant_id", "timeline_id"]
     )
-    .expect("failed to define a metric")
-});
+    .expect("failed to define a metric");
+}
 
 ///
 /// A virtual file descriptor. You can use this just like std::fs::File, but internally

pageserver/src/walingest.rs

@@ -30,8 +30,11 @@ use anyhow::Result;
 use bytes::{Buf, Bytes, BytesMut};
 use tracing::*;
 
+use std::collections::HashMap;
+
 use crate::pgdatadir_mapping::*;
 use crate::reltag::{RelTag, SlruKind};
+use crate::repository::Repository;
 use crate::walrecord::*;
 use postgres_ffi::nonrelfile_utils::mx_offset_to_member_segment;
 use postgres_ffi::xlog_utils::*;
@@ -41,15 +44,17 @@ 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, R: Repository> {
+    timeline: &'a DatadirTimeline<R>,
 
     checkpoint: CheckPoint,
     checkpoint_modified: bool,
+
+    relsize_cache: HashMap<RelTag, BlockNumber>,
 }
 
-impl<'a, T: DatadirTimeline> WalIngest<'a, T> {
-    pub fn new(timeline: &T, startpoint: Lsn) -> Result<WalIngest<T>> {
+impl<'a, R: Repository> WalIngest<'a, R> {
+    pub fn new(timeline: &DatadirTimeline<R>, startpoint: Lsn) -> Result<WalIngest<R>> {
         // 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)?;
@@ -60,27 +65,26 @@ impl<'a, T: DatadirTimeline> WalIngest<'a, T> {
             timeline,
             checkpoint,
             checkpoint_modified: false,
+            relsize_cache: HashMap::new(),
         })
     }
 
     ///
     /// Decode a PostgreSQL WAL record and store it in the repository, in the given timeline.
     ///
-    /// This function updates `lsn` field of `DatadirModification`
     ///
     /// Helper function to parse a WAL record and call the Timeline's PUT functions for all the
     /// relations/pages that the record affects.
     ///
     pub fn ingest_record(
         &mut self,
+        timeline: &DatadirTimeline<R>,
         recdata: Bytes,
         lsn: Lsn,
-        modification: &mut DatadirModification<T>,
-        decoded: &mut DecodedWALRecord,
     ) -> Result<()> {
-        modification.lsn = lsn;
-        decode_wal_record(recdata, decoded).context("failed decoding wal record")?;
+        let mut modification = timeline.begin_modification(lsn);
 
+        let mut decoded = decode_wal_record(recdata).context("failed decoding wal record")?;
         let mut buf = decoded.record.clone();
         buf.advance(decoded.main_data_offset);
 
@@ -94,7 +98,7 @@ impl<'a, T: DatadirTimeline> WalIngest<'a, T> {
         if decoded.xl_rmid == pg_constants::RM_HEAP_ID
             || decoded.xl_rmid == pg_constants::RM_HEAP2_ID
         {
-            self.ingest_heapam_record(&mut buf, modification, decoded)?;
+            self.ingest_heapam_record(&mut buf, &mut modification, &mut decoded)?;
         }
         // Handle other special record types
         if decoded.xl_rmid == pg_constants::RM_SMGR_ID
@@ -102,19 +106,19 @@ impl<'a, T: DatadirTimeline> WalIngest<'a, T> {
                 == pg_constants::XLOG_SMGR_CREATE
         {
             let create = XlSmgrCreate::decode(&mut buf);
-            self.ingest_xlog_smgr_create(modification, &create)?;
+            self.ingest_xlog_smgr_create(&mut modification, &create)?;
         } else if decoded.xl_rmid == pg_constants::RM_SMGR_ID
             && (decoded.xl_info & pg_constants::XLR_RMGR_INFO_MASK)
                 == pg_constants::XLOG_SMGR_TRUNCATE
         {
             let truncate = XlSmgrTruncate::decode(&mut buf);
-            self.ingest_xlog_smgr_truncate(modification, &truncate)?;
+            self.ingest_xlog_smgr_truncate(&mut modification, &truncate)?;
         } else if decoded.xl_rmid == pg_constants::RM_DBASE_ID {
             if (decoded.xl_info & pg_constants::XLR_RMGR_INFO_MASK)
                 == pg_constants::XLOG_DBASE_CREATE
             {
                 let createdb = XlCreateDatabase::decode(&mut buf);
-                self.ingest_xlog_dbase_create(modification, &createdb)?;
+                self.ingest_xlog_dbase_create(&mut modification, &createdb)?;
             } else if (decoded.xl_info & pg_constants::XLR_RMGR_INFO_MASK)
                 == pg_constants::XLOG_DBASE_DROP
             {
@@ -133,7 +137,7 @@ impl<'a, T: DatadirTimeline> WalIngest<'a, T> {
                 let segno = pageno / pg_constants::SLRU_PAGES_PER_SEGMENT;
                 let rpageno = pageno % pg_constants::SLRU_PAGES_PER_SEGMENT;
                 self.put_slru_page_image(
-                    modification,
+                    &mut modification,
                     SlruKind::Clog,
                     segno,
                     rpageno,
@@ -142,7 +146,7 @@ impl<'a, T: DatadirTimeline> WalIngest<'a, T> {
             } else {
                 assert!(info == pg_constants::CLOG_TRUNCATE);
                 let xlrec = XlClogTruncate::decode(&mut buf);
-                self.ingest_clog_truncate_record(modification, &xlrec)?;
+                self.ingest_clog_truncate_record(&mut modification, &xlrec)?;
             }
         } else if decoded.xl_rmid == pg_constants::RM_XACT_ID {
             let info = decoded.xl_info & pg_constants::XLOG_XACT_OPMASK;
@@ -150,7 +154,7 @@ impl<'a, T: DatadirTimeline> WalIngest<'a, T> {
                 let parsed_xact =
                     XlXactParsedRecord::decode(&mut buf, decoded.xl_xid, decoded.xl_info);
                 self.ingest_xact_record(
-                    modification,
+                    &mut modification,
                     &parsed_xact,
                     info == pg_constants::XLOG_XACT_COMMIT,
                 )?;
@@ -160,7 +164,7 @@ impl<'a, T: DatadirTimeline> WalIngest<'a, T> {
                 let parsed_xact =
                     XlXactParsedRecord::decode(&mut buf, decoded.xl_xid, decoded.xl_info);
                 self.ingest_xact_record(
-                    modification,
+                    &mut modification,
                     &parsed_xact,
                     info == pg_constants::XLOG_XACT_COMMIT_PREPARED,
                 )?;
@@ -183,7 +187,7 @@ impl<'a, T: DatadirTimeline> WalIngest<'a, T> {
                 let segno = pageno / pg_constants::SLRU_PAGES_PER_SEGMENT;
                 let rpageno = pageno % pg_constants::SLRU_PAGES_PER_SEGMENT;
                 self.put_slru_page_image(
-                    modification,
+                    &mut modification,
                     SlruKind::MultiXactOffsets,
                     segno,
                     rpageno,
@@ -194,7 +198,7 @@ impl<'a, T: DatadirTimeline> WalIngest<'a, T> {
                 let segno = pageno / pg_constants::SLRU_PAGES_PER_SEGMENT;
                 let rpageno = pageno % pg_constants::SLRU_PAGES_PER_SEGMENT;
                 self.put_slru_page_image(
-                    modification,
+                    &mut modification,
                     SlruKind::MultiXactMembers,
                     segno,
                     rpageno,
@@ -202,14 +206,14 @@ impl<'a, T: DatadirTimeline> WalIngest<'a, T> {
                 )?;
             } else if info == pg_constants::XLOG_MULTIXACT_CREATE_ID {
                 let xlrec = XlMultiXactCreate::decode(&mut buf);
-                self.ingest_multixact_create_record(modification, &xlrec)?;
+                self.ingest_multixact_create_record(&mut modification, &xlrec)?;
             } else if info == pg_constants::XLOG_MULTIXACT_TRUNCATE_ID {
                 let xlrec = XlMultiXactTruncate::decode(&mut buf);
-                self.ingest_multixact_truncate_record(modification, &xlrec)?;
+                self.ingest_multixact_truncate_record(&mut modification, &xlrec)?;
             }
         } else if decoded.xl_rmid == pg_constants::RM_RELMAP_ID {
             let xlrec = XlRelmapUpdate::decode(&mut buf);
-            self.ingest_relmap_page(modification, &xlrec, decoded)?;
+            self.ingest_relmap_page(&mut modification, &xlrec, &decoded)?;
         } else if decoded.xl_rmid == pg_constants::RM_XLOG_ID {
             let info = decoded.xl_info & pg_constants::XLR_RMGR_INFO_MASK;
             if info == pg_constants::XLOG_NEXTOID {
@@ -244,7 +248,7 @@ impl<'a, T: DatadirTimeline> WalIngest<'a, T> {
         // Iterate through all the blocks that the record modifies, and
         // "put" a separate copy of the record for each block.
         for blk in decoded.blocks.iter() {
-            self.ingest_decoded_block(modification, lsn, decoded, blk)?;
+            self.ingest_decoded_block(&mut modification, lsn, &decoded, blk)?;
         }
 
         // If checkpoint data was updated, store the new version in the repository
@@ -264,7 +268,7 @@ impl<'a, T: DatadirTimeline> WalIngest<'a, T> {
 
     fn ingest_decoded_block(
         &mut self,
-        modification: &mut DatadirModification<T>,
+        modification: &mut DatadirModification<R>,
         lsn: Lsn,
         decoded: &DecodedWALRecord,
         blk: &DecodedBkpBlock,
@@ -324,7 +328,7 @@ impl<'a, T: DatadirTimeline> WalIngest<'a, T> {
     fn ingest_heapam_record(
         &mut self,
         buf: &mut Bytes,
-        modification: &mut DatadirModification<T>,
+        modification: &mut DatadirModification<R>,
         decoded: &mut DecodedWALRecord,
     ) -> Result<()> {
         // Handle VM bit updates that are implicitly part of heap records.
@@ -405,7 +409,7 @@ impl<'a, T: DatadirTimeline> WalIngest<'a, T> {
             // replaying it would fail to find the previous image of the page, because
             // it doesn't exist. So check if the VM page(s) exist, and skip the WAL
             // record if it doesn't.
-            let vm_size = self.get_relsize(vm_rel, modification.lsn)?;
+            let vm_size = self.get_relsize(vm_rel)?;
             if let Some(blknum) = new_vm_blk {
                 if blknum >= vm_size {
                     new_vm_blk = None;
@@ -468,7 +472,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<R>,
         rec: &XlCreateDatabase,
     ) -> Result<()> {
         let db_id = rec.db_id;
@@ -535,7 +539,7 @@ impl<'a, T: DatadirTimeline> WalIngest<'a, T> {
 
     fn ingest_xlog_smgr_create(
         &mut self,
-        modification: &mut DatadirModification<T>,
+        modification: &mut DatadirModification<R>,
         rec: &XlSmgrCreate,
     ) -> Result<()> {
         let rel = RelTag {
@@ -553,7 +557,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<R>,
         rec: &XlSmgrTruncate,
     ) -> Result<()> {
         let spcnode = rec.rnode.spcnode;
@@ -618,7 +622,7 @@ impl<'a, T: DatadirTimeline> WalIngest<'a, T> {
     ///
     fn ingest_xact_record(
         &mut self,
-        modification: &mut DatadirModification<T>,
+        modification: &mut DatadirModification<R>,
         parsed: &XlXactParsedRecord,
         is_commit: bool,
     ) -> Result<()> {
@@ -687,7 +691,7 @@ impl<'a, T: DatadirTimeline> WalIngest<'a, T> {
 
     fn ingest_clog_truncate_record(
         &mut self,
-        modification: &mut DatadirModification<T>,
+        modification: &mut DatadirModification<R>,
         xlrec: &XlClogTruncate,
     ) -> Result<()> {
         info!(
@@ -745,7 +749,7 @@ impl<'a, T: DatadirTimeline> WalIngest<'a, T> {
 
     fn ingest_multixact_create_record(
         &mut self,
-        modification: &mut DatadirModification<T>,
+        modification: &mut DatadirModification<R>,
         xlrec: &XlMultiXactCreate,
     ) -> Result<()> {
         // Create WAL record for updating the multixact-offsets page
@@ -824,7 +828,7 @@ impl<'a, T: DatadirTimeline> WalIngest<'a, T> {
 
     fn ingest_multixact_truncate_record(
         &mut self,
-        modification: &mut DatadirModification<T>,
+        modification: &mut DatadirModification<R>,
         xlrec: &XlMultiXactTruncate,
     ) -> Result<()> {
         self.checkpoint.oldestMulti = xlrec.end_trunc_off;
@@ -858,7 +862,7 @@ impl<'a, T: DatadirTimeline> WalIngest<'a, T> {
 
     fn ingest_relmap_page(
         &mut self,
-        modification: &mut DatadirModification<T>,
+        modification: &mut DatadirModification<R>,
         xlrec: &XlRelmapUpdate,
         decoded: &DecodedWALRecord,
     ) -> Result<()> {
@@ -874,16 +878,17 @@ impl<'a, T: DatadirTimeline> WalIngest<'a, T> {
 
     fn put_rel_creation(
         &mut self,
-        modification: &mut DatadirModification<T>,
+        modification: &mut DatadirModification<R>,
         rel: RelTag,
     ) -> Result<()> {
+        self.relsize_cache.insert(rel, 0);
         modification.put_rel_creation(rel, 0)?;
         Ok(())
     }
 
     fn put_rel_page_image(
         &mut self,
-        modification: &mut DatadirModification<T>,
+        modification: &mut DatadirModification<R>,
         rel: RelTag,
         blknum: BlockNumber,
         img: Bytes,
@@ -895,7 +900,7 @@ impl<'a, T: DatadirTimeline> WalIngest<'a, T> {
 
     fn put_rel_wal_record(
         &mut self,
-        modification: &mut DatadirModification<T>,
+        modification: &mut DatadirModification<R>,
         rel: RelTag,
         blknum: BlockNumber,
         rec: ZenithWalRecord,
@@ -907,49 +912,63 @@ impl<'a, T: DatadirTimeline> WalIngest<'a, T> {
 
     fn put_rel_truncation(
         &mut self,
-        modification: &mut DatadirModification<T>,
+        modification: &mut DatadirModification<R>,
         rel: RelTag,
         nblocks: BlockNumber,
     ) -> Result<()> {
         modification.put_rel_truncation(rel, nblocks)?;
+        self.relsize_cache.insert(rel, nblocks);
         Ok(())
     }
 
     fn put_rel_drop(
         &mut self,
-        modification: &mut DatadirModification<T>,
+        modification: &mut DatadirModification<R>,
         rel: RelTag,
     ) -> Result<()> {
         modification.put_rel_drop(rel)?;
+        self.relsize_cache.remove(&rel);
         Ok(())
     }
 
-    fn get_relsize(&mut self, rel: RelTag, lsn: Lsn) -> Result<BlockNumber> {
-        let nblocks = if !self.timeline.get_rel_exists(rel, lsn)? {
-            0
+    fn get_relsize(&mut self, rel: RelTag) -> Result<BlockNumber> {
+        if let Some(nblocks) = self.relsize_cache.get(&rel) {
+            Ok(*nblocks)
         } else {
-            self.timeline.get_rel_size(rel, lsn)?
-        };
-        Ok(nblocks)
+            let last_lsn = self.timeline.get_last_record_lsn();
+            let nblocks = if !self.timeline.get_rel_exists(rel, last_lsn)? {
+                0
+            } else {
+                self.timeline.get_rel_size(rel, last_lsn)?
+            };
+            self.relsize_cache.insert(rel, nblocks);
+            Ok(nblocks)
+        }
     }
 
     fn handle_rel_extend(
         &mut self,
-        modification: &mut DatadirModification<T>,
+        modification: &mut DatadirModification<R>,
         rel: RelTag,
         blknum: BlockNumber,
     ) -> Result<()> {
         let new_nblocks = blknum + 1;
-        // Check if the relation exists. We implicitly create relations on first
-        // record.
-        // TODO: would be nice if to be more explicit about it
-        let last_lsn = modification.lsn;
-        let old_nblocks = if !self.timeline.get_rel_exists(rel, last_lsn)? {
-            // create it with 0 size initially, the logic below will extend it
-            modification.put_rel_creation(rel, 0)?;
-            0
+        let old_nblocks = if let Some(nblocks) = self.relsize_cache.get(&rel) {
+            *nblocks
         } else {
-            self.timeline.get_rel_size(rel, last_lsn)?
+            // Check if the relation exists. We implicitly create relations on first
+            // record.
+            // TODO: would be nice if to be more explicit about it
+            let last_lsn = self.timeline.get_last_record_lsn();
+            let nblocks = if !self.timeline.get_rel_exists(rel, last_lsn)? {
+                // create it with 0 size initially, the logic below will extend it
+                modification.put_rel_creation(rel, 0)?;
+                0
+            } else {
+                self.timeline.get_rel_size(rel, last_lsn)?
+            };
+            self.relsize_cache.insert(rel, nblocks);
+            nblocks
         };
 
         if new_nblocks > old_nblocks {
@@ -960,13 +979,14 @@ impl<'a, T: DatadirTimeline> WalIngest<'a, T> {
             for gap_blknum in old_nblocks..blknum {
                 modification.put_rel_page_image(rel, gap_blknum, ZERO_PAGE.clone())?;
             }
+            self.relsize_cache.insert(rel, new_nblocks);
         }
         Ok(())
     }
 
     fn put_slru_page_image(
         &mut self,
-        modification: &mut DatadirModification<T>,
+        modification: &mut DatadirModification<R>,
         kind: SlruKind,
         segno: u32,
         blknum: BlockNumber,
@@ -979,7 +999,7 @@ impl<'a, T: DatadirTimeline> WalIngest<'a, T> {
 
     fn handle_slru_extend(
         &mut self,
-        modification: &mut DatadirModification<T>,
+        modification: &mut DatadirModification<R>,
         kind: SlruKind,
         segno: u32,
         blknum: BlockNumber,
@@ -1032,7 +1052,6 @@ mod tests {
     use super::*;
     use crate::pgdatadir_mapping::create_test_timeline;
     use crate::repository::repo_harness::*;
-    use crate::repository::Timeline;
     use postgres_ffi::pg_constants;
 
     /// Arbitrary relation tag, for testing.
@@ -1043,13 +1062,13 @@ mod tests {
         forknum: 0,
     };
 
-    fn assert_current_logical_size<T: Timeline>(_timeline: &T, _lsn: Lsn) {
+    fn assert_current_logical_size<R: Repository>(_timeline: &DatadirTimeline<R>, _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<R: Repository>(tline: &DatadirTimeline<R>) -> Result<WalIngest<R>> {
         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
@@ -1063,7 +1082,7 @@ mod tests {
     fn test_relsize() -> Result<()> {
         let repo = RepoHarness::create("test_relsize")?.load();
         let tline = create_test_timeline(repo, TIMELINE_ID)?;
-        let mut walingest = init_walingest_test(&*tline)?;
+        let mut walingest = init_walingest_test(&tline)?;
 
         let mut m = tline.begin_modification(Lsn(0x20));
         walingest.put_rel_creation(&mut m, TESTREL_A)?;
@@ -1079,7 +1098,7 @@ mod tests {
         walingest.put_rel_page_image(&mut m, TESTREL_A, 2, TEST_IMG("foo blk 2 at 5"))?;
         m.commit()?;
 
-        assert_current_logical_size(&*tline, Lsn(0x50));
+        assert_current_logical_size(&tline, Lsn(0x50));
 
         // The relation was created at LSN 2, not visible at LSN 1 yet.
         assert_eq!(tline.get_rel_exists(TESTREL_A, Lsn(0x10))?, false);
@@ -1126,7 +1145,7 @@ mod tests {
         let mut m = tline.begin_modification(Lsn(0x60));
         walingest.put_rel_truncation(&mut m, TESTREL_A, 2)?;
         m.commit()?;
-        assert_current_logical_size(&*tline, Lsn(0x60));
+        assert_current_logical_size(&tline, Lsn(0x60));
 
         // Check reported size and contents after truncation
         assert_eq!(tline.get_rel_size(TESTREL_A, Lsn(0x60))?, 2);
@@ -1191,7 +1210,7 @@ mod tests {
     fn test_drop_extend() -> Result<()> {
         let repo = RepoHarness::create("test_drop_extend")?.load();
         let tline = create_test_timeline(repo, TIMELINE_ID)?;
-        let mut walingest = init_walingest_test(&*tline)?;
+        let mut walingest = init_walingest_test(&tline)?;
 
         let mut m = tline.begin_modification(Lsn(0x20));
         walingest.put_rel_page_image(&mut m, TESTREL_A, 0, TEST_IMG("foo blk 0 at 2"))?;
@@ -1231,7 +1250,7 @@ mod tests {
     fn test_truncate_extend() -> Result<()> {
         let repo = RepoHarness::create("test_truncate_extend")?.load();
         let tline = create_test_timeline(repo, TIMELINE_ID)?;
-        let mut walingest = init_walingest_test(&*tline)?;
+        let mut walingest = init_walingest_test(&tline)?;
 
         // Create a 20 MB relation (the size is arbitrary)
         let relsize = 20 * 1024 * 1024 / 8192;
@@ -1319,7 +1338,7 @@ mod tests {
     fn test_large_rel() -> Result<()> {
         let repo = RepoHarness::create("test_large_rel")?.load();
         let tline = create_test_timeline(repo, TIMELINE_ID)?;
-        let mut walingest = init_walingest_test(&*tline)?;
+        let mut walingest = init_walingest_test(&tline)?;
 
         let mut lsn = 0x10;
         for blknum in 0..pg_constants::RELSEG_SIZE + 1 {
@@ -1330,7 +1349,7 @@ mod tests {
             m.commit()?;
         }
 
-        assert_current_logical_size(&*tline, Lsn(lsn));
+        assert_current_logical_size(&tline, Lsn(lsn));
 
         assert_eq!(
             tline.get_rel_size(TESTREL_A, Lsn(lsn))?,
@@ -1346,7 +1365,7 @@ mod tests {
             tline.get_rel_size(TESTREL_A, Lsn(lsn))?,
             pg_constants::RELSEG_SIZE
         );
-        assert_current_logical_size(&*tline, Lsn(lsn));
+        assert_current_logical_size(&tline, Lsn(lsn));
 
         // Truncate another block
         lsn += 0x10;
@@ -1357,7 +1376,7 @@ mod tests {
             tline.get_rel_size(TESTREL_A, Lsn(lsn))?,
             pg_constants::RELSEG_SIZE - 1
         );
-        assert_current_logical_size(&*tline, Lsn(lsn));
+        assert_current_logical_size(&tline, Lsn(lsn));
 
         // Truncate to 1500, and then truncate all the way down to 0, one block at a time
         // This tests the behavior at segment boundaries
@@ -1374,7 +1393,7 @@ mod tests {
 
             size -= 1;
         }
-        assert_current_logical_size(&*tline, Lsn(lsn));
+        assert_current_logical_size(&tline, Lsn(lsn));
 
         Ok(())
     }

pageserver/src/walreceiver.rs

@@ -26,6 +26,7 @@ mod walreceiver_connection;
 use anyhow::{ensure, Context};
 use etcd_broker::Client;
 use itertools::Itertools;
+use once_cell::sync::Lazy;
 use std::cell::Cell;
 use std::collections::{hash_map, HashMap, HashSet};
 use std::future::Future;
@@ -35,13 +36,14 @@ use std::thread_local;
 use std::time::Duration;
 use tokio::{
     select,
-    sync::{mpsc, watch},
+    sync::{mpsc, watch, RwLock},
     task::JoinHandle,
 };
 use tracing::*;
 use url::Url;
 
 use crate::config::PageServerConf;
+use crate::http::models::WalReceiverEntry;
 use crate::tenant_mgr::{self, LocalTimelineUpdate, TenantState};
 use crate::thread_mgr::{self, ThreadKind};
 use utils::zid::{ZTenantId, ZTenantTimelineId, ZTimelineId};
@@ -53,6 +55,23 @@ thread_local! {
     pub(crate) static IS_WAL_RECEIVER: Cell<bool> = Cell::new(false);
 }
 
+/// WAL receiver state for sharing with the outside world.
+/// Only entries for timelines currently available in pageserver are stored.
+static WAL_RECEIVER_ENTRIES: Lazy<RwLock<HashMap<ZTenantTimelineId, WalReceiverEntry>>> =
+    Lazy::new(|| RwLock::new(HashMap::new()));
+
+/// Gets the public WAL streaming entry for a certain timeline.
+pub async fn get_wal_receiver_entry(
+    tenant_id: ZTenantId,
+    timeline_id: ZTimelineId,
+) -> Option<WalReceiverEntry> {
+    WAL_RECEIVER_ENTRIES
+        .read()
+        .await
+        .get(&ZTenantTimelineId::new(tenant_id, timeline_id))
+        .cloned()
+}
+
 /// Sets up the main WAL receiver thread that manages the rest of the subtasks inside of it, per timeline.
 /// See comments in [`wal_receiver_main_thread_loop_step`] for more details on per timeline activities.
 pub fn init_wal_receiver_main_thread(
@@ -66,7 +85,7 @@ pub fn init_wal_receiver_main_thread(
     );
     let broker_prefix = &conf.broker_etcd_prefix;
     info!(
-        "Starting wal receiver main thread, etcd endpoints: {}",
+        "Starting wal receiver main thread, etdc endpoints: {}",
         etcd_endpoints.iter().map(Url::to_string).join(", ")
     );
 
@@ -262,10 +281,13 @@ async fn wal_receiver_main_thread_loop_step<'a>(
                         }
                         None => warn!("Timeline {id} does not have a tenant entry in wal receiver main thread"),
                     };
-                    if let Err(e) = join_confirmation_sender.send(()) {
-                        warn!("cannot send wal_receiver shutdown confirmation {e}")
-                    } else {
-                        info!("confirm walreceiver shutdown for {id}");
+                    {
+                        WAL_RECEIVER_ENTRIES.write().await.remove(&id);
+                        if let Err(e) = join_confirmation_sender.send(()) {
+                            warn!("cannot send wal_receiver shutdown confirmation {e}")
+                        } else {
+                            info!("confirm walreceiver shutdown for {id}");
+                        }
                     }
                 }
                 // Timeline got attached, retrieve all necessary information to start its broker loop and maintain this loop endlessly.
@@ -300,6 +322,17 @@ async fn wal_receiver_main_thread_loop_step<'a>(
                             }
                         };
 
+                    {
+                        WAL_RECEIVER_ENTRIES.write().await.insert(
+                            id,
+                            WalReceiverEntry {
+                                wal_producer_connstr: None,
+                                last_received_msg_lsn: None,
+                                last_received_msg_ts: None,
+                            },
+                        );
+                    }
+
                     vacant_connection_manager_entry.insert(
                         connection_manager::spawn_connection_manager_task(
                             id,

pageserver/src/walreceiver/connection_manager.rs

@@ -17,7 +17,7 @@ use std::{
 };
 
 use anyhow::Context;
-use chrono::{NaiveDateTime, Utc};
+use chrono::{DateTime, Local, NaiveDateTime, Utc};
 use etcd_broker::{
     subscription_key::SubscriptionKey, subscription_value::SkTimelineInfo, BrokerSubscription,
     BrokerUpdate, Client,
@@ -25,25 +25,21 @@ use etcd_broker::{
 use tokio::select;
 use tracing::*;
 
-use crate::{
-    exponential_backoff,
-    repository::{Repository, Timeline},
-    DEFAULT_BASE_BACKOFF_SECONDS, DEFAULT_MAX_BACKOFF_SECONDS,
-};
-use crate::{RepositoryImpl, TimelineImpl};
+use crate::DatadirTimelineImpl;
 use utils::{
     lsn::Lsn,
+    pq_proto::ReplicationFeedback,
     zid::{NodeId, ZTenantTimelineId},
 };
 
-use super::{walreceiver_connection::WalConnectionStatus, TaskEvent, TaskHandle};
+use super::{TaskEvent, TaskHandle};
 
 /// Spawns the loop to take care of the timeline's WAL streaming connection.
 pub(super) fn spawn_connection_manager_task(
     id: ZTenantTimelineId,
     broker_loop_prefix: String,
     mut client: Client,
-    local_timeline: Arc<TimelineImpl>,
+    local_timeline: Arc<DatadirTimelineImpl>,
     wal_connect_timeout: Duration,
     lagging_wal_timeout: Duration,
     max_lsn_wal_lag: NonZeroU64,
@@ -113,26 +109,21 @@ async fn connection_manager_loop_step(
                 }
             } => {
                 let wal_connection = walreceiver_state.wal_connection.as_mut().expect("Should have a connection, as checked by the corresponding select! guard");
-                match wal_connection_update {
+                match &wal_connection_update {
                     TaskEvent::Started => {
+                        wal_connection.latest_connection_update = Utc::now().naive_utc();
                         *walreceiver_state.wal_connection_attempts.entry(wal_connection.sk_id).or_insert(0) += 1;
                     },
-                    TaskEvent::NewEvent(status) => {
-                        if status.has_received_wal {
-                            // Reset connection attempts here only, we know that safekeeper is healthy
-                            // because it can send us a WAL update.
-                            walreceiver_state.wal_connection_attempts.remove(&wal_connection.sk_id);
-                        }
-                        wal_connection.status = status;
+                    TaskEvent::NewEvent(replication_feedback) => {
+                        wal_connection.latest_connection_update = DateTime::<Local>::from(replication_feedback.ps_replytime).naive_utc();
+                        // reset connection attempts here only, the only place where both nodes
+                        // explicitly confirmn with replication feedback that they are connected to each other
+                        walreceiver_state.wal_connection_attempts.remove(&wal_connection.sk_id);
                     },
                     TaskEvent::End(end_result) => {
                         match end_result {
                             Ok(()) => debug!("WAL receiving task finished"),
-                            Err(e) => {
-                                warn!("WAL receiving task failed: {e}");
-                                // If the task failed, set the connection attempts to at least 1, to try other safekeepers.
-                                let _ = *walreceiver_state.wal_connection_attempts.entry(wal_connection.sk_id).or_insert(1);
-                            }
+                            Err(e) => warn!("WAL receiving task failed: {e}"),
                         };
                         walreceiver_state.wal_connection = None;
                     },
@@ -176,7 +167,7 @@ async fn connection_manager_loop_step(
             walreceiver_state
                 .change_connection(
                     new_candidate.safekeeper_id,
-                    new_candidate.wal_source_connstr,
+                    new_candidate.wal_producer_connstr,
                 )
                 .await
         }
@@ -238,11 +229,23 @@ async fn subscribe_for_timeline_updates(
     }
 }
 
+const DEFAULT_BASE_BACKOFF_SECONDS: f64 = 2.0;
+const DEFAULT_MAX_BACKOFF_SECONDS: f64 = 60.0;
+
+async fn exponential_backoff(n: u32, base: f64, max_seconds: f64) {
+    if n == 0 {
+        return;
+    }
+    let seconds_to_wait = base.powf(f64::from(n) - 1.0).min(max_seconds);
+    info!("Backoff: waiting {seconds_to_wait} seconds before proceeding with the task");
+    tokio::time::sleep(Duration::from_secs_f64(seconds_to_wait)).await;
+}
+
 /// All data that's needed to run endless broker loop and keep the WAL streaming connection alive, if possible.
 struct WalreceiverState {
     id: ZTenantTimelineId,
     /// Use pageserver data about the timeline to filter out some of the safekeepers.
-    local_timeline: Arc<TimelineImpl>,
+    local_timeline: Arc<DatadirTimelineImpl>,
     /// 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.
@@ -261,21 +264,10 @@ struct WalreceiverState {
 struct WalConnection {
     /// Current safekeeper pageserver is connected to for WAL streaming.
     sk_id: NodeId,
-    /// Status of the connection.
-    status: WalConnectionStatus,
+    /// Connection task start time or the timestamp of a latest connection message received.
+    latest_connection_update: NaiveDateTime,
     /// WAL streaming task handle.
-    connection_task: TaskHandle<WalConnectionStatus>,
-    /// Have we discovered that other safekeeper has more recent WAL than we do?
-    discovered_new_wal: Option<NewCommittedWAL>,
-}
-
-/// Notion of a new committed WAL, which exists on other safekeeper.
-#[derive(Debug, Clone, Copy)]
-struct NewCommittedWAL {
-    /// LSN of the new committed WAL.
-    lsn: Lsn,
-    /// When we discovered that the new committed WAL exists on other safekeeper.
-    discovered_at: NaiveDateTime,
+    connection_task: TaskHandle<ReplicationFeedback>,
 }
 
 /// Data about the timeline to connect to, received from etcd.
@@ -291,7 +283,7 @@ struct EtcdSkTimeline {
 impl WalreceiverState {
     fn new(
         id: ZTenantTimelineId,
-        local_timeline: Arc<<RepositoryImpl as Repository>::Timeline>,
+        local_timeline: Arc<DatadirTimelineImpl>,
         wal_connect_timeout: Duration,
         lagging_wal_timeout: Duration,
         max_lsn_wal_lag: NonZeroU64,
@@ -309,7 +301,7 @@ impl WalreceiverState {
     }
 
     /// Shuts down the current connection (if any) and immediately starts another one with the given connection string.
-    async fn change_connection(&mut self, new_sk_id: NodeId, new_wal_source_connstr: String) {
+    async fn change_connection(&mut self, new_sk_id: NodeId, new_wal_producer_connstr: String) {
         if let Some(old_connection) = self.wal_connection.take() {
             old_connection.connection_task.shutdown().await
         }
@@ -331,7 +323,7 @@ impl WalreceiverState {
                 .await;
                 super::walreceiver_connection::handle_walreceiver_connection(
                     id,
-                    &new_wal_source_connstr,
+                    &new_wal_producer_connstr,
                     events_sender.as_ref(),
                     cancellation,
                     connect_timeout,
@@ -342,19 +334,10 @@ impl WalreceiverState {
             .instrument(info_span!("walreceiver_connection", id = %id))
         });
 
-        let now = Utc::now().naive_utc();
         self.wal_connection = Some(WalConnection {
             sk_id: new_sk_id,
-            status: WalConnectionStatus {
-                is_connected: false,
-                has_received_wal: false,
-                latest_connection_update: now,
-                latest_wal_update: now,
-                streaming_lsn: None,
-                commit_lsn: None,
-            },
+            latest_connection_update: Utc::now().naive_utc(),
             connection_task: connection_handle,
-            discovered_new_wal: None,
         });
     }
 
@@ -385,16 +368,14 @@ impl WalreceiverState {
     /// Cleans up stale etcd records and checks the rest for the new connection candidate.
     /// Returns a new candidate, if the current state is absent or somewhat lagging, `None` otherwise.
     /// The current rules for approving new candidates:
-    /// * pick a candidate different from the connected safekeeper with biggest `commit_lsn` and lowest failed connection attemps
+    /// * pick from the input data from etcd for currently connected safekeeper (if any)
+    /// * out of the rest input entries, pick one with biggest `commit_lsn` that's after than pageserver's latest Lsn for the timeline
     /// * if there's no such entry, no new candidate found, abort
-    /// * otherwise check if the candidate is much better than the current one
-    ///
-    /// To understand exact rules for determining if the candidate is better than the current one, refer to this function's implementation.
-    /// General rules are following:
-    /// * if connected safekeeper is not present, pick the candidate
-    /// * if we haven't received any updates for some time, pick the candidate
-    /// * if the candidate commit_lsn is much higher than the current one, pick the candidate
-    /// * if connected safekeeper stopped sending us new WAL which is available on other safekeeper, pick the candidate
+    /// * check the current connection time data for staleness, reconnect if stale
+    /// * otherwise, check if etcd updates contain currently connected safekeeper
+    ///     * if not, that means no WAL updates happened after certain time (either none since the connection time or none since the last event after the connection)
+    ///       Reconnect if the time exceeds the threshold.
+    ///     * if there's one, compare its Lsn with the other candidate's, reconnect if candidate's over threshold
     ///
     /// This way we ensure to keep up with the most up-to-date safekeeper and don't try to jump from one safekeeper to another too frequently.
     /// Both thresholds are configured per tenant.
@@ -405,140 +386,65 @@ impl WalreceiverState {
             Some(existing_wal_connection) => {
                 let connected_sk_node = existing_wal_connection.sk_id;
 
-                let (new_sk_id, new_safekeeper_etcd_data, new_wal_source_connstr) =
+                let (new_sk_id, new_safekeeper_etcd_data, new_wal_producer_connstr) =
                     self.select_connection_candidate(Some(connected_sk_node))?;
 
                 let now = Utc::now().naive_utc();
                 if let Ok(latest_interaciton) =
-                    (now - existing_wal_connection.status.latest_connection_update).to_std()
+                    (now - existing_wal_connection.latest_connection_update).to_std()
                 {
-                    // Drop connection if we haven't received keepalive message for a while.
-                    if latest_interaciton > self.wal_connect_timeout {
+                    if latest_interaciton > self.lagging_wal_timeout {
                         return Some(NewWalConnectionCandidate {
                             safekeeper_id: new_sk_id,
-                            wal_source_connstr: new_wal_source_connstr,
-                            reason: ReconnectReason::NoKeepAlives {
-                                last_keep_alive: Some(
-                                    existing_wal_connection.status.latest_connection_update,
+                            wal_producer_connstr: new_wal_producer_connstr,
+                            reason: ReconnectReason::NoWalTimeout {
+                                last_wal_interaction: Some(
+                                    existing_wal_connection.latest_connection_update,
                                 ),
                                 check_time: now,
-                                threshold: self.wal_connect_timeout,
+                                threshold: self.lagging_wal_timeout,
                             },
                         });
                     }
                 }
 
-                if !existing_wal_connection.status.is_connected {
-                    // We haven't connected yet and we shouldn't switch until connection timeout (condition above).
-                    return None;
-                }
-
-                if let Some(current_commit_lsn) = existing_wal_connection.status.commit_lsn {
-                    let new_commit_lsn = new_safekeeper_etcd_data.commit_lsn.unwrap_or(Lsn(0));
-                    // Check if the new candidate has much more WAL than the current one.
-                    match new_commit_lsn.0.checked_sub(current_commit_lsn.0) {
-                        Some(new_sk_lsn_advantage) => {
-                            if new_sk_lsn_advantage >= self.max_lsn_wal_lag.get() {
-                                return Some(NewWalConnectionCandidate {
-                                    safekeeper_id: new_sk_id,
-                                    wal_source_connstr: new_wal_source_connstr,
-                                    reason: ReconnectReason::LaggingWal {
-                                        current_commit_lsn,
-                                        new_commit_lsn,
-                                        threshold: self.max_lsn_wal_lag,
-                                    },
-                                });
+                match self.wal_stream_candidates.get(&connected_sk_node) {
+                    Some(current_connection_etcd_data) => {
+                        let new_lsn = new_safekeeper_etcd_data.commit_lsn.unwrap_or(Lsn(0));
+                        let current_lsn = current_connection_etcd_data
+                            .timeline
+                            .commit_lsn
+                            .unwrap_or(Lsn(0));
+                        match new_lsn.0.checked_sub(current_lsn.0)
+                            {
+                                Some(new_sk_lsn_advantage) => {
+                                    if new_sk_lsn_advantage >= self.max_lsn_wal_lag.get() {
+                                        return Some(
+                                            NewWalConnectionCandidate {
+                                                safekeeper_id: new_sk_id,
+                                                wal_producer_connstr: new_wal_producer_connstr,
+                                                reason: ReconnectReason::LaggingWal { current_lsn, new_lsn, threshold: self.max_lsn_wal_lag },
+                                            });
+                                    }
+                                }
+                                None => debug!("Best SK candidate has its commit Lsn behind the current timeline's latest consistent Lsn"),
                             }
-                        }
-                        None => debug!(
-                            "Best SK candidate has its commit_lsn behind connected SK's commit_lsn"
-                        ),
                     }
-                }
-
-                let current_lsn = match existing_wal_connection.status.streaming_lsn {
-                    Some(lsn) => lsn,
-                    None => self.local_timeline.get_last_record_lsn(),
-                };
-                let current_commit_lsn = existing_wal_connection
-                    .status
-                    .commit_lsn
-                    .unwrap_or(current_lsn);
-                let candidate_commit_lsn = new_safekeeper_etcd_data.commit_lsn.unwrap_or(Lsn(0));
-
-                // Keep discovered_new_wal only if connected safekeeper has not caught up yet.
-                let mut discovered_new_wal = existing_wal_connection
-                    .discovered_new_wal
-                    .filter(|new_wal| new_wal.lsn > current_commit_lsn);
-
-                if discovered_new_wal.is_none() {
-                    // Check if the new candidate has more WAL than the current one.
-                    // If the new candidate has more WAL than the current one, we consider switching to the new candidate.
-                    discovered_new_wal = if candidate_commit_lsn > current_commit_lsn {
-                        trace!(
-                            "New candidate has commit_lsn {}, higher than current_commit_lsn {}",
-                            candidate_commit_lsn,
-                            current_commit_lsn
-                        );
-                        Some(NewCommittedWAL {
-                            lsn: candidate_commit_lsn,
-                            discovered_at: Utc::now().naive_utc(),
+                    None => {
+                        return Some(NewWalConnectionCandidate {
+                            safekeeper_id: new_sk_id,
+                            wal_producer_connstr: new_wal_producer_connstr,
+                            reason: ReconnectReason::NoEtcdDataForExistingConnection,
                         })
-                    } else {
-                        None
-                    };
-                }
-
-                let waiting_for_new_lsn_since = if current_lsn < current_commit_lsn {
-                    // Connected safekeeper has more WAL, but we haven't received updates for some time.
-                    trace!(
-                        "Connected safekeeper has more WAL, but we haven't received updates for {:?}. current_lsn: {}, current_commit_lsn: {}",
-                        (now - existing_wal_connection.status.latest_wal_update).to_std(),
-                        current_lsn,
-                        current_commit_lsn
-                    );
-                    Some(existing_wal_connection.status.latest_wal_update)
-                } else {
-                    discovered_new_wal.as_ref().map(|new_wal| {
-                        // We know that new WAL is available on other safekeeper, but connected safekeeper don't have it.
-                        new_wal
-                            .discovered_at
-                            .max(existing_wal_connection.status.latest_wal_update)
-                    })
-                };
-
-                // If we haven't received any WAL updates for a while and candidate has more WAL, switch to it.
-                if let Some(waiting_for_new_lsn_since) = waiting_for_new_lsn_since {
-                    if let Ok(waiting_for_new_wal) = (now - waiting_for_new_lsn_since).to_std() {
-                        if candidate_commit_lsn > current_commit_lsn
-                            && waiting_for_new_wal > self.lagging_wal_timeout
-                        {
-                            return Some(NewWalConnectionCandidate {
-                                safekeeper_id: new_sk_id,
-                                wal_source_connstr: new_wal_source_connstr,
-                                reason: ReconnectReason::NoWalTimeout {
-                                    current_lsn,
-                                    current_commit_lsn,
-                                    candidate_commit_lsn,
-                                    last_wal_interaction: Some(
-                                        existing_wal_connection.status.latest_wal_update,
-                                    ),
-                                    check_time: now,
-                                    threshold: self.lagging_wal_timeout,
-                                },
-                            });
-                        }
                     }
                 }
-
-                self.wal_connection.as_mut().unwrap().discovered_new_wal = discovered_new_wal;
             }
             None => {
-                let (new_sk_id, _, new_wal_source_connstr) =
+                let (new_sk_id, _, new_wal_producer_connstr) =
                     self.select_connection_candidate(None)?;
                 return Some(NewWalConnectionCandidate {
                     safekeeper_id: new_sk_id,
-                    wal_source_connstr: new_wal_source_connstr,
+                    wal_producer_connstr: new_wal_producer_connstr,
                     reason: ReconnectReason::NoExistingConnection,
                 });
             }
@@ -551,7 +457,7 @@ impl WalreceiverState {
     /// Optionally, omits the given node, to support gracefully switching from a healthy safekeeper to another.
     ///
     /// The candidate that is chosen:
-    /// * has fewest connection attempts from pageserver to safekeeper node (reset every time we receive a WAL message from the node)
+    /// * has fewest connection attempts from pageserver to safekeeper node (reset every time the WAL replication feedback is sent)
     /// * has greatest data Lsn among the ones that are left
     ///
     /// NOTE:
@@ -590,13 +496,14 @@ impl WalreceiverState {
             .max_by_key(|(_, info, _)| info.commit_lsn)
     }
 
-    /// Returns a list of safekeepers that have valid info and ready for connection.
     fn applicable_connection_candidates(
         &self,
     ) -> impl Iterator<Item = (NodeId, &SkTimelineInfo, String)> {
         self.wal_stream_candidates
             .iter()
-            .filter(|(_, info)| info.timeline.commit_lsn.is_some())
+            .filter(|(_, etcd_info)| {
+                etcd_info.timeline.commit_lsn > Some(self.local_timeline.get_last_record_lsn())
+            })
             .filter_map(|(sk_id, etcd_info)| {
                 let info = &etcd_info.timeline;
                 match wal_stream_connection_string(
@@ -612,7 +519,6 @@ impl WalreceiverState {
             })
     }
 
-    /// Remove candidates which haven't sent etcd updates for a while.
     fn cleanup_old_candidates(&mut self) {
         let mut node_ids_to_remove = Vec::with_capacity(self.wal_stream_candidates.len());
 
@@ -639,7 +545,7 @@ impl WalreceiverState {
 #[derive(Debug, PartialEq, Eq)]
 struct NewWalConnectionCandidate {
     safekeeper_id: NodeId,
-    wal_source_connstr: String,
+    wal_producer_connstr: String,
     reason: ReconnectReason,
 }
 
@@ -647,24 +553,17 @@ struct NewWalConnectionCandidate {
 #[derive(Debug, PartialEq, Eq)]
 enum ReconnectReason {
     NoExistingConnection,
+    NoEtcdDataForExistingConnection,
     LaggingWal {
-        current_commit_lsn: Lsn,
-        new_commit_lsn: Lsn,
+        current_lsn: Lsn,
+        new_lsn: Lsn,
         threshold: NonZeroU64,
     },
     NoWalTimeout {
-        current_lsn: Lsn,
-        current_commit_lsn: Lsn,
-        candidate_commit_lsn: Lsn,
         last_wal_interaction: Option<NaiveDateTime>,
         check_time: NaiveDateTime,
         threshold: Duration,
     },
-    NoKeepAlives {
-        last_keep_alive: Option<NaiveDateTime>,
-        check_time: NaiveDateTime,
-        threshold: Duration,
-    },
 }
 
 fn wal_stream_connection_string(
@@ -688,6 +587,7 @@ fn wal_stream_connection_string(
 
 #[cfg(test)]
 mod tests {
+    use std::time::SystemTime;
 
     use crate::repository::{
         repo_harness::{RepoHarness, TIMELINE_ID},
@@ -765,7 +665,7 @@ mod tests {
                         backup_lsn: None,
                         remote_consistent_lsn: None,
                         peer_horizon_lsn: None,
-                        safekeeper_connstr: None,
+                        safekeeper_connstr: Some(DUMMY_SAFEKEEPER_CONNSTR.to_string()),
                     },
                     etcd_version: 0,
                     latest_update: delay_over_threshold,
@@ -791,26 +691,22 @@ mod tests {
         let connected_sk_id = NodeId(0);
         let current_lsn = 100_000;
 
-        let connection_status = WalConnectionStatus {
-            is_connected: true,
-            has_received_wal: true,
-            latest_connection_update: now,
-            latest_wal_update: now,
-            commit_lsn: Some(Lsn(current_lsn)),
-            streaming_lsn: Some(Lsn(current_lsn)),
-        };
-
         state.max_lsn_wal_lag = NonZeroU64::new(100).unwrap();
         state.wal_connection = Some(WalConnection {
             sk_id: connected_sk_id,
-            status: connection_status.clone(),
+            latest_connection_update: now,
             connection_task: TaskHandle::spawn(move |sender, _| async move {
                 sender
-                    .send(TaskEvent::NewEvent(connection_status.clone()))
+                    .send(TaskEvent::NewEvent(ReplicationFeedback {
+                        current_timeline_size: 1,
+                        ps_writelsn: 1,
+                        ps_applylsn: current_lsn,
+                        ps_flushlsn: 1,
+                        ps_replytime: SystemTime::now(),
+                    }))
                     .ok();
                 Ok(())
             }),
-            discovered_new_wal: None,
         });
         state.wal_stream_candidates = HashMap::from([
             (
@@ -906,7 +802,7 @@ mod tests {
             "Should select new safekeeper due to missing connection, even if there's also a lag in the wal over the threshold"
         );
         assert!(only_candidate
-            .wal_source_connstr
+            .wal_producer_connstr
             .contains(DUMMY_SAFEKEEPER_CONNSTR));
 
         let selected_lsn = 100_000;
@@ -971,7 +867,7 @@ mod tests {
             "Should select new safekeeper due to missing connection, even if there's also a lag in the wal over the threshold"
         );
         assert!(biggest_wal_candidate
-            .wal_source_connstr
+            .wal_producer_connstr
             .contains(DUMMY_SAFEKEEPER_CONNSTR));
 
         Ok(())
@@ -1035,6 +931,65 @@ mod tests {
         Ok(())
     }
 
+    #[tokio::test]
+    async fn connection_no_etcd_data_candidate() -> anyhow::Result<()> {
+        let harness = RepoHarness::create("connection_no_etcd_data_candidate")?;
+        let mut state = dummy_state(&harness);
+
+        let now = Utc::now().naive_utc();
+        let current_lsn = Lsn(100_000).align();
+        let connected_sk_id = NodeId(0);
+        let other_sk_id = NodeId(connected_sk_id.0 + 1);
+
+        state.wal_connection = Some(WalConnection {
+            sk_id: connected_sk_id,
+            latest_connection_update: now,
+            connection_task: TaskHandle::spawn(move |sender, _| async move {
+                sender
+                    .send(TaskEvent::NewEvent(ReplicationFeedback {
+                        current_timeline_size: 1,
+                        ps_writelsn: current_lsn.0,
+                        ps_applylsn: 1,
+                        ps_flushlsn: 1,
+                        ps_replytime: SystemTime::now(),
+                    }))
+                    .ok();
+                Ok(())
+            }),
+        });
+        state.wal_stream_candidates = HashMap::from([(
+            other_sk_id,
+            EtcdSkTimeline {
+                timeline: SkTimelineInfo {
+                    last_log_term: None,
+                    flush_lsn: None,
+                    commit_lsn: Some(Lsn(1 + state.max_lsn_wal_lag.get())),
+                    backup_lsn: None,
+                    remote_consistent_lsn: None,
+                    peer_horizon_lsn: None,
+                    safekeeper_connstr: Some(DUMMY_SAFEKEEPER_CONNSTR.to_string()),
+                },
+                etcd_version: 0,
+                latest_update: now,
+            },
+        )]);
+
+        let only_candidate = state
+            .next_connection_candidate()
+            .expect("Expected one candidate selected out of the only data option, but got none");
+        assert_eq!(only_candidate.safekeeper_id, other_sk_id);
+        assert_eq!(
+            only_candidate.reason,
+            ReconnectReason::NoEtcdDataForExistingConnection,
+            "Should select new safekeeper due to missing etcd data, even if there's an existing connection with this safekeeper"
+        );
+        assert!(only_candidate
+            .wal_producer_connstr
+            .contains(DUMMY_SAFEKEEPER_CONNSTR));
+
+        Ok(())
+    }
+
     #[tokio::test]
     async fn lsn_wal_over_threshhold_current_candidate() -> anyhow::Result<()> {
         let harness = RepoHarness::create("lsn_wal_over_threshcurrent_candidate")?;
@@ -1045,25 +1000,21 @@ mod tests {
         let connected_sk_id = NodeId(0);
         let new_lsn = Lsn(current_lsn.0 + state.max_lsn_wal_lag.get() + 1);
 
-        let connection_status = WalConnectionStatus {
-            is_connected: true,
-            has_received_wal: true,
-            latest_connection_update: now,
-            latest_wal_update: now,
-            commit_lsn: Some(current_lsn),
-            streaming_lsn: Some(current_lsn),
-        };
-
         state.wal_connection = Some(WalConnection {
             sk_id: connected_sk_id,
-            status: connection_status.clone(),
+            latest_connection_update: now,
             connection_task: TaskHandle::spawn(move |sender, _| async move {
                 sender
-                    .send(TaskEvent::NewEvent(connection_status.clone()))
+                    .send(TaskEvent::NewEvent(ReplicationFeedback {
+                        current_timeline_size: 1,
+                        ps_writelsn: current_lsn.0,
+                        ps_applylsn: 1,
+                        ps_flushlsn: 1,
+                        ps_replytime: SystemTime::now(),
+                    }))
                     .ok();
                 Ok(())
             }),
-            discovered_new_wal: None,
         });
         state.wal_stream_candidates = HashMap::from([
             (
@@ -1108,49 +1059,45 @@ mod tests {
         assert_eq!(
             over_threshcurrent_candidate.reason,
             ReconnectReason::LaggingWal {
-                current_commit_lsn: current_lsn,
-                new_commit_lsn: new_lsn,
+                current_lsn,
+                new_lsn,
                 threshold: state.max_lsn_wal_lag
             },
             "Should select bigger WAL safekeeper if it starts to lag enough"
         );
         assert!(over_threshcurrent_candidate
-            .wal_source_connstr
+            .wal_producer_connstr
             .contains("advanced by Lsn safekeeper"));
 
         Ok(())
     }
 
     #[tokio::test]
-    async fn timeout_connection_threshhold_current_candidate() -> anyhow::Result<()> {
-        let harness = RepoHarness::create("timeout_connection_threshhold_current_candidate")?;
+    async fn timeout_wal_over_threshhold_current_candidate() -> anyhow::Result<()> {
+        let harness = RepoHarness::create("timeout_wal_over_threshhold_current_candidate")?;
         let mut state = dummy_state(&harness);
         let current_lsn = Lsn(100_000).align();
         let now = Utc::now().naive_utc();
 
-        let wal_connect_timeout = chrono::Duration::from_std(state.wal_connect_timeout)?;
+        let lagging_wal_timeout = chrono::Duration::from_std(state.lagging_wal_timeout)?;
         let time_over_threshold =
-            Utc::now().naive_utc() - wal_connect_timeout - wal_connect_timeout;
-
-        let connection_status = WalConnectionStatus {
-            is_connected: true,
-            has_received_wal: true,
-            latest_connection_update: time_over_threshold,
-            latest_wal_update: time_over_threshold,
-            commit_lsn: Some(current_lsn),
-            streaming_lsn: Some(current_lsn),
-        };
+            Utc::now().naive_utc() - lagging_wal_timeout - lagging_wal_timeout;
 
         state.wal_connection = Some(WalConnection {
             sk_id: NodeId(1),
-            status: connection_status.clone(),
+            latest_connection_update: time_over_threshold,
             connection_task: TaskHandle::spawn(move |sender, _| async move {
                 sender
-                    .send(TaskEvent::NewEvent(connection_status.clone()))
+                    .send(TaskEvent::NewEvent(ReplicationFeedback {
+                        current_timeline_size: 1,
+                        ps_writelsn: current_lsn.0,
+                        ps_applylsn: 1,
+                        ps_flushlsn: 1,
+                        ps_replytime: SystemTime::now(),
+                    }))
                     .ok();
                 Ok(())
             }),
-            discovered_new_wal: None,
         });
         state.wal_stream_candidates = HashMap::from([(
             NodeId(0),
@@ -1175,52 +1122,38 @@ mod tests {
 
         assert_eq!(over_threshcurrent_candidate.safekeeper_id, NodeId(0));
         match over_threshcurrent_candidate.reason {
-            ReconnectReason::NoKeepAlives {
-                last_keep_alive,
+            ReconnectReason::NoWalTimeout {
+                last_wal_interaction,
                 threshold,
                 ..
             } => {
-                assert_eq!(last_keep_alive, Some(time_over_threshold));
+                assert_eq!(last_wal_interaction, Some(time_over_threshold));
                 assert_eq!(threshold, state.lagging_wal_timeout);
             }
             unexpected => panic!("Unexpected reason: {unexpected:?}"),
         }
         assert!(over_threshcurrent_candidate
-            .wal_source_connstr
+            .wal_producer_connstr
             .contains(DUMMY_SAFEKEEPER_CONNSTR));
 
         Ok(())
     }
 
     #[tokio::test]
-    async fn timeout_wal_over_threshhold_current_candidate() -> anyhow::Result<()> {
-        let harness = RepoHarness::create("timeout_wal_over_threshhold_current_candidate")?;
+    async fn timeout_connection_over_threshhold_current_candidate() -> anyhow::Result<()> {
+        let harness = RepoHarness::create("timeout_connection_over_threshhold_current_candidate")?;
         let mut state = dummy_state(&harness);
         let current_lsn = Lsn(100_000).align();
-        let new_lsn = Lsn(100_100).align();
         let now = Utc::now().naive_utc();
 
         let lagging_wal_timeout = chrono::Duration::from_std(state.lagging_wal_timeout)?;
         let time_over_threshold =
             Utc::now().naive_utc() - lagging_wal_timeout - lagging_wal_timeout;
 
-        let connection_status = WalConnectionStatus {
-            is_connected: true,
-            has_received_wal: true,
-            latest_connection_update: now,
-            latest_wal_update: time_over_threshold,
-            commit_lsn: Some(current_lsn),
-            streaming_lsn: Some(current_lsn),
-        };
-
         state.wal_connection = Some(WalConnection {
             sk_id: NodeId(1),
-            status: connection_status,
+            latest_connection_update: time_over_threshold,
             connection_task: TaskHandle::spawn(move |_, _| async move { Ok(()) }),
-            discovered_new_wal: Some(NewCommittedWAL {
-                discovered_at: time_over_threshold,
-                lsn: new_lsn,
-            }),
         });
         state.wal_stream_candidates = HashMap::from([(
             NodeId(0),
@@ -1228,7 +1161,7 @@ mod tests {
                 timeline: SkTimelineInfo {
                     last_log_term: None,
                     flush_lsn: None,
-                    commit_lsn: Some(new_lsn),
+                    commit_lsn: Some(current_lsn),
                     backup_lsn: None,
                     remote_consistent_lsn: None,
                     peer_horizon_lsn: None,
@@ -1246,23 +1179,17 @@ mod tests {
         assert_eq!(over_threshcurrent_candidate.safekeeper_id, NodeId(0));
         match over_threshcurrent_candidate.reason {
             ReconnectReason::NoWalTimeout {
-                current_lsn,
-                current_commit_lsn,
-                candidate_commit_lsn,
                 last_wal_interaction,
                 threshold,
                 ..
             } => {
-                assert_eq!(current_lsn, current_lsn);
-                assert_eq!(current_commit_lsn, current_lsn);
-                assert_eq!(candidate_commit_lsn, new_lsn);
                 assert_eq!(last_wal_interaction, Some(time_over_threshold));
                 assert_eq!(threshold, state.lagging_wal_timeout);
             }
             unexpected => panic!("Unexpected reason: {unexpected:?}"),
         }
         assert!(over_threshcurrent_candidate
-            .wal_source_connstr
+            .wal_producer_connstr
             .contains(DUMMY_SAFEKEEPER_CONNSTR));
 
         Ok(())
@@ -1276,13 +1203,16 @@ mod tests {
                 tenant_id: harness.tenant_id,
                 timeline_id: TIMELINE_ID,
             },
-            local_timeline: harness
-                .load()
-                .create_empty_timeline(TIMELINE_ID, Lsn(0))
-                .expect("Failed to create an empty timeline for dummy wal connection manager"),
+            local_timeline: Arc::new(DatadirTimelineImpl::new(
+                harness
+                    .load()
+                    .create_empty_timeline(TIMELINE_ID, Lsn(0))
+                    .expect("Failed to create an empty timeline for dummy wal connection manager"),
+                10_000,
+            )),
             wal_connect_timeout: Duration::from_secs(1),
             lagging_wal_timeout: Duration::from_secs(1),
-            max_lsn_wal_lag: NonZeroU64::new(1024 * 1024).unwrap(),
+            max_lsn_wal_lag: NonZeroU64::new(1).unwrap(),
             wal_connection: None,
             wal_stream_candidates: HashMap::new(),
             wal_connection_attempts: HashMap::new(),

pageserver/src/walreceiver/walreceiver_connection.rs

@@ -8,57 +8,37 @@ use std::{
 
 use anyhow::{bail, ensure, Context};
 use bytes::BytesMut;
-use chrono::{NaiveDateTime, Utc};
 use fail::fail_point;
-use futures::StreamExt;
 use postgres::{SimpleQueryMessage, SimpleQueryRow};
 use postgres_protocol::message::backend::ReplicationMessage;
 use postgres_types::PgLsn;
 use tokio::{pin, select, sync::watch, time};
 use tokio_postgres::{replication::ReplicationStream, Client};
+use tokio_stream::StreamExt;
 use tracing::{debug, error, info, info_span, trace, warn, Instrument};
 
 use super::TaskEvent;
 use crate::{
-    layered_repository::WalReceiverInfo,
-    pgdatadir_mapping::DatadirTimeline,
+    http::models::WalReceiverEntry,
     repository::{Repository, Timeline},
     tenant_mgr,
     walingest::WalIngest,
-    walrecord::DecodedWALRecord,
 };
 use postgres_ffi::waldecoder::WalStreamDecoder;
 use utils::{lsn::Lsn, pq_proto::ReplicationFeedback, zid::ZTenantTimelineId};
 
-/// Status of the connection.
-#[derive(Debug, Clone)]
-pub struct WalConnectionStatus {
-    /// If we were able to initiate a postgres connection, this means that safekeeper process is at least running.
-    pub is_connected: bool,
-    /// Defines a healthy connection as one on which we have received at least some WAL bytes.
-    pub has_received_wal: bool,
-    /// Connection establishment time or the timestamp of a latest connection message received.
-    pub latest_connection_update: NaiveDateTime,
-    /// Time of the latest WAL message received.
-    pub latest_wal_update: NaiveDateTime,
-    /// Latest WAL update contained WAL up to this LSN. Next WAL message with start from that LSN.
-    pub streaming_lsn: Option<Lsn>,
-    /// Latest commit_lsn received from the safekeeper. Can be zero if no message has been received yet.
-    pub commit_lsn: Option<Lsn>,
-}
-
-/// Open a connection to the given safekeeper and receive WAL, sending back progress
-/// messages as we go.
+/// Opens a conneciton to the given wal producer and streams the WAL, sending progress messages during streaming.
 pub async fn handle_walreceiver_connection(
     id: ZTenantTimelineId,
-    wal_source_connstr: &str,
-    events_sender: &watch::Sender<TaskEvent<WalConnectionStatus>>,
+    wal_producer_connstr: &str,
+    events_sender: &watch::Sender<TaskEvent<ReplicationFeedback>>,
     mut cancellation: watch::Receiver<()>,
     connect_timeout: Duration,
 ) -> anyhow::Result<()> {
     // Connect to the database in replication mode.
-    info!("connecting to {wal_source_connstr}");
-    let connect_cfg = format!("{wal_source_connstr} application_name=pageserver replication=true");
+    info!("connecting to {wal_producer_connstr}");
+    let connect_cfg =
+        format!("{wal_producer_connstr} application_name=pageserver replication=true");
 
     let (mut replication_client, connection) = time::timeout(
         connect_timeout,
@@ -67,26 +47,12 @@ pub async fn handle_walreceiver_connection(
     .await
     .context("Timed out while waiting for walreceiver connection to open")?
     .context("Failed to open walreceiver conection")?;
-
-    info!("connected!");
-    let mut connection_status = WalConnectionStatus {
-        is_connected: true,
-        has_received_wal: false,
-        latest_connection_update: Utc::now().naive_utc(),
-        latest_wal_update: Utc::now().naive_utc(),
-        streaming_lsn: None,
-        commit_lsn: None,
-    };
-    if let Err(e) = events_sender.send(TaskEvent::NewEvent(connection_status.clone())) {
-        warn!("Wal connection event listener dropped right after connection init, aborting the connection: {e}");
-        return Ok(());
-    }
-
     // The connection object performs the actual communication with the database,
     // so spawn it off to run on its own.
     let mut connection_cancellation = cancellation.clone();
     tokio::spawn(
         async move {
+            info!("connected!");
             select! {
                     connection_result = connection => match connection_result{
                             Ok(()) => info!("Walreceiver db connection closed"),
@@ -116,14 +82,6 @@ pub async fn handle_walreceiver_connection(
 
     let identify = identify_system(&mut replication_client).await?;
     info!("{identify:?}");
-
-    connection_status.latest_connection_update = Utc::now().naive_utc();
-    if let Err(e) = events_sender.send(TaskEvent::NewEvent(connection_status.clone())) {
-        warn!("Wal connection event listener dropped after IDENTIFY_SYSTEM, aborting the connection: {e}");
-        return Ok(());
-    }
-
-    // NB: this is a flush_lsn, not a commit_lsn.
     let end_of_wal = Lsn::from(u64::from(identify.xlogpos));
     let mut caught_up = false;
     let ZTenantTimelineId {
@@ -158,7 +116,7 @@ pub async fn handle_walreceiver_connection(
     // There might be some padding after the last full record, skip it.
     startpoint += startpoint.calc_padding(8u32);
 
-    info!("last_record_lsn {last_rec_lsn} starting replication from {startpoint}, safekeeper is at {end_of_wal}...");
+    info!("last_record_lsn {last_rec_lsn} starting replication from {startpoint}, server is at {end_of_wal}...");
 
     let query = format!("START_REPLICATION PHYSICAL {startpoint}");
 
@@ -180,33 +138,6 @@ pub async fn handle_walreceiver_connection(
         }
     } {
         let replication_message = replication_message?;
-        let now = Utc::now().naive_utc();
-
-        // Update the connection status before processing the message. If the message processing
-        // fails (e.g. in walingest), we still want to know latests LSNs from the safekeeper.
-        match &replication_message {
-            ReplicationMessage::XLogData(xlog_data) => {
-                connection_status.latest_connection_update = now;
-                connection_status.commit_lsn = Some(Lsn::from(xlog_data.wal_end()));
-                connection_status.streaming_lsn = Some(Lsn::from(
-                    xlog_data.wal_start() + xlog_data.data().len() as u64,
-                ));
-                if !xlog_data.data().is_empty() {
-                    connection_status.latest_wal_update = now;
-                    connection_status.has_received_wal = true;
-                }
-            }
-            ReplicationMessage::PrimaryKeepAlive(keepalive) => {
-                connection_status.latest_connection_update = now;
-                connection_status.commit_lsn = Some(Lsn::from(keepalive.wal_end()));
-            }
-            &_ => {}
-        };
-        if let Err(e) = events_sender.send(TaskEvent::NewEvent(connection_status.clone())) {
-            warn!("Wal connection event listener dropped, aborting the connection: {e}");
-            return Ok(());
-        }
-
         let status_update = match replication_message {
             ReplicationMessage::XLogData(xlog_data) => {
                 // Pass the WAL data to the decoder, and see if we can decode
@@ -219,25 +150,19 @@ pub async fn handle_walreceiver_connection(
 
                 waldecoder.feed_bytes(data);
 
-                {
-                    let mut decoded = DecodedWALRecord::default();
-                    let mut modification = timeline.begin_modification(endlsn);
-                    while let Some((lsn, recdata)) = waldecoder.poll_decode()? {
-                        // let _enter = info_span!("processing record", lsn = %lsn).entered();
+                while let Some((lsn, recdata)) = waldecoder.poll_decode()? {
+                    let _enter = info_span!("processing record", lsn = %lsn).entered();
 
-                        // It is important to deal with the aligned records as lsn in getPage@LSN is
-                        // aligned and can be several bytes bigger. Without this alignment we are
-                        // at risk of hitting a deadlock.
-                        ensure!(lsn.is_aligned());
+                    // It is important to deal with the aligned records as lsn in getPage@LSN is
+                    // aligned and can be several bytes bigger. Without this alignment we are
+                    // at risk of hitting a deadlock.
+                    ensure!(lsn.is_aligned());
 
-                        walingest
-                            .ingest_record(recdata, lsn, &mut modification, &mut decoded)
-                            .context("could not ingest record at {lsn}")?;
+                    walingest.ingest_record(&timeline, recdata, lsn)?;
 
-                        fail_point!("walreceiver-after-ingest");
+                    fail_point!("walreceiver-after-ingest");
 
-                        last_rec_lsn = lsn;
-                    }
+                    last_rec_lsn = lsn;
                 }
 
                 if !caught_up && endlsn >= end_of_wal {
@@ -245,6 +170,16 @@ pub async fn handle_walreceiver_connection(
                     caught_up = true;
                 }
 
+                let timeline_to_check = Arc::clone(&timeline.tline);
+                tokio::task::spawn_blocking(move || timeline_to_check.check_checkpoint_distance())
+                    .await
+                    .with_context(|| {
+                        format!("Spawned checkpoint check task panicked for timeline {id}")
+                    })?
+                    .with_context(|| {
+                        format!("Failed to check checkpoint distance for timeline {id}")
+                    })?;
+
                 Some(endlsn)
             }
 
@@ -265,12 +200,6 @@ pub async fn handle_walreceiver_connection(
             _ => None,
         };
 
-        let timeline_to_check = Arc::clone(&timeline);
-        tokio::task::spawn_blocking(move || timeline_to_check.check_checkpoint_distance())
-            .await
-            .with_context(|| format!("Spawned checkpoint check task panicked for timeline {id}"))?
-            .with_context(|| format!("Failed to check checkpoint distance for timeline {id}"))?;
-
         if let Some(last_lsn) = status_update {
             let remote_index = repo.get_remote_index();
             let timeline_remote_consistent_lsn = remote_index
@@ -289,22 +218,27 @@ pub async fn handle_walreceiver_connection(
             // The last LSN we processed. It is not guaranteed to survive pageserver crash.
             let write_lsn = u64::from(last_lsn);
             // `disk_consistent_lsn` is the LSN at which page server guarantees local persistence of all received data
-            let flush_lsn = u64::from(timeline.get_disk_consistent_lsn());
+            let flush_lsn = u64::from(timeline.tline.get_disk_consistent_lsn());
             // The last LSN that is synced to remote storage and is guaranteed to survive pageserver crash
             // Used by safekeepers to remove WAL preceding `remote_consistent_lsn`.
             let apply_lsn = u64::from(timeline_remote_consistent_lsn);
             let ts = SystemTime::now();
 
-            // Update the status about what we just received. This is shown in the mgmt API.
-            let last_received_wal = WalReceiverInfo {
-                wal_source_connstr: wal_source_connstr.to_owned(),
-                last_received_msg_lsn: last_lsn,
-                last_received_msg_ts: ts
-                    .duration_since(SystemTime::UNIX_EPOCH)
-                    .expect("Received message time should be before UNIX EPOCH!")
-                    .as_micros(),
-            };
-            *timeline.last_received_wal.lock().unwrap() = Some(last_received_wal);
+            // Update the current WAL receiver's data stored inside the global hash table `WAL_RECEIVERS`
+            {
+                super::WAL_RECEIVER_ENTRIES.write().await.insert(
+                    id,
+                    WalReceiverEntry {
+                        wal_producer_connstr: Some(wal_producer_connstr.to_owned()),
+                        last_received_msg_lsn: Some(last_lsn),
+                        last_received_msg_ts: Some(
+                            ts.duration_since(SystemTime::UNIX_EPOCH)
+                                .expect("Received message time should be before UNIX EPOCH!")
+                                .as_micros(),
+                        ),
+                    },
+                );
+            }
 
             // Send zenith feedback message.
             // Regular standby_status_update fields are put into this message.
@@ -324,6 +258,10 @@ pub async fn handle_walreceiver_connection(
                 .as_mut()
                 .zenith_status_update(data.len() as u64, &data)
                 .await?;
+            if let Err(e) = events_sender.send(TaskEvent::NewEvent(zenith_status_update)) {
+                warn!("Wal connection event listener dropped, aborting the connection: {e}");
+                return Ok(());
+            }
         }
     }
 

pageserver/src/walrecord.rs

@@ -96,7 +96,6 @@ impl DecodedBkpBlock {
     }
 }
 
-#[derive(Default)]
 pub struct DecodedWALRecord {
     pub xl_xid: TransactionId,
     pub xl_info: u8,
@@ -506,17 +505,7 @@ impl XlMultiXactTruncate {
 //      block data
 //      ...
 //      main data
-//
-//
-// For performance reasons, the caller provides the DecodedWALRecord struct and the function just fills it in.
-// It would be more natural for this function to return a DecodedWALRecord as return value,
-// but reusing the caller-supplied struct avoids an allocation.
-// This code is in the hot path for digesting incoming WAL, and is very performance sensitive.
-//
-pub fn decode_wal_record(
-    record: Bytes,
-    decoded: &mut DecodedWALRecord,
-) -> Result<(), DeserializeError> {
+pub fn decode_wal_record(record: Bytes) -> Result<DecodedWALRecord, DeserializeError> {
     let mut rnode_spcnode: u32 = 0;
     let mut rnode_dbnode: u32 = 0;
     let mut rnode_relnode: u32 = 0;
@@ -545,7 +534,7 @@ pub fn decode_wal_record(
     let mut blocks_total_len: u32 = 0;
     let mut main_data_len = 0;
     let mut datatotal: u32 = 0;
-    decoded.blocks.clear();
+    let mut blocks: Vec<DecodedBkpBlock> = Vec::new();
 
     // 2. Decode the headers.
     // XLogRecordBlockHeaders if any,
@@ -724,7 +713,7 @@ pub fn decode_wal_record(
                     blk.blkno
                 );
 
-                decoded.blocks.push(blk);
+                blocks.push(blk);
             }
 
             _ => {
@@ -735,7 +724,7 @@ pub fn decode_wal_record(
 
     // 3. Decode blocks.
     let mut ptr = record.len() - buf.remaining();
-    for blk in decoded.blocks.iter_mut() {
+    for blk in blocks.iter_mut() {
         if blk.has_image {
             blk.bimg_offset = ptr as u32;
             ptr += blk.bimg_len as usize;
@@ -755,13 +744,14 @@ pub fn decode_wal_record(
         assert_eq!(buf.remaining(), main_data_len as usize);
     }
 
-    decoded.xl_xid = xlogrec.xl_xid;
-    decoded.xl_info = xlogrec.xl_info;
-    decoded.xl_rmid = xlogrec.xl_rmid;
-    decoded.record = record;
-    decoded.main_data_offset = main_data_offset;
-
-    Ok(())
+    Ok(DecodedWALRecord {
+        xl_xid: xlogrec.xl_xid,
+        xl_info: xlogrec.xl_info,
+        xl_rmid: xlogrec.xl_rmid,
+        record,
+        blocks,
+        main_data_offset,
+    })
 }
 
 ///

pageserver/src/walredo.rs

@@ -20,8 +20,8 @@
 //!
 use byteorder::{ByteOrder, LittleEndian};
 use bytes::{BufMut, Bytes, BytesMut};
+use lazy_static::lazy_static;
 use nix::poll::*;
-use once_cell::sync::Lazy;
 use serde::Serialize;
 use std::fs;
 use std::fs::OpenOptions;
@@ -105,27 +105,21 @@ impl crate::walredo::WalRedoManager for DummyRedoManager {
 // We collect the time spent in actual WAL redo ('redo'), and time waiting
 // for access to the postgres process ('wait') since there is only one for
 // each tenant.
-
-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")
-});
-
-static WAL_REDO_WAIT_TIME: Lazy<Histogram> = Lazy::new(|| {
-    register_histogram!(
+lazy_static! {
+    static ref WAL_REDO_TIME: Histogram =
+        register_histogram!("pageserver_wal_redo_seconds", "Time spent on WAL redo")
+            .expect("failed to define a metric");
+    static ref WAL_REDO_WAIT_TIME: Histogram = register_histogram!(
         "pageserver_wal_redo_wait_seconds",
         "Time spent waiting for access to the WAL redo process"
     )
-    .expect("failed to define a metric")
-});
-
-static WAL_REDO_RECORD_COUNTER: Lazy<IntCounter> = Lazy::new(|| {
-    register_int_counter!(
+    .expect("failed to define a metric");
+    static ref WAL_REDO_RECORD_COUNTER: IntCounter = register_int_counter!(
         "pageserver_replayed_wal_records_total",
         "Number of WAL records replayed in WAL redo process"
     )
-    .unwrap()
-});
+    .unwrap();
+}
 
 ///
 /// This is the real implementation that uses a Postgres process to

proxy/Cargo.toml

@@ -7,7 +7,6 @@ edition = "2021"
 anyhow = "1.0"
 async-trait = "0.1"
 base64 = "0.13.0"
-bstr = "0.2.17"
 bytes = { version = "1.0.1", features = ['serde'] }
 clap = "3.0"
 futures = "0.3.13"
@@ -15,7 +14,7 @@ hashbrown = "0.11.2"
 hex = "0.4.3"
 hmac = "0.12.1"
 hyper = "0.14"
-once_cell = "1.13.0"
+lazy_static = "1.4.0"
 md5 = "0.7.0"
 parking_lot = "0.12"
 pin-project-lite = "0.2.7"

proxy/src/auth.rs

@@ -1,18 +1,15 @@
 //! Client authentication mechanisms.
 
 pub mod backend;
-pub use backend::{BackendType, DatabaseInfo};
+pub use backend::DatabaseInfo;
 
 mod credentials;
 pub use credentials::ClientCredentials;
 
-mod password_hack;
-use password_hack::PasswordHackPayload;
-
 mod flow;
 pub use flow::*;
 
-use crate::error::UserFacingError;
+use crate::{error::UserFacingError, waiters};
 use std::io;
 use thiserror::Error;
 
@@ -22,54 +19,52 @@ pub type Result<T> = std::result::Result<T, AuthError>;
 /// Common authentication error.
 #[derive(Debug, Error)]
 pub enum AuthErrorImpl {
-    // This will be dropped in the future.
+    /// Authentication error reported by the console.
     #[error(transparent)]
-    Legacy(#[from] backend::LegacyAuthError),
+    Console(#[from] backend::AuthError),
 
     #[error(transparent)]
-    Link(#[from] backend::LinkAuthError),
+    GetAuthInfo(#[from] backend::console::ConsoleAuthError),
 
-    #[error(transparent)]
-    GetAuthInfo(#[from] backend::GetAuthInfoError),
-
-    #[error(transparent)]
-    WakeCompute(#[from] backend::WakeComputeError),
-
-    /// SASL protocol errors (includes [SCRAM](crate::scram)).
     #[error(transparent)]
     Sasl(#[from] crate::sasl::Error),
 
-    #[error("Unsupported authentication method: {0}")]
-    BadAuthMethod(Box<str>),
+    /// For passwords that couldn't be processed by [`backend::legacy_console::parse_password`].
+    #[error("Malformed password message")]
+    MalformedPassword,
 
-    #[error("Malformed password message: {0}")]
-    MalformedPassword(&'static str),
-
-    #[error(
-        "Project name is not specified. \
-        Either please upgrade the postgres client library (libpq) for SNI support \
-        or pass the project name as a parameter: '&options=project%3D<project-name>'. \
-        See more at https://neon.tech/sni"
-    )]
-    MissingProjectName,
-
-    /// Errors produced by e.g. [`crate::stream::PqStream`].
+    /// Errors produced by [`crate::stream::PqStream`].
     #[error(transparent)]
     Io(#[from] io::Error),
 }
 
+impl AuthErrorImpl {
+    pub fn auth_failed(msg: impl Into<String>) -> Self {
+        Self::Console(backend::AuthError::auth_failed(msg))
+    }
+}
+
+impl From<waiters::RegisterError> for AuthErrorImpl {
+    fn from(e: waiters::RegisterError) -> Self {
+        Self::Console(backend::AuthError::from(e))
+    }
+}
+
+impl From<waiters::WaitError> for AuthErrorImpl {
+    fn from(e: waiters::WaitError) -> Self {
+        Self::Console(backend::AuthError::from(e))
+    }
+}
+
 #[derive(Debug, Error)]
 #[error(transparent)]
 pub struct AuthError(Box<AuthErrorImpl>);
 
-impl AuthError {
-    pub fn bad_auth_method(name: impl Into<Box<str>>) -> Self {
-        AuthErrorImpl::BadAuthMethod(name.into()).into()
-    }
-}
-
-impl<E: Into<AuthErrorImpl>> From<E> for AuthError {
-    fn from(e: E) -> Self {
+impl<T> From<T> for AuthError
+where
+    AuthErrorImpl: From<T>,
+{
+    fn from(e: T) -> Self {
         Self(Box::new(e.into()))
     }
 }
@@ -78,14 +73,10 @@ impl UserFacingError for AuthError {
     fn to_string_client(&self) -> String {
         use AuthErrorImpl::*;
         match self.0.as_ref() {
-            Legacy(e) => e.to_string_client(),
-            Link(e) => e.to_string_client(),
+            Console(e) => e.to_string_client(),
             GetAuthInfo(e) => e.to_string_client(),
-            WakeCompute(e) => e.to_string_client(),
             Sasl(e) => e.to_string_client(),
-            BadAuthMethod(_) => self.to_string(),
-            MalformedPassword(_) => self.to_string(),
-            MissingProjectName => self.to_string(),
+            MalformedPassword => self.to_string(),
             _ => "Internal error".to_string(),
         }
     }

proxy/src/auth/backend.rs

@@ -1,26 +1,26 @@
+mod legacy_console;
+mod link;
 mod postgres;
 
-mod link;
-pub use link::LinkAuthError;
+pub mod console;
 
-mod console;
-pub use console::{GetAuthInfoError, WakeComputeError};
-
-mod legacy_console;
-pub use legacy_console::LegacyAuthError;
+pub use legacy_console::{AuthError, AuthErrorImpl};
 
+use super::ClientCredentials;
 use crate::{
-    auth::{self, AuthFlow, ClientCredentials},
-    compute, config, mgmt,
+    compute,
+    config::{AuthBackendType, ProxyConfig},
+    mgmt,
     stream::PqStream,
     waiters::{self, Waiter, Waiters},
 };
-
-use once_cell::sync::Lazy;
+use lazy_static::lazy_static;
 use serde::{Deserialize, Serialize};
 use tokio::io::{AsyncRead, AsyncWrite};
 
-static CPLANE_WAITERS: Lazy<Waiters<mgmt::ComputeReady>> = Lazy::new(Default::default);
+lazy_static! {
+    static ref CPLANE_WAITERS: Waiters<mgmt::ComputeReady> = Default::default();
+}
 
 /// Give caller an opportunity to wait for the cloud's reply.
 pub async fn with_waiter<R, T, E>(
@@ -78,158 +78,32 @@ impl From<DatabaseInfo> for tokio_postgres::Config {
     }
 }
 
-/// This type serves two purposes:
-///
-/// * When `T` is `()`, it's just a regular auth backend selector
-///   which we use in [`crate::config::ProxyConfig`].
-///
-/// * However, when we substitute `T` with [`ClientCredentials`],
-///   this helps us provide the credentials only to those auth
-///   backends which require them for the authentication process.
-#[derive(Debug, Clone, Copy, PartialEq)]
-pub enum BackendType<T> {
-    /// Legacy Cloud API (V1) + link auth.
-    LegacyConsole(T),
-    /// Current Cloud API (V2).
-    Console(T),
-    /// Local mock of Cloud API (V2).
-    Postgres(T),
-    /// Authentication via a web browser.
-    Link,
-}
-
-impl<T> BackendType<T> {
-    /// Very similar to [`std::option::Option::map`].
-    /// Maps [`BackendType<T>`] to [`BackendType<R>`] by applying
-    /// a function to a contained value.
-    pub fn map<R>(self, f: impl FnOnce(T) -> R) -> BackendType<R> {
-        use BackendType::*;
-        match self {
-            LegacyConsole(x) => LegacyConsole(f(x)),
-            Console(x) => Console(f(x)),
-            Postgres(x) => Postgres(f(x)),
-            Link => Link,
+pub(super) async fn handle_user(
+    config: &ProxyConfig,
+    client: &mut PqStream<impl AsyncRead + AsyncWrite + Unpin + Send>,
+    creds: ClientCredentials,
+) -> super::Result<compute::NodeInfo> {
+    use AuthBackendType::*;
+    match config.auth_backend {
+        LegacyConsole => {
+            legacy_console::handle_user(
+                &config.auth_endpoint,
+                &config.auth_link_uri,
+                client,
+                &creds,
+            )
+            .await
         }
-    }
-}
-
-impl<T, E> BackendType<Result<T, E>> {
-    /// Very similar to [`std::option::Option::transpose`].
-    /// This is most useful for error handling.
-    pub fn transpose(self) -> Result<BackendType<T>, E> {
-        use BackendType::*;
-        match self {
-            LegacyConsole(x) => x.map(LegacyConsole),
-            Console(x) => x.map(Console),
-            Postgres(x) => x.map(Postgres),
-            Link => Ok(Link),
-        }
-    }
-}
-
-impl BackendType<ClientCredentials> {
-    /// Authenticate the client via the requested backend, possibly using credentials.
-    pub async fn authenticate(
-        mut self,
-        urls: &config::AuthUrls,
-        client: &mut PqStream<impl AsyncRead + AsyncWrite + Unpin + Send>,
-    ) -> super::Result<compute::NodeInfo> {
-        use BackendType::*;
-
-        if let Console(creds) | Postgres(creds) = &mut self {
-            // If there's no project so far, that entails that client doesn't
-            // support SNI or other means of passing the project name.
-            // We now expect to see a very specific payload in the place of password.
-            if creds.project().is_none() {
-                let payload = AuthFlow::new(client)
-                    .begin(auth::PasswordHack)
-                    .await?
-                    .authenticate()
-                    .await?;
-
-                // Finally we may finish the initialization of `creds`.
-                // TODO: add missing type safety to ClientCredentials.
-                creds.project = Some(payload.project);
-
-                let mut config = match &self {
-                    Console(creds) => {
-                        console::Api::new(&urls.auth_endpoint, creds)
-                            .wake_compute()
-                            .await?
-                    }
-                    Postgres(creds) => {
-                        postgres::Api::new(&urls.auth_endpoint, creds)
-                            .wake_compute()
-                            .await?
-                    }
-                    _ => unreachable!("see the patterns above"),
-                };
-
-                // We should use a password from payload as well.
-                config.password(payload.password);
-
-                return Ok(compute::NodeInfo {
-                    reported_auth_ok: false,
-                    config,
-                });
-            }
-        }
-
-        match self {
-            LegacyConsole(creds) => {
-                legacy_console::handle_user(
-                    &urls.auth_endpoint,
-                    &urls.auth_link_uri,
-                    &creds,
-                    client,
-                )
+        Console => {
+            console::Api::new(&config.auth_endpoint, &creds)?
+                .handle_user(client)
                 .await
-            }
-            Console(creds) => {
-                console::Api::new(&urls.auth_endpoint, &creds)
-                    .handle_user(client)
-                    .await
-            }
-            Postgres(creds) => {
-                postgres::Api::new(&urls.auth_endpoint, &creds)
-                    .handle_user(client)
-                    .await
-            }
-            // NOTE: this auth backend doesn't use client credentials.
-            Link => link::handle_user(&urls.auth_link_uri, client).await,
-        }
-    }
-}
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-
-    #[test]
-    fn test_backend_type_map() {
-        let values = [
-            BackendType::LegacyConsole(0),
-            BackendType::Console(0),
-            BackendType::Postgres(0),
-            BackendType::Link,
-        ];
-
-        for value in values {
-            assert_eq!(value.map(|x| x), value);
-        }
-    }
-
-    #[test]
-    fn test_backend_type_transpose() {
-        let values = [
-            BackendType::LegacyConsole(Ok::<_, ()>(0)),
-            BackendType::Console(Ok(0)),
-            BackendType::Postgres(Ok(0)),
-            BackendType::Link,
-        ];
-
-        for value in values {
-            assert_eq!(value.map(Result::unwrap), value.transpose().unwrap());
         }
+        Postgres => {
+            postgres::Api::new(&config.auth_endpoint, &creds)?
+                .handle_user(client)
+                .await
+        }
+        Link => link::handle_user(&config.auth_link_uri, client).await,
     }
 }

proxy/src/auth/backend/console.rs

@@ -1,23 +1,34 @@
 //! Cloud API V2.
 
 use crate::{
-    auth::{self, AuthFlow, ClientCredentials},
-    compute::{self, ComputeConnCfg},
-    error::{io_error, UserFacingError},
+    auth::{self, AuthFlow, ClientCredentials, DatabaseInfo},
+    compute,
+    error::UserFacingError,
     scram,
     stream::PqStream,
     url::ApiUrl,
 };
 use serde::{Deserialize, Serialize};
-use std::future::Future;
+use std::{future::Future, io};
 use thiserror::Error;
 use tokio::io::{AsyncRead, AsyncWrite};
+use utils::pq_proto::{BeMessage as Be, BeParameterStatusMessage};
 
-const REQUEST_FAILED: &str = "Console request failed";
+pub type Result<T> = std::result::Result<T, ConsoleAuthError>;
 
 #[derive(Debug, Error)]
-pub enum TransportError {
-    #[error("Console responded with a malformed JSON: {0}")]
+pub enum ConsoleAuthError {
+    #[error(transparent)]
+    BadProjectName(#[from] auth::credentials::ClientCredsParseError),
+
+    // We shouldn't include the actual secret here.
+    #[error("Bad authentication secret")]
+    BadSecret,
+
+    #[error("Console responded with a malformed compute address: '{0}'")]
+    BadComputeAddress(String),
+
+    #[error("Console responded with a malformed JSON: '{0}'")]
     BadResponse(#[from] serde_json::Error),
 
     /// HTTP status (other than 200) returned by the console.
@@ -28,72 +39,19 @@ pub enum TransportError {
     Io(#[from] std::io::Error),
 }
 
-impl UserFacingError for TransportError {
+impl UserFacingError for ConsoleAuthError {
     fn to_string_client(&self) -> String {
-        use TransportError::*;
+        use ConsoleAuthError::*;
         match self {
-            HttpStatus(_) => self.to_string(),
-            _ => REQUEST_FAILED.to_owned(),
+            BadProjectName(e) => e.to_string_client(),
+            _ => "Internal error".to_string(),
         }
     }
 }
 
-// Helps eliminate graceless `.map_err` calls without introducing another ctor.
-impl From<reqwest::Error> for TransportError {
-    fn from(e: reqwest::Error) -> Self {
-        io_error(e).into()
-    }
-}
-
-#[derive(Debug, Error)]
-pub enum GetAuthInfoError {
-    // We shouldn't include the actual secret here.
-    #[error("Console responded with a malformed auth secret")]
-    BadSecret,
-
-    #[error(transparent)]
-    Transport(TransportError),
-}
-
-impl UserFacingError for GetAuthInfoError {
-    fn to_string_client(&self) -> String {
-        use GetAuthInfoError::*;
-        match self {
-            BadSecret => REQUEST_FAILED.to_owned(),
-            Transport(e) => e.to_string_client(),
-        }
-    }
-}
-
-impl<E: Into<TransportError>> From<E> for GetAuthInfoError {
-    fn from(e: E) -> Self {
-        Self::Transport(e.into())
-    }
-}
-
-#[derive(Debug, Error)]
-pub enum WakeComputeError {
-    // We shouldn't show users the address even if it's broken.
-    #[error("Console responded with a malformed compute address: {0}")]
-    BadComputeAddress(String),
-
-    #[error(transparent)]
-    Transport(TransportError),
-}
-
-impl UserFacingError for WakeComputeError {
-    fn to_string_client(&self) -> String {
-        use WakeComputeError::*;
-        match self {
-            BadComputeAddress(_) => REQUEST_FAILED.to_owned(),
-            Transport(e) => e.to_string_client(),
-        }
-    }
-}
-
-impl<E: Into<TransportError>> From<E> for WakeComputeError {
-    fn from(e: E) -> Self {
-        Self::Transport(e.into())
+impl From<&auth::credentials::ClientCredsParseError> for ConsoleAuthError {
+    fn from(e: &auth::credentials::ClientCredsParseError) -> Self {
+        ConsoleAuthError::BadProjectName(e.clone())
     }
 }
 
@@ -126,8 +84,8 @@ pub(super) struct Api<'a> {
 
 impl<'a> Api<'a> {
     /// Construct an API object containing the auth parameters.
-    pub(super) fn new(endpoint: &'a ApiUrl, creds: &'a ClientCredentials) -> Self {
-        Self { endpoint, creds }
+    pub(super) fn new(endpoint: &'a ApiUrl, creds: &'a ClientCredentials) -> Result<Self> {
+        Ok(Self { endpoint, creds })
     }
 
     /// Authenticate the existing user or throw an error.
@@ -138,59 +96,57 @@ impl<'a> Api<'a> {
         handle_user(client, &self, Self::get_auth_info, Self::wake_compute).await
     }
 
-    async fn get_auth_info(&self) -> Result<AuthInfo, GetAuthInfoError> {
+    async fn get_auth_info(&self) -> Result<AuthInfo> {
         let mut url = self.endpoint.clone();
         url.path_segments_mut().push("proxy_get_role_secret");
         url.query_pairs_mut()
-            .append_pair("project", self.creds.project().expect("impossible"))
+            .append_pair("project", self.creds.project_name.as_ref()?)
             .append_pair("role", &self.creds.user);
 
         // TODO: use a proper logger
         println!("cplane request: {url}");
 
-        let resp = reqwest::get(url.into_inner()).await?;
+        let resp = reqwest::get(url.into_inner()).await.map_err(io_error)?;
         if !resp.status().is_success() {
-            return Err(TransportError::HttpStatus(resp.status()).into());
+            return Err(ConsoleAuthError::HttpStatus(resp.status()));
         }
 
-        let response: GetRoleSecretResponse = serde_json::from_str(&resp.text().await?)?;
+        let response: GetRoleSecretResponse =
+            serde_json::from_str(&resp.text().await.map_err(io_error)?)?;
 
-        scram::ServerSecret::parse(&response.role_secret)
+        scram::ServerSecret::parse(response.role_secret.as_str())
             .map(AuthInfo::Scram)
-            .ok_or(GetAuthInfoError::BadSecret)
+            .ok_or(ConsoleAuthError::BadSecret)
     }
 
     /// Wake up the compute node and return the corresponding connection info.
-    pub(super) async fn wake_compute(&self) -> Result<ComputeConnCfg, WakeComputeError> {
+    async fn wake_compute(&self) -> Result<DatabaseInfo> {
         let mut url = self.endpoint.clone();
         url.path_segments_mut().push("proxy_wake_compute");
-        url.query_pairs_mut()
-            .append_pair("project", self.creds.project().expect("impossible"));
+        let project_name = self.creds.project_name.as_ref()?;
+        url.query_pairs_mut().append_pair("project", project_name);
 
         // TODO: use a proper logger
         println!("cplane request: {url}");
 
-        let resp = reqwest::get(url.into_inner()).await?;
+        let resp = reqwest::get(url.into_inner()).await.map_err(io_error)?;
         if !resp.status().is_success() {
-            return Err(TransportError::HttpStatus(resp.status()).into());
+            return Err(ConsoleAuthError::HttpStatus(resp.status()));
         }
 
-        let response: GetWakeComputeResponse = serde_json::from_str(&resp.text().await?)?;
+        let response: GetWakeComputeResponse =
+            serde_json::from_str(&resp.text().await.map_err(io_error)?)?;
 
-        // Unfortunately, ownership won't let us use `Option::ok_or` here.
-        let (host, port) = match parse_host_port(&response.address) {
-            None => return Err(WakeComputeError::BadComputeAddress(response.address)),
-            Some(x) => x,
-        };
+        let (host, port) = parse_host_port(&response.address)
+            .ok_or(ConsoleAuthError::BadComputeAddress(response.address))?;
 
-        let mut config = ComputeConnCfg::new();
-        config
-            .host(host)
-            .port(port)
-            .dbname(&self.creds.dbname)
-            .user(&self.creds.user);
-
-        Ok(config)
+        Ok(DatabaseInfo {
+            host,
+            port,
+            dbname: self.creds.dbname.to_owned(),
+            user: self.creds.user.to_owned(),
+            password: None,
+        })
     }
 }
 
@@ -203,8 +159,8 @@ pub(super) async fn handle_user<'a, Endpoint, GetAuthInfo, WakeCompute>(
     wake_compute: impl FnOnce(&'a Endpoint) -> WakeCompute,
 ) -> auth::Result<compute::NodeInfo>
 where
-    GetAuthInfo: Future<Output = Result<AuthInfo, GetAuthInfoError>>,
-    WakeCompute: Future<Output = Result<ComputeConnCfg, WakeComputeError>>,
+    GetAuthInfo: Future<Output = Result<AuthInfo>>,
+    WakeCompute: Future<Output = Result<DatabaseInfo>>,
 {
     let auth_info = get_auth_info(endpoint).await?;
 
@@ -212,7 +168,7 @@ where
     let scram_keys = match auth_info {
         AuthInfo::Md5(_) => {
             // TODO: decide if we should support MD5 in api v2
-            return Err(auth::AuthError::bad_auth_method("MD5"));
+            return Err(auth::AuthErrorImpl::auth_failed("MD5 is not supported").into());
         }
         AuthInfo::Scram(secret) => {
             let scram = auth::Scram(&secret);
@@ -223,18 +179,48 @@ where
         }
     };
 
-    let mut config = wake_compute(endpoint).await?;
-    if let Some(keys) = scram_keys {
-        config.auth_keys(tokio_postgres::config::AuthKeys::ScramSha256(keys));
-    }
+    client
+        .write_message_noflush(&Be::AuthenticationOk)?
+        .write_message_noflush(&BeParameterStatusMessage::encoding())?;
 
     Ok(compute::NodeInfo {
-        reported_auth_ok: false,
-        config,
+        db_info: wake_compute(endpoint).await?,
+        scram_keys,
     })
 }
 
-fn parse_host_port(input: &str) -> Option<(&str, u16)> {
-    let (host, port) = input.split_once(':')?;
-    Some((host, port.parse().ok()?))
+/// Upcast (almost) any error into an opaque [`io::Error`].
+pub(super) fn io_error(e: impl Into<Box<dyn std::error::Error + Send + Sync>>) -> io::Error {
+    io::Error::new(io::ErrorKind::Other, e)
+}
+
+fn parse_host_port(input: &str) -> Option<(String, u16)> {
+    let (host, port) = input.split_once(':')?;
+    Some((host.to_owned(), port.parse().ok()?))
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use serde_json::json;
+
+    #[test]
+    fn parse_db_info() -> anyhow::Result<()> {
+        let _: DatabaseInfo = serde_json::from_value(json!({
+            "host": "localhost",
+            "port": 5432,
+            "dbname": "postgres",
+            "user": "john_doe",
+            "password": "password",
+        }))?;
+
+        let _: DatabaseInfo = serde_json::from_value(json!({
+            "host": "localhost",
+            "port": 5432,
+            "dbname": "postgres",
+            "user": "john_doe",
+        }))?;
+
+        Ok(())
+    }
 }

proxy/src/auth/backend/legacy_console.rs

@@ -11,10 +11,10 @@ use crate::{
 use serde::{Deserialize, Serialize};
 use thiserror::Error;
 use tokio::io::{AsyncRead, AsyncWrite};
-use utils::pq_proto::BeMessage as Be;
+use utils::pq_proto::{BeMessage as Be, BeParameterStatusMessage};
 
 #[derive(Debug, Error)]
-pub enum LegacyAuthError {
+pub enum AuthErrorImpl {
     /// Authentication error reported by the console.
     #[error("Authentication failed: {0}")]
     AuthFailed(String),
@@ -24,7 +24,7 @@ pub enum LegacyAuthError {
     HttpStatus(reqwest::StatusCode),
 
     #[error("Console responded with a malformed JSON: {0}")]
-    BadResponse(#[from] serde_json::Error),
+    MalformedResponse(#[from] serde_json::Error),
 
     #[error(transparent)]
     Transport(#[from] reqwest::Error),
@@ -36,10 +36,30 @@ pub enum LegacyAuthError {
     WaiterWait(#[from] waiters::WaitError),
 }
 
-impl UserFacingError for LegacyAuthError {
+#[derive(Debug, Error)]
+#[error(transparent)]
+pub struct AuthError(Box<AuthErrorImpl>);
+
+impl AuthError {
+    /// Smart constructor for authentication error reported by `mgmt`.
+    pub fn auth_failed(msg: impl Into<String>) -> Self {
+        Self(Box::new(AuthErrorImpl::AuthFailed(msg.into())))
+    }
+}
+
+impl<T> From<T> for AuthError
+where
+    AuthErrorImpl: From<T>,
+{
+    fn from(e: T) -> Self {
+        Self(Box::new(e.into()))
+    }
+}
+
+impl UserFacingError for AuthError {
     fn to_string_client(&self) -> String {
-        use LegacyAuthError::*;
-        match self {
+        use AuthErrorImpl::*;
+        match self.0.as_ref() {
             AuthFailed(_) | HttpStatus(_) => self.to_string(),
             _ => "Internal error".to_string(),
         }
@@ -56,19 +76,13 @@ enum ProxyAuthResponse {
     NotReady { ready: bool }, // TODO: get rid of `ready`
 }
 
-impl ClientCredentials {
-    fn is_existing_user(&self) -> bool {
-        self.user.ends_with("@zenith")
-    }
-}
-
 async fn authenticate_proxy_client(
     auth_endpoint: &reqwest::Url,
     creds: &ClientCredentials,
     md5_response: &str,
     salt: &[u8; 4],
     psql_session_id: &str,
-) -> Result<DatabaseInfo, LegacyAuthError> {
+) -> Result<DatabaseInfo, AuthError> {
     let mut url = auth_endpoint.clone();
     url.query_pairs_mut()
         .append_pair("login", &creds.user)
@@ -82,17 +96,17 @@ async fn authenticate_proxy_client(
         // TODO: leverage `reqwest::Client` to reuse connections
         let resp = reqwest::get(url).await?;
         if !resp.status().is_success() {
-            return Err(LegacyAuthError::HttpStatus(resp.status()));
+            return Err(AuthErrorImpl::HttpStatus(resp.status()).into());
         }
 
-        let auth_info = serde_json::from_str(resp.text().await?.as_str())?;
-        println!("got auth info: {:?}", auth_info);
+        let auth_info: ProxyAuthResponse = serde_json::from_str(resp.text().await?.as_str())?;
+        println!("got auth info: #{:?}", auth_info);
 
         use ProxyAuthResponse::*;
         let db_info = match auth_info {
             Ready { conn_info } => conn_info,
-            Error { error } => return Err(LegacyAuthError::AuthFailed(error)),
-            NotReady { .. } => waiter.await?.map_err(LegacyAuthError::AuthFailed)?,
+            Error { error } => return Err(AuthErrorImpl::AuthFailed(error).into()),
+            NotReady { .. } => waiter.await?.map_err(AuthErrorImpl::AuthFailed)?,
         };
 
         Ok(db_info)
@@ -104,7 +118,7 @@ async fn handle_existing_user(
     auth_endpoint: &reqwest::Url,
     client: &mut PqStream<impl AsyncRead + AsyncWrite + Unpin + Send>,
     creds: &ClientCredentials,
-) -> auth::Result<compute::NodeInfo> {
+) -> Result<compute::NodeInfo, auth::AuthError> {
     let psql_session_id = super::link::new_psql_session_id();
     let md5_salt = rand::random();
 
@@ -114,9 +128,7 @@ async fn handle_existing_user(
 
     // Read client's password hash
     let msg = client.read_password_message().await?;
-    let md5_response = parse_password(&msg).ok_or(auth::AuthErrorImpl::MalformedPassword(
-        "the password should be a valid null-terminated utf-8 string",
-    ))?;
+    let md5_response = parse_password(&msg).ok_or(auth::AuthErrorImpl::MalformedPassword)?;
 
     let db_info = authenticate_proxy_client(
         auth_endpoint,
@@ -127,17 +139,21 @@ async fn handle_existing_user(
     )
     .await?;
 
+    client
+        .write_message_noflush(&Be::AuthenticationOk)?
+        .write_message_noflush(&BeParameterStatusMessage::encoding())?;
+
     Ok(compute::NodeInfo {
-        reported_auth_ok: false,
-        config: db_info.into(),
+        db_info,
+        scram_keys: None,
     })
 }
 
 pub async fn handle_user(
     auth_endpoint: &reqwest::Url,
     auth_link_uri: &reqwest::Url,
-    creds: &ClientCredentials,
     client: &mut PqStream<impl AsyncRead + AsyncWrite + Unpin + Send>,
+    creds: &ClientCredentials,
 ) -> auth::Result<compute::NodeInfo> {
     if creds.is_existing_user() {
         handle_existing_user(auth_endpoint, client, creds).await
@@ -185,24 +201,4 @@ mod tests {
         .unwrap();
         assert!(matches!(auth, ProxyAuthResponse::NotReady { .. }));
     }
-
-    #[test]
-    fn parse_db_info() -> anyhow::Result<()> {
-        let _: DatabaseInfo = serde_json::from_value(json!({
-            "host": "localhost",
-            "port": 5432,
-            "dbname": "postgres",
-            "user": "john_doe",
-            "password": "password",
-        }))?;
-
-        let _: DatabaseInfo = serde_json::from_value(json!({
-            "host": "localhost",
-            "port": 5432,
-            "dbname": "postgres",
-            "user": "john_doe",
-        }))?;
-
-        Ok(())
-    }
 }

proxy/src/auth/backend/link.rs

@@ -1,34 +1,7 @@
-use crate::{auth, compute, error::UserFacingError, stream::PqStream, waiters};
-use thiserror::Error;
+use crate::{auth, compute, stream::PqStream};
 use tokio::io::{AsyncRead, AsyncWrite};
 use utils::pq_proto::{BeMessage as Be, BeParameterStatusMessage};
 
-#[derive(Debug, Error)]
-pub enum LinkAuthError {
-    /// Authentication error reported by the console.
-    #[error("Authentication failed: {0}")]
-    AuthFailed(String),
-
-    #[error(transparent)]
-    WaiterRegister(#[from] waiters::RegisterError),
-
-    #[error(transparent)]
-    WaiterWait(#[from] waiters::WaitError),
-
-    #[error(transparent)]
-    Io(#[from] std::io::Error),
-}
-
-impl UserFacingError for LinkAuthError {
-    fn to_string_client(&self) -> String {
-        use LinkAuthError::*;
-        match self {
-            AuthFailed(_) => self.to_string(),
-            _ => "Internal error".to_string(),
-        }
-    }
-}
-
 fn hello_message(redirect_uri: &str, session_id: &str) -> String {
     format!(
         concat![
@@ -61,14 +34,14 @@ pub async fn handle_user(
             .await?;
 
         // Wait for web console response (see `mgmt`)
-        waiter.await?.map_err(LinkAuthError::AuthFailed)
+        waiter.await?.map_err(auth::AuthErrorImpl::auth_failed)
     })
     .await?;
 
     client.write_message_noflush(&Be::NoticeResponse("Connecting to database."))?;
 
     Ok(compute::NodeInfo {
-        reported_auth_ok: true,
-        config: db_info.into(),
+        db_info,
+        scram_keys: None,
     })
 }

proxy/src/auth/backend/postgres.rs

@@ -3,12 +3,10 @@
 use crate::{
     auth::{
         self,
-        backend::console::{self, AuthInfo, GetAuthInfoError, TransportError, WakeComputeError},
-        ClientCredentials,
+        backend::console::{self, io_error, AuthInfo, Result},
+        ClientCredentials, DatabaseInfo,
     },
-    compute::{self, ComputeConnCfg},
-    error::io_error,
-    scram,
+    compute, scram,
     stream::PqStream,
     url::ApiUrl,
 };
@@ -20,17 +18,10 @@ pub(super) struct Api<'a> {
     creds: &'a ClientCredentials,
 }
 
-// Helps eliminate graceless `.map_err` calls without introducing another ctor.
-impl From<tokio_postgres::Error> for TransportError {
-    fn from(e: tokio_postgres::Error) -> Self {
-        io_error(e).into()
-    }
-}
-
 impl<'a> Api<'a> {
     /// Construct an API object containing the auth parameters.
-    pub(super) fn new(endpoint: &'a ApiUrl, creds: &'a ClientCredentials) -> Self {
-        Self { endpoint, creds }
+    pub(super) fn new(endpoint: &'a ApiUrl, creds: &'a ClientCredentials) -> Result<Self> {
+        Ok(Self { endpoint, creds })
     }
 
     /// Authenticate the existing user or throw an error.
@@ -43,16 +34,21 @@ impl<'a> Api<'a> {
     }
 
     /// This implementation fetches the auth info from a local postgres instance.
-    async fn get_auth_info(&self) -> Result<AuthInfo, GetAuthInfoError> {
+    async fn get_auth_info(&self) -> Result<AuthInfo> {
         // Perhaps we could persist this connection, but then we'd have to
         // write more code for reopening it if it got closed, which doesn't
         // seem worth it.
         let (client, connection) =
-            tokio_postgres::connect(self.endpoint.as_str(), tokio_postgres::NoTls).await?;
+            tokio_postgres::connect(self.endpoint.as_str(), tokio_postgres::NoTls)
+                .await
+                .map_err(io_error)?;
 
         tokio::spawn(connection);
         let query = "select rolpassword from pg_catalog.pg_authid where rolname = $1";
-        let rows = client.query(query, &[&self.creds.user]).await?;
+        let rows = client
+            .query(query, &[&self.creds.user])
+            .await
+            .map_err(io_error)?;
 
         match &rows[..] {
             // We can't get a secret if there's no such user.
@@ -60,10 +56,7 @@ impl<'a> Api<'a> {
 
             // We shouldn't get more than one row anyway.
             [row, ..] => {
-                let entry = row
-                    .try_get("rolpassword")
-                    .map_err(|e| io_error(format!("failed to read user's password: {e}")))?;
-
+                let entry = row.try_get(0).map_err(io_error)?;
                 scram::ServerSecret::parse(entry)
                     .map(AuthInfo::Scram)
                     .or_else(|| {
@@ -76,20 +69,20 @@ impl<'a> Api<'a> {
                         }))
                     })
                     // Putting the secret into this message is a security hazard!
-                    .ok_or(GetAuthInfoError::BadSecret)
+                    .ok_or(console::ConsoleAuthError::BadSecret)
             }
         }
     }
 
     /// We don't need to wake anything locally, so we just return the connection info.
-    pub(super) async fn wake_compute(&self) -> Result<ComputeConnCfg, WakeComputeError> {
-        let mut config = ComputeConnCfg::new();
-        config
-            .host(self.endpoint.host_str().unwrap_or("localhost"))
-            .port(self.endpoint.port().unwrap_or(5432))
-            .dbname(&self.creds.dbname)
-            .user(&self.creds.user);
-
-        Ok(config)
+    async fn wake_compute(&self) -> Result<DatabaseInfo> {
+        Ok(DatabaseInfo {
+            // TODO: handle that near CLI params parsing
+            host: self.endpoint.host_str().unwrap_or("localhost").to_owned(),
+            port: self.endpoint.port().unwrap_or(5432),
+            dbname: self.creds.dbname.to_owned(),
+            user: self.creds.user.to_owned(),
+            password: None,
+        })
     }
 }

proxy/src/auth/credentials.rs

@@ -1,25 +1,39 @@
 //! User credentials used in authentication.
 
+use crate::compute;
+use crate::config::ProxyConfig;
 use crate::error::UserFacingError;
+use crate::stream::PqStream;
+use std::collections::HashMap;
 use thiserror::Error;
-use utils::pq_proto::StartupMessageParams;
+use tokio::io::{AsyncRead, AsyncWrite};
 
 #[derive(Debug, Error, PartialEq, Eq, Clone)]
 pub enum ClientCredsParseError {
-    #[error("Parameter '{0}' is missing in startup packet.")]
+    #[error("Parameter `{0}` is missing in startup packet.")]
     MissingKey(&'static str),
 
+    #[error(
+        "Project name is not specified. \
+        EITHER please upgrade the postgres client library (libpq) for SNI support \
+        OR pass the project name as a parameter: '&options=project%3D<project-name>'."
+    )]
+    MissingSNIAndProjectName,
+
     #[error("Inconsistent project name inferred from SNI ('{0}') and project option ('{1}').")]
-    InconsistentProjectNames(String, String),
+    InconsistentProjectNameAndSNI(String, String),
+
+    #[error("Common name is not set.")]
+    CommonNameNotSet,
 
     #[error(
         "SNI ('{1}') inconsistently formatted with respect to common name ('{0}'). \
-        SNI should be formatted as '<project-name>.{0}'."
+        SNI should be formatted as '<project-name>.<common-name>'."
     )]
-    InconsistentSni(String, String),
+    InconsistentCommonNameAndSNI(String, String),
 
-    #[error("Project name ('{0}') must contain only alphanumeric characters and hyphen.")]
-    MalformedProjectName(String),
+    #[error("Project name ('{0}') must contain only alphanumeric characters and hyphens ('-').")]
+    ProjectNameContainsIllegalChars(String),
 }
 
 impl UserFacingError for ClientCredsParseError {}
@@ -30,171 +44,286 @@ impl UserFacingError for ClientCredsParseError {}
 pub struct ClientCredentials {
     pub user: String,
     pub dbname: String,
-    pub project: Option<String>,
+    pub project_name: Result<String, ClientCredsParseError>,
 }
 
 impl ClientCredentials {
-    pub fn project(&self) -> Option<&str> {
-        self.project.as_deref()
+    pub fn is_existing_user(&self) -> bool {
+        // This logic will likely change in the future.
+        self.user.ends_with("@zenith")
     }
-}
 
-impl ClientCredentials {
     pub fn parse(
-        mut options: StartupMessageParams,
-        sni: Option<&str>,
+        mut options: HashMap<String, String>,
+        sni_data: Option<&str>,
         common_name: Option<&str>,
     ) -> Result<Self, ClientCredsParseError> {
-        use ClientCredsParseError::*;
+        let mut get_param = |key| {
+            options
+                .remove(key)
+                .ok_or(ClientCredsParseError::MissingKey(key))
+        };
 
-        // Some parameters are absolutely necessary, others not so much.
-        let mut get_param = |key| options.remove(key).ok_or(MissingKey(key));
-
-        // Some parameters are stored in the startup message.
         let user = get_param("user")?;
         let dbname = get_param("database")?;
-        let project_a = get_param("project").ok();
-
-        // Alternative project name is in fact a subdomain from SNI.
-        // NOTE: we do not consider SNI if `common_name` is missing.
-        let project_b = sni
-            .zip(common_name)
-            .map(|(sni, cn)| {
-                // TODO: what if SNI is present but just a common name?
-                subdomain_from_sni(sni, cn)
-                    .ok_or_else(|| InconsistentSni(sni.to_owned(), cn.to_owned()))
-            })
-            .transpose()?;
-
-        let project = match (project_a, project_b) {
-            // Invariant: if we have both project name variants, they should match.
-            (Some(a), Some(b)) if a != b => Some(Err(InconsistentProjectNames(a, b))),
-            (a, b) => a.or(b).map(|name| {
-                // Invariant: project name may not contain certain characters.
-                check_project_name(name).map_err(MalformedProjectName)
-            }),
-        }
-        .transpose()?;
+        let project_name = get_param("project").ok();
+        let project_name = get_project_name(sni_data, common_name, project_name.as_deref());
 
         Ok(Self {
             user,
             dbname,
-            project,
+            project_name,
         })
     }
-}
 
-fn check_project_name(name: String) -> Result<String, String> {
-    if name.chars().all(|c| c.is_alphanumeric() || c == '-') {
-        Ok(name)
-    } else {
-        Err(name)
+    /// Use credentials to authenticate the user.
+    pub async fn authenticate(
+        self,
+        config: &ProxyConfig,
+        client: &mut PqStream<impl AsyncRead + AsyncWrite + Unpin + Send>,
+    ) -> super::Result<compute::NodeInfo> {
+        // This method is just a convenient facade for `handle_user`
+        super::backend::handle_user(config, client, self).await
     }
 }
 
-fn subdomain_from_sni(sni: &str, common_name: &str) -> Option<String> {
-    sni.strip_suffix(common_name)?
-        .strip_suffix('.')
-        .map(str::to_owned)
+/// Inferring project name from sni_data.
+fn project_name_from_sni_data(
+    sni_data: &str,
+    common_name: &str,
+) -> Result<String, ClientCredsParseError> {
+    let common_name_with_dot = format!(".{common_name}");
+    // check that ".{common_name_with_dot}" is the actual suffix in sni_data
+    if !sni_data.ends_with(&common_name_with_dot) {
+        return Err(ClientCredsParseError::InconsistentCommonNameAndSNI(
+            common_name.to_string(),
+            sni_data.to_string(),
+        ));
+    }
+    // return sni_data without the common name suffix.
+    Ok(sni_data
+        .strip_suffix(&common_name_with_dot)
+        .unwrap()
+        .to_string())
 }
 
 #[cfg(test)]
-mod tests {
+mod tests_for_project_name_from_sni_data {
     use super::*;
 
-    fn make_options<'a, const N: usize>(pairs: [(&'a str, &'a str); N]) -> StartupMessageParams {
-        StartupMessageParams::from(pairs.map(|(k, v)| (k.to_owned(), v.to_owned())))
+    #[test]
+    fn passing() {
+        let target_project_name = "my-project-123";
+        let common_name = "localtest.me";
+        let sni_data = format!("{target_project_name}.{common_name}");
+        assert_eq!(
+            project_name_from_sni_data(&sni_data, common_name),
+            Ok(target_project_name.to_string())
+        );
     }
 
     #[test]
-    #[ignore = "TODO: fix how database is handled"]
-    fn parse_bare_minimum() -> anyhow::Result<()> {
-        // According to postgresql, only `user` should be required.
-        let options = make_options([("user", "john_doe")]);
-
-        // TODO: check that `creds.dbname` is None.
-        let creds = ClientCredentials::parse(options, None, None)?;
-        assert_eq!(creds.user, "john_doe");
-
-        Ok(())
-    }
-
-    #[test]
-    fn parse_missing_project() -> anyhow::Result<()> {
-        let options = make_options([("user", "john_doe"), ("database", "world")]);
-
-        let creds = ClientCredentials::parse(options, None, None)?;
-        assert_eq!(creds.user, "john_doe");
-        assert_eq!(creds.dbname, "world");
-        assert_eq!(creds.project, None);
-
-        Ok(())
-    }
-
-    #[test]
-    fn parse_project_from_sni() -> anyhow::Result<()> {
-        let options = make_options([("user", "john_doe"), ("database", "world")]);
-
-        let sni = Some("foo.localhost");
-        let common_name = Some("localhost");
-
-        let creds = ClientCredentials::parse(options, sni, common_name)?;
-        assert_eq!(creds.user, "john_doe");
-        assert_eq!(creds.dbname, "world");
-        assert_eq!(creds.project.as_deref(), Some("foo"));
-
-        Ok(())
-    }
-
-    #[test]
-    fn parse_project_from_options() -> anyhow::Result<()> {
-        let options = make_options([
-            ("user", "john_doe"),
-            ("database", "world"),
-            ("project", "bar"),
-        ]);
-
-        let creds = ClientCredentials::parse(options, None, None)?;
-        assert_eq!(creds.user, "john_doe");
-        assert_eq!(creds.dbname, "world");
-        assert_eq!(creds.project.as_deref(), Some("bar"));
-
-        Ok(())
-    }
-
-    #[test]
-    fn parse_projects_identical() -> anyhow::Result<()> {
-        let options = make_options([
-            ("user", "john_doe"),
-            ("database", "world"),
-            ("project", "baz"),
-        ]);
-
-        let sni = Some("baz.localhost");
-        let common_name = Some("localhost");
-
-        let creds = ClientCredentials::parse(options, sni, common_name)?;
-        assert_eq!(creds.user, "john_doe");
-        assert_eq!(creds.dbname, "world");
-        assert_eq!(creds.project.as_deref(), Some("baz"));
-
-        Ok(())
-    }
-
-    #[test]
-    fn parse_projects_different() {
-        let options = make_options([
-            ("user", "john_doe"),
-            ("database", "world"),
-            ("project", "first"),
-        ]);
-
-        let sni = Some("second.localhost");
-        let common_name = Some("localhost");
-
-        assert!(matches!(
-            ClientCredentials::parse(options, sni, common_name).expect_err("should fail"),
-            ClientCredsParseError::InconsistentProjectNames(_, _)
-        ));
+    fn throws_inconsistent_common_name_and_sni_data() {
+        let target_project_name = "my-project-123";
+        let common_name = "localtest.me";
+        let wrong_suffix = "wrongtest.me";
+        assert_eq!(common_name.len(), wrong_suffix.len());
+        let wrong_common_name = format!("wrong{wrong_suffix}");
+        let sni_data = format!("{target_project_name}.{wrong_common_name}");
+        assert_eq!(
+            project_name_from_sni_data(&sni_data, common_name),
+            Err(ClientCredsParseError::InconsistentCommonNameAndSNI(
+                common_name.to_string(),
+                sni_data
+            ))
+        );
+    }
+}
+
+/// Determine project name from SNI or from project_name parameter from options argument.
+fn get_project_name(
+    sni_data: Option<&str>,
+    common_name: Option<&str>,
+    project_name: Option<&str>,
+) -> Result<String, ClientCredsParseError> {
+    // determine the project name from sni_data if it exists, otherwise from project_name.
+    let ret = match sni_data {
+        Some(sni_data) => {
+            let common_name = common_name.ok_or(ClientCredsParseError::CommonNameNotSet)?;
+            let project_name_from_sni = project_name_from_sni_data(sni_data, common_name)?;
+            // check invariant: project name from options and from sni should match
+            if let Some(project_name) = &project_name {
+                if !project_name_from_sni.eq(project_name) {
+                    return Err(ClientCredsParseError::InconsistentProjectNameAndSNI(
+                        project_name_from_sni,
+                        project_name.to_string(),
+                    ));
+                }
+            }
+            project_name_from_sni
+        }
+        None => project_name
+            .ok_or(ClientCredsParseError::MissingSNIAndProjectName)?
+            .to_string(),
+    };
+
+    // check formatting invariant: project name must contain only alphanumeric characters and hyphens.
+    if !ret.chars().all(|x: char| x.is_alphanumeric() || x == '-') {
+        return Err(ClientCredsParseError::ProjectNameContainsIllegalChars(ret));
+    }
+
+    Ok(ret)
+}
+
+#[cfg(test)]
+mod tests_for_project_name_only {
+    use super::*;
+
+    #[test]
+    fn passing_from_sni_data_only() {
+        let target_project_name = "my-project-123";
+        let common_name = "localtest.me";
+        let sni_data = format!("{target_project_name}.{common_name}");
+        assert_eq!(
+            get_project_name(Some(&sni_data), Some(common_name), None),
+            Ok(target_project_name.to_string())
+        );
+    }
+
+    #[test]
+    fn throws_project_name_contains_illegal_chars_from_sni_data_only() {
+        let project_name_prefix = "my-project";
+        let project_name_suffix = "123";
+        let common_name = "localtest.me";
+
+        for illegal_char_id in 0..256 {
+            let illegal_char = char::from_u32(illegal_char_id).unwrap();
+            if !(illegal_char.is_alphanumeric() || illegal_char == '-')
+                && illegal_char.to_string().len() == 1
+            {
+                let target_project_name =
+                    format!("{project_name_prefix}{illegal_char}{project_name_suffix}");
+                let sni_data = format!("{target_project_name}.{common_name}");
+                assert_eq!(
+                    get_project_name(Some(&sni_data), Some(common_name), None),
+                    Err(ClientCredsParseError::ProjectNameContainsIllegalChars(
+                        target_project_name
+                    ))
+                );
+            }
+        }
+    }
+
+    #[test]
+    fn passing_from_project_name_only() {
+        let target_project_name = "my-project-123";
+        let common_names = [Some("localtest.me"), None];
+        for common_name in common_names {
+            assert_eq!(
+                get_project_name(None, common_name, Some(target_project_name)),
+                Ok(target_project_name.to_string())
+            );
+        }
+    }
+
+    #[test]
+    fn throws_project_name_contains_illegal_chars_from_project_name_only() {
+        let project_name_prefix = "my-project";
+        let project_name_suffix = "123";
+        let common_names = [Some("localtest.me"), None];
+
+        for common_name in common_names {
+            for illegal_char_id in 0..256 {
+                let illegal_char: char = char::from_u32(illegal_char_id).unwrap();
+                if !(illegal_char.is_alphanumeric() || illegal_char == '-')
+                    && illegal_char.to_string().len() == 1
+                {
+                    let target_project_name =
+                        format!("{project_name_prefix}{illegal_char}{project_name_suffix}");
+                    assert_eq!(
+                        get_project_name(None, common_name, Some(&target_project_name)),
+                        Err(ClientCredsParseError::ProjectNameContainsIllegalChars(
+                            target_project_name
+                        ))
+                    );
+                }
+            }
+        }
+    }
+
+    #[test]
+    fn passing_from_sni_data_and_project_name() {
+        let target_project_name = "my-project-123";
+        let common_name = "localtest.me";
+        let sni_data = format!("{target_project_name}.{common_name}");
+        assert_eq!(
+            get_project_name(
+                Some(&sni_data),
+                Some(common_name),
+                Some(target_project_name)
+            ),
+            Ok(target_project_name.to_string())
+        );
+    }
+
+    #[test]
+    fn throws_inconsistent_project_name_and_sni() {
+        let project_name_param = "my-project-123";
+        let wrong_project_name = "not-my-project-123";
+        let common_name = "localtest.me";
+        let sni_data = format!("{wrong_project_name}.{common_name}");
+        assert_eq!(
+            get_project_name(Some(&sni_data), Some(common_name), Some(project_name_param)),
+            Err(ClientCredsParseError::InconsistentProjectNameAndSNI(
+                wrong_project_name.to_string(),
+                project_name_param.to_string()
+            ))
+        );
+    }
+
+    #[test]
+    fn throws_common_name_not_set() {
+        let target_project_name = "my-project-123";
+        let wrong_project_name = "not-my-project-123";
+        let common_name = "localtest.me";
+        let sni_datas = [
+            Some(format!("{wrong_project_name}.{common_name}")),
+            Some(format!("{target_project_name}.{common_name}")),
+        ];
+        let project_names = [None, Some(target_project_name)];
+        for sni_data in sni_datas {
+            for project_name_param in project_names {
+                assert_eq!(
+                    get_project_name(sni_data.as_deref(), None, project_name_param),
+                    Err(ClientCredsParseError::CommonNameNotSet)
+                );
+            }
+        }
+    }
+
+    #[test]
+    fn throws_inconsistent_common_name_and_sni_data() {
+        let target_project_name = "my-project-123";
+        let wrong_project_name = "not-my-project-123";
+        let common_name = "localtest.me";
+        let wrong_suffix = "wrongtest.me";
+        assert_eq!(common_name.len(), wrong_suffix.len());
+        let wrong_common_name = format!("wrong{wrong_suffix}");
+        let sni_datas = [
+            Some(format!("{wrong_project_name}.{wrong_common_name}")),
+            Some(format!("{target_project_name}.{wrong_common_name}")),
+        ];
+        let project_names = [None, Some(target_project_name)];
+        for project_name_param in project_names {
+            for sni_data in &sni_datas {
+                assert_eq!(
+                    get_project_name(sni_data.as_deref(), Some(common_name), project_name_param),
+                    Err(ClientCredsParseError::InconsistentCommonNameAndSNI(
+                        common_name.to_string(),
+                        sni_data.clone().unwrap().to_string()
+                    ))
+                );
+            }
+        }
     }
 }

proxy/src/auth/flow.rs

@@ -1,7 +1,8 @@
 //! Main authentication flow.
 
-use super::{AuthErrorImpl, PasswordHackPayload};
-use crate::{sasl, scram, stream::PqStream};
+use super::AuthErrorImpl;
+use crate::stream::PqStream;
+use crate::{sasl, scram};
 use std::io;
 use tokio::io::{AsyncRead, AsyncWrite};
 use utils::pq_proto::{BeAuthenticationSaslMessage, BeMessage, BeMessage as Be};
@@ -26,17 +27,6 @@ impl AuthMethod for Scram<'_> {
     }
 }
 
-/// Use an ad hoc auth flow (for clients which don't support SNI) proposed in
-/// <https://github.com/neondatabase/cloud/issues/1620#issuecomment-1165332290>.
-pub struct PasswordHack;
-
-impl AuthMethod for PasswordHack {
-    #[inline(always)]
-    fn first_message(&self) -> BeMessage<'_> {
-        Be::AuthenticationCleartextPassword
-    }
-}
-
 /// This wrapper for [`PqStream`] performs client authentication.
 #[must_use]
 pub struct AuthFlow<'a, Stream, State> {
@@ -67,37 +57,17 @@ impl<'a, S: AsyncWrite + Unpin> AuthFlow<'a, S, Begin> {
     }
 }
 
-impl<S: AsyncRead + AsyncWrite + Unpin> AuthFlow<'_, S, PasswordHack> {
-    /// Perform user authentication. Raise an error in case authentication failed.
-    pub async fn authenticate(self) -> super::Result<PasswordHackPayload> {
-        let msg = self.stream.read_password_message().await?;
-        let password = msg
-            .strip_suffix(&[0])
-            .ok_or(AuthErrorImpl::MalformedPassword("missing terminator"))?;
-
-        let payload = PasswordHackPayload::parse(password)
-            // If we ended up here and the payload is malformed, it means that
-            // the user neither enabled SNI nor resorted to any other method
-            // for passing the project name we rely on. We should show them
-            // the most helpful error message and point to the documentation.
-            .ok_or(AuthErrorImpl::MissingProjectName)?;
-
-        Ok(payload)
-    }
-}
-
 /// Stream wrapper for handling [SCRAM](crate::scram) auth.
 impl<S: AsyncRead + AsyncWrite + Unpin> AuthFlow<'_, S, Scram<'_>> {
     /// Perform user authentication. Raise an error in case authentication failed.
     pub async fn authenticate(self) -> super::Result<scram::ScramKey> {
         // Initial client message contains the chosen auth method's name.
         let msg = self.stream.read_password_message().await?;
-        let sasl = sasl::FirstMessage::parse(&msg)
-            .ok_or(AuthErrorImpl::MalformedPassword("bad sasl message"))?;
+        let sasl = sasl::FirstMessage::parse(&msg).ok_or(AuthErrorImpl::MalformedPassword)?;
 
         // Currently, the only supported SASL method is SCRAM.
         if !scram::METHODS.contains(&sasl.method) {
-            return Err(super::AuthError::bad_auth_method(sasl.method));
+            return Err(AuthErrorImpl::auth_failed("method not supported").into());
         }
 
         let secret = self.state.0;

proxy/src/auth/password_hack.rs

@@ -1,46 +0,0 @@
-//! Payload for ad hoc authentication method for clients that don't support SNI.
-//! See the `impl` for [`super::backend::BackendType<ClientCredentials>`].
-//! Read more: <https://github.com/neondatabase/cloud/issues/1620#issuecomment-1165332290>.
-//! UPDATE (Mon Aug  8 13:20:34 UTC 2022): the payload format has been simplified.
-
-use bstr::ByteSlice;
-
-pub struct PasswordHackPayload {
-    pub project: String,
-    pub password: Vec<u8>,
-}
-
-impl PasswordHackPayload {
-    pub fn parse(bytes: &[u8]) -> Option<Self> {
-        // The format is `project=<utf-8>;<password-bytes>`.
-        let mut iter = bytes.strip_prefix(b"project=")?.splitn_str(2, ";");
-        let project = iter.next()?.to_str().ok()?.to_owned();
-        let password = iter.next()?.to_owned();
-
-        Some(Self { project, password })
-    }
-}
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-
-    #[test]
-    fn parse_password_hack_payload() {
-        let bytes = b"";
-        assert!(PasswordHackPayload::parse(bytes).is_none());
-
-        let bytes = b"project=";
-        assert!(PasswordHackPayload::parse(bytes).is_none());
-
-        let bytes = b"project=;";
-        let payload = PasswordHackPayload::parse(bytes).expect("parsing failed");
-        assert_eq!(payload.project, "");
-        assert_eq!(payload.password, b"");
-
-        let bytes = b"project=foobar;pass;word";
-        let payload = PasswordHackPayload::parse(bytes).expect("parsing failed");
-        assert_eq!(payload.project, "foobar");
-        assert_eq!(payload.password, b"pass;word");
-    }
-}

proxy/src/compute.rs

@@ -1,6 +1,8 @@
-use crate::{cancellation::CancelClosure, error::UserFacingError};
-use futures::TryFutureExt;
-use std::{io, net::SocketAddr};
+use crate::auth::DatabaseInfo;
+use crate::cancellation::CancelClosure;
+use crate::error::UserFacingError;
+use std::io;
+use std::net::SocketAddr;
 use thiserror::Error;
 use tokio::net::TcpStream;
 use tokio_postgres::NoTls;
@@ -19,96 +21,44 @@ pub enum ConnectionError {
     FailedToFetchPgVersion,
 }
 
-impl UserFacingError for ConnectionError {
-    fn to_string_client(&self) -> String {
-        use ConnectionError::*;
-        match self {
-            // This helps us drop irrelevant library-specific prefixes.
-            // TODO: propagate severity level and other parameters.
-            Postgres(err) => match err.as_db_error() {
-                Some(err) => err.message().to_string(),
-                None => err.to_string(),
-            },
-            other => other.to_string(),
-        }
-    }
-}
+impl UserFacingError for ConnectionError {}
+
+/// PostgreSQL version as [`String`].
+pub type Version = String;
 
 /// A pair of `ClientKey` & `ServerKey` for `SCRAM-SHA-256`.
 pub type ScramKeys = tokio_postgres::config::ScramKeys<32>;
 
-pub type ComputeConnCfg = tokio_postgres::Config;
-
-/// Various compute node info for establishing connection etc.
+/// Compute node connection params.
 pub struct NodeInfo {
-    /// Did we send [`utils::pq_proto::BeMessage::AuthenticationOk`]?
-    pub reported_auth_ok: bool,
-    /// Compute node connection params.
-    pub config: tokio_postgres::Config,
+    pub db_info: DatabaseInfo,
+    pub scram_keys: Option<ScramKeys>,
 }
 
 impl NodeInfo {
     async fn connect_raw(&self) -> io::Result<(SocketAddr, TcpStream)> {
-        use tokio_postgres::config::Host;
+        let host_port = (self.db_info.host.as_str(), self.db_info.port);
+        let socket = TcpStream::connect(host_port).await?;
+        let socket_addr = socket.peer_addr()?;
+        socket2::SockRef::from(&socket).set_keepalive(true)?;
 
-        let connect_once = |host, port| {
-            TcpStream::connect((host, port)).and_then(|socket| async {
-                let socket_addr = socket.peer_addr()?;
-                // This prevents load balancer from severing the connection.
-                socket2::SockRef::from(&socket).set_keepalive(true)?;
-                Ok((socket_addr, socket))
-            })
-        };
-
-        // We can't reuse connection establishing logic from `tokio_postgres` here,
-        // because it has no means for extracting the underlying socket which we
-        // require for our business.
-        let mut connection_error = None;
-        let ports = self.config.get_ports();
-        for (i, host) in self.config.get_hosts().iter().enumerate() {
-            let port = ports.get(i).or_else(|| ports.get(0)).unwrap_or(&5432);
-            let host = match host {
-                Host::Tcp(host) => host.as_str(),
-                Host::Unix(_) => continue, // unix sockets are not welcome here
-            };
-
-            // TODO: maybe we should add a timeout.
-            match connect_once(host, *port).await {
-                Ok(socket) => return Ok(socket),
-                Err(err) => {
-                    // We can't throw an error here, as there might be more hosts to try.
-                    println!("failed to connect to compute `{host}:{port}`: {err}");
-                    connection_error = Some(err);
-                }
-            }
-        }
-
-        Err(connection_error.unwrap_or_else(|| {
-            io::Error::new(
-                io::ErrorKind::Other,
-                format!("couldn't connect: bad compute config: {:?}", self.config),
-            )
-        }))
+        Ok((socket_addr, socket))
     }
-}
 
-pub struct PostgresConnection {
-    /// Socket connected to a compute node.
-    pub stream: TcpStream,
-    /// PostgreSQL version of this instance.
-    pub version: String,
-}
-
-impl NodeInfo {
     /// Connect to a corresponding compute node.
-    pub async fn connect(&self) -> Result<(PostgresConnection, CancelClosure), ConnectionError> {
-        let (socket_addr, mut stream) = self
+    pub async fn connect(self) -> Result<(TcpStream, Version, CancelClosure), ConnectionError> {
+        let (socket_addr, mut socket) = self
             .connect_raw()
             .await
             .map_err(|_| ConnectionError::FailedToConnectToCompute)?;
 
+        let mut config = tokio_postgres::Config::from(self.db_info);
+        if let Some(scram_keys) = self.scram_keys {
+            config.auth_keys(tokio_postgres::config::AuthKeys::ScramSha256(scram_keys));
+        }
+
         // TODO: establish a secure connection to the DB
-        let (client, conn) = self.config.connect_raw(&mut stream, NoTls).await?;
+        let (client, conn) = config.connect_raw(&mut socket, NoTls).await?;
         let version = conn
             .parameter("server_version")
             .ok_or(ConnectionError::FailedToFetchPgVersion)?
@@ -116,8 +66,6 @@ impl NodeInfo {
 
         let cancel_closure = CancelClosure::new(socket_addr, client.cancel_token());
 
-        let db = PostgresConnection { stream, version };
-
-        Ok((db, cancel_closure))
+        Ok((socket, version, cancel_closure))
     }
 }