assume primary is the default

Signed-off-by: Alex Chi Z <chi@neon.tech>

.github/actionlint.yml

@@ -33,3 +33,5 @@ config-variables:
   - NEON_PROD_AWS_ACCOUNT_ID
   - AWS_ECR_REGION
   - BENCHMARK_LARGE_OLTP_PROJECTID
+  - SLACK_ON_CALL_DEVPROD_STREAM
+  - SLACK_RUST_CHANNEL_ID

.github/workflows/_meta.yml

@@ -19,7 +19,7 @@ on:
         description: "Tag of the last compute release"
         value: ${{ jobs.tags.outputs.compute }}
       run-kind:
-        description: "The kind of run we're currently in. Will be one of `pr`, `push-main`, `storage-rc`, `storage-release`, `proxy-rc`, `proxy-release`, `compute-rc`, `compute-release` or `merge_queue`"
+        description: "The kind of run we're currently in. Will be one of `push-main`, `storage-release`, `compute-release`, `proxy-release`, `storage-rc-pr`, `compute-rc-pr`,  `proxy-rc-pr`, `pr`, or `workflow-dispatch`"
         value: ${{ jobs.tags.outputs.run-kind }}
 
 permissions: {}
@@ -55,6 +55,7 @@ jobs:
               || (inputs.github-event-name == 'pull_request' && github.base_ref == 'release-compute') && 'compute-rc-pr'
               || (inputs.github-event-name == 'pull_request' && github.base_ref == 'release-proxy')   && 'proxy-rc-pr'
               || (inputs.github-event-name == 'pull_request')                                         && 'pr'
+              || (inputs.github-event-name == 'workflow_dispatch')                                    && 'workflow-dispatch'
               || 'unknown'
             }}
         run: |
@@ -85,6 +86,9 @@ jobs:
             BUILD_AND_TEST_RUN_ID=$(gh run list -b $CURRENT_BRANCH -c $CURRENT_SHA -w 'Build and Test' -L 1 --json databaseId --jq '.[].databaseId')
             echo "tag=$BUILD_AND_TEST_RUN_ID" | tee -a $GITHUB_OUTPUT
             ;;
+          workflow-dispatch)
+            echo "tag=$GITHUB_RUN_ID" | tee -a $GITHUB_OUTPUT
+            ;;
           *)
             echo "Unexpected RUN_KIND ('${RUN_KIND}'), failing to assign build-tag!"
             exit 1

.github/workflows/build_and_test.yml

@@ -1061,7 +1061,7 @@ jobs:
           exit 1
 
   deploy:
-    needs: [ check-permissions, push-neon-image-prod, push-compute-image-prod, meta, build-and-test-locally, trigger-custom-extensions-build-and-wait ]
+    needs: [ check-permissions, push-neon-image-dev, push-compute-image-dev, push-neon-image-prod, push-compute-image-prod, meta, build-and-test-locally, trigger-custom-extensions-build-and-wait ]
     # `!failure() && !cancelled()` is required because the workflow depends on the job that can be skipped: `push-neon-image-prod` and `push-compute-image-prod`
     if: ${{ contains(fromJSON('["push-main", "storage-release", "proxy-release", "compute-release"]'), needs.meta.outputs.run-kind) && !failure() && !cancelled() }}
     permissions:
@@ -1175,7 +1175,7 @@ jobs:
               -f deployPgSniRouter=false \
               -f deployProxy=false \
               -f deployStorage=true \
-              -f deployStorageBroker=true \
+              -f deployStorageBroker=false \
               -f deployStorageController=true \
               -f branch=main \
               -f dockerTag=${{needs.meta.outputs.build-tag}} \
@@ -1183,7 +1183,7 @@ jobs:
 
             gh workflow --repo neondatabase/infra run deploy-prod.yml --ref main \
               -f deployStorage=true \
-              -f deployStorageBroker=true \
+              -f deployStorageBroker=false \
               -f deployStorageController=true \
               -f branch=main \
               -f dockerTag=${{needs.meta.outputs.build-tag}}
@@ -1231,7 +1231,7 @@ jobs:
           payload: |
             channel: ${{ vars.SLACK_STORAGE_CHANNEL_ID }}
             text: |
-              🔴 @oncall-storage: deploy job on release branch had unexpected status "${{ needs.deploy.result }}" <${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}|GitHub Run>.
+              🔴 <!subteam^S06CJ87UMNY|@oncall-storage>: deploy job on release branch had unexpected status "${{ needs.deploy.result }}" <${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}|GitHub Run>.
 
   # The job runs on `release` branch and copies compatibility data and Neon artifact from the last *release PR* to the latest directory
   promote-compatibility-data:

.github/workflows/cargo-deny.yml

@@ -7,7 +7,7 @@ on:
         required: false
         type: string
   schedule:
-    - cron: '0 0 * * *'
+    - cron: '0 10 * * *'
 
 jobs:
   cargo-deny:
@@ -50,8 +50,9 @@ jobs:
           method: chat.postMessage
           token: ${{ secrets.SLACK_BOT_TOKEN }}
           payload: |
-            channel: ${{ vars.SLACK_CICD_CHANNEL_ID }}
+            channel: ${{ vars.SLACK_ON_CALL_DEVPROD_STREAM }}
             text: |
               Periodic cargo-deny on ${{ matrix.ref }}: ${{ job.status }}
               <${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}|GitHub Run>
-              Pinging @oncall-devprod.
+              Fixing the problem should be fairly straight forward from the logs. If not, <#${{ vars.SLACK_RUST_CHANNEL_ID }}> is there to help.
+              Pinging <!subteam^S0838JPSH32|@oncall-devprod>.

.github/workflows/periodic_pagebench.yml

@@ -3,12 +3,12 @@ name: Periodic pagebench performance test on dedicated EC2 machine in eu-central
 on:
   schedule:
     # * is a special character in YAML so you have to quote this string
-    #          ┌───────────── minute (0 - 59)
-    #          │ ┌───────────── hour (0 - 23)
-    #          │ │ ┌───────────── day of the month (1 - 31)
-    #          │ │ │ ┌───────────── month (1 - 12 or JAN-DEC)
-    #          │ │ │ │ ┌───────────── day of the week (0 - 6 or SUN-SAT)
-    - cron:  '0 18 * * *' # Runs at 6 PM UTC every day
+    #        ┌───────────── minute (0 - 59)
+    #        │   ┌───────────── hour (0 - 23)
+    #        │   │ ┌───────────── day of the month (1 - 31)
+    #        │   │ │ ┌───────────── month (1 - 12 or JAN-DEC)
+    #        │   │ │ │ ┌───────────── day of the week (0 - 6 or SUN-SAT)
+    - cron: '0 */3 * * *' # Runs every 3 hours
   workflow_dispatch: # Allows manual triggering of the workflow
     inputs:
       commit_hash:

CODEOWNERS

@@ -1,8 +1,8 @@
 # Autoscaling
 /libs/vm_monitor/ @neondatabase/autoscaling
 
-# DevProd
-/.github/ @neondatabase/developer-productivity
+# DevProd & PerfCorr
+/.github/ @neondatabase/developer-productivity @neondatabase/performance-correctness
 
 # Compute
 /pgxn/ @neondatabase/compute

Cargo.lock

@@ -1127,9 +1127,9 @@ checksum = "37b2a672a2cb129a2e41c10b1224bb368f9f37a2b16b612598138befd7b37eb5"
 
 [[package]]
 name = "cc"
-version = "1.1.30"
+version = "1.2.16"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b16803a61b81d9eabb7eae2588776c4c1e584b738ede45fdbb4c972cec1e9945"
+checksum = "be714c154be609ec7f5dad223a33bf1482fff90472de28f7362806e6d4832b8c"
 dependencies = [
  "jobserver",
  "libc",
@@ -2848,6 +2848,7 @@ dependencies = [
  "anyhow",
  "bytes",
  "fail",
+ "futures",
  "hyper 0.14.30",
  "itertools 0.10.5",
  "jemalloc_pprof",
@@ -2861,6 +2862,7 @@ dependencies = [
  "serde_path_to_error",
  "thiserror 1.0.69",
  "tokio",
+ "tokio-rustls 0.26.0",
  "tokio-stream",
  "tokio-util",
  "tracing",
@@ -4062,7 +4064,7 @@ dependencies = [
  "opentelemetry-http",
  "opentelemetry-proto",
  "opentelemetry_sdk",
- "prost",
+ "prost 0.13.3",
  "reqwest",
  "thiserror 1.0.69",
 ]
@@ -4075,7 +4077,7 @@ checksum = "a6e05acbfada5ec79023c85368af14abd0b307c015e9064d249b2a950ef459a6"
 dependencies = [
  "opentelemetry",
  "opentelemetry_sdk",
- "prost",
+ "prost 0.13.3",
  "tonic",
 ]
 
@@ -4189,6 +4191,7 @@ dependencies = [
  "pageserver_api",
  "pageserver_client",
  "rand 0.8.5",
+ "reqwest",
  "serde",
  "serde_json",
  "tokio",
@@ -4278,6 +4281,9 @@ dependencies = [
  "remote_storage",
  "reqwest",
  "rpds",
+ "rustls 0.23.18",
+ "rustls-pemfile 2.1.1",
+ "rustls-pki-types",
  "scopeguard",
  "send-future",
  "serde",
@@ -4296,6 +4302,7 @@ dependencies = [
  "tokio-epoll-uring",
  "tokio-io-timeout",
  "tokio-postgres",
+ "tokio-rustls 0.26.0",
  "tokio-stream",
  "tokio-tar",
  "tokio-util",
@@ -4303,6 +4310,7 @@ dependencies = [
  "tracing",
  "url",
  "utils",
+ "uuid",
  "wal_decoder",
  "walkdir",
  "workspace_hack",
@@ -4829,8 +4837,10 @@ dependencies = [
  "nix 0.26.4",
  "once_cell",
  "parking_lot 0.12.1",
- "protobuf",
- "protobuf-codegen-pure",
+ "prost 0.12.6",
+ "prost-build 0.12.6",
+ "prost-derive 0.12.6",
+ "sha2",
  "smallvec",
  "symbolic-demangle",
  "tempfile",
@@ -4849,7 +4859,7 @@ dependencies = [
  "inferno 0.12.0",
  "num",
  "paste",
- "prost",
+ "prost 0.13.3",
 ]
 
 [[package]]
@@ -4942,6 +4952,16 @@ dependencies = [
  "thiserror 1.0.69",
 ]
 
+[[package]]
+name = "prost"
+version = "0.12.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "deb1435c188b76130da55f17a466d252ff7b1418b2ad3e037d127b94e3411f29"
+dependencies = [
+ "bytes",
+ "prost-derive 0.12.6",
+]
+
 [[package]]
 name = "prost"
 version = "0.13.3"
@@ -4949,7 +4969,28 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "7b0487d90e047de87f984913713b85c601c05609aad5b0df4b4573fbf69aa13f"
 dependencies = [
  "bytes",
- "prost-derive",
+ "prost-derive 0.13.3",
+]
+
+[[package]]
+name = "prost-build"
+version = "0.12.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "22505a5c94da8e3b7c2996394d1c933236c4d743e81a410bcca4e6989fc066a4"
+dependencies = [
+ "bytes",
+ "heck",
+ "itertools 0.10.5",
+ "log",
+ "multimap",
+ "once_cell",
+ "petgraph",
+ "prettyplease",
+ "prost 0.12.6",
+ "prost-types 0.12.6",
+ "regex",
+ "syn 2.0.90",
+ "tempfile",
 ]
 
 [[package]]
@@ -4966,13 +5007,26 @@ dependencies = [
  "once_cell",
  "petgraph",
  "prettyplease",
- "prost",
- "prost-types",
+ "prost 0.13.3",
+ "prost-types 0.13.3",
  "regex",
  "syn 2.0.90",
  "tempfile",
 ]
 
+[[package]]
+name = "prost-derive"
+version = "0.12.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "81bddcdb20abf9501610992b6759a4c888aef7d1a7247ef75e2404275ac24af1"
+dependencies = [
+ "anyhow",
+ "itertools 0.10.5",
+ "proc-macro2",
+ "quote",
+ "syn 2.0.90",
+]
+
 [[package]]
 name = "prost-derive"
 version = "0.13.3"
@@ -4986,38 +5040,22 @@ dependencies = [
  "syn 2.0.90",
 ]
 
+[[package]]
+name = "prost-types"
+version = "0.12.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9091c90b0a32608e984ff2fa4091273cbdd755d54935c51d520887f4a1dbd5b0"
+dependencies = [
+ "prost 0.12.6",
+]
+
 [[package]]
 name = "prost-types"
 version = "0.13.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "4759aa0d3a6232fb8dbdb97b61de2c20047c68aca932c7ed76da9d788508d670"
 dependencies = [
- "prost",
-]
-
-[[package]]
-name = "protobuf"
-version = "2.28.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "106dd99e98437432fed6519dedecfade6a06a73bb7b2a1e019fdd2bee5778d94"
-
-[[package]]
-name = "protobuf-codegen"
-version = "2.28.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "033460afb75cf755fcfc16dfaed20b86468082a2ea24e05ac35ab4a099a017d6"
-dependencies = [
- "protobuf",
-]
-
-[[package]]
-name = "protobuf-codegen-pure"
-version = "2.28.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "95a29399fc94bcd3eeaa951c715f7bea69409b2445356b00519740bcd6ddd865"
-dependencies = [
- "protobuf",
- "protobuf-codegen",
+ "prost 0.13.3",
 ]
 
 [[package]]
@@ -5626,16 +5664,16 @@ dependencies = [
 
 [[package]]
 name = "ring"
-version = "0.17.6"
+version = "0.17.13"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "684d5e6e18f669ccebf64a92236bb7db9a34f07be010e3627368182027180866"
+checksum = "70ac5d832aa16abd7d1def883a8545280c20a60f523a370aa3a9617c2b8550ee"
 dependencies = [
  "cc",
+ "cfg-if",
  "getrandom 0.2.11",
  "libc",
- "spin",
  "untrusted",
- "windows-sys 0.48.0",
+ "windows-sys 0.52.0",
 ]
 
 [[package]]
@@ -5877,9 +5915,9 @@ dependencies = [
 
 [[package]]
 name = "rustls-pki-types"
-version = "1.10.0"
+version = "1.11.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "16f1201b3c9a7ee8039bcadc17b7e605e2945b27eee7631788c1bd2b0643674b"
+checksum = "917ce264624a4b4db1c364dcc35bfca9ded014d0a958cd47ad3e960e988ea51c"
 
 [[package]]
 name = "rustls-webpki"
@@ -6565,7 +6603,7 @@ dependencies = [
  "metrics",
  "once_cell",
  "parking_lot 0.12.1",
- "prost",
+ "prost 0.13.3",
  "rustls 0.23.18",
  "tokio",
  "tonic",
@@ -6583,6 +6621,7 @@ dependencies = [
  "bytes",
  "chrono",
  "clap",
+ "clashmap",
  "control_plane",
  "cron",
  "diesel",
@@ -7337,7 +7376,7 @@ dependencies = [
  "hyper-util",
  "percent-encoding",
  "pin-project",
- "prost",
+ "prost 0.13.3",
  "rustls-native-certs 0.8.0",
  "rustls-pemfile 2.1.1",
  "tokio",
@@ -7357,8 +7396,8 @@ checksum = "9557ce109ea773b399c9b9e5dca39294110b74f1f342cb347a80d1fce8c26a11"
 dependencies = [
  "prettyplease",
  "proc-macro2",
- "prost-build",
- "prost-types",
+ "prost-build 0.13.3",
+ "prost-types 0.13.3",
  "quote",
  "syn 2.0.90",
 ]
@@ -7891,7 +7930,7 @@ dependencies = [
  "pageserver_api",
  "postgres_ffi",
  "pprof",
- "prost",
+ "prost 0.13.3",
  "remote_storage",
  "serde",
  "serde_json",
@@ -8348,6 +8387,7 @@ dependencies = [
  "hyper-util",
  "indexmap 1.9.3",
  "indexmap 2.0.1",
+ "itertools 0.10.5",
  "itertools 0.12.1",
  "lazy_static",
  "libc",
@@ -8366,7 +8406,7 @@ dependencies = [
  "parquet",
  "prettyplease",
  "proc-macro2",
- "prost",
+ "prost 0.13.3",
  "quote",
  "rand 0.8.5",
  "regex",

Cargo.toml

@@ -139,7 +139,7 @@ parquet = { version = "53", default-features = false, features = ["zstd"] }
 parquet_derive = "53"
 pbkdf2 = { version = "0.12.1", features = ["simple", "std"] }
 pin-project-lite = "0.2"
-pprof = { version = "0.14", features = ["criterion", "flamegraph", "frame-pointer", "protobuf", "protobuf-codec"] }
+pprof = { version = "0.14", features = ["criterion", "flamegraph", "frame-pointer", "prost-codec"] }
 procfs = "0.16"
 prometheus = {version = "0.13", default-features=false, features = ["process"]} # removes protobuf dependency
 prost = "0.13"
@@ -155,6 +155,7 @@ rpds = "0.13"
 rustc-hash = "1.1.0"
 rustls = { version = "0.23.16", default-features = false }
 rustls-pemfile = "2"
+rustls-pki-types = "1.11"
 scopeguard = "1.1"
 sysinfo = "0.29.2"
 sd-notify = "0.4.1"

compute/compute-node.Dockerfile

@@ -1484,7 +1484,7 @@ WORKDIR /ext-src
 COPY compute/patches/pg_duckdb_v031.patch .
 COPY compute/patches/duckdb_v120.patch .
 # pg_duckdb build requires source dir to be a git repo to get submodules
-# allow neon_superuser to execute some functions that in pg_duckdb are available to superuser only: 
+# allow neon_superuser to execute some functions that in pg_duckdb are available to superuser only:
 # - extension management function duckdb.install_extension()
 # - access to duckdb.extensions table and its sequence
 RUN git clone --depth 1 --branch v0.3.1 https://github.com/duckdb/pg_duckdb.git pg_duckdb-src && \
@@ -1499,8 +1499,8 @@ ARG PG_VERSION
 COPY --from=pg_duckdb-src /ext-src/ /ext-src/
 WORKDIR /ext-src/pg_duckdb-src
 RUN make install -j $(getconf _NPROCESSORS_ONLN) && \
-    echo 'trusted = true' >> /usr/local/pgsql/share/extension/pg_duckdb.control 
-        
+    echo 'trusted = true' >> /usr/local/pgsql/share/extension/pg_duckdb.control
+
 #########################################################################################
 #
 # Layer "pg_repack"
@@ -1758,15 +1758,15 @@ ARG TARGETARCH
 # test_runner/regress/test_compute_metrics.py
 # See comment on the top of the file regading `echo`, `-e` and `\n`
 RUN if [ "$TARGETARCH" = "amd64" ]; then\
-        postgres_exporter_sha256='027e75dda7af621237ff8f5ac66b78a40b0093595f06768612b92b1374bd3105';\
+        postgres_exporter_sha256='59aa4a7bb0f7d361f5e05732f5ed8c03cc08f78449cef5856eadec33a627694b';\
         pgbouncer_exporter_sha256='c9f7cf8dcff44f0472057e9bf52613d93f3ffbc381ad7547a959daa63c5e84ac';\
         sql_exporter_sha256='38e439732bbf6e28ca4a94d7bc3686d3fa1abdb0050773d5617a9efdb9e64d08';\
     else\
-        postgres_exporter_sha256='131a376d25778ff9701a4c81f703f179e0b58db5c2c496e66fa43f8179484786';\
+        postgres_exporter_sha256='d1dedea97f56c6d965837bfd1fbb3e35a3b4a4556f8cccee8bd513d8ee086124';\
         pgbouncer_exporter_sha256='217c4afd7e6492ae904055bc14fe603552cf9bac458c063407e991d68c519da3';\
         sql_exporter_sha256='11918b00be6e2c3a67564adfdb2414fdcbb15a5db76ea17d1d1a944237a893c6';\
     fi\
-    && curl -sL https://github.com/prometheus-community/postgres_exporter/releases/download/v0.16.0/postgres_exporter-0.16.0.linux-${TARGETARCH}.tar.gz\
+    && curl -sL https://github.com/prometheus-community/postgres_exporter/releases/download/v0.17.1/postgres_exporter-0.17.1.linux-${TARGETARCH}.tar.gz\
      | tar xzf - --strip-components=1 -C.\
     && curl -sL https://github.com/prometheus-community/pgbouncer_exporter/releases/download/v0.10.2/pgbouncer_exporter-0.10.2.linux-${TARGETARCH}.tar.gz\
      | tar xzf - --strip-components=1 -C.\
@@ -1980,12 +1980,10 @@ COPY --from=sql_exporter_preprocessor --chmod=0644 /home/nonroot/compute/etc/neo
 RUN echo '/usr/local/lib' >> /etc/ld.so.conf && /sbin/ldconfig
 
 # rsyslog config permissions
-RUN chown postgres:postgres /etc/rsyslog.conf && \
-    touch /etc/compute_rsyslog.conf && \
-    chown -R postgres:postgres /etc/compute_rsyslog.conf && \
-    # directory for rsyslogd pid file
-    mkdir /var/run/rsyslogd && \
-    chown -R postgres:postgres /var/run/rsyslogd
+# directory for rsyslogd pid file
+RUN mkdir /var/run/rsyslogd && \
+    chown -R postgres:postgres /var/run/rsyslogd && \
+    chown -R postgres:postgres /etc/rsyslog.d/
 
 
 ENV LANG=en_US.utf8

compute/etc/neon_collector.jsonnet

@@ -29,6 +29,7 @@
     import 'sql_exporter/lfc_approximate_working_set_size.libsonnet',
     import 'sql_exporter/lfc_approximate_working_set_size_windows.libsonnet',
     import 'sql_exporter/lfc_cache_size_limit.libsonnet',
+    import 'sql_exporter/lfc_chunk_size.libsonnet',
     import 'sql_exporter/lfc_hits.libsonnet',
     import 'sql_exporter/lfc_misses.libsonnet',
     import 'sql_exporter/lfc_used.libsonnet',

compute/etc/sql_exporter/db_total_size.sql

@@ -1 +1,5 @@
-SELECT sum(pg_database_size(datname)) AS total FROM pg_database;
+SELECT sum(pg_database_size(datname)) AS total
+FROM pg_database
+-- Ignore invalid databases, as we will likely have problems with
+-- getting their size from the Pageserver.
+WHERE datconnlimit != -2;

compute/etc/sql_exporter/lfc_chunk_size.libsonnet

@@ -0,0 +1,10 @@
+{
+  metric_name: 'lfc_chunk_size',
+  type: 'gauge',
+  help: 'LFC chunk size, measured in 8KiB pages',
+  key_labels: null,
+  values: [
+    'lfc_chunk_size_pages',
+  ],
+  query: importstr 'sql_exporter/lfc_chunk_size.sql',
+}

compute/etc/sql_exporter/lfc_chunk_size.sql

@@ -0,0 +1 @@
+SELECT lfc_value AS lfc_chunk_size_pages FROM neon.neon_lfc_stats WHERE lfc_key = 'file_cache_chunk_size_pages';

compute/etc/sql_exporter/pg_stats_userdb.sql

@@ -1,10 +1,20 @@
 -- We export stats for 10 non-system databases. Without this limit it is too
 -- easy to abuse the system by creating lots of databases.
 
-SELECT pg_database_size(datname) AS db_size, deadlocks, tup_inserted AS inserted,
-  tup_updated AS updated, tup_deleted AS deleted, datname
+SELECT pg_database_size(datname) AS db_size,
+  deadlocks,
+  tup_inserted AS inserted,
+  tup_updated AS updated,
+  tup_deleted AS deleted,
+  datname
 FROM pg_stat_database
 WHERE datname IN (
   SELECT datname FROM pg_database
-  WHERE datname <> 'postgres' AND NOT datistemplate ORDER BY oid LIMIT 10
+  -- Ignore invalid databases, as we will likely have problems with
+  -- getting their size from the Pageserver.
+  WHERE datconnlimit != -2
+    AND datname <> 'postgres'
+    AND NOT datistemplate
+  ORDER BY oid
+  LIMIT 10
 );

compute/vm-image-spec-bookworm.yaml

@@ -42,7 +42,7 @@ commands:
   - name: rsyslogd
     user: postgres
     sysvInitAction: respawn
-    shell: '/usr/sbin/rsyslogd -i /var/run/rsyslogd/rsyslogd.pid -f /etc/compute_rsyslog.conf'
+    shell: '/usr/sbin/rsyslogd -n -i /var/run/rsyslogd/rsyslogd.pid -f /etc/compute_rsyslog.conf'
 shutdownHook: |
   su -p postgres --session-command '/usr/local/bin/pg_ctl stop -D /var/db/postgres/compute/pgdata -m fast --wait -t 10'
 files:
@@ -78,6 +78,7 @@ files:
   - filename: compute_rsyslog.conf
     content: |
       *.*    /dev/null
+      $IncludeConfig /etc/rsyslog.d/*.conf
 build: |
   # Build cgroup-tools
   #
@@ -141,9 +142,11 @@ merge: |
   RUN set -e \
       && chmod 0644 /etc/cgconfig.conf
 
+
   COPY compute_rsyslog.conf /etc/compute_rsyslog.conf
-  RUN set -e \
-      && chmod 0644 /etc/compute_rsyslog.conf
+  RUN chmod 0666 /etc/compute_rsyslog.conf
+  RUN chmod 0666 /var/log/
+
 
   COPY --from=libcgroup-builder /libcgroup-install/bin/*  /usr/bin/
   COPY --from=libcgroup-builder /libcgroup-install/lib/*  /usr/lib/

compute/vm-image-spec-bullseye.yaml

@@ -42,7 +42,7 @@ commands:
   - name: rsyslogd
     user: postgres
     sysvInitAction: respawn
-    shell: '/usr/sbin/rsyslogd -i /var/run/rsyslogd/rsyslogd.pid -f /etc/compute_rsyslog.conf'
+    shell: '/usr/sbin/rsyslogd -n -i /var/run/rsyslogd/rsyslogd.pid -f /etc/compute_rsyslog.conf'
 shutdownHook: |
   su -p postgres --session-command '/usr/local/bin/pg_ctl stop -D /var/db/postgres/compute/pgdata -m fast --wait -t 10'
 files:
@@ -78,6 +78,7 @@ files:
   - filename: compute_rsyslog.conf
     content: |
       *.*    /dev/null
+      $IncludeConfig /etc/rsyslog.d/*.conf
 build: |
   # Build cgroup-tools
   #
@@ -138,8 +139,9 @@ merge: |
       && chmod 0644 /etc/cgconfig.conf
 
   COPY compute_rsyslog.conf /etc/compute_rsyslog.conf
-  RUN set -e \
-      && chmod 0644 /etc/compute_rsyslog.conf
+  RUN chmod 0666 /etc/compute_rsyslog.conf
+  RUN chmod 0666 /var/log/
+
 
   COPY --from=libcgroup-builder /libcgroup-install/bin/*  /usr/bin/
   COPY --from=libcgroup-builder /libcgroup-install/lib/*  /usr/lib/

compute_tools/src/compute.rs

@@ -37,7 +37,7 @@ use crate::logger::startup_context_from_env;
 use crate::lsn_lease::launch_lsn_lease_bg_task_for_static;
 use crate::monitor::launch_monitor;
 use crate::pg_helpers::*;
-use crate::rsyslog::configure_and_start_rsyslog;
+use crate::rsyslog::configure_audit_rsyslog;
 use crate::spec::*;
 use crate::swap::resize_swap;
 use crate::sync_sk::{check_if_synced, ping_safekeeper};
@@ -297,79 +297,6 @@ struct StartVmMonitorResult {
     vm_monitor: Option<tokio::task::JoinHandle<Result<()>>>,
 }
 
-pub(crate) fn construct_superuser_query(spec: &ComputeSpec) -> String {
-    let roles = spec
-        .cluster
-        .roles
-        .iter()
-        .map(|r| escape_literal(&r.name))
-        .collect::<Vec<_>>();
-
-    let dbs = spec
-        .cluster
-        .databases
-        .iter()
-        .map(|db| escape_literal(&db.name))
-        .collect::<Vec<_>>();
-
-    let roles_decl = if roles.is_empty() {
-        String::from("roles text[] := NULL;")
-    } else {
-        format!(
-            r#"
-               roles text[] := ARRAY(SELECT rolname
-                                     FROM pg_catalog.pg_roles
-                                     WHERE rolname IN ({}));"#,
-            roles.join(", ")
-        )
-    };
-
-    let database_decl = if dbs.is_empty() {
-        String::from("dbs text[] := NULL;")
-    } else {
-        format!(
-            r#"
-               dbs text[] := ARRAY(SELECT datname
-                                   FROM pg_catalog.pg_database
-                                   WHERE datname IN ({}));"#,
-            dbs.join(", ")
-        )
-    };
-
-    // ALL PRIVILEGES grants CREATE, CONNECT, and TEMPORARY on all databases
-    // (see https://www.postgresql.org/docs/current/ddl-priv.html)
-    let query = format!(
-        r#"
-            DO $$
-                DECLARE
-                    r text;
-                    {}
-                    {}
-                BEGIN
-                    IF NOT EXISTS (
-                        SELECT FROM pg_catalog.pg_roles WHERE rolname = 'neon_superuser')
-                    THEN
-                        CREATE ROLE neon_superuser CREATEDB CREATEROLE NOLOGIN REPLICATION BYPASSRLS IN ROLE pg_read_all_data, pg_write_all_data;
-                        IF array_length(roles, 1) IS NOT NULL THEN
-                            EXECUTE format('GRANT neon_superuser TO %s',
-                                           array_to_string(ARRAY(SELECT quote_ident(x) FROM unnest(roles) as x), ', '));
-                            FOREACH r IN ARRAY roles LOOP
-                                EXECUTE format('ALTER ROLE %s CREATEROLE CREATEDB', quote_ident(r));
-                            END LOOP;
-                        END IF;
-                        IF array_length(dbs, 1) IS NOT NULL THEN
-                            EXECUTE format('GRANT ALL PRIVILEGES ON DATABASE %s TO neon_superuser',
-                                           array_to_string(ARRAY(SELECT quote_ident(x) FROM unnest(dbs) as x), ', '));
-                        END IF;
-                    END IF;
-                END
-            $$;"#,
-        roles_decl, database_decl,
-    );
-
-    query
-}
-
 impl ComputeNode {
     pub fn new(
         params: ComputeNodeParams,
@@ -689,7 +616,7 @@ impl ComputeNode {
             let log_directory_path = Path::new(&self.params.pgdata).join("log");
             // TODO: make this more robust
             // now rsyslog starts once and there is no monitoring or restart if it fails
-            configure_and_start_rsyslog(
+            configure_audit_rsyslog(
                 log_directory_path.to_str().unwrap(),
                 "hipaa",
                 &remote_endpoint,
@@ -928,6 +855,12 @@ impl ComputeNode {
             info!("Storage auth token not set");
         }
 
+        if let Some(spec) = &compute_state.pspec {
+            config.application_name(&format!("compute_ctl-{}", spec.spec.mode.to_type_str()));
+        } else {
+            config.application_name("compute_ctl");
+        }
+
         // Connect to pageserver
         let mut client = config.connect(NoTls)?;
         let pageserver_connect_micros = start_time.elapsed().as_micros() as u64;

compute_tools/src/config.rs

@@ -99,6 +99,7 @@ pub fn write_postgres_conf(
         writeln!(file, "lc_numeric='C.UTF-8'")?;
     }
 
+    writeln!(file, "neon.endpoint_type={}", spec.mode.to_type_str())?;
     match spec.mode {
         ComputeMode::Primary => {}
         ComputeMode::Static(lsn) => {

compute_tools/src/extension_server.rs

@@ -202,8 +202,24 @@ pub async fn download_extension(
     // move contents of the libdir / sharedir in unzipped archive to the correct local paths
     for paths in [sharedir_paths, libdir_paths] {
         let (zip_dir, real_dir) = paths;
+
+        let dir = match std::fs::read_dir(&zip_dir) {
+            Ok(dir) => dir,
+            Err(e) => match e.kind() {
+                // In the event of a SQL-only extension, there would be nothing
+                // to move from the lib/ directory, so note that in the log and
+                // move on.
+                std::io::ErrorKind::NotFound => {
+                    info!("nothing to move from {}", zip_dir);
+                    continue;
+                }
+                _ => return Err(anyhow::anyhow!(e)),
+            },
+        };
+
         info!("mv {zip_dir:?}/*  {real_dir:?}");
-        for file in std::fs::read_dir(zip_dir)? {
+
+        for file in dir {
             let old_file = file?.path();
             let new_file =
                 Path::new(&real_dir).join(old_file.file_name().context("error parsing file")?);
@@ -253,27 +269,31 @@ pub fn create_control_files(remote_extensions: &RemoteExtSpec, pgbin: &str) {
     }
 }
 
-// Do request to extension storage proxy, i.e.
+// Do request to extension storage proxy, e.g.,
 // curl http://pg-ext-s3-gateway/latest/v15/extensions/anon.tar.zst
-// using HHTP GET
-// and return the response body as bytes
-//
+// using HTTP GET and return the response body as bytes.
 async fn download_extension_tar(ext_remote_storage: &str, ext_path: &str) -> Result<Bytes> {
     let uri = format!("{}/{}", ext_remote_storage, ext_path);
+    let filename = Path::new(ext_path)
+        .file_name()
+        .unwrap_or_else(|| std::ffi::OsStr::new("unknown"))
+        .to_str()
+        .unwrap_or("unknown")
+        .to_string();
 
-    info!("Download extension {} from uri {}", ext_path, uri);
+    info!("Downloading extension file '{}' from uri {}", filename, uri);
 
     match do_extension_server_request(&uri).await {
         Ok(resp) => {
             info!("Successfully downloaded remote extension data {}", ext_path);
             REMOTE_EXT_REQUESTS_TOTAL
-                .with_label_values(&[&StatusCode::OK.to_string()])
+                .with_label_values(&[&StatusCode::OK.to_string(), &filename])
                 .inc();
             Ok(resp)
         }
         Err((msg, status)) => {
             REMOTE_EXT_REQUESTS_TOTAL
-                .with_label_values(&[&status])
+                .with_label_values(&[&status, &filename])
                 .inc();
             bail!(msg);
         }

compute_tools/src/metrics.rs

@@ -54,9 +54,7 @@ pub(crate) static REMOTE_EXT_REQUESTS_TOTAL: Lazy<IntCounterVec> = Lazy::new(||
     register_int_counter_vec!(
         "compute_ctl_remote_ext_requests_total",
         "Total number of requests made by compute_ctl to download extensions from S3 proxy by status",
-        // Do not use any labels like extension name yet.
-        // We can add them later if needed.
-        &["http_status"]
+        &["http_status", "filename"]
     )
     .expect("failed to define a metric")
 });

compute_tools/src/pg_helpers.rs

@@ -186,15 +186,40 @@ impl DatabaseExt for Database {
 /// Postgres SQL queries and DATABASE_URL.
 pub trait Escaping {
     fn pg_quote(&self) -> String;
+    fn pg_quote_dollar(&self) -> (String, String);
 }
 
 impl Escaping for PgIdent {
     /// This is intended to mimic Postgres quote_ident(), but for simplicity it
     /// always quotes provided string with `""` and escapes every `"`.
     /// **Not idempotent**, i.e. if string is already escaped it will be escaped again.
+    /// N.B. it's not useful for escaping identifiers that are used inside WHERE
+    /// clause, use `escape_literal()` instead.
     fn pg_quote(&self) -> String {
-        let result = format!("\"{}\"", self.replace('"', "\"\""));
-        result
+        format!("\"{}\"", self.replace('"', "\"\""))
+    }
+
+    /// This helper is intended to be used for dollar-escaping strings for usage
+    /// inside PL/pgSQL procedures. In addition to dollar-escaping the string,
+    /// it also returns a tag that is intended to be used inside the outer
+    /// PL/pgSQL procedure. If you do not need an outer tag, just discard it.
+    /// Here we somewhat mimic the logic of Postgres' `pg_get_functiondef()`,
+    /// <https://github.com/postgres/postgres/blob/8b49392b270b4ac0b9f5c210e2a503546841e832/src/backend/utils/adt/ruleutils.c#L2924>
+    fn pg_quote_dollar(&self) -> (String, String) {
+        let mut tag: String = "".to_string();
+        let mut outer_tag = "x".to_string();
+
+        // Find the first suitable tag that is not present in the string.
+        // Postgres' max role/DB name length is 63 bytes, so even in the
+        // worst case it won't take long.
+        while self.contains(&format!("${tag}$")) || self.contains(&format!("${outer_tag}$")) {
+            tag += "x";
+            outer_tag = tag.clone() + "x";
+        }
+
+        let escaped = format!("${tag}${self}${tag}$");
+
+        (escaped, outer_tag)
     }
 }
 
@@ -226,10 +251,13 @@ pub async fn get_existing_dbs_async(
     // invalid state. See:
     //   https://github.com/postgres/postgres/commit/a4b4cc1d60f7e8ccfcc8ff8cb80c28ee411ad9a9
     let rowstream = client
+        // We use a subquery instead of a fancy `datdba::regrole::text AS owner`,
+        // because the latter automatically wraps the result in double quotes,
+        // if the role name contains special characters.
         .query_raw::<str, &String, &[String; 0]>(
             "SELECT
                 datname AS name,
-                datdba::regrole::text AS owner,
+                (SELECT rolname FROM pg_roles WHERE oid = datdba) AS owner,
                 NOT datallowconn AS restrict_conn,
                 datconnlimit = - 2 AS invalid
             FROM

compute_tools/src/rsyslog.rs

@@ -21,8 +21,7 @@ fn get_rsyslog_pid() -> Option<String> {
     }
 }
 
-// Start rsyslogd with the specified configuration file
-// If it is already running - restart it.
+// Restart rsyslogd to apply the new configuration.
 // This is necessary, because there is no other way to reload the rsyslog configuration.
 //
 // Rsyslogd shouldn't lose any messages, because of the restart,
@@ -30,42 +29,26 @@ fn get_rsyslog_pid() -> Option<String> {
 // and will continue reading from that position.
 // TODO: test it properly
 //
-fn start_rsyslog(rsyslog_conf_path: &str) -> Result<()> {
-    let old_pid = get_rsyslog_pid();
-    if let Some(pid) = old_pid {
-        info!("rsyslogd is already running with pid: {}, restart it", pid);
-        // kill it to restart
-        let _ = Command::new("pkill")
-            .arg("rsyslogd")
-            .output()
-            .context("Failed to stop rsyslogd")?;
-    }
+fn restart_rsyslog() -> Result<()> {
+    let old_pid = get_rsyslog_pid().context("rsyslogd is not running")?;
+    info!("rsyslogd is running with pid: {}, restart it", old_pid);
 
-    let _ = Command::new("/usr/sbin/rsyslogd")
-        .arg("-f")
-        .arg(rsyslog_conf_path)
-        .arg("-i")
-        .arg("/var/run/rsyslogd/rsyslogd.pid")
+    // kill it to restart
+    let _ = Command::new("pkill")
+        .arg("rsyslogd")
         .output()
-        .context("Failed to start rsyslogd")?;
-
-    // Check that rsyslogd is running
-    if let Some(pid) = get_rsyslog_pid() {
-        info!("rsyslogd started successfully with pid: {}", pid);
-    } else {
-        return Err(anyhow::anyhow!("Failed to start rsyslogd"));
-    }
+        .context("Failed to stop rsyslogd")?;
 
     Ok(())
 }
 
-pub fn configure_and_start_rsyslog(
+pub fn configure_audit_rsyslog(
     log_directory: &str,
     tag: &str,
     remote_endpoint: &str,
 ) -> Result<()> {
     let config_content: String = format!(
-        include_str!("config_template/compute_rsyslog_template.conf"),
+        include_str!("config_template/compute_audit_rsyslog_template.conf"),
         log_directory = log_directory,
         tag = tag,
         remote_endpoint = remote_endpoint
@@ -73,7 +56,7 @@ pub fn configure_and_start_rsyslog(
 
     info!("rsyslog config_content: {}", config_content);
 
-    let rsyslog_conf_path = "/etc/compute_rsyslog.conf";
+    let rsyslog_conf_path = "/etc/rsyslog.d/compute_audit_rsyslog.conf";
     let mut file = OpenOptions::new()
         .create(true)
         .write(true)
@@ -82,10 +65,13 @@ pub fn configure_and_start_rsyslog(
 
     file.write_all(config_content.as_bytes())?;
 
-    info!("rsyslog configuration added successfully. Starting rsyslogd");
+    info!(
+        "rsyslog configuration file {} added successfully. Starting rsyslogd",
+        rsyslog_conf_path
+    );
 
     // start the service, using the configuration
-    start_rsyslog(rsyslog_conf_path)?;
+    restart_rsyslog()?;
 
     Ok(())
 }

compute_tools/src/spec_apply.rs

@@ -13,16 +13,17 @@ use tokio_postgres::Client;
 use tokio_postgres::error::SqlState;
 use tracing::{Instrument, debug, error, info, info_span, instrument, warn};
 
-use crate::compute::{ComputeNode, ComputeState, construct_superuser_query};
+use crate::compute::{ComputeNode, ComputeState};
 use crate::pg_helpers::{
-    DatabaseExt, Escaping, GenericOptionsSearch, RoleExt, escape_literal, get_existing_dbs_async,
+    DatabaseExt, Escaping, GenericOptionsSearch, RoleExt, get_existing_dbs_async,
     get_existing_roles_async,
 };
 use crate::spec_apply::ApplySpecPhase::{
-    CreateAndAlterDatabases, CreateAndAlterRoles, CreateAvailabilityCheck, CreatePgauditExtension,
-    CreatePgauditlogtofileExtension, CreateSchemaNeon, CreateSuperUser, DisablePostgresDBPgAudit,
-    DropInvalidDatabases, DropRoles, FinalizeDropLogicalSubscriptions, HandleNeonExtension,
-    HandleOtherExtensions, RenameAndDeleteDatabases, RenameRoles, RunInEachDatabase,
+    CreateAndAlterDatabases, CreateAndAlterRoles, CreateAvailabilityCheck, CreateNeonSuperuser,
+    CreatePgauditExtension, CreatePgauditlogtofileExtension, CreateSchemaNeon,
+    DisablePostgresDBPgAudit, DropInvalidDatabases, DropRoles, FinalizeDropLogicalSubscriptions,
+    HandleNeonExtension, HandleOtherExtensions, RenameAndDeleteDatabases, RenameRoles,
+    RunInEachDatabase,
 };
 use crate::spec_apply::PerDatabasePhase::{
     ChangeSchemaPerms, DeleteDBRoleReferences, DropLogicalSubscriptions, HandleAnonExtension,
@@ -187,7 +188,7 @@ impl ComputeNode {
             }
 
             for phase in [
-                CreateSuperUser,
+                CreateNeonSuperuser,
                 DropInvalidDatabases,
                 RenameRoles,
                 CreateAndAlterRoles,
@@ -468,7 +469,7 @@ pub enum PerDatabasePhase {
 
 #[derive(Clone, Debug)]
 pub enum ApplySpecPhase {
-    CreateSuperUser,
+    CreateNeonSuperuser,
     DropInvalidDatabases,
     RenameRoles,
     CreateAndAlterRoles,
@@ -595,14 +596,10 @@ async fn get_operations<'a>(
     apply_spec_phase: &'a ApplySpecPhase,
 ) -> Result<Box<dyn Iterator<Item = Operation> + 'a + Send>> {
     match apply_spec_phase {
-        ApplySpecPhase::CreateSuperUser => {
-            let query = construct_superuser_query(spec);
-
-            Ok(Box::new(once(Operation {
-                query,
-                comment: None,
-            })))
-        }
+        ApplySpecPhase::CreateNeonSuperuser => Ok(Box::new(once(Operation {
+            query: include_str!("sql/create_neon_superuser.sql").to_string(),
+            comment: None,
+        }))),
         ApplySpecPhase::DropInvalidDatabases => {
             let mut ctx = ctx.write().await;
             let databases = &mut ctx.dbs;
@@ -736,14 +733,15 @@ async fn get_operations<'a>(
                         // We do not check whether the DB exists or not,
                         // Postgres will take care of it for us
                         "delete_db" => {
+                            let (db_name, outer_tag) = op.name.pg_quote_dollar();
                             // In Postgres we can't drop a database if it is a template.
                             // So we need to unset the template flag first, but it could
                             // be a retry, so we could've already dropped the database.
                             // Check that database exists first to make it idempotent.
                             let unset_template_query: String = format!(
                                 include_str!("sql/unset_template_for_drop_dbs.sql"),
-                                datname_str = escape_literal(&op.name),
-                                datname = &op.name.pg_quote()
+                                datname = db_name,
+                                outer_tag = outer_tag,
                             );
 
                             // Use FORCE to drop database even if there are active connections.
@@ -850,6 +848,8 @@ async fn get_operations<'a>(
                                 comment: None,
                             },
                             Operation {
+                                // ALL PRIVILEGES grants CREATE, CONNECT, and TEMPORARY on the database
+                                // (see https://www.postgresql.org/docs/current/ddl-priv.html)
                                 query: format!(
                                     "GRANT ALL PRIVILEGES ON DATABASE {} TO neon_superuser",
                                     db.name.pg_quote()
@@ -909,9 +909,11 @@ async fn get_operations<'a>(
                 PerDatabasePhase::DropLogicalSubscriptions => {
                     match &db {
                         DB::UserDB(db) => {
+                            let (db_name, outer_tag) = db.name.pg_quote_dollar();
                             let drop_subscription_query: String = format!(
                                 include_str!("sql/drop_subscriptions.sql"),
-                                datname_str = escape_literal(&db.name),
+                                datname_str = db_name,
+                                outer_tag = outer_tag,
                             );
 
                             let operations = vec![Operation {
@@ -950,6 +952,7 @@ async fn get_operations<'a>(
                                     DB::SystemDB => PgIdent::from("cloud_admin").pg_quote(),
                                     DB::UserDB(db) => db.owner.pg_quote(),
                                 };
+                                let (escaped_role, outer_tag) = op.name.pg_quote_dollar();
 
                                 Some(vec![
                                     // This will reassign all dependent objects to the db owner
@@ -964,7 +967,9 @@ async fn get_operations<'a>(
                                     Operation {
                                         query: format!(
                                             include_str!("sql/pre_drop_role_revoke_privileges.sql"),
-                                            role_name = quoted,
+                                            // N.B. this has to be properly dollar-escaped with `pg_quote_dollar()`
+                                            role_name = escaped_role,
+                                            outer_tag = outer_tag,
                                         ),
                                         comment: None,
                                     },
@@ -989,12 +994,14 @@ async fn get_operations<'a>(
                         DB::SystemDB => return Ok(Box::new(empty())),
                         DB::UserDB(db) => db,
                     };
+                    let (db_owner, outer_tag) = db.owner.pg_quote_dollar();
 
                     let operations = vec![
                         Operation {
                             query: format!(
                                 include_str!("sql/set_public_schema_owner.sql"),
-                                db_owner = db.owner.pg_quote()
+                                db_owner = db_owner,
+                                outer_tag = outer_tag,
                             ),
                             comment: None,
                         },

compute_tools/src/sql/create_neon_superuser.sql

@@ -0,0 +1,8 @@
+DO $$
+    BEGIN
+        IF NOT EXISTS (SELECT FROM pg_catalog.pg_roles WHERE rolname = 'neon_superuser')
+        THEN
+            CREATE ROLE neon_superuser CREATEDB CREATEROLE NOLOGIN REPLICATION BYPASSRLS IN ROLE pg_read_all_data, pg_write_all_data;
+        END IF;
+    END
+$$;

compute_tools/src/sql/drop_subscriptions.sql

@@ -1,4 +1,4 @@
-DO $$
+DO ${outer_tag}$
 DECLARE
     subname TEXT;
 BEGIN
@@ -9,4 +9,4 @@ BEGIN
         EXECUTE format('DROP SUBSCRIPTION %I;', subname);
     END LOOP;
 END;
-$$;
+${outer_tag}$;

compute_tools/src/sql/pre_drop_role_revoke_privileges.sql

@@ -1,8 +1,7 @@
-SET SESSION ROLE neon_superuser;
-
-DO $$
+DO ${outer_tag}$
 DECLARE
     schema TEXT;
+    grantor TEXT;
     revoke_query TEXT;
 BEGIN
     FOR schema IN
@@ -15,14 +14,25 @@ BEGIN
         -- ii) it's easy to add more schemas to the list if needed.
         WHERE schema_name IN ('public')
     LOOP
-        revoke_query := format(
-            'REVOKE ALL PRIVILEGES ON ALL TABLES IN SCHEMA %I FROM {role_name} GRANTED BY neon_superuser;',
-            schema
-        );
+        FOR grantor IN EXECUTE
+            format(
+                'SELECT DISTINCT rtg.grantor FROM information_schema.role_table_grants AS rtg WHERE grantee = %s',
+                -- N.B. this has to be properly dollar-escaped with `pg_quote_dollar()`
+                quote_literal({role_name})
+            )
+        LOOP
+            EXECUTE format('SET LOCAL ROLE %I', grantor);
 
-        EXECUTE revoke_query;
+            revoke_query := format(
+                'REVOKE ALL PRIVILEGES ON ALL TABLES IN SCHEMA %I FROM %I GRANTED BY %I',
+                schema,
+                -- N.B. this has to be properly dollar-escaped with `pg_quote_dollar()`
+                {role_name},
+                grantor
+            );
+
+            EXECUTE revoke_query;
+        END LOOP;
     END LOOP;
 END;
-$$;
-
-RESET ROLE;
+${outer_tag}$;

compute_tools/src/sql/set_public_schema_owner.sql

@@ -1,5 +1,4 @@
-DO
-$$
+DO ${outer_tag}$
     DECLARE
         schema_owner TEXT;
     BEGIN
@@ -16,8 +15,8 @@ $$
 
             IF schema_owner = 'cloud_admin' OR schema_owner = 'zenith_admin'
             THEN
-                ALTER SCHEMA public OWNER TO {db_owner};
+                EXECUTE format('ALTER SCHEMA public OWNER TO %I', {db_owner});
             END IF;
         END IF;
     END
-$$;
+${outer_tag}$;

compute_tools/src/sql/unset_template_for_drop_dbs.sql

@@ -1,12 +1,12 @@
-DO $$
+DO ${outer_tag}$
     BEGIN
         IF EXISTS(
             SELECT 1
             FROM pg_catalog.pg_database
-            WHERE datname = {datname_str}
+            WHERE datname = {datname}
         )
         THEN
-            ALTER DATABASE {datname} is_template false;
+            EXECUTE format('ALTER DATABASE %I is_template false', {datname});
         END IF;
     END
-$$;
+${outer_tag}$;

compute_tools/tests/pg_helpers_tests.rs

@@ -61,6 +61,23 @@ test.escaping = 'here''s a backslash \\ and a quote '' and a double-quote " hoor
         assert_eq!(ident.pg_quote(), "\"\"\"name\"\";\\n select 1;\"");
     }
 
+    #[test]
+    fn ident_pg_quote_dollar() {
+        let test_cases = vec![
+            ("name", ("$$name$$", "x")),
+            ("name$$", ("$x$name$$$x$", "xx")),
+            ("name$$$", ("$x$name$$$$x$", "xx")),
+            ("name$$$$", ("$x$name$$$$$x$", "xx")),
+            ("name$x$", ("$xx$name$x$$xx$", "xxx")),
+        ];
+
+        for (input, expected) in test_cases {
+            let (escaped, tag) = PgIdent::from(input).pg_quote_dollar();
+            assert_eq!(escaped, expected.0);
+            assert_eq!(tag, expected.1);
+        }
+    }
+
     #[test]
     fn generic_options_search() {
         let generic_options: GenericOptions = Some(vec![

control_plane/src/bin/neon_local.rs

@@ -963,6 +963,7 @@ fn handle_init(args: &InitCmdArgs) -> anyhow::Result<LocalEnv> {
                         id: pageserver_id,
                         listen_pg_addr: format!("127.0.0.1:{pg_port}"),
                         listen_http_addr: format!("127.0.0.1:{http_port}"),
+                        listen_https_addr: None,
                         pg_auth_type: AuthType::Trust,
                         http_auth_type: AuthType::Trust,
                         other: Default::default(),
@@ -977,6 +978,7 @@ fn handle_init(args: &InitCmdArgs) -> anyhow::Result<LocalEnv> {
             default_tenant_id: TenantId::from_array(std::array::from_fn(|_| 0)),
             storage_controller: None,
             control_plane_compute_hook_api: None,
+            generate_local_ssl_certs: false,
         }
     };
 

control_plane/src/local_env.rs

@@ -81,6 +81,10 @@ pub struct LocalEnv {
     // but deserialization into a generic toml object as `toml::Value::try_from` fails with an error.
     // https://toml.io/en/v1.0.0 does not contain a concept of "a table inside another table".
     pub branch_name_mappings: HashMap<String, Vec<(TenantId, TimelineId)>>,
+
+    /// Flag to generate SSL certificates for components that need it.
+    /// Also generates root CA certificate that is used to sign all other certificates.
+    pub generate_local_ssl_certs: bool,
 }
 
 /// On-disk state stored in `.neon/config`.
@@ -102,6 +106,10 @@ pub struct OnDiskConfig {
     pub control_plane_api: Option<Url>,
     pub control_plane_compute_hook_api: Option<Url>,
     branch_name_mappings: HashMap<String, Vec<(TenantId, TimelineId)>>,
+    // Note: skip serializing because in compat tests old storage controller fails
+    // to load new config file. May be removed after this field is in release branch.
+    #[serde(skip_serializing_if = "std::ops::Not::not")]
+    pub generate_local_ssl_certs: bool,
 }
 
 fn fail_if_pageservers_field_specified<'de, D>(_: D) -> Result<Vec<PageServerConf>, D::Error>
@@ -129,6 +137,7 @@ pub struct NeonLocalInitConf {
     pub safekeepers: Vec<SafekeeperConf>,
     pub control_plane_api: Option<Url>,
     pub control_plane_compute_hook_api: Option<Option<Url>>,
+    pub generate_local_ssl_certs: bool,
 }
 
 /// Broker config for cluster internal communication.
@@ -166,7 +175,10 @@ pub struct NeonStorageControllerConf {
     #[serde(with = "humantime_serde")]
     pub long_reconcile_threshold: Option<Duration>,
 
-    pub load_safekeepers: bool,
+    #[serde(default)]
+    pub use_https_pageserver_api: bool,
+
+    pub timelines_onto_safekeepers: bool,
 }
 
 impl NeonStorageControllerConf {
@@ -190,7 +202,8 @@ impl Default for NeonStorageControllerConf {
             max_secondary_lag_bytes: None,
             heartbeat_interval: Self::DEFAULT_HEARTBEAT_INTERVAL,
             long_reconcile_threshold: None,
-            load_safekeepers: true,
+            use_https_pageserver_api: false,
+            timelines_onto_safekeepers: false,
         }
     }
 }
@@ -220,6 +233,7 @@ pub struct PageServerConf {
     pub id: NodeId,
     pub listen_pg_addr: String,
     pub listen_http_addr: String,
+    pub listen_https_addr: Option<String>,
     pub pg_auth_type: AuthType,
     pub http_auth_type: AuthType,
     pub no_sync: bool,
@@ -231,6 +245,7 @@ impl Default for PageServerConf {
             id: NodeId(0),
             listen_pg_addr: String::new(),
             listen_http_addr: String::new(),
+            listen_https_addr: None,
             pg_auth_type: AuthType::Trust,
             http_auth_type: AuthType::Trust,
             no_sync: false,
@@ -246,6 +261,7 @@ pub struct NeonLocalInitPageserverConf {
     pub id: NodeId,
     pub listen_pg_addr: String,
     pub listen_http_addr: String,
+    pub listen_https_addr: Option<String>,
     pub pg_auth_type: AuthType,
     pub http_auth_type: AuthType,
     #[serde(default, skip_serializing_if = "std::ops::Not::not")]
@@ -260,6 +276,7 @@ impl From<&NeonLocalInitPageserverConf> for PageServerConf {
             id,
             listen_pg_addr,
             listen_http_addr,
+            listen_https_addr,
             pg_auth_type,
             http_auth_type,
             no_sync,
@@ -269,6 +286,7 @@ impl From<&NeonLocalInitPageserverConf> for PageServerConf {
             id: *id,
             listen_pg_addr: listen_pg_addr.clone(),
             listen_http_addr: listen_http_addr.clone(),
+            listen_https_addr: listen_https_addr.clone(),
             pg_auth_type: *pg_auth_type,
             http_auth_type: *http_auth_type,
             no_sync: *no_sync,
@@ -413,6 +431,41 @@ impl LocalEnv {
         }
     }
 
+    pub fn ssl_ca_cert_path(&self) -> Option<PathBuf> {
+        if self.generate_local_ssl_certs {
+            Some(self.base_data_dir.join("rootCA.crt"))
+        } else {
+            None
+        }
+    }
+
+    pub fn ssl_ca_key_path(&self) -> Option<PathBuf> {
+        if self.generate_local_ssl_certs {
+            Some(self.base_data_dir.join("rootCA.key"))
+        } else {
+            None
+        }
+    }
+
+    pub fn generate_ssl_ca_cert(&self) -> anyhow::Result<()> {
+        let cert_path = self.ssl_ca_cert_path().unwrap();
+        let key_path = self.ssl_ca_key_path().unwrap();
+        if !fs::exists(cert_path.as_path())? {
+            generate_ssl_ca_cert(cert_path.as_path(), key_path.as_path())?;
+        }
+        Ok(())
+    }
+
+    pub fn generate_ssl_cert(&self, cert_path: &Path, key_path: &Path) -> anyhow::Result<()> {
+        self.generate_ssl_ca_cert()?;
+        generate_ssl_cert(
+            cert_path,
+            key_path,
+            self.ssl_ca_cert_path().unwrap().as_path(),
+            self.ssl_ca_key_path().unwrap().as_path(),
+        )
+    }
+
     /// Inspect the base data directory and extract the instance id and instance directory path
     /// for all storage controller instances
     pub async fn storage_controller_instances(&self) -> std::io::Result<Vec<(u8, PathBuf)>> {
@@ -522,6 +575,7 @@ impl LocalEnv {
                 control_plane_api,
                 control_plane_compute_hook_api,
                 branch_name_mappings,
+                generate_local_ssl_certs,
             } = on_disk_config;
             LocalEnv {
                 base_data_dir: repopath.to_owned(),
@@ -536,6 +590,7 @@ impl LocalEnv {
                 control_plane_api: control_plane_api.unwrap(),
                 control_plane_compute_hook_api,
                 branch_name_mappings,
+                generate_local_ssl_certs,
             }
         };
 
@@ -571,6 +626,7 @@ impl LocalEnv {
                 struct PageserverConfigTomlSubset {
                     listen_pg_addr: String,
                     listen_http_addr: String,
+                    listen_https_addr: Option<String>,
                     pg_auth_type: AuthType,
                     http_auth_type: AuthType,
                     #[serde(default)]
@@ -595,6 +651,7 @@ impl LocalEnv {
                 let PageserverConfigTomlSubset {
                     listen_pg_addr,
                     listen_http_addr,
+                    listen_https_addr,
                     pg_auth_type,
                     http_auth_type,
                     no_sync,
@@ -612,6 +669,7 @@ impl LocalEnv {
                     },
                     listen_pg_addr,
                     listen_http_addr,
+                    listen_https_addr,
                     pg_auth_type,
                     http_auth_type,
                     no_sync,
@@ -639,6 +697,7 @@ impl LocalEnv {
                 control_plane_api: Some(self.control_plane_api.clone()),
                 control_plane_compute_hook_api: self.control_plane_compute_hook_api.clone(),
                 branch_name_mappings: self.branch_name_mappings.clone(),
+                generate_local_ssl_certs: self.generate_local_ssl_certs,
             },
         )
     }
@@ -721,6 +780,7 @@ impl LocalEnv {
             safekeepers,
             control_plane_api,
             control_plane_compute_hook_api,
+            generate_local_ssl_certs,
         } = conf;
 
         // Find postgres binaries.
@@ -769,8 +829,13 @@ impl LocalEnv {
             control_plane_api: control_plane_api.unwrap(),
             control_plane_compute_hook_api: control_plane_compute_hook_api.unwrap_or_default(),
             branch_name_mappings: Default::default(),
+            generate_local_ssl_certs,
         };
 
+        if generate_local_ssl_certs {
+            env.generate_ssl_ca_cert()?;
+        }
+
         // create endpoints dir
         fs::create_dir_all(env.endpoints_path())?;
 
@@ -854,3 +919,80 @@ fn generate_auth_keys(private_key_path: &Path, public_key_path: &Path) -> anyhow
     }
     Ok(())
 }
+
+fn generate_ssl_ca_cert(cert_path: &Path, key_path: &Path) -> anyhow::Result<()> {
+    // openssl req -x509 -newkey rsa:2048 -nodes -subj "/CN=Neon Local CA" -days 36500 \
+    // -out rootCA.crt -keyout rootCA.key
+    let keygen_output = Command::new("openssl")
+        .args([
+            "req", "-x509", "-newkey", "rsa:2048", "-nodes", "-days", "36500",
+        ])
+        .args(["-subj", "/CN=Neon Local CA"])
+        .args(["-out", cert_path.to_str().unwrap()])
+        .args(["-keyout", key_path.to_str().unwrap()])
+        .output()
+        .context("failed to generate CA certificate")?;
+    if !keygen_output.status.success() {
+        bail!(
+            "openssl failed: '{}'",
+            String::from_utf8_lossy(&keygen_output.stderr)
+        );
+    }
+    Ok(())
+}
+
+fn generate_ssl_cert(
+    cert_path: &Path,
+    key_path: &Path,
+    ca_cert_path: &Path,
+    ca_key_path: &Path,
+) -> anyhow::Result<()> {
+    // Generate Certificate Signing Request (CSR).
+    let mut csr_path = cert_path.to_path_buf();
+    csr_path.set_extension(".csr");
+
+    // openssl req -new -nodes -newkey rsa:2048 -keyout server.key -out server.csr \
+    // -subj "/CN=localhost" -addext "subjectAltName=DNS:localhost,IP:127.0.0.1"
+    let keygen_output = Command::new("openssl")
+        .args(["req", "-new", "-nodes"])
+        .args(["-newkey", "rsa:2048"])
+        .args(["-subj", "/CN=localhost"])
+        .args(["-addext", "subjectAltName=DNS:localhost,IP:127.0.0.1"])
+        .args(["-keyout", key_path.to_str().unwrap()])
+        .args(["-out", csr_path.to_str().unwrap()])
+        .output()
+        .context("failed to generate CSR")?;
+    if !keygen_output.status.success() {
+        bail!(
+            "openssl failed: '{}'",
+            String::from_utf8_lossy(&keygen_output.stderr)
+        );
+    }
+
+    // Sign CSR with CA key.
+    //
+    // openssl x509 -req -in server.csr -CA rootCA.crt -CAkey rootCA.key -CAcreateserial \
+    // -out server.crt -days 36500 -copy_extensions copyall
+    let keygen_output = Command::new("openssl")
+        .args(["x509", "-req"])
+        .args(["-in", csr_path.to_str().unwrap()])
+        .args(["-CA", ca_cert_path.to_str().unwrap()])
+        .args(["-CAkey", ca_key_path.to_str().unwrap()])
+        .arg("-CAcreateserial")
+        .args(["-out", cert_path.to_str().unwrap()])
+        .args(["-days", "36500"])
+        .args(["-copy_extensions", "copyall"])
+        .output()
+        .context("failed to sign CSR")?;
+    if !keygen_output.status.success() {
+        bail!(
+            "openssl failed: '{}'",
+            String::from_utf8_lossy(&keygen_output.stderr)
+        );
+    }
+
+    // Remove CSR file as it's not needed anymore.
+    fs::remove_file(csr_path)?;
+
+    Ok(())
+}

control_plane/src/pageserver.rs

@@ -21,6 +21,7 @@ use pageserver_api::shard::TenantShardId;
 use pageserver_client::mgmt_api;
 use postgres_backend::AuthType;
 use postgres_connection::{PgConnectionConfig, parse_host_port};
+use reqwest::Certificate;
 use utils::auth::{Claims, Scope};
 use utils::id::{NodeId, TenantId, TimelineId};
 use utils::lsn::Lsn;
@@ -49,12 +50,29 @@ impl PageServerNode {
         let (host, port) =
             parse_host_port(&conf.listen_pg_addr).expect("Unable to parse listen_pg_addr");
         let port = port.unwrap_or(5432);
+
+        let ssl_ca_cert = env.ssl_ca_cert_path().map(|ssl_ca_file| {
+            let buf = std::fs::read(ssl_ca_file).expect("SSL root CA file should exist");
+            Certificate::from_pem(&buf).expect("CA certificate should be valid")
+        });
+
+        let endpoint = if env.storage_controller.use_https_pageserver_api {
+            format!(
+                "https://{}",
+                conf.listen_https_addr.as_ref().expect(
+                    "listen https address should be specified if use_https_pageserver_api is on"
+                )
+            )
+        } else {
+            format!("http://{}", conf.listen_http_addr)
+        };
+
         Self {
             pg_connection_config: PgConnectionConfig::new_host_port(host, port),
             conf: conf.clone(),
             env: env.clone(),
             http_client: mgmt_api::Client::new(
-                format!("http://{}", conf.listen_http_addr),
+                endpoint,
                 {
                     match conf.http_auth_type {
                         AuthType::Trust => None,
@@ -65,7 +83,9 @@ impl PageServerNode {
                     }
                 }
                 .as_deref(),
-            ),
+                ssl_ca_cert,
+            )
+            .expect("Client constructs with no errors"),
         }
     }
 
@@ -220,6 +240,13 @@ impl PageServerNode {
             .context("write identity toml")?;
         drop(identity_toml);
 
+        if self.env.generate_local_ssl_certs {
+            self.env.generate_ssl_cert(
+                datadir.join("server.crt").as_path(),
+                datadir.join("server.key").as_path(),
+            )?;
+        }
+
         // TODO: invoke a TBD config-check command to validate that pageserver will start with the written config
 
         // Write metadata file, used by pageserver on startup to register itself with
@@ -230,6 +257,15 @@ impl PageServerNode {
             parse_host_port(&self.conf.listen_http_addr).expect("Unable to parse listen_http_addr");
         let http_port = http_port.unwrap_or(9898);
 
+        let https_port = match self.conf.listen_https_addr.as_ref() {
+            Some(https_addr) => {
+                let (_https_host, https_port) =
+                    parse_host_port(https_addr).expect("Unable to parse listen_https_addr");
+                Some(https_port.unwrap_or(9899))
+            }
+            None => None,
+        };
+
         // Intentionally hand-craft JSON: this acts as an implicit format compat test
         // in case the pageserver-side structure is edited, and reflects the real life
         // situation: the metadata is written by some other script.
@@ -240,6 +276,7 @@ impl PageServerNode {
                 postgres_port: self.pg_connection_config.port(),
                 http_host: "localhost".to_string(),
                 http_port,
+                https_port,
                 other: HashMap::from([(
                     "availability_zone_id".to_string(),
                     serde_json::json!(az_id),

control_plane/src/storage_controller.rs

@@ -12,13 +12,10 @@ use hyper0::Uri;
 use nix::unistd::Pid;
 use pageserver_api::controller_api::{
     NodeConfigureRequest, NodeDescribeResponse, NodeRegisterRequest, TenantCreateRequest,
-    TenantCreateResponse, TenantLocateResponse, TenantShardMigrateRequest,
-    TenantShardMigrateResponse,
+    TenantCreateResponse, TenantLocateResponse,
 };
-use pageserver_api::models::{
-    TenantShardSplitRequest, TenantShardSplitResponse, TimelineCreateRequest, TimelineInfo,
-};
-use pageserver_api::shard::{ShardStripeSize, TenantShardId};
+use pageserver_api::models::{TimelineCreateRequest, TimelineInfo};
+use pageserver_api::shard::TenantShardId;
 use pageserver_client::mgmt_api::ResponseErrorMessageExt;
 use postgres_backend::AuthType;
 use reqwest::Method;
@@ -537,8 +534,12 @@ impl StorageController {
             args.push("--start-as-candidate".to_string());
         }
 
-        if self.config.load_safekeepers {
-            args.push("--load-safekeepers".to_string());
+        if self.config.use_https_pageserver_api {
+            args.push("--use-https-pageserver-api".to_string());
+        }
+
+        if let Some(ssl_ca_file) = self.env.ssl_ca_cert_path() {
+            args.push(format!("--ssl-ca-file={}", ssl_ca_file.to_str().unwrap()));
         }
 
         if let Some(private_key) = &self.private_key {
@@ -583,6 +584,10 @@ impl StorageController {
             self.env.base_data_dir.display()
         ));
 
+        if self.config.timelines_onto_safekeepers {
+            args.push("--timelines-onto-safekeepers".to_string());
+        }
+
         background_process::start_process(
             COMMAND,
             &instance_dir,
@@ -829,41 +834,6 @@ impl StorageController {
         .await
     }
 
-    #[instrument(skip(self))]
-    pub async fn tenant_migrate(
-        &self,
-        tenant_shard_id: TenantShardId,
-        node_id: NodeId,
-    ) -> anyhow::Result<TenantShardMigrateResponse> {
-        self.dispatch(
-            Method::PUT,
-            format!("control/v1/tenant/{tenant_shard_id}/migrate"),
-            Some(TenantShardMigrateRequest {
-                node_id,
-                migration_config: None,
-            }),
-        )
-        .await
-    }
-
-    #[instrument(skip(self), fields(%tenant_id, %new_shard_count))]
-    pub async fn tenant_split(
-        &self,
-        tenant_id: TenantId,
-        new_shard_count: u8,
-        new_stripe_size: Option<ShardStripeSize>,
-    ) -> anyhow::Result<TenantShardSplitResponse> {
-        self.dispatch(
-            Method::PUT,
-            format!("control/v1/tenant/{tenant_id}/shard_split"),
-            Some(TenantShardSplitRequest {
-                new_shard_count,
-                new_stripe_size,
-            }),
-        )
-        .await
-    }
-
     #[instrument(skip_all, fields(node_id=%req.node_id))]
     pub async fn node_register(&self, req: NodeRegisterRequest) -> anyhow::Result<()> {
         self.dispatch::<_, ()>(Method::POST, "control/v1/node".to_string(), Some(req))

control_plane/storcon_cli/src/main.rs

@@ -1,16 +1,17 @@
 use std::collections::{HashMap, HashSet};
+use std::path::PathBuf;
 use std::str::FromStr;
 use std::time::Duration;
 
 use clap::{Parser, Subcommand};
 use futures::StreamExt;
 use pageserver_api::controller_api::{
-    AvailabilityZone, NodeAvailabilityWrapper, NodeConfigureRequest, NodeDescribeResponse,
-    NodeRegisterRequest, NodeSchedulingPolicy, NodeShardResponse, PlacementPolicy,
-    SafekeeperDescribeResponse, SafekeeperSchedulingPolicyRequest, ShardSchedulingPolicy,
-    ShardsPreferredAzsRequest, ShardsPreferredAzsResponse, SkSchedulingPolicy, TenantCreateRequest,
-    TenantDescribeResponse, TenantPolicyRequest, TenantShardMigrateRequest,
-    TenantShardMigrateResponse,
+    AvailabilityZone, MigrationConfig, NodeAvailabilityWrapper, NodeConfigureRequest,
+    NodeDescribeResponse, NodeRegisterRequest, NodeSchedulingPolicy, NodeShardResponse,
+    PlacementPolicy, SafekeeperDescribeResponse, SafekeeperSchedulingPolicyRequest,
+    ShardSchedulingPolicy, ShardsPreferredAzsRequest, ShardsPreferredAzsResponse,
+    SkSchedulingPolicy, TenantCreateRequest, TenantDescribeResponse, TenantPolicyRequest,
+    TenantShardMigrateRequest, TenantShardMigrateResponse,
 };
 use pageserver_api::models::{
     EvictionPolicy, EvictionPolicyLayerAccessThreshold, LocationConfigSecondary, ShardParameters,
@@ -112,6 +113,15 @@ enum Command {
         tenant_shard_id: TenantShardId,
         #[arg(long)]
         node: NodeId,
+        #[arg(long, default_value_t = true, action = clap::ArgAction::Set)]
+        prewarm: bool,
+        #[arg(long, default_value_t = false, action = clap::ArgAction::Set)]
+        override_scheduler: bool,
+    },
+    /// Watch the location of a tenant shard evolve, e.g. while expecting it to migrate
+    TenantShardWatch {
+        #[arg(long)]
+        tenant_shard_id: TenantShardId,
     },
     /// Migrate the secondary location for a tenant shard to a specific pageserver.
     TenantShardMigrateSecondary {
@@ -269,6 +279,10 @@ struct Cli {
     /// a token with both scopes to use with this tool.
     jwt: Option<String>,
 
+    #[arg(long)]
+    /// Trusted root CA certificate to use in https APIs.
+    ssl_ca_file: Option<PathBuf>,
+
     #[command(subcommand)]
     command: Command,
 }
@@ -379,9 +393,17 @@ async fn main() -> anyhow::Result<()> {
 
     let storcon_client = Client::new(cli.api.clone(), cli.jwt.clone());
 
+    let ssl_ca_cert = match &cli.ssl_ca_file {
+        Some(ssl_ca_file) => {
+            let buf = tokio::fs::read(ssl_ca_file).await?;
+            Some(reqwest::Certificate::from_pem(&buf)?)
+        }
+        None => None,
+    };
+
     let mut trimmed = cli.api.to_string();
     trimmed.pop();
-    let vps_client = mgmt_api::Client::new(trimmed, cli.jwt.as_deref());
+    let vps_client = mgmt_api::Client::new(trimmed, cli.jwt.as_deref(), ssl_ca_cert)?;
 
     match cli.command {
         Command::NodeRegister {
@@ -619,19 +641,43 @@ async fn main() -> anyhow::Result<()> {
         Command::TenantShardMigrate {
             tenant_shard_id,
             node,
+            prewarm,
+            override_scheduler,
         } => {
-            let req = TenantShardMigrateRequest {
-                node_id: node,
-                migration_config: None,
+            let migration_config = MigrationConfig {
+                prewarm,
+                override_scheduler,
+                ..Default::default()
             };
 
-            storcon_client
+            let req = TenantShardMigrateRequest {
+                node_id: node,
+                origin_node_id: None,
+                migration_config,
+            };
+
+            match storcon_client
                 .dispatch::<TenantShardMigrateRequest, TenantShardMigrateResponse>(
                     Method::PUT,
                     format!("control/v1/tenant/{tenant_shard_id}/migrate"),
                     Some(req),
                 )
-                .await?;
+                .await
+            {
+                Err(mgmt_api::Error::ApiError(StatusCode::PRECONDITION_FAILED, msg)) => {
+                    anyhow::bail!(
+                        "Migration to {node} rejected, may require `--force` ({}) ",
+                        msg
+                    );
+                }
+                Err(e) => return Err(e.into()),
+                Ok(_) => {}
+            }
+
+            watch_tenant_shard(storcon_client, tenant_shard_id, Some(node)).await?;
+        }
+        Command::TenantShardWatch { tenant_shard_id } => {
+            watch_tenant_shard(storcon_client, tenant_shard_id, None).await?;
         }
         Command::TenantShardMigrateSecondary {
             tenant_shard_id,
@@ -639,7 +685,8 @@ async fn main() -> anyhow::Result<()> {
         } => {
             let req = TenantShardMigrateRequest {
                 node_id: node,
-                migration_config: None,
+                origin_node_id: None,
+                migration_config: MigrationConfig::default(),
             };
 
             storcon_client
@@ -1105,7 +1152,8 @@ async fn main() -> anyhow::Result<()> {
                                 format!("control/v1/tenant/{}/migrate", mv.tenant_shard_id),
                                 Some(TenantShardMigrateRequest {
                                     node_id: mv.to,
-                                    migration_config: None,
+                                    origin_node_id: Some(mv.from),
+                                    migration_config: MigrationConfig::default(),
                                 }),
                             )
                             .await
@@ -1284,3 +1332,68 @@ async fn main() -> anyhow::Result<()> {
 
     Ok(())
 }
+
+static WATCH_INTERVAL: Duration = Duration::from_secs(5);
+
+async fn watch_tenant_shard(
+    storcon_client: Client,
+    tenant_shard_id: TenantShardId,
+    until_migrated_to: Option<NodeId>,
+) -> anyhow::Result<()> {
+    if let Some(until_migrated_to) = until_migrated_to {
+        println!(
+            "Waiting for tenant shard {} to be migrated to node {}",
+            tenant_shard_id, until_migrated_to
+        );
+    }
+
+    loop {
+        let desc = storcon_client
+            .dispatch::<(), TenantDescribeResponse>(
+                Method::GET,
+                format!("control/v1/tenant/{}", tenant_shard_id.tenant_id),
+                None,
+            )
+            .await?;
+
+        // Output the current state of the tenant shard
+        let shard = desc
+            .shards
+            .iter()
+            .find(|s| s.tenant_shard_id == tenant_shard_id)
+            .ok_or(anyhow::anyhow!("Tenant shard not found"))?;
+        let summary = format!(
+            "attached: {} secondary: {} {}",
+            shard
+                .node_attached
+                .map(|n| format!("{}", n))
+                .unwrap_or("none".to_string()),
+            shard
+                .node_secondary
+                .iter()
+                .map(|n| n.to_string())
+                .collect::<Vec<_>>()
+                .join(","),
+            if shard.is_reconciling {
+                "(reconciler active)"
+            } else {
+                "(reconciler idle)"
+            }
+        );
+        println!("{}", summary);
+
+        // Maybe drop out if we finished migration
+        if let Some(until_migrated_to) = until_migrated_to {
+            if shard.node_attached == Some(until_migrated_to) && !shard.is_reconciling {
+                println!(
+                    "Tenant shard {} is now on node {}",
+                    tenant_shard_id, until_migrated_to
+                );
+                break;
+            }
+        }
+
+        tokio::time::sleep(WATCH_INTERVAL).await;
+    }
+    Ok(())
+}

deny.toml

@@ -27,6 +27,14 @@ yanked = "warn"
 id = "RUSTSEC-2023-0071"
 reason = "the marvin attack only affects private key decryption, not public key signature verification"
 
+[[advisories.ignore]]
+id = "RUSTSEC-2024-0436"
+reason = "The paste crate is a build-only dependency with no runtime components. It is unlikely to have any security impact."
+
+[[advisories.ignore]]
+id = "RUSTSEC-2025-0014"
+reason = "The humantime is widely used and is not easy to replace right now. It is unmaintained, but it has no known vulnerabilities to care about. #11179"
+
 # This section is considered when running `cargo deny check licenses`
 # More documentation for the licenses section can be found here:
 # https://embarkstudios.github.io/cargo-deny/checks/licenses/cfg.html

docs/rfcs/041-rel-sparse-keyspace.md

@@ -0,0 +1,201 @@
+# Sparse Keyspace for Relation Directories
+
+## Summary
+
+This is an RFC describing a new storage strategy for storing relation directories.
+
+## Motivation
+
+Postgres maintains a directory structure for databases and relations. In Neon, we store these information
+by serializing the directory data in a single key (see `pgdatadir_mapping.rs`).
+
+```rust
+// DbDir:
+// 00 00000000 00000000 00000000 00   00000000
+
+// RelDir:
+// 00 SPCNODE  DBNODE   00000000 00   00000001 (Postgres never uses relfilenode 0)
+```
+
+We have a dedicated structure on the ingestion path to serialize the relation directory into this single key.
+
+```rust
+#[derive(Debug, Serialize, Deserialize, Default)]
+pub(crate) struct RelDirectory {
+    // Set of relations that exist. (relfilenode, forknum)
+    //
+    // TODO: Store it as a btree or radix tree or something else that spans multiple
+    // key-value pairs, if you have a lot of relations
+    pub(crate) rels: HashSet<(Oid, u8)>,
+}
+```
+
+The current codebase has the following three access patterns for the relation directory.
+
+1. Check if a relation exists.
+2. List all relations.
+3. Create/drop a relation.
+
+For (1), we currently have to get the reldir key, deserialize it, and check whether the relation exists in the
+hash set. For (2), we get the reldir key and the hash set. For (3), we need first to get
+and deserialize the key, add the new relation record to the hash set, and then serialize it and write it back.
+
+If we have 100k relations in a database, we would have a 100k-large hash set. Then, every
+relation created and dropped would have deserialized and serialized this 100k-large hash set. This makes the
+relation create/drop process to be quadratic. When we check if a relation exists in the ingestion path,
+we would have to deserialize this super big 100k-large key before checking if a single relation exists.
+
+In this RFC, we will propose a new way to store the reldir data in the sparse keyspace and propose how
+to seamlessly migrate users to use the new keyspace.
+
+The PoC patch is implemented in [PR10316](https://github.com/neondatabase/neon/pull/10316).
+
+## Key Mapping
+
+We will use the recently introduced sparse keyspace to store actual data. Sparse keyspace was proposed in
+[038-aux-file-v2.md](038-aux-file-v2.md). The original reldir has one single value of `HashSet<(Oid, u8)>`
+for each of the databases (identified as `spcnode, dbnode`). We encode the `Oid` (`relnode, forknum`),
+into the key.
+
+```plain
+(REL_DIR_KEY_PREFIX, spcnode, dbnode, relnode, forknum, 1) -> deleted
+(REL_DIR_KEY_PREFIX, spcnode, dbnode, relnode, forknum, 1) -> exists
+```
+
+Assume all reldir data are stored in this new keyspace; the 3 reldir operations we mentioned before can be
+implemented as follows.
+
+1. Check if a relation exists: check if the key maps to "exists".
+2. List all relations: scan the sprase keyspace over the `rel_dir_key_prefix`. Extract relnode and forknum from the key.
+3. Create/drop a relation: write "exists" or "deleted" to the corresponding key of the relation. The delete tombstone will
+   be removed during image layer generation upon compaction.
+
+Note that "exists" and "deleted" will be encoded as a single byte as two variants of an enum.
+The mapping is implemented as `rel_tag_sparse_key` in the PoC patch.
+
+## Changes to Sparse Keyspace
+
+Previously, we only used sparse keyspaces for the aux files, which did not carry over when branching. The reldir
+information needs to be preserved from the parent branch to the child branch. Therefore, the read path needs
+to be updated accordingly to accommodate such "inherited sparse keys". This is done in
+[PR#10313](https://github.com/neondatabase/neon/pull/10313).
+
+## Coexistence of the Old and New Keyspaces
+
+Migrating to the new keyspace will be done gradually: when we flip a config item to enable the new reldir keyspace, the
+ingestion path will start to write to the new keyspace and the old reldir data will be kept in the old one. The read
+path needs to combine the data from both keyspaces.
+
+Theoretically, we could do a rewrite at the startup time that scans all relation directories and copies that data into the
+new keyspace. However, this could take a long time, especially if we have thousands of tenants doing the migration
+process simultaneously after the pageserver restarts. Therefore, we propose the coexistence strategy so that the
+migration can happen seamlessly and imposes no potential downtime for the user.
+
+With the coexistence assumption, the 3 reldir operations will be implemented as follows:
+
+1. Check if a relation exists
+   - Check the new keyspace if the key maps to any value. If it maps to "exists" or "deleted", directly
+    return it to the user.
+   - Otherwise, deserialize the old reldir key and get the result.
+2. List all relations: scan the sparse keyspace over the `rel_dir_key_prefix` and deserialize the old reldir key.
+   Combine them to obtain the final result.
+3. Create/drop a relation: write "exists" or "deleted" to the corresponding key of the relation into the new keyspace.
+   - We assume no overwrite of relations will happen (i.e., the user won't create a relation at the same Oid). This will be implemented as a runtime check.
+   - For relation creation, we add `sparse_reldir_tableX -> exists` to the keyspace.
+   - For relation drop, we first check if the relation is recorded in the old keyspace. If yes, we deserialize the old reldir key,
+    remove the relation, and then write it back. Otherwise, we put `sparse_reldir_tableX -> deleted` to the keyspace.
+   - The delete tombstone will be removed during image layer generation upon compaction.
+
+This process ensures that the transition will not introduce any downtime and all new updates are written to the new keyspace. The total
+amount of data in the storage would be `O(relations_modifications)` and we can guarantee `O(current_relations)` after compaction.
+There could be some relations that exist in the old reldir key for a long time. Refer to the "Full Migration" section on how to deal
+with them. Plus, for relation modifications, it will have `O(old_relations)` complexity until we do the full migration, which gives
+us `O(1)` complexity after fully opt-in the sparse keyspace.
+
+The process also implies that a relation will only exists either in the old reldir key or in the new sparse keyspace. It is not possible
+to have a table to be recorded in the old reldir key while later having a delete tombstone for it in the sparse keyspace at any LSN.
+
+We will introduce a config item and an index_part record to record the current status of the migration process.
+
+- Config item `enable_reldir_v2`: controls whether the ingestion path writes the reldir info into the new keyspace.
+- `index_part.json` field `reldir_v2_status`: whether the timeline has written any key into the new reldir keyspace.
+
+If `enable_reldir_v2` is set to `true` and the timeline ingests the first key into the new reldir keyspace, it will update
+`index_part.json` to set `reldir_v2_status` to `Status::Migrating`. Even if `enable_reldir_v2` gets flipped back to
+`false` (i.e., when the pageserver restarts and such config isn't persisted), the read/write path will still
+read/write to the new keyspace to avoid data inconsistency. This also indicates that the migration is one-way only:
+once v2 is enabled, the user cannot go back to v1.
+
+## Next Steps
+
+### Full Migration
+
+This won't be implemented in the project's first phase but might be implemented in the future. Having both v1 and
+v2 existing in the system would force us to keep the code to deserialize the old reldir key forever. To entirely deprecate this
+code path, we must ensure the timeline has no old reldir data.
+
+We can trigger a special image layer generation process at the gc-horizon. The generated image layers will cover several keyspaces:
+the old reldir key in each of the databases, and the new reldir sparse keyspace. It will remove the old reldir key while
+copying them into the corresponding keys in the sparse keyspace in the resulting image. This special process happens in
+the background during compaction. For example, assume this special process is triggered at LSN 0/180. The `create_image_layers`
+process discovers the following keys at this LSN.
+
+```plain
+db1/reldir_key -> (table 1, table 2, table 3)
+...db1 rel keys
+db2/reldir_key -> (table 4, table 5, table 6)
+...db2 rel keys
+sparse_reldir_db2_table7 -> exists
+sparse_reldir_db1_table8 -> deleted
+```
+
+It will generate the following keys:
+
+```plain
+db1/reldir_key -> () # we have to keep the key because it is part of `collect_keyspace`.
+...db1 rel keys
+db2/reldir_key -> ()
+...db2 rel keys
+
+-- start image layer for the sparse keyspace at sparse_reldir_prefix at LSN 0/180
+sparse_reldir_db1_table1 -> exists
+sparse_reldir_db1_table2 -> exists
+sparse_reldir_db1_table3 -> exists
+sparse_reldir_db2_table4 -> exists
+sparse_reldir_db2_table5 -> exists
+sparse_reldir_db2_table6 -> exists
+sparse_reldir_db2_table7 -> exists
+-- end image layer for the sparse keyspace at sparse_reldir_prefix+1
+
+# The `sparse_reldir_db1_table8` key gets dropped as part of the image layer generation code for the sparse keyspace.
+# Note that the read path will stop reading if a key is not found in the image layer covering the key range so there
+# are no correctness issue.
+```
+
+We must verify that no pending modifications to the old reldir exists in the delta/image layers above the gc-horizon before
+we start this process (We can do a vectored read to get the full key history of the old reldir key and ensure there are no more images
+above the gc-horizon). Otherwise, it will violate the property that "a relation will only exists either in the old reldir key or
+in the new sparse keyspace". After we run this migration process, we can mark `reldir_v2_status` in the `index_part.json` to
+`Status::Migrated`, and the read path won't need to read from the old reldir anymore. Once the status is set to `Migrated`, we
+don't need to add the key into `collect_keyspace` and therefore all of them will be removed from all future image layers.
+
+The migration process can be proactively triggered across all attached/detached tenants to help us fully remove the old reldir code.
+
+### Consolidate Relation Size Keys
+
+We have relsize at the end of all relation nodes.
+
+```plain
+// RelSize:
+// 00 SPCNODE  DBNODE   RELNODE  FORK FFFFFFFF
+```
+
+This means that computing logical size requires us to do several single-key gets across the keyspace,
+potentially requiring downloading many layer files. We could consolidate them into a single
+keyspace, improving logical size calculation performance.
+
+### Migrate DBDir Keys
+
+We assume the number of databases created by the users will be small, and therefore, the current way
+of storing the database directory would be acceptable. In the future, we could also migrate DBDir keys into
+the sparse keyspace to support large amount of databases.

libs/compute_api/src/spec.rs

@@ -272,6 +272,18 @@ pub enum ComputeMode {
     Replica,
 }
 
+impl ComputeMode {
+    /// Convert the compute mode to a string that can be used to identify the type of compute,
+    /// which means that if it's a static compute, the LSN will not be included.
+    pub fn to_type_str(&self) -> &'static str {
+        match self {
+            ComputeMode::Primary => "primary",
+            ComputeMode::Static(_) => "static",
+            ComputeMode::Replica => "replica",
+        }
+    }
+}
+
 /// Log level for audit logging
 /// Disabled, log, hipaa
 /// Default is Disabled

libs/http-utils/Cargo.toml

@@ -8,6 +8,7 @@ license.workspace = true
 anyhow.workspace = true
 bytes.workspace = true
 fail.workspace = true
+futures.workspace = true
 hyper0.workspace = true
 itertools.workspace = true
 jemalloc_pprof.workspace = true
@@ -21,6 +22,7 @@ serde_path_to_error.workspace = true
 thiserror.workspace = true
 tracing.workspace = true
 tokio.workspace = true
+tokio-rustls.workspace = true
 tokio-util.workspace = true
 url.workspace = true
 uuid.workspace = true

libs/http-utils/src/endpoint.rs

@@ -399,12 +399,10 @@ pub async fn profile_cpu_handler(req: Request<Body>) -> Result<Response<Body>, A
     // Return the report in the requested format.
     match format {
         Format::Pprof => {
-            let mut body = Vec::new();
-            report
+            let body = report
                 .pprof()
                 .map_err(|err| ApiError::InternalServerError(err.into()))?
-                .write_to_vec(&mut body)
-                .map_err(|err| ApiError::InternalServerError(err.into()))?;
+                .encode_to_vec();
 
             Response::builder()
                 .status(200)

libs/http-utils/src/lib.rs

@@ -3,9 +3,10 @@ pub mod error;
 pub mod failpoints;
 pub mod json;
 pub mod request;
+pub mod server;
 
 extern crate hyper0 as hyper;
 
 /// Current fast way to apply simple http routing in various Neon binaries.
 /// Re-exported for sake of uniform approach, that could be later replaced with better alternatives, if needed.
-pub use routerify::{RouterBuilder, RouterService, ext::RequestExt};
+pub use routerify::{RequestServiceBuilder, RouterBuilder, RouterService, ext::RequestExt};

libs/http-utils/src/server.rs

@@ -0,0 +1,155 @@
+use std::{error::Error, sync::Arc};
+
+use futures::StreamExt;
+use futures::stream::FuturesUnordered;
+use hyper0::Body;
+use hyper0::server::conn::Http;
+use routerify::{RequestService, RequestServiceBuilder};
+use tokio::io::{AsyncRead, AsyncWrite};
+use tokio_rustls::TlsAcceptor;
+use tokio_util::sync::CancellationToken;
+use tracing::{error, info};
+
+use crate::error::ApiError;
+
+/// A simple HTTP server over hyper library.
+/// You may want to use it instead of [`hyper0::server::Server`] because:
+/// 1. hyper0's Server was removed from hyper v1.
+///    It's recommended to replace hyepr0's Server with a manual loop, which is done here.
+/// 2. hyper0's Server doesn't support TLS out of the box, and there is no way
+///    to support it efficiently with the Accept trait that hyper0's Server uses.
+///    That's one of the reasons why it was removed from v1.
+///    <https://github.com/hyperium/hyper/blob/115339d3df50f20c8717680aa35f48858e9a6205/docs/ROADMAP.md#higher-level-client-and-server-problems>
+pub struct Server {
+    request_service: Arc<RequestServiceBuilder<Body, ApiError>>,
+    listener: tokio::net::TcpListener,
+    tls_acceptor: Option<TlsAcceptor>,
+}
+
+impl Server {
+    pub fn new(
+        request_service: Arc<RequestServiceBuilder<Body, ApiError>>,
+        listener: std::net::TcpListener,
+        tls_acceptor: Option<TlsAcceptor>,
+    ) -> anyhow::Result<Self> {
+        // Note: caller of from_std is responsible for setting nonblocking mode.
+        listener.set_nonblocking(true)?;
+        let listener = tokio::net::TcpListener::from_std(listener)?;
+
+        Ok(Self {
+            request_service,
+            listener,
+            tls_acceptor,
+        })
+    }
+
+    pub async fn serve(self, cancel: CancellationToken) -> anyhow::Result<()> {
+        fn suppress_io_error(err: &std::io::Error) -> bool {
+            use std::io::ErrorKind::*;
+            matches!(err.kind(), ConnectionReset | ConnectionAborted | BrokenPipe)
+        }
+        fn suppress_hyper_error(err: &hyper0::Error) -> bool {
+            if err.is_incomplete_message() || err.is_closed() || err.is_timeout() {
+                return true;
+            }
+            if let Some(inner) = err.source() {
+                if let Some(io) = inner.downcast_ref::<std::io::Error>() {
+                    return suppress_io_error(io);
+                }
+            }
+            false
+        }
+
+        let mut connections = FuturesUnordered::new();
+        loop {
+            tokio::select! {
+                stream = self.listener.accept() => {
+                    let (tcp_stream, remote_addr) = match stream {
+                        Ok(stream) => stream,
+                        Err(err) => {
+                            if !suppress_io_error(&err) {
+                                info!("Failed to accept TCP connection: {err:#}");
+                            }
+                            continue;
+                        }
+                    };
+
+                    let service = self.request_service.build(remote_addr);
+                    let tls_acceptor = self.tls_acceptor.clone();
+                    let cancel = cancel.clone();
+
+                    connections.push(tokio::spawn(
+                        async move {
+                            match tls_acceptor {
+                                Some(tls_acceptor) => {
+                                    // Handle HTTPS connection.
+                                    let tls_stream = tokio::select! {
+                                        tls_stream = tls_acceptor.accept(tcp_stream) => tls_stream,
+                                        _ = cancel.cancelled() => return,
+                                    };
+                                    let tls_stream = match tls_stream {
+                                        Ok(tls_stream) => tls_stream,
+                                        Err(err) => {
+                                            if !suppress_io_error(&err) {
+                                                info!("Failed to accept TLS connection: {err:#}");
+                                            }
+                                            return;
+                                        }
+                                    };
+                                    if let Err(err) = Self::serve_connection(tls_stream, service, cancel).await {
+                                        if !suppress_hyper_error(&err) {
+                                            info!("Failed to serve HTTPS connection: {err:#}");
+                                        }
+                                    }
+                                }
+                                None => {
+                                    // Handle HTTP connection.
+                                    if let Err(err) = Self::serve_connection(tcp_stream, service, cancel).await {
+                                        if !suppress_hyper_error(&err) {
+                                            info!("Failed to serve HTTP connection: {err:#}");
+                                        }
+                                    }
+                                }
+                            };
+                        }));
+                 }
+                Some(conn) = connections.next() => {
+                    if let Err(err) = conn {
+                        error!("Connection panicked: {err:#}");
+                    }
+                }
+                _ = cancel.cancelled() => {
+                    // Wait for graceful shutdown of all connections.
+                    while let Some(conn) = connections.next().await {
+                        if let Err(err) = conn {
+                            error!("Connection panicked: {err:#}");
+                        }
+                    }
+                    break;
+                }
+            }
+        }
+        Ok(())
+    }
+
+    /// Serves HTTP connection with graceful shutdown.
+    async fn serve_connection<I>(
+        io: I,
+        service: RequestService<Body, ApiError>,
+        cancel: CancellationToken,
+    ) -> Result<(), hyper0::Error>
+    where
+        I: AsyncRead + AsyncWrite + Unpin + Send + 'static,
+    {
+        let mut conn = Http::new().serve_connection(io, service).with_upgrades();
+
+        tokio::select! {
+            res = &mut conn => res,
+            _ = cancel.cancelled() => {
+                Pin::new(&mut conn).graceful_shutdown();
+                // Note: connection should still be awaited for graceful shutdown to complete.
+                conn.await
+            }
+        }
+    }
+}

libs/pageserver_api/src/config.rs

@@ -35,6 +35,7 @@ pub struct NodeMetadata {
     pub postgres_port: u16,
     pub http_host: String,
     pub http_port: u16,
+    pub https_port: Option<u16>,
 
     // Deployment tools may write fields to the metadata file beyond what we
     // use in this type: this type intentionally only names fields that require.
@@ -57,6 +58,9 @@ pub struct ConfigToml {
     // types mapped 1:1 into the runtime PageServerConfig type
     pub listen_pg_addr: String,
     pub listen_http_addr: String,
+    pub listen_https_addr: Option<String>,
+    pub ssl_key_file: Utf8PathBuf,
+    pub ssl_cert_file: Utf8PathBuf,
     pub availability_zone: Option<String>,
     #[serde(with = "humantime_serde")]
     pub wait_lsn_timeout: Duration,
@@ -123,6 +127,10 @@ pub struct ConfigToml {
     pub enable_read_path_debugging: Option<bool>,
     #[serde(skip_serializing_if = "Option::is_none")]
     pub validate_wal_contiguity: Option<bool>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub load_previous_heatmap: Option<bool>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub generate_unarchival_heatmap: Option<bool>,
 }
 
 #[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
@@ -417,6 +425,9 @@ pub mod defaults {
 
     pub const DEFAULT_WAL_RECEIVER_PROTOCOL: utils::postgres_client::PostgresClientProtocol =
         utils::postgres_client::PostgresClientProtocol::Vanilla;
+
+    pub const DEFAULT_SSL_KEY_FILE: &str = "server.key";
+    pub const DEFAULT_SSL_CERT_FILE: &str = "server.crt";
 }
 
 impl Default for ConfigToml {
@@ -426,6 +437,9 @@ impl Default for ConfigToml {
         Self {
             listen_pg_addr: (DEFAULT_PG_LISTEN_ADDR.to_string()),
             listen_http_addr: (DEFAULT_HTTP_LISTEN_ADDR.to_string()),
+            listen_https_addr: (None),
+            ssl_key_file: Utf8PathBuf::from(DEFAULT_SSL_KEY_FILE),
+            ssl_cert_file: Utf8PathBuf::from(DEFAULT_SSL_CERT_FILE),
             availability_zone: (None),
             wait_lsn_timeout: (humantime::parse_duration(DEFAULT_WAIT_LSN_TIMEOUT)
                 .expect("cannot parse default wait lsn timeout")),
@@ -523,6 +537,8 @@ impl Default for ConfigToml {
                 None
             },
             validate_wal_contiguity: None,
+            load_previous_heatmap: None,
+            generate_unarchival_heatmap: None,
         }
     }
 }

libs/pageserver_api/src/config/tests.rs

@@ -16,6 +16,30 @@ fn test_node_metadata_v1_backward_compatibilty() {
             postgres_port: 23,
             http_host: "localhost".to_string(),
             http_port: 42,
+            https_port: None,
+            other: HashMap::new(),
+        }
+    )
+}
+
+#[test]
+fn test_node_metadata_v2_backward_compatibilty() {
+    let v2 = serde_json::to_vec(&serde_json::json!({
+        "host": "localhost",
+        "port": 23,
+        "http_host": "localhost",
+        "http_port": 42,
+        "https_port": 123,
+    }));
+
+    assert_eq!(
+        serde_json::from_slice::<NodeMetadata>(&v2.unwrap()).unwrap(),
+        NodeMetadata {
+            postgres_host: "localhost".to_string(),
+            postgres_port: 23,
+            http_host: "localhost".to_string(),
+            http_port: 42,
+            https_port: Some(123),
             other: HashMap::new(),
         }
     )

libs/pageserver_api/src/controller_api.rs

@@ -182,20 +182,66 @@ pub struct TenantDescribeResponseShard {
 #[derive(Serialize, Deserialize, Debug)]
 pub struct TenantShardMigrateRequest {
     pub node_id: NodeId,
+
+    /// Optionally, callers may specify the node they are migrating _from_, and the server will
+    /// reject the request if the shard is no longer attached there: this enables writing safer
+    /// clients that don't risk fighting with some other movement of the shard.
     #[serde(default)]
-    pub migration_config: Option<MigrationConfig>,
+    pub origin_node_id: Option<NodeId>,
+
+    #[serde(default)]
+    pub migration_config: MigrationConfig,
 }
 
-#[derive(Serialize, Deserialize, Debug)]
+#[derive(Serialize, Deserialize, Debug, PartialEq, Eq)]
 pub struct MigrationConfig {
+    /// If true, the migration will be executed even if it is to a location with a sub-optimal scheduling
+    /// score: this is usually not what you want, and if you use this then you'll also need to set the
+    /// tenant's scheduling policy to Essential or Pause to avoid the optimiser reverting your migration.
+    ///
+    /// Default: false
+    #[serde(default)]
+    pub override_scheduler: bool,
+
+    /// If true, the migration will be done gracefully by creating a secondary location first and
+    /// waiting for it to warm up before cutting over.  If false, if there is no existing secondary
+    /// location at the destination, the tenant will be migrated immediately.  If the tenant's data
+    /// can't be downloaded within [`Self::secondary_warmup_timeout`], then the migration will go
+    /// ahead but run with a cold cache that can severely reduce performance until it warms up.
+    ///
+    /// When doing a graceful migration, the migration API returns as soon as it is started.
+    ///
+    /// Default: true
+    #[serde(default = "default_prewarm")]
+    pub prewarm: bool,
+
+    /// For non-prewarm migrations which will immediately enter a cutover to the new node: how long to wait
+    /// overall for secondary warmup before cutting over
     #[serde(default)]
     #[serde(with = "humantime_serde")]
     pub secondary_warmup_timeout: Option<Duration>,
+    /// For non-prewarm migrations which will immediately enter a cutover to the new node: how long to wait
+    /// within each secondary download poll call to pageserver.
     #[serde(default)]
     #[serde(with = "humantime_serde")]
     pub secondary_download_request_timeout: Option<Duration>,
 }
 
+fn default_prewarm() -> bool {
+    true
+}
+
+impl Default for MigrationConfig {
+    fn default() -> Self {
+        Self {
+            override_scheduler: false,
+            prewarm: default_prewarm(),
+            secondary_warmup_timeout: None,
+            secondary_download_request_timeout: None,
+        }
+    }
+}
+
 #[derive(Serialize, Clone, Debug)]
 #[serde(into = "NodeAvailabilityWrapper")]
 pub enum NodeAvailability {
@@ -443,6 +489,7 @@ pub struct SafekeeperDescribeResponse {
     pub host: String,
     pub port: i32,
     pub http_port: i32,
+    pub https_port: Option<i32>,
     pub availability_zone_id: String,
     pub scheduling_policy: SkSchedulingPolicy,
 }
@@ -487,4 +534,43 @@ mod test {
             err
         );
     }
+
+    /// Check that a minimal migrate request with no config results in the expected default settings
+    #[test]
+    fn test_migrate_request_decode_defaults() {
+        let json = r#"{
+            "node_id": 123
+        }"#;
+
+        let request: TenantShardMigrateRequest = serde_json::from_str(json).unwrap();
+        assert_eq!(request.node_id, NodeId(123));
+        assert_eq!(request.origin_node_id, None);
+        assert!(!request.migration_config.override_scheduler);
+        assert!(request.migration_config.prewarm);
+        assert_eq!(request.migration_config.secondary_warmup_timeout, None);
+        assert_eq!(
+            request.migration_config.secondary_download_request_timeout,
+            None
+        );
+    }
+
+    /// Check that a partially specified migration config results in the expected default settings
+    #[test]
+    fn test_migration_config_decode_defaults() {
+        // Specify just one field of the config
+        let json = r#"{
+        }"#;
+
+        let config: MigrationConfig = serde_json::from_str(json).unwrap();
+
+        // Check each field's expected default value
+        assert!(!config.override_scheduler);
+        assert!(config.prewarm);
+        assert_eq!(config.secondary_warmup_timeout, None);
+        assert_eq!(config.secondary_download_request_timeout, None);
+        assert_eq!(config.secondary_warmup_timeout, None);
+
+        // Consistency check that the Default impl agrees with our serde defaults
+        assert_eq!(MigrationConfig::default(), config);
+    }
 }

libs/pageserver_api/src/models.rs

@@ -274,6 +274,31 @@ pub struct TimelineCreateRequest {
     pub mode: TimelineCreateRequestMode,
 }
 
+/// Storage controller specific extensions to [`TimelineInfo`].
+#[derive(Serialize, Deserialize, Clone)]
+pub struct TimelineCreateResponseStorcon {
+    #[serde(flatten)]
+    pub timeline_info: TimelineInfo,
+
+    pub safekeepers: Option<SafekeepersInfo>,
+}
+
+/// Safekeepers as returned in timeline creation request to storcon or pushed to
+/// cplane in the post migration hook.
+#[derive(Serialize, Deserialize, Clone)]
+pub struct SafekeepersInfo {
+    pub tenant_id: TenantId,
+    pub timeline_id: TimelineId,
+    pub generation: u32,
+    pub safekeepers: Vec<SafekeeperInfo>,
+}
+
+#[derive(Serialize, Deserialize, Clone)]
+pub struct SafekeeperInfo {
+    pub id: NodeId,
+    pub hostname: String,
+}
+
 #[derive(Serialize, Deserialize, Clone)]
 #[serde(untagged)]
 pub enum TimelineCreateRequestMode {
@@ -1146,6 +1171,15 @@ pub struct TimelineArchivalConfigRequest {
     pub state: TimelineArchivalState,
 }
 
+#[derive(Serialize, Deserialize, PartialEq, Eq, Clone)]
+pub struct TimelinePatchIndexPartRequest {
+    pub rel_size_migration: Option<RelSizeMigration>,
+    pub gc_compaction_last_completed_lsn: Option<Lsn>,
+    pub applied_gc_cutoff_lsn: Option<Lsn>,
+    #[serde(default)]
+    pub force_index_update: bool,
+}
+
 #[derive(Debug, Serialize, Deserialize, Clone)]
 pub struct TimelinesInfoAndOffloaded {
     pub timelines: Vec<TimelineInfo>,
@@ -1442,8 +1476,14 @@ pub struct TenantScanRemoteStorageResponse {
 #[derive(Serialize, Deserialize, Debug, Clone)]
 #[serde(rename_all = "snake_case")]
 pub enum TenantSorting {
+    /// Total size of layers on local disk for all timelines in a shard.
     ResidentSize,
+    /// The logical size of the largest timeline within a _tenant_ (not shard). Only tracked on
+    /// shard 0, contains the sum across all shards.
     MaxLogicalSize,
+    /// The logical size of the largest timeline within a _tenant_ (not shard), divided by number of
+    /// shards. Only tracked on shard 0, and estimates the per-shard logical size.
+    MaxLogicalSizePerShard,
 }
 
 impl Default for TenantSorting {
@@ -1473,14 +1513,20 @@ pub struct TopTenantShardsRequest {
 pub struct TopTenantShardItem {
     pub id: TenantShardId,
 
-    /// Total size of layers on local disk for all timelines in this tenant
+    /// Total size of layers on local disk for all timelines in this shard.
     pub resident_size: u64,
 
-    /// Total size of layers in remote storage for all timelines in this tenant
+    /// Total size of layers in remote storage for all timelines in this shard.
     pub physical_size: u64,
 
-    /// The largest logical size of a timeline within this tenant
+    /// The largest logical size of a timeline within this _tenant_ (not shard). This is only
+    /// tracked on shard 0, and contains the sum of the logical size across all shards.
     pub max_logical_size: u64,
+
+    /// The largest logical size of a timeline within this _tenant_ (not shard) divided by number of
+    /// shards. This is only tracked on shard 0, and is only an estimate as we divide it evenly by
+    /// shard count, rounded up.
+    pub max_logical_size_per_shard: u64,
 }
 
 #[derive(Serialize, Deserialize, Debug, Default)]

libs/postgres_ffi/src/lib.rs

@@ -396,6 +396,14 @@ pub mod waldecoder {
             self.lsn + self.inputbuf.remaining() as u64
         }
 
+        /// Returns the LSN up to which the WAL decoder has processed.
+        ///
+        /// If [`Self::poll_decode`] returned a record, then this will return
+        /// the end LSN of said record.
+        pub fn lsn(&self) -> Lsn {
+            self.lsn
+        }
+
         pub fn feed_bytes(&mut self, buf: &[u8]) {
             self.inputbuf.extend_from_slice(buf);
         }

libs/proxy/tokio-postgres2/src/client.rs

@@ -284,6 +284,18 @@ impl Client {
         simple_query::batch_execute(self.inner(), query).await
     }
 
+    pub async fn discard_all(&self) -> Result<ReadyForQueryStatus, Error> {
+        // clear the prepared statements that are about to be nuked from the postgres session
+        {
+            let mut typeinfo = self.inner.cached_typeinfo.lock();
+            typeinfo.typeinfo = None;
+            typeinfo.typeinfo_composite = None;
+            typeinfo.typeinfo_enum = None;
+        }
+
+        self.batch_execute("discard all").await
+    }
+
     /// Begins a new database transaction.
     ///
     /// The transaction will roll back by default - use the `commit` method to commit it.

libs/safekeeper_api/src/membership.rs

@@ -131,6 +131,14 @@ impl Configuration {
         }
     }
 
+    pub fn new(members: MemberSet) -> Self {
+        Configuration {
+            generation: INITIAL_GENERATION,
+            members,
+            new_members: None,
+        }
+    }
+
     /// Is `sk_id` member of the configuration?
     pub fn contains(&self, sk_id: NodeId) -> bool {
         self.members.contains(sk_id) || self.new_members.as_ref().is_some_and(|m| m.contains(sk_id))

libs/safekeeper_api/src/models.rs

@@ -18,7 +18,7 @@ pub struct SafekeeperStatus {
     pub id: NodeId,
 }
 
-#[derive(Serialize, Deserialize)]
+#[derive(Serialize, Deserialize, Clone)]
 pub struct TimelineCreateRequest {
     pub tenant_id: TenantId,
     pub timeline_id: TimelineId,
@@ -283,7 +283,7 @@ pub struct SafekeeperUtilization {
 }
 
 /// pull_timeline request body.
-#[derive(Debug, Deserialize, Serialize)]
+#[derive(Debug, Clone, Deserialize, Serialize)]
 pub struct PullTimelineRequest {
     pub tenant_id: TenantId,
     pub timeline_id: TimelineId,

pageserver/Cargo.toml

@@ -48,6 +48,9 @@ pprof.workspace = true
 rand.workspace = true
 range-set-blaze = { version = "0.1.16", features = ["alloc"] }
 regex.workspace = true
+rustls-pemfile.workspace = true
+rustls-pki-types.workspace = true
+rustls.workspace = true
 scopeguard.workspace = true
 send-future.workspace = true
 serde.workspace = true
@@ -62,6 +65,7 @@ tokio = { workspace = true, features = ["process", "sync", "fs", "rt", "io-util"
 tokio-epoll-uring.workspace = true
 tokio-io-timeout.workspace = true
 tokio-postgres.workspace = true
+tokio-rustls.workspace = true
 tokio-stream.workspace = true
 tokio-util.workspace = true
 toml_edit = { workspace = true, features = [ "serde" ] }
@@ -98,6 +102,7 @@ criterion.workspace = true
 hex-literal.workspace = true
 tokio = { workspace = true, features = ["process", "sync", "fs", "rt", "io-util", "time", "test-util"] }
 indoc.workspace = true
+uuid.workspace = true
 
 [[bench]]
 name = "bench_layer_map"
@@ -115,6 +120,10 @@ harness = false
 name = "upload_queue"
 harness = false
 
+[[bench]]
+name = "bench_metrics"
+harness = false
+
 [[bin]]
 name = "test_helper_slow_client_reads"
 required-features = [ "testing" ]

pageserver/benches/bench_metrics.rs

@@ -0,0 +1,366 @@
+use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
+use utils::id::{TenantId, TimelineId};
+
+//
+// Demonstrates that repeat label values lookup is a multicore scalability bottleneck
+// that is worth avoiding.
+//
+criterion_group!(
+    label_values,
+    label_values::bench_naive_usage,
+    label_values::bench_cache_label_values_lookup
+);
+mod label_values {
+    use super::*;
+
+    pub fn bench_naive_usage(c: &mut Criterion) {
+        let mut g = c.benchmark_group("label_values__naive_usage");
+
+        for ntimelines in [1, 4, 8] {
+            g.bench_with_input(
+                BenchmarkId::new("ntimelines", ntimelines),
+                &ntimelines,
+                |b, ntimelines| {
+                    b.iter_custom(|iters| {
+                        let barrier = std::sync::Barrier::new(*ntimelines + 1);
+
+                        let timelines = (0..*ntimelines)
+                            .map(|_| {
+                                (
+                                    TenantId::generate().to_string(),
+                                    "0000".to_string(),
+                                    TimelineId::generate().to_string(),
+                                )
+                            })
+                            .collect::<Vec<_>>();
+
+                        let metric_vec = metrics::UIntGaugeVec::new(
+                            metrics::opts!("testmetric", "testhelp"),
+                            &["tenant_id", "shard_id", "timeline_id"],
+                        )
+                        .unwrap();
+
+                        std::thread::scope(|s| {
+                            for (tenant_id, shard_id, timeline_id) in &timelines {
+                                s.spawn(|| {
+                                    barrier.wait();
+                                    for _ in 0..iters {
+                                        metric_vec
+                                            .with_label_values(&[tenant_id, shard_id, timeline_id])
+                                            .inc();
+                                    }
+                                    barrier.wait();
+                                });
+                            }
+                            barrier.wait();
+                            let start = std::time::Instant::now();
+                            barrier.wait();
+                            start.elapsed()
+                        })
+                    })
+                },
+            );
+        }
+        g.finish();
+    }
+
+    pub fn bench_cache_label_values_lookup(c: &mut Criterion) {
+        let mut g = c.benchmark_group("label_values__cache_label_values_lookup");
+
+        for ntimelines in [1, 4, 8] {
+            g.bench_with_input(
+                BenchmarkId::new("ntimelines", ntimelines),
+                &ntimelines,
+                |b, ntimelines| {
+                    b.iter_custom(|iters| {
+                        let barrier = std::sync::Barrier::new(*ntimelines + 1);
+
+                        let timelines = (0..*ntimelines)
+                            .map(|_| {
+                                (
+                                    TenantId::generate().to_string(),
+                                    "0000".to_string(),
+                                    TimelineId::generate().to_string(),
+                                )
+                            })
+                            .collect::<Vec<_>>();
+
+                        let metric_vec = metrics::UIntGaugeVec::new(
+                            metrics::opts!("testmetric", "testhelp"),
+                            &["tenant_id", "shard_id", "timeline_id"],
+                        )
+                        .unwrap();
+
+                        std::thread::scope(|s| {
+                            for (tenant_id, shard_id, timeline_id) in &timelines {
+                                s.spawn(|| {
+                                    let metric = metric_vec.with_label_values(&[
+                                        tenant_id,
+                                        shard_id,
+                                        timeline_id,
+                                    ]);
+                                    barrier.wait();
+                                    for _ in 0..iters {
+                                        metric.inc();
+                                    }
+                                    barrier.wait();
+                                });
+                            }
+                            barrier.wait();
+                            let start = std::time::Instant::now();
+                            barrier.wait();
+                            start.elapsed()
+                        })
+                    })
+                },
+            );
+        }
+        g.finish();
+    }
+}
+
+//
+// Demonstrates that even a single metric can be a scalability bottleneck
+// if multiple threads in it concurrently but there's nothing we can do
+// about it without changing the metrics framework to use e.g. sharded counte atomics.
+//
+criterion_group!(
+    single_metric_multicore_scalability,
+    single_metric_multicore_scalability::bench,
+);
+mod single_metric_multicore_scalability {
+    use super::*;
+
+    pub fn bench(c: &mut Criterion) {
+        let mut g = c.benchmark_group("single_metric_multicore_scalability");
+
+        for nthreads in [1, 4, 8] {
+            g.bench_with_input(
+                BenchmarkId::new("nthreads", nthreads),
+                &nthreads,
+                |b, nthreads| {
+                    b.iter_custom(|iters| {
+                        let barrier = std::sync::Barrier::new(*nthreads + 1);
+
+                        let metric = metrics::UIntGauge::new("testmetric", "testhelp").unwrap();
+
+                        std::thread::scope(|s| {
+                            for _ in 0..*nthreads {
+                                s.spawn(|| {
+                                    barrier.wait();
+                                    for _ in 0..iters {
+                                        metric.inc();
+                                    }
+                                    barrier.wait();
+                                });
+                            }
+                            barrier.wait();
+                            let start = std::time::Instant::now();
+                            barrier.wait();
+                            start.elapsed()
+                        })
+                    })
+                },
+            );
+        }
+        g.finish();
+    }
+}
+
+//
+// Demonstrates that even if we cache label value, the propagation of such a cached metric value
+// by Clone'ing it is a scalability bottleneck.
+// The reason is that it's an Arc internally and thus there's contention on the reference count atomics.
+//
+// We can avoid that by having long-lived references per thread (= indirection).
+//
+criterion_group!(
+    propagation_of_cached_label_value,
+    propagation_of_cached_label_value::bench_naive,
+    propagation_of_cached_label_value::bench_long_lived_reference_per_thread,
+);
+mod propagation_of_cached_label_value {
+    use std::sync::Arc;
+
+    use super::*;
+
+    pub fn bench_naive(c: &mut Criterion) {
+        let mut g = c.benchmark_group("propagation_of_cached_label_value__naive");
+
+        for nthreads in [1, 4, 8] {
+            g.bench_with_input(
+                BenchmarkId::new("nthreads", nthreads),
+                &nthreads,
+                |b, nthreads| {
+                    b.iter_custom(|iters| {
+                        let barrier = std::sync::Barrier::new(*nthreads + 1);
+
+                        let metric = metrics::UIntGauge::new("testmetric", "testhelp").unwrap();
+
+                        std::thread::scope(|s| {
+                            for _ in 0..*nthreads {
+                                s.spawn(|| {
+                                    barrier.wait();
+                                    for _ in 0..iters {
+                                        // propagating the metric means we'd clone it into the child RequestContext
+                                        let propagated = metric.clone();
+                                        // simulate some work
+                                        criterion::black_box(propagated);
+                                    }
+                                    barrier.wait();
+                                });
+                            }
+                            barrier.wait();
+                            let start = std::time::Instant::now();
+                            barrier.wait();
+                            start.elapsed()
+                        })
+                    })
+                },
+            );
+        }
+        g.finish();
+    }
+
+    pub fn bench_long_lived_reference_per_thread(c: &mut Criterion) {
+        let mut g =
+            c.benchmark_group("propagation_of_cached_label_value__long_lived_reference_per_thread");
+
+        for nthreads in [1, 4, 8] {
+            g.bench_with_input(
+                BenchmarkId::new("nthreads", nthreads),
+                &nthreads,
+                |b, nthreads| {
+                    b.iter_custom(|iters| {
+                        let barrier = std::sync::Barrier::new(*nthreads + 1);
+
+                        let metric = metrics::UIntGauge::new("testmetric", "testhelp").unwrap();
+
+                        std::thread::scope(|s| {
+                            for _ in 0..*nthreads {
+                                s.spawn(|| {
+                                    // This is the technique.
+                                    let this_threads_metric_reference = Arc::new(metric.clone());
+
+                                    barrier.wait();
+                                    for _ in 0..iters {
+                                        // propagating the metric means we'd clone it into the child RequestContext
+                                        let propagated = Arc::clone(&this_threads_metric_reference);
+                                        // simulate some work (include the pointer chase!)
+                                        criterion::black_box(&*propagated);
+                                    }
+                                    barrier.wait();
+                                });
+                            }
+                            barrier.wait();
+                            let start = std::time::Instant::now();
+                            barrier.wait();
+                            start.elapsed()
+                        })
+                    })
+                },
+            );
+        }
+    }
+}
+
+criterion_main!(
+    label_values,
+    single_metric_multicore_scalability,
+    propagation_of_cached_label_value
+);
+
+/*
+RUST_BACKTRACE=full cargo bench --bench bench_metrics --  --discard-baseline --noplot
+
+Results on an im4gn.2xlarge instance
+
+label_values__naive_usage/ntimelines/1 time:   [178.71 ns 178.74 ns 178.76 ns]
+label_values__naive_usage/ntimelines/4 time:   [532.94 ns 539.59 ns 546.31 ns]
+label_values__naive_usage/ntimelines/8 time:   [1.1082 µs 1.1109 µs 1.1135 µs]
+label_values__cache_label_values_lookup/ntimelines/1 time:   [6.4116 ns 6.4119 ns 6.4123 ns]
+label_values__cache_label_values_lookup/ntimelines/4 time:   [6.3482 ns 6.3819 ns 6.4079 ns]
+label_values__cache_label_values_lookup/ntimelines/8 time:   [6.4213 ns 6.5279 ns 6.6293 ns]
+single_metric_multicore_scalability/nthreads/1 time:   [6.0102 ns 6.0104 ns 6.0106 ns]
+single_metric_multicore_scalability/nthreads/4 time:   [38.127 ns 38.275 ns 38.416 ns]
+single_metric_multicore_scalability/nthreads/8 time:   [73.698 ns 74.882 ns 75.864 ns]
+propagation_of_cached_label_value__naive/nthreads/1 time:   [14.424 ns 14.425 ns 14.426 ns]
+propagation_of_cached_label_value__naive/nthreads/4 time:   [100.71 ns 102.53 ns 104.35 ns]
+propagation_of_cached_label_value__naive/nthreads/8 time:   [211.50 ns 214.44 ns 216.87 ns]
+propagation_of_cached_label_value__long_lived_reference_per_thread/nthreads/1 time:   [14.135 ns 14.147 ns 14.160 ns]
+propagation_of_cached_label_value__long_lived_reference_per_thread/nthreads/4 time:   [14.243 ns 14.255 ns 14.268 ns]
+propagation_of_cached_label_value__long_lived_reference_per_thread/nthreads/8 time:   [14.470 ns 14.682 ns 14.895 ns]
+
+Results on an i3en.3xlarge instance
+
+label_values__naive_usage/ntimelines/1      time:   [117.32 ns 117.53 ns 117.74 ns]
+label_values__naive_usage/ntimelines/4      time:   [736.58 ns 741.12 ns 745.61 ns]
+label_values__naive_usage/ntimelines/8      time:   [1.4513 µs 1.4596 µs 1.4665 µs]
+label_values__cache_label_values_lookup/ntimelines/1      time:   [8.0964 ns 8.0979 ns 8.0995 ns]
+label_values__cache_label_values_lookup/ntimelines/4      time:   [8.1620 ns 8.2912 ns 8.4491 ns]
+label_values__cache_label_values_lookup/ntimelines/8      time:   [14.148 ns 14.237 ns 14.324 ns]
+single_metric_multicore_scalability/nthreads/1      time:   [8.0993 ns 8.1013 ns 8.1046 ns]
+single_metric_multicore_scalability/nthreads/4      time:   [80.039 ns 80.672 ns 81.297 ns]
+single_metric_multicore_scalability/nthreads/8      time:   [153.58 ns 154.23 ns 154.90 ns]
+propagation_of_cached_label_value__naive/nthreads/1     time:   [13.924 ns 13.926 ns 13.928 ns]
+propagation_of_cached_label_value__naive/nthreads/4     time:   [143.66 ns 145.27 ns 146.59 ns]
+propagation_of_cached_label_value__naive/nthreads/8     time:   [296.51 ns 297.90 ns 299.30 ns]
+propagation_of_cached_label_value__long_lived_reference_per_thread/nthreads/1     time:   [14.013 ns 14.149 ns 14.308 ns]
+propagation_of_cached_label_value__long_lived_reference_per_thread/nthreads/4     time:   [14.311 ns 14.625 ns 14.984 ns]
+propagation_of_cached_label_value__long_lived_reference_per_thread/nthreads/8     time:   [25.981 ns 26.227 ns 26.476 ns]
+
+Results on an Standard L16s v3 (16 vcpus, 128 GiB memory)  Intel(R) Xeon(R) Platinum 8370C CPU @ 2.80GHz
+
+label_values__naive_usage/ntimelines/1      time:   [101.63 ns 101.84 ns 102.06 ns]
+label_values__naive_usage/ntimelines/4      time:   [417.55 ns 424.73 ns 432.63 ns]
+label_values__naive_usage/ntimelines/8      time:   [874.91 ns 889.51 ns 904.25 ns]
+label_values__cache_label_values_lookup/ntimelines/1      time:   [5.7724 ns 5.7760 ns 5.7804 ns]
+label_values__cache_label_values_lookup/ntimelines/4      time:   [7.8878 ns 7.9401 ns 8.0034 ns]
+label_values__cache_label_values_lookup/ntimelines/8      time:   [7.2621 ns 7.6354 ns 8.0337 ns]
+single_metric_multicore_scalability/nthreads/1      time:   [5.7710 ns 5.7744 ns 5.7785 ns]
+single_metric_multicore_scalability/nthreads/4      time:   [66.629 ns 66.994 ns 67.336 ns]
+single_metric_multicore_scalability/nthreads/8      time:   [130.85 ns 131.98 ns 132.91 ns]
+propagation_of_cached_label_value__naive/nthreads/1     time:   [11.540 ns 11.546 ns 11.553 ns]
+propagation_of_cached_label_value__naive/nthreads/4     time:   [131.22 ns 131.90 ns 132.56 ns]
+propagation_of_cached_label_value__naive/nthreads/8     time:   [260.99 ns 262.75 ns 264.26 ns]
+propagation_of_cached_label_value__long_lived_reference_per_thread/nthreads/1     time:   [11.544 ns 11.550 ns 11.557 ns]
+propagation_of_cached_label_value__long_lived_reference_per_thread/nthreads/4     time:   [11.568 ns 11.642 ns 11.763 ns]
+propagation_of_cached_label_value__long_lived_reference_per_thread/nthreads/8     time:   [13.416 ns 14.121 ns 14.886 ns
+
+Results on an M4 MAX MacBook Pro   Total Number of Cores:	14 (10 performance and 4 efficiency)
+
+label_values__naive_usage/ntimelines/1      time:   [52.711 ns 53.026 ns 53.381 ns]
+label_values__naive_usage/ntimelines/4      time:   [323.99 ns 330.40 ns 337.53 ns]
+label_values__naive_usage/ntimelines/8      time:   [1.1615 µs 1.1998 µs 1.2399 µs]
+label_values__cache_label_values_lookup/ntimelines/1      time:   [1.6635 ns 1.6715 ns 1.6809 ns]
+label_values__cache_label_values_lookup/ntimelines/4      time:   [1.7786 ns 1.7876 ns 1.8028 ns]
+label_values__cache_label_values_lookup/ntimelines/8      time:   [1.8195 ns 1.8371 ns 1.8665 ns]
+single_metric_multicore_scalability/nthreads/1      time:   [1.7764 ns 1.7909 ns 1.8079 ns]
+single_metric_multicore_scalability/nthreads/4      time:   [33.875 ns 34.868 ns 35.923 ns]
+single_metric_multicore_scalability/nthreads/8      time:   [226.85 ns 235.30 ns 244.18 ns]
+propagation_of_cached_label_value__naive/nthreads/1     time:   [3.4337 ns 3.4491 ns 3.4660 ns]
+propagation_of_cached_label_value__naive/nthreads/4     time:   [69.486 ns 71.937 ns 74.472 ns]
+propagation_of_cached_label_value__naive/nthreads/8     time:   [434.87 ns 456.47 ns 477.84 ns]
+propagation_of_cached_label_value__long_lived_reference_per_thread/nthreads/1     time:   [3.3767 ns 3.3974 ns 3.4220 ns]
+propagation_of_cached_label_value__long_lived_reference_per_thread/nthreads/4     time:   [3.6105 ns 4.2355 ns 5.1463 ns]
+propagation_of_cached_label_value__long_lived_reference_per_thread/nthreads/8     time:   [4.0889 ns 4.9714 ns 6.0779 ns]
+
+Results on a Hetzner AX102 AMD Ryzen 9 7950X3D 16-Core Processor
+
+label_values__naive_usage/ntimelines/1      time:   [64.510 ns 64.559 ns 64.610 ns]
+label_values__naive_usage/ntimelines/4      time:   [309.71 ns 326.09 ns 342.32 ns]
+label_values__naive_usage/ntimelines/8      time:   [776.92 ns 819.35 ns 856.93 ns]
+label_values__cache_label_values_lookup/ntimelines/1      time:   [1.2855 ns 1.2943 ns 1.3021 ns]
+label_values__cache_label_values_lookup/ntimelines/4      time:   [1.3865 ns 1.4139 ns 1.4441 ns]
+label_values__cache_label_values_lookup/ntimelines/8      time:   [1.5311 ns 1.5669 ns 1.6046 ns]
+single_metric_multicore_scalability/nthreads/1      time:   [1.1927 ns 1.1981 ns 1.2049 ns]
+single_metric_multicore_scalability/nthreads/4      time:   [24.346 ns 25.439 ns 26.634 ns]
+single_metric_multicore_scalability/nthreads/8      time:   [58.666 ns 60.137 ns 61.486 ns]
+propagation_of_cached_label_value__naive/nthreads/1     time:   [2.7067 ns 2.7238 ns 2.7402 ns]
+propagation_of_cached_label_value__naive/nthreads/4     time:   [62.723 ns 66.214 ns 69.787 ns]
+propagation_of_cached_label_value__naive/nthreads/8     time:   [164.24 ns 170.10 ns 175.68 ns]
+propagation_of_cached_label_value__long_lived_reference_per_thread/nthreads/1     time:   [2.2915 ns 2.2960 ns 2.3012 ns]
+propagation_of_cached_label_value__long_lived_reference_per_thread/nthreads/4     time:   [2.5726 ns 2.6158 ns 2.6624 ns]
+propagation_of_cached_label_value__long_lived_reference_per_thread/nthreads/8     time:   [2.7068 ns 2.8243 ns 2.9824 ns]
+
+*/

pageserver/client/src/mgmt_api.rs

@@ -7,7 +7,7 @@ use http_utils::error::HttpErrorBody;
 use pageserver_api::models::*;
 use pageserver_api::shard::TenantShardId;
 pub use reqwest::Body as ReqwestBody;
-use reqwest::{IntoUrl, Method, StatusCode};
+use reqwest::{Certificate, IntoUrl, Method, StatusCode};
 use utils::id::{TenantId, TimelineId};
 use utils::lsn::Lsn;
 
@@ -38,6 +38,9 @@ pub enum Error {
 
     #[error("Cancelled")]
     Cancelled,
+
+    #[error("create client: {0}{}", .0.source().map(|e| format!(": {e}")).unwrap_or_default())]
+    CreateClient(reqwest::Error),
 }
 
 pub type Result<T> = std::result::Result<T, Error>;
@@ -69,8 +72,17 @@ pub enum ForceAwaitLogicalSize {
 }
 
 impl Client {
-    pub fn new(mgmt_api_endpoint: String, jwt: Option<&str>) -> Self {
-        Self::from_client(reqwest::Client::new(), mgmt_api_endpoint, jwt)
+    pub fn new(
+        mgmt_api_endpoint: String,
+        jwt: Option<&str>,
+        ssl_ca_cert: Option<Certificate>,
+    ) -> Result<Self> {
+        let mut http_client = reqwest::Client::builder();
+        if let Some(ssl_ca_cert) = ssl_ca_cert {
+            http_client = http_client.add_root_certificate(ssl_ca_cert);
+        }
+        let http_client = http_client.build().map_err(Error::CreateClient)?;
+        Ok(Self::from_client(http_client, mgmt_api_endpoint, jwt))
     }
 
     pub fn from_client(
@@ -101,12 +113,10 @@ impl Client {
         debug_assert!(path.starts_with('/'));
         let uri = format!("{}{}", self.mgmt_api_endpoint, path);
 
-        let req = self.client.request(Method::GET, uri);
-        let req = if let Some(value) = &self.authorization_header {
-            req.header(reqwest::header::AUTHORIZATION, value)
-        } else {
-            req
-        };
+        let mut req = self.client.request(Method::GET, uri);
+        if let Some(value) = &self.authorization_header {
+            req = req.header(reqwest::header::AUTHORIZATION, value);
+        }
         req.send().await.map_err(Error::ReceiveBody)
     }
 

pageserver/pagebench/Cargo.toml

@@ -15,6 +15,7 @@ hdrhistogram.workspace = true
 humantime.workspace = true
 humantime-serde.workspace = true
 rand.workspace = true
+reqwest.workspace=true
 serde.workspace = true
 serde_json.workspace = true
 tracing.workspace = true

pageserver/pagebench/src/cmd/aux_files.rs

@@ -36,7 +36,8 @@ async fn main_impl(args: Args) -> anyhow::Result<()> {
     let mgmt_api_client = Arc::new(pageserver_client::mgmt_api::Client::new(
         args.mgmt_api_endpoint.clone(),
         args.pageserver_jwt.as_deref(),
-    ));
+        None, // TODO: support ssl_ca_file for https APIs in pagebench.
+    )?);
 
     // discover targets
     let timelines: Vec<TenantTimelineId> = crate::util::cli::targets::discover(

pageserver/pagebench/src/cmd/basebackup.rs

@@ -77,7 +77,8 @@ async fn main_impl(
     let mgmt_api_client = Arc::new(pageserver_client::mgmt_api::Client::new(
         args.mgmt_api_endpoint.clone(),
         args.pageserver_jwt.as_deref(),
-    ));
+        None, // TODO: support ssl_ca_file for https APIs in pagebench.
+    )?);
 
     // discover targets
     let timelines: Vec<TenantTimelineId> = crate::util::cli::targets::discover(

pageserver/pagebench/src/cmd/getpage_latest_lsn.rs

@@ -125,7 +125,8 @@ async fn main_impl(
     let mgmt_api_client = Arc::new(pageserver_client::mgmt_api::Client::new(
         args.mgmt_api_endpoint.clone(),
         args.pageserver_jwt.as_deref(),
-    ));
+        None, // TODO: support ssl_ca_file for https APIs in pagebench.
+    )?);
 
     if let Some(engine_str) = &args.set_io_engine {
         mgmt_api_client.put_io_engine(engine_str).await?;

pageserver/pagebench/src/cmd/ondemand_download_churn.rs

@@ -83,7 +83,8 @@ async fn main_impl(args: Args) -> anyhow::Result<()> {
     let mgmt_api_client = Arc::new(pageserver_client::mgmt_api::Client::new(
         args.mgmt_api_endpoint.clone(),
         args.pageserver_jwt.as_deref(),
-    ));
+        None, // TODO: support ssl_ca_file for https APIs in pagebench.
+    )?);
 
     if let Some(engine_str) = &args.set_io_engine {
         mgmt_api_client.put_io_engine(engine_str).await?;

pageserver/pagebench/src/cmd/trigger_initial_size_calculation.rs

@@ -40,7 +40,8 @@ async fn main_impl(args: Args) -> anyhow::Result<()> {
     let mgmt_api_client = Arc::new(pageserver_client::mgmt_api::Client::new(
         args.mgmt_api_endpoint.clone(),
         args.pageserver_jwt.as_deref(),
-    ));
+        None, // TODO: support ssl_ca_file for https APIs in pagebench.
+    )?);
 
     // discover targets
     let timelines: Vec<TenantTimelineId> = crate::util::cli::targets::discover(

pageserver/src/bin/pageserver.rs

@@ -25,11 +25,12 @@ use pageserver::task_mgr::{
 };
 use pageserver::tenant::{TenantSharedResources, mgr, secondary};
 use pageserver::{
-    CancellableTask, ConsumptionMetricsTasks, HttpEndpointListener, http, page_cache, page_service,
-    task_mgr, virtual_file,
+    CancellableTask, ConsumptionMetricsTasks, HttpEndpointListener, HttpsEndpointListener, http,
+    page_cache, page_service, task_mgr, virtual_file,
 };
 use postgres_backend::AuthType;
 use remote_storage::GenericRemoteStorage;
+use rustls_pki_types::{CertificateDer, PrivateKeyDer};
 use tokio::signal::unix::SignalKind;
 use tokio::time::Instant;
 use tokio_util::sync::CancellationToken;
@@ -343,8 +344,15 @@ fn start_pageserver(
     info!("Starting pageserver http handler on {http_addr}");
     let http_listener = tcp_listener::bind(http_addr)?;
 
-    let pg_addr = &conf.listen_pg_addr;
+    let https_listener = match conf.listen_https_addr.as_ref() {
+        Some(https_addr) => {
+            info!("Starting pageserver https handler on {https_addr}");
+            Some(tcp_listener::bind(https_addr)?)
+        }
+        None => None,
+    };
 
+    let pg_addr = &conf.listen_pg_addr;
     info!("Starting pageserver pg protocol handler on {pg_addr}");
     let pageserver_listener = tcp_listener::bind(pg_addr)?;
 
@@ -575,9 +583,8 @@ fn start_pageserver(
 
     // Start up the service to handle HTTP mgmt API request. We created the
     // listener earlier already.
-    let http_endpoint_listener = {
+    let (http_endpoint_listener, https_endpoint_listener) = {
         let _rt_guard = MGMT_REQUEST_RUNTIME.enter(); // for hyper
-        let cancel = CancellationToken::new();
 
         let router_state = Arc::new(
             http::routes::State::new(
@@ -592,22 +599,51 @@ fn start_pageserver(
             )
             .context("Failed to initialize router state")?,
         );
+
         let router = http::make_router(router_state, launch_ts, http_auth.clone())?
             .build()
             .map_err(|err| anyhow!(err))?;
-        let service = http_utils::RouterService::new(router).unwrap();
-        let server = hyper0::Server::from_tcp(http_listener)?
-            .serve(service)
-            .with_graceful_shutdown({
-                let cancel = cancel.clone();
-                async move { cancel.clone().cancelled().await }
-            });
 
-        let task = MGMT_REQUEST_RUNTIME.spawn(task_mgr::exit_on_panic_or_error(
-            "http endpoint listener",
-            server,
-        ));
-        HttpEndpointListener(CancellableTask { task, cancel })
+        let service =
+            Arc::new(http_utils::RequestServiceBuilder::new(router).map_err(|err| anyhow!(err))?);
+
+        let http_task = {
+            let server =
+                http_utils::server::Server::new(Arc::clone(&service), http_listener, None)?;
+            let cancel = CancellationToken::new();
+
+            let task = MGMT_REQUEST_RUNTIME.spawn(task_mgr::exit_on_panic_or_error(
+                "http endpoint listener",
+                server.serve(cancel.clone()),
+            ));
+            HttpEndpointListener(CancellableTask { task, cancel })
+        };
+
+        let https_task = match https_listener {
+            Some(https_listener) => {
+                let certs = load_certs(&conf.ssl_cert_file)?;
+                let key = load_private_key(&conf.ssl_key_file)?;
+
+                let server_config = rustls::ServerConfig::builder()
+                    .with_no_client_auth()
+                    .with_single_cert(certs, key)?;
+
+                let tls_acceptor = tokio_rustls::TlsAcceptor::from(Arc::new(server_config));
+
+                let server =
+                    http_utils::server::Server::new(service, https_listener, Some(tls_acceptor))?;
+                let cancel = CancellationToken::new();
+
+                let task = MGMT_REQUEST_RUNTIME.spawn(task_mgr::exit_on_panic_or_error(
+                    "https endpoint listener",
+                    server.serve(cancel.clone()),
+                ));
+                Some(HttpsEndpointListener(CancellableTask { task, cancel }))
+            }
+            None => None,
+        };
+
+        (http_task, https_task)
     };
 
     let consumption_metrics_tasks = {
@@ -683,6 +719,7 @@ fn start_pageserver(
         shutdown_pageserver.cancel();
         pageserver::shutdown_pageserver(
             http_endpoint_listener,
+            https_endpoint_listener,
             page_service,
             consumption_metrics_tasks,
             disk_usage_eviction_task,
@@ -697,6 +734,25 @@ fn start_pageserver(
     })
 }
 
+fn load_certs(filename: &Utf8Path) -> std::io::Result<Vec<CertificateDer<'static>>> {
+    let file = std::fs::File::open(filename)?;
+    let mut reader = std::io::BufReader::new(file);
+
+    rustls_pemfile::certs(&mut reader).collect()
+}
+
+fn load_private_key(filename: &Utf8Path) -> anyhow::Result<PrivateKeyDer<'static>> {
+    let file = std::fs::File::open(filename)?;
+    let mut reader = std::io::BufReader::new(file);
+
+    let key = rustls_pemfile::private_key(&mut reader)?;
+
+    key.ok_or(anyhow::anyhow!(
+        "no private key found in {}",
+        filename.as_str(),
+    ))
+}
+
 async fn create_remote_storage_client(
     conf: &'static PageServerConf,
 ) -> anyhow::Result<GenericRemoteStorage> {

pageserver/src/config.rs

@@ -53,6 +53,11 @@ pub struct PageServerConf {
     pub listen_pg_addr: String,
     /// Example (default): 127.0.0.1:9898
     pub listen_http_addr: String,
+    /// Example: 127.0.0.1:9899
+    pub listen_https_addr: Option<String>,
+
+    pub ssl_key_file: Utf8PathBuf,
+    pub ssl_cert_file: Utf8PathBuf,
 
     /// Current availability zone. Used for traffic metrics.
     pub availability_zone: Option<String>,
@@ -194,6 +199,13 @@ pub struct PageServerConf {
     /// Interpreted protocol feature: if enabled, validate that the logical WAL received from
     /// safekeepers does not have gaps.
     pub validate_wal_contiguity: bool,
+
+    /// When set, the previously written to disk heatmap is loaded on tenant attach and used
+    /// to avoid clobbering the heatmap from new, cold, attached locations.
+    pub load_previous_heatmap: bool,
+
+    /// When set, include visible layers in the next uploaded heatmaps of an unarchived timeline.
+    pub generate_unarchival_heatmap: bool,
 }
 
 /// Token for authentication to safekeepers
@@ -310,6 +322,9 @@ impl PageServerConf {
         let pageserver_api::config::ConfigToml {
             listen_pg_addr,
             listen_http_addr,
+            listen_https_addr,
+            ssl_key_file,
+            ssl_cert_file,
             availability_zone,
             wait_lsn_timeout,
             wal_redo_timeout,
@@ -358,6 +373,8 @@ impl PageServerConf {
             get_vectored_concurrent_io,
             enable_read_path_debugging,
             validate_wal_contiguity,
+            load_previous_heatmap,
+            generate_unarchival_heatmap,
         } = config_toml;
 
         let mut conf = PageServerConf {
@@ -366,6 +383,9 @@ impl PageServerConf {
             // ------------------------------------------------------------
             listen_pg_addr,
             listen_http_addr,
+            listen_https_addr,
+            ssl_key_file,
+            ssl_cert_file,
             availability_zone,
             wait_lsn_timeout,
             wal_redo_timeout,
@@ -447,6 +467,8 @@ impl PageServerConf {
             no_sync: no_sync.unwrap_or(false),
             enable_read_path_debugging: enable_read_path_debugging.unwrap_or(false),
             validate_wal_contiguity: validate_wal_contiguity.unwrap_or(false),
+            load_previous_heatmap: load_previous_heatmap.unwrap_or(true),
+            generate_unarchival_heatmap: generate_unarchival_heatmap.unwrap_or(true),
         };
 
         // ------------------------------------------------------------
@@ -480,7 +502,9 @@ impl PageServerConf {
     #[cfg(test)]
     pub fn test_repo_dir(test_name: &str) -> Utf8PathBuf {
         let test_output_dir = std::env::var("TEST_OUTPUT").unwrap_or("../tmp_check".into());
-        Utf8PathBuf::from(format!("{test_output_dir}/test_{test_name}"))
+
+        let test_id = uuid::Uuid::new_v4();
+        Utf8PathBuf::from(format!("{test_output_dir}/test_{test_name}_{test_id}"))
     }
 
     pub fn dummy_conf(repo_dir: Utf8PathBuf) -> Self {
@@ -493,6 +517,8 @@ impl PageServerConf {
             metric_collection_interval: Duration::from_secs(60),
             synthetic_size_calculation_interval: Duration::from_secs(60),
             background_task_maximum_delay: Duration::ZERO,
+            load_previous_heatmap: Some(true),
+            generate_unarchival_heatmap: Some(true),
             ..Default::default()
         };
         PageServerConf::parse_and_validate(NodeId(0), config_toml, &repo_dir).unwrap()

pageserver/src/context.rs

@@ -89,16 +89,112 @@
 //! [`RequestContext`] argument. Functions in the middle of the call chain
 //! only need to pass it on.
 
-use crate::task_mgr::TaskKind;
+use std::sync::Arc;
+
+use once_cell::sync::Lazy;
+use tracing::warn;
+use utils::{id::TimelineId, shard::TenantShardId};
+
+use crate::{
+    metrics::{StorageIoSizeMetrics, TimelineMetrics},
+    task_mgr::TaskKind,
+    tenant::Timeline,
+};
 
 // The main structure of this module, see module-level comment.
-#[derive(Debug)]
 pub struct RequestContext {
     task_kind: TaskKind,
     download_behavior: DownloadBehavior,
     access_stats_behavior: AccessStatsBehavior,
     page_content_kind: PageContentKind,
     read_path_debug: bool,
+    scope: Scope,
+}
+
+#[derive(Clone)]
+pub(crate) enum Scope {
+    Global {
+        io_size_metrics: &'static crate::metrics::StorageIoSizeMetrics,
+    },
+    SecondaryTenant {
+        io_size_metrics: &'static crate::metrics::StorageIoSizeMetrics,
+    },
+    SecondaryTimeline {
+        io_size_metrics: crate::metrics::StorageIoSizeMetrics,
+    },
+    Timeline {
+        // We wrap the `Arc<TimelineMetrics>`s inside another Arc to avoid child
+        // context creation contending for the ref counters of the Arc<TimelineMetrics>,
+        // which are shared among all tasks that operate on the timeline, especially
+        // concurrent page_service connections.
+        #[allow(clippy::redundant_allocation)]
+        arc_arc: Arc<Arc<TimelineMetrics>>,
+    },
+    #[cfg(test)]
+    UnitTest {
+        io_size_metrics: &'static crate::metrics::StorageIoSizeMetrics,
+    },
+}
+
+static GLOBAL_IO_SIZE_METRICS: Lazy<crate::metrics::StorageIoSizeMetrics> =
+    Lazy::new(|| crate::metrics::StorageIoSizeMetrics::new("*", "*", "*"));
+
+impl Scope {
+    pub(crate) fn new_global() -> Self {
+        Scope::Global {
+            io_size_metrics: &GLOBAL_IO_SIZE_METRICS,
+        }
+    }
+    /// NB: this allocates, so, use only at relatively long-lived roots, e.g., at start
+    /// of a compaction iteration.
+    pub(crate) fn new_timeline(timeline: &Timeline) -> Self {
+        Scope::Timeline {
+            arc_arc: Arc::new(Arc::clone(&timeline.metrics)),
+        }
+    }
+    pub(crate) fn new_page_service_pagestream(
+        timeline_handle: &crate::tenant::timeline::handle::Handle<
+            crate::page_service::TenantManagerTypes,
+        >,
+    ) -> Self {
+        Scope::Timeline {
+            arc_arc: Arc::clone(&timeline_handle.metrics),
+        }
+    }
+    pub(crate) fn new_secondary_timeline(
+        tenant_shard_id: &TenantShardId,
+        timeline_id: &TimelineId,
+    ) -> Self {
+        // TODO(https://github.com/neondatabase/neon/issues/11156): secondary timelines have no infrastructure for metrics lifecycle.
+
+        let tenant_id = tenant_shard_id.tenant_id.to_string();
+        let shard_id = tenant_shard_id.shard_slug().to_string();
+        let timeline_id = timeline_id.to_string();
+
+        let io_size_metrics =
+            crate::metrics::StorageIoSizeMetrics::new(&tenant_id, &shard_id, &timeline_id);
+        Scope::SecondaryTimeline { io_size_metrics }
+    }
+    pub(crate) fn new_secondary_tenant(_tenant_shard_id: &TenantShardId) -> Self {
+        // Before propagating metrics via RequestContext, the labels were inferred from file path.
+        // The only user of VirtualFile at tenant scope is the heatmap download & read.
+        // The inferred labels for the path of the heatmap file on local disk were that of the global metric (*,*,*).
+        // Thus, we do the same here, and extend that for anything secondary-tenant scoped.
+        //
+        // If we want to have (tenant_id, shard_id, '*') labels for secondary tenants in the future,
+        // we will need to think about the metric lifecycle, i.e., remove them during secondary tenant shutdown,
+        // like we do for attached timelines. (We don't have attached-tenant-scoped usage of VirtualFile
+        // at this point, so, we were able to completely side-step tenant-scoped stuff there).
+        Scope::SecondaryTenant {
+            io_size_metrics: &GLOBAL_IO_SIZE_METRICS,
+        }
+    }
+    #[cfg(test)]
+    pub(crate) fn new_unit_test() -> Self {
+        Scope::UnitTest {
+            io_size_metrics: &GLOBAL_IO_SIZE_METRICS,
+        }
+    }
 }
 
 /// The kind of access to the page cache.
@@ -157,6 +253,7 @@ impl RequestContextBuilder {
                 access_stats_behavior: AccessStatsBehavior::Update,
                 page_content_kind: PageContentKind::Unknown,
                 read_path_debug: false,
+                scope: Scope::new_global(),
             },
         }
     }
@@ -171,10 +268,16 @@ impl RequestContextBuilder {
                 access_stats_behavior: original.access_stats_behavior,
                 page_content_kind: original.page_content_kind,
                 read_path_debug: original.read_path_debug,
+                scope: original.scope.clone(),
             },
         }
     }
 
+    pub fn task_kind(mut self, k: TaskKind) -> Self {
+        self.inner.task_kind = k;
+        self
+    }
+
     /// Configure the DownloadBehavior of the context: whether to
     /// download missing layers, and/or warn on the download.
     pub fn download_behavior(mut self, b: DownloadBehavior) -> Self {
@@ -199,6 +302,11 @@ impl RequestContextBuilder {
         self
     }
 
+    pub(crate) fn scope(mut self, s: Scope) -> Self {
+        self.inner.scope = s;
+        self
+    }
+
     pub fn build(self) -> RequestContext {
         self.inner
     }
@@ -281,7 +389,50 @@ impl RequestContext {
     }
 
     fn child_impl(&self, task_kind: TaskKind, download_behavior: DownloadBehavior) -> Self {
-        Self::new(task_kind, download_behavior)
+        RequestContextBuilder::extend(self)
+            .task_kind(task_kind)
+            .download_behavior(download_behavior)
+            .build()
+    }
+
+    pub fn with_scope_timeline(&self, timeline: &Arc<Timeline>) -> Self {
+        RequestContextBuilder::extend(self)
+            .scope(Scope::new_timeline(timeline))
+            .build()
+    }
+
+    pub(crate) fn with_scope_page_service_pagestream(
+        &self,
+        timeline_handle: &crate::tenant::timeline::handle::Handle<
+            crate::page_service::TenantManagerTypes,
+        >,
+    ) -> Self {
+        RequestContextBuilder::extend(self)
+            .scope(Scope::new_page_service_pagestream(timeline_handle))
+            .build()
+    }
+
+    pub fn with_scope_secondary_timeline(
+        &self,
+        tenant_shard_id: &TenantShardId,
+        timeline_id: &TimelineId,
+    ) -> Self {
+        RequestContextBuilder::extend(self)
+            .scope(Scope::new_secondary_timeline(tenant_shard_id, timeline_id))
+            .build()
+    }
+
+    pub fn with_scope_secondary_tenant(&self, tenant_shard_id: &TenantShardId) -> Self {
+        RequestContextBuilder::extend(self)
+            .scope(Scope::new_secondary_tenant(tenant_shard_id))
+            .build()
+    }
+
+    #[cfg(test)]
+    pub fn with_scope_unit_test(&self) -> Self {
+        RequestContextBuilder::new(TaskKind::UnitTest)
+            .scope(Scope::new_unit_test())
+            .build()
     }
 
     pub fn task_kind(&self) -> TaskKind {
@@ -303,4 +454,38 @@ impl RequestContext {
     pub(crate) fn read_path_debug(&self) -> bool {
         self.read_path_debug
     }
+
+    pub(crate) fn io_size_metrics(&self) -> &StorageIoSizeMetrics {
+        match &self.scope {
+            Scope::Global { io_size_metrics } => {
+                let is_unit_test = cfg!(test);
+                let is_regress_test_build = cfg!(feature = "testing");
+                if is_unit_test || is_regress_test_build {
+                    panic!("all VirtualFile instances are timeline-scoped");
+                } else {
+                    use once_cell::sync::Lazy;
+                    use std::sync::Mutex;
+                    use std::time::Duration;
+                    use utils::rate_limit::RateLimit;
+                    static LIMIT: Lazy<Mutex<RateLimit>> =
+                        Lazy::new(|| Mutex::new(RateLimit::new(Duration::from_secs(1))));
+                    let mut guard = LIMIT.lock().unwrap();
+                    guard.call2(|rate_limit_stats| {
+                        warn!(
+                            %rate_limit_stats,
+                            backtrace=%std::backtrace::Backtrace::force_capture(),
+                            "all VirtualFile instances are timeline-scoped",
+                        );
+                    });
+
+                    io_size_metrics
+                }
+            }
+            Scope::Timeline { arc_arc } => &arc_arc.storage_io_size,
+            Scope::SecondaryTimeline { io_size_metrics } => io_size_metrics,
+            Scope::SecondaryTenant { io_size_metrics } => io_size_metrics,
+            #[cfg(test)]
+            Scope::UnitTest { io_size_metrics } => io_size_metrics,
+        }
+    }
 }

pageserver/src/controller_upcall_client.rs

@@ -181,7 +181,7 @@ impl ControlPlaneGenerationsApi for ControllerUpcallClient {
                         listen_pg_port: m.postgres_port,
                         listen_http_addr: m.http_host,
                         listen_http_port: m.http_port,
-                        listen_https_port: None, // TODO: Support https.
+                        listen_https_port: m.https_port,
                         availability_zone_id: az_id.expect("Checked above"),
                     })
                 }

pageserver/src/http/routes.rs

@@ -37,7 +37,8 @@ use pageserver_api::models::{
     TenantShardSplitResponse, TenantSorting, TenantState, TenantWaitLsnRequest,
     TimelineArchivalConfigRequest, TimelineCreateRequest, TimelineCreateRequestMode,
     TimelineCreateRequestModeImportPgdata, TimelineGcRequest, TimelineInfo,
-    TimelinesInfoAndOffloaded, TopTenantShardItem, TopTenantShardsRequest, TopTenantShardsResponse,
+    TimelinePatchIndexPartRequest, TimelinesInfoAndOffloaded, TopTenantShardItem,
+    TopTenantShardsRequest, TopTenantShardsResponse,
 };
 use pageserver_api::shard::{ShardCount, TenantShardId};
 use remote_storage::{DownloadError, GenericRemoteStorage, TimeTravelError};
@@ -54,6 +55,7 @@ use utils::id::{TenantId, TimelineId};
 use utils::lsn::Lsn;
 
 use crate::config::PageServerConf;
+use crate::context;
 use crate::context::{DownloadBehavior, RequestContext, RequestContextBuilder};
 use crate::deletion_queue::DeletionQueueClient;
 use crate::pgdatadir_mapping::LsnForTimestamp;
@@ -63,6 +65,7 @@ use crate::tenant::mgr::{
     GetActiveTenantError, GetTenantError, TenantManager, TenantMapError, TenantMapInsertError,
     TenantSlot, TenantSlotError, TenantSlotUpsertError, TenantStateError, UpsertLocationError,
 };
+use crate::tenant::remote_timeline_client::index::GcCompactionState;
 use crate::tenant::remote_timeline_client::{
     download_index_part, list_remote_tenant_shards, list_remote_timelines,
 };
@@ -858,6 +861,75 @@ async fn timeline_archival_config_handler(
     json_response(StatusCode::OK, ())
 }
 
+/// This API is used to patch the index part of a timeline. You must ensure such patches are safe to apply. Use this API as an emergency
+/// measure only.
+///
+/// Some examples of safe patches:
+/// - Increase the gc_cutoff and gc_compaction_cutoff to a larger value in case of a bug that didn't bump the cutoff and cause read errors.
+/// - Force set the index part to use reldir v2 (migrating/migrated).
+///
+/// Some examples of unsafe patches:
+/// - Force set the index part from v2 to v1 (legacy). This will cause the code path to ignore anything written to the new keyspace and cause
+///   errors.
+/// - Decrease the gc_cutoff without validating the data really exists. It will cause read errors in the background.
+async fn timeline_patch_index_part_handler(
+    mut request: Request<Body>,
+    _cancel: CancellationToken,
+) -> Result<Response<Body>, ApiError> {
+    let tenant_shard_id: TenantShardId = parse_request_param(&request, "tenant_shard_id")?;
+    let timeline_id: TimelineId = parse_request_param(&request, "timeline_id")?;
+
+    let request_data: TimelinePatchIndexPartRequest = json_request(&mut request).await?;
+    check_permission(&request, None)?; // require global permission for this request
+    let state = get_state(&request);
+
+    async {
+        let timeline =
+            active_timeline_of_active_tenant(&state.tenant_manager, tenant_shard_id, timeline_id)
+                .await?;
+
+        if let Some(rel_size_migration) = request_data.rel_size_migration {
+            timeline
+                .update_rel_size_v2_status(rel_size_migration)
+                .map_err(ApiError::InternalServerError)?;
+        }
+
+        if let Some(gc_compaction_last_completed_lsn) =
+            request_data.gc_compaction_last_completed_lsn
+        {
+            timeline
+                .update_gc_compaction_state(GcCompactionState {
+                    last_completed_lsn: gc_compaction_last_completed_lsn,
+                })
+                .map_err(ApiError::InternalServerError)?;
+        }
+
+        if let Some(applied_gc_cutoff_lsn) = request_data.applied_gc_cutoff_lsn {
+            {
+                let guard = timeline.applied_gc_cutoff_lsn.lock_for_write();
+                guard.store_and_unlock(applied_gc_cutoff_lsn);
+            }
+        }
+
+        if request_data.force_index_update {
+            timeline
+                .remote_client
+                .force_schedule_index_upload()
+                .context("force schedule index upload")
+                .map_err(ApiError::InternalServerError)?;
+        }
+
+        Ok::<_, ApiError>(())
+    }
+    .instrument(info_span!("timeline_patch_index_part",
+                tenant_id = %tenant_shard_id.tenant_id,
+                shard_id = %tenant_shard_id.shard_slug(),
+                %timeline_id))
+    .await?;
+
+    json_response(StatusCode::OK, ())
+}
+
 async fn timeline_detail_handler(
     request: Request<Body>,
     _cancel: CancellationToken,
@@ -882,12 +954,13 @@ async fn timeline_detail_handler(
         tenant.wait_to_become_active(ACTIVE_TENANT_TIMEOUT).await?;
 
         let timeline = tenant.get_timeline(timeline_id, false)?;
+        let ctx = &ctx.with_scope_timeline(&timeline);
 
         let timeline_info = build_timeline_info(
             &timeline,
             include_non_incremental_logical_size.unwrap_or(false),
             force_await_initial_logical_size.unwrap_or(false),
-            &ctx,
+            ctx,
         )
         .await
         .context("get local timeline info")
@@ -931,7 +1004,8 @@ async fn get_lsn_by_timestamp_handler(
     let timeline =
         active_timeline_of_active_tenant(&state.tenant_manager, tenant_shard_id, timeline_id)
             .await?;
-    let ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Download);
+    let ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Download)
+        .with_scope_timeline(&timeline);
     let result = timeline
         .find_lsn_for_timestamp(timestamp_pg, &cancel, &ctx)
         .await?;
@@ -1003,7 +1077,8 @@ async fn get_timestamp_of_lsn_handler(
     let timeline =
         active_timeline_of_active_tenant(&state.tenant_manager, tenant_shard_id, timeline_id)
             .await?;
-    let ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Download);
+    let ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Download)
+        .with_scope_timeline(&timeline);
     let result = timeline.get_timestamp_for_lsn(lsn, &ctx).await?;
 
     match result {
@@ -1358,7 +1433,8 @@ async fn timeline_layer_scan_disposable_keys(
         active_timeline_of_active_tenant(&state.tenant_manager, tenant_shard_id, timeline_id)
             .await?;
 
-    let ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Download);
+    let ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Download)
+        .with_scope_timeline(&timeline);
 
     let guard = timeline.layers.read().await;
     let Some(layer) = guard.try_get_from_key(&layer_name.clone().into()) else {
@@ -1444,7 +1520,8 @@ async fn timeline_download_heatmap_layers_handler(
     let timeline =
         active_timeline_of_active_tenant(&state.tenant_manager, tenant_shard_id, timeline_id)
             .await?;
-    let ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Download);
+    let ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Download)
+        .with_scope_timeline(&timeline);
 
     let max_concurrency = get_config(&request)
         .remote_storage_config
@@ -1492,7 +1569,8 @@ async fn layer_download_handler(
     let timeline =
         active_timeline_of_active_tenant(&state.tenant_manager, tenant_shard_id, timeline_id)
             .await?;
-    let ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Download);
+    let ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Download)
+        .with_scope_timeline(&timeline);
     let downloaded = timeline
         .download_layer(&layer_name, &ctx)
         .await
@@ -2228,8 +2306,8 @@ async fn timeline_compact_handler(
         .unwrap_or(false);
 
     async {
-        let ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Download);
         let timeline = active_timeline_of_active_tenant(&state.tenant_manager, tenant_shard_id, timeline_id).await?;
+        let ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Download).with_scope_timeline(&timeline);
         if scheduled {
             let tenant = state
                 .tenant_manager
@@ -2336,8 +2414,8 @@ async fn timeline_checkpoint_handler(
         parse_query_param::<_, bool>(&request, "wait_until_uploaded")?.unwrap_or(false);
 
     async {
-        let ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Download);
         let timeline = active_timeline_of_active_tenant(&state.tenant_manager, tenant_shard_id, timeline_id).await?;
+        let ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Download).with_scope_timeline(&timeline);
         if wait_until_flushed {
             timeline.freeze_and_flush().await
         } else {
@@ -2392,7 +2470,8 @@ async fn timeline_download_remote_layers_handler_post(
     let timeline =
         active_timeline_of_active_tenant(&state.tenant_manager, tenant_shard_id, timeline_id)
             .await?;
-    let ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Download);
+    let ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Download)
+        .with_scope_timeline(&timeline);
     match timeline.spawn_download_all_remote_layers(body, &ctx).await {
         Ok(st) => json_response(StatusCode::ACCEPTED, st),
         Err(st) => json_response(StatusCode::CONFLICT, st),
@@ -2475,6 +2554,7 @@ async fn timeline_detach_ancestor_handler(
         tracing::info!("all timeline upload queues are drained");
 
         let timeline = tenant.get_timeline(timeline_id, true)?;
+        let ctx = &ctx.with_scope_timeline(&timeline);
 
         let progress = timeline
             .prepare_to_detach_from_ancestor(&tenant, options, ctx)
@@ -2581,8 +2661,9 @@ async fn getpage_at_lsn_handler_inner(
     async {
         let ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Download);
         // Enable read path debugging
-        let ctx = RequestContextBuilder::extend(&ctx).read_path_debug(true).build();
         let timeline = active_timeline_of_active_tenant(&state.tenant_manager, tenant_shard_id, timeline_id).await?;
+        let ctx = RequestContextBuilder::extend(&ctx).read_path_debug(true)
+        .scope(context::Scope::new_timeline(&timeline)).build();
 
         // Use last_record_lsn if no lsn is provided
         let lsn = lsn.unwrap_or_else(|| timeline.get_last_record_lsn());
@@ -2616,8 +2697,8 @@ async fn timeline_collect_keyspace(
     let at_lsn: Option<Lsn> = parse_query_param(&request, "at_lsn")?;
 
     async {
-        let ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Download);
         let timeline = active_timeline_of_active_tenant(&state.tenant_manager, tenant_shard_id, timeline_id).await?;
+        let ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Download).with_scope_timeline(&timeline);
         let at_lsn = at_lsn.unwrap_or_else(|| timeline.get_last_record_lsn());
         let (dense_ks, sparse_ks) = timeline
             .collect_keyspace(at_lsn, &ctx)
@@ -3142,6 +3223,7 @@ async fn post_top_tenants(
         match order_by {
             TenantSorting::ResidentSize => sizes.resident_size,
             TenantSorting::MaxLogicalSize => sizes.max_logical_size,
+            TenantSorting::MaxLogicalSizePerShard => sizes.max_logical_size_per_shard,
         }
     }
 
@@ -3254,7 +3336,7 @@ async fn put_tenant_timeline_import_basebackup(
 
         tenant.wait_to_become_active(ACTIVE_TENANT_TIMEOUT).await?;
 
-        let timeline = tenant
+        let (timeline, timeline_ctx) = tenant
             .create_empty_timeline(timeline_id, base_lsn, pg_version, &ctx)
             .map_err(ApiError::InternalServerError)
             .await?;
@@ -3273,7 +3355,13 @@ async fn put_tenant_timeline_import_basebackup(
         info!("importing basebackup");
 
         timeline
-            .import_basebackup_from_tar(tenant.clone(), &mut body, base_lsn, broker_client, &ctx)
+            .import_basebackup_from_tar(
+                tenant.clone(),
+                &mut body,
+                base_lsn,
+                broker_client,
+                &timeline_ctx,
+            )
             .await
             .map_err(ApiError::InternalServerError)?;
 
@@ -3313,6 +3401,7 @@ async fn put_tenant_timeline_import_wal(
         let state = get_state(&request);
 
         let timeline = active_timeline_of_active_tenant(&state.tenant_manager, TenantShardId::unsharded(tenant_id), timeline_id).await?;
+        let ctx = RequestContextBuilder::extend(&ctx).scope(context::Scope::new_timeline(&timeline)).build();
 
         let mut body = StreamReader::new(request.into_body().map(|res| {
             res.map_err(|error| {
@@ -3629,6 +3718,10 @@ pub fn make_router(
             "/v1/tenant/:tenant_shard_id/timeline/:timeline_id/get_timestamp_of_lsn",
             |r| api_handler(r, get_timestamp_of_lsn_handler),
         )
+        .post(
+            "/v1/tenant/:tenant_shard_id/timeline/:timeline_id/patch_index_part",
+            |r| api_handler(r, timeline_patch_index_part_handler),
+        )
         .post(
             "/v1/tenant/:tenant_shard_id/timeline/:timeline_id/lsn_lease",
             |r| api_handler(r, lsn_lease_handler),

pageserver/src/lib.rs

@@ -64,6 +64,7 @@ pub struct CancellableTask {
     pub cancel: CancellationToken,
 }
 pub struct HttpEndpointListener(pub CancellableTask);
+pub struct HttpsEndpointListener(pub CancellableTask);
 pub struct ConsumptionMetricsTasks(pub CancellableTask);
 pub struct DiskUsageEvictionTask(pub CancellableTask);
 impl CancellableTask {
@@ -77,6 +78,7 @@ impl CancellableTask {
 #[allow(clippy::too_many_arguments)]
 pub async fn shutdown_pageserver(
     http_listener: HttpEndpointListener,
+    https_listener: Option<HttpsEndpointListener>,
     page_service: page_service::Listener,
     consumption_metrics_worker: ConsumptionMetricsTasks,
     disk_usage_eviction_task: Option<DiskUsageEvictionTask>,
@@ -213,6 +215,15 @@ pub async fn shutdown_pageserver(
     )
     .await;
 
+    if let Some(https_listener) = https_listener {
+        timed(
+            https_listener.0.shutdown(),
+            "shutdown https",
+            Duration::from_secs(1),
+        )
+        .await;
+    }
+
     // Shut down the HTTP endpoint last, so that you can still check the server's
     // status while it's shutting down.
     // FIXME: We should probably stop accepting commands like attach/detach earlier.

pageserver/src/metrics.rs

@@ -143,6 +143,29 @@ pub(crate) static LAYERS_PER_READ_GLOBAL: Lazy<Histogram> = Lazy::new(|| {
     .expect("failed to define a metric")
 });
 
+pub(crate) static LAYERS_PER_READ_BATCH_GLOBAL: Lazy<Histogram> = Lazy::new(|| {
+    register_histogram!(
+        "pageserver_layers_per_read_batch_global",
+        "Layers visited to serve a single read batch (read amplification), regardless of number of reads.",
+        vec![
+            1.0, 2.0, 4.0, 8.0, 16.0, 32.0, 64.0, 128.0, 256.0, 512.0, 1024.0
+        ],
+    )
+    .expect("failed to define a metric")
+});
+
+pub(crate) static LAYERS_PER_READ_AMORTIZED_GLOBAL: Lazy<Histogram> = Lazy::new(|| {
+    register_histogram!(
+        "pageserver_layers_per_read_amortized_global",
+        "Layers visited to serve a single read (read amplification). Amortized across a batch: \
+            all visited layers are divided by number of reads.",
+        vec![
+            1.0, 2.0, 4.0, 8.0, 16.0, 32.0, 64.0, 128.0, 256.0, 512.0, 1024.0
+        ],
+    )
+    .expect("failed to define a metric")
+});
+
 pub(crate) static DELTAS_PER_READ_GLOBAL: Lazy<Histogram> = Lazy::new(|| {
     // We expect this to be low because of Postgres checkpoints. Let's see if that holds.
     register_histogram!(
@@ -1204,11 +1227,24 @@ impl StorageIoTime {
 
 pub(crate) static STORAGE_IO_TIME_METRIC: Lazy<StorageIoTime> = Lazy::new(StorageIoTime::new);
 
-const STORAGE_IO_SIZE_OPERATIONS: &[&str] = &["read", "write"];
+#[derive(Clone, Copy)]
+#[repr(usize)]
+enum StorageIoSizeOperation {
+    Read,
+    Write,
+}
+
+impl StorageIoSizeOperation {
+    const VARIANTS: &'static [&'static str] = &["read", "write"];
+
+    fn as_str(&self) -> &'static str {
+        Self::VARIANTS[*self as usize]
+    }
+}
 
 // Needed for the https://neonprod.grafana.net/d/5uK9tHL4k/picking-tenant-for-relocation?orgId=1
-pub(crate) static STORAGE_IO_SIZE: Lazy<IntGaugeVec> = Lazy::new(|| {
-    register_int_gauge_vec!(
+static STORAGE_IO_SIZE: Lazy<UIntGaugeVec> = Lazy::new(|| {
+    register_uint_gauge_vec!(
         "pageserver_io_operations_bytes_total",
         "Total amount of bytes read/written in IO operations",
         &["operation", "tenant_id", "shard_id", "timeline_id"]
@@ -1216,6 +1252,34 @@ pub(crate) static STORAGE_IO_SIZE: Lazy<IntGaugeVec> = Lazy::new(|| {
     .expect("failed to define a metric")
 });
 
+#[derive(Clone, Debug)]
+pub(crate) struct StorageIoSizeMetrics {
+    pub read: UIntGauge,
+    pub write: UIntGauge,
+}
+
+impl StorageIoSizeMetrics {
+    pub(crate) fn new(tenant_id: &str, shard_id: &str, timeline_id: &str) -> Self {
+        let read = STORAGE_IO_SIZE
+            .get_metric_with_label_values(&[
+                StorageIoSizeOperation::Read.as_str(),
+                tenant_id,
+                shard_id,
+                timeline_id,
+            ])
+            .unwrap();
+        let write = STORAGE_IO_SIZE
+            .get_metric_with_label_values(&[
+                StorageIoSizeOperation::Write.as_str(),
+                tenant_id,
+                shard_id,
+                timeline_id,
+            ])
+            .unwrap();
+        Self { read, write }
+    }
+}
+
 #[cfg(not(test))]
 pub(crate) mod virtual_file_descriptor_cache {
     use super::*;
@@ -2798,6 +2862,7 @@ pub(crate) struct TimelineMetrics {
     /// Number of valid LSN leases.
     pub valid_lsn_lease_count_gauge: UIntGauge,
     pub wal_records_received: IntCounter,
+    pub storage_io_size: StorageIoSizeMetrics,
     shutdown: std::sync::atomic::AtomicBool,
 }
 
@@ -2933,6 +2998,8 @@ impl TimelineMetrics {
             .get_metric_with_label_values(&[&tenant_id, &shard_id, &timeline_id])
             .unwrap();
 
+        let storage_io_size = StorageIoSizeMetrics::new(&tenant_id, &shard_id, &timeline_id);
+
         TimelineMetrics {
             tenant_id,
             shard_id,
@@ -2962,6 +3029,7 @@ impl TimelineMetrics {
             evictions_with_low_residence_duration: std::sync::RwLock::new(
                 evictions_with_low_residence_duration,
             ),
+            storage_io_size,
             valid_lsn_lease_count_gauge,
             wal_records_received,
             shutdown: std::sync::atomic::AtomicBool::default(),
@@ -3152,7 +3220,7 @@ impl TimelineMetrics {
             ]);
         }
 
-        for op in STORAGE_IO_SIZE_OPERATIONS {
+        for op in StorageIoSizeOperation::VARIANTS {
             let _ = STORAGE_IO_SIZE.remove_label_values(&[op, tenant_id, shard_id, timeline_id]);
         }
 
@@ -4074,6 +4142,8 @@ pub fn preinitialize_metrics(conf: &'static PageServerConf) {
     // histograms
     [
         &LAYERS_PER_READ_GLOBAL,
+        &LAYERS_PER_READ_BATCH_GLOBAL,
+        &LAYERS_PER_READ_AMORTIZED_GLOBAL,
         &DELTAS_PER_READ_GLOBAL,
         &WAIT_LSN_TIME,
         &WAL_REDO_TIME,

pageserver/src/page_service.rs

@@ -56,6 +56,7 @@ use crate::config::PageServerConf;
 use crate::context::{DownloadBehavior, RequestContext};
 use crate::metrics::{
     self, COMPUTE_COMMANDS_COUNTERS, ComputeCommandKind, LIVE_CONNECTIONS, SmgrOpTimer,
+    TimelineMetrics,
 };
 use crate::pgdatadir_mapping::Version;
 use crate::span::{
@@ -423,6 +424,9 @@ impl timeline::handle::Types for TenantManagerTypes {
 
 pub(crate) struct TenantManagerCacheItem {
     pub(crate) timeline: Arc<Timeline>,
+    // allow() for cheap propagation through RequestContext inside a task
+    #[allow(clippy::redundant_allocation)]
+    pub(crate) metrics: Arc<Arc<TimelineMetrics>>,
     #[allow(dead_code)] // we store it to keep the gate open
     pub(crate) gate_guard: GateGuard,
 }
@@ -506,8 +510,11 @@ impl timeline::handle::TenantManager<TenantManagerTypes> for TenantManagerWrappe
             }
         };
 
+        let metrics = Arc::new(Arc::clone(&timeline.metrics));
+
         Ok(TenantManagerCacheItem {
             timeline,
+            metrics,
             gate_guard,
         })
     }
@@ -1238,6 +1245,14 @@ impl PageServerHandler {
         ),
         QueryError,
     > {
+        macro_rules! upgrade_handle_and_set_context {
+            ($shard:ident) => {{
+                let weak_handle = &$shard;
+                let handle = weak_handle.upgrade()?;
+                let ctx = ctx.with_scope_page_service_pagestream(&handle);
+                (handle, ctx)
+            }};
+        }
         Ok(match batch {
             BatchedFeMessage::Exists {
                 span,
@@ -1246,9 +1261,10 @@ impl PageServerHandler {
                 req,
             } => {
                 fail::fail_point!("ps::handle-pagerequest-message::exists");
+                let (shard, ctx) = upgrade_handle_and_set_context!(shard);
                 (
                     vec![
-                        self.handle_get_rel_exists_request(&*shard.upgrade()?, &req, ctx)
+                        self.handle_get_rel_exists_request(&shard, &req, &ctx)
                             .instrument(span.clone())
                             .await
                             .map(|msg| (msg, timer))
@@ -1264,9 +1280,10 @@ impl PageServerHandler {
                 req,
             } => {
                 fail::fail_point!("ps::handle-pagerequest-message::nblocks");
+                let (shard, ctx) = upgrade_handle_and_set_context!(shard);
                 (
                     vec![
-                        self.handle_get_nblocks_request(&*shard.upgrade()?, &req, ctx)
+                        self.handle_get_nblocks_request(&shard, &req, &ctx)
                             .instrument(span.clone())
                             .await
                             .map(|msg| (msg, timer))
@@ -1282,17 +1299,18 @@ impl PageServerHandler {
                 pages,
             } => {
                 fail::fail_point!("ps::handle-pagerequest-message::getpage");
+                let (shard, ctx) = upgrade_handle_and_set_context!(shard);
                 (
                     {
                         let npages = pages.len();
                         trace!(npages, "handling getpage request");
                         let res = self
                             .handle_get_page_at_lsn_request_batched(
-                                &*shard.upgrade()?,
+                                &shard,
                                 effective_request_lsn,
                                 pages,
                                 io_concurrency,
-                                ctx,
+                                &ctx,
                             )
                             .instrument(span.clone())
                             .await;
@@ -1309,9 +1327,10 @@ impl PageServerHandler {
                 req,
             } => {
                 fail::fail_point!("ps::handle-pagerequest-message::dbsize");
+                let (shard, ctx) = upgrade_handle_and_set_context!(shard);
                 (
                     vec![
-                        self.handle_db_size_request(&*shard.upgrade()?, &req, ctx)
+                        self.handle_db_size_request(&shard, &req, &ctx)
                             .instrument(span.clone())
                             .await
                             .map(|msg| (msg, timer))
@@ -1327,9 +1346,10 @@ impl PageServerHandler {
                 req,
             } => {
                 fail::fail_point!("ps::handle-pagerequest-message::slrusegment");
+                let (shard, ctx) = upgrade_handle_and_set_context!(shard);
                 (
                     vec![
-                        self.handle_get_slru_segment_request(&*shard.upgrade()?, &req, ctx)
+                        self.handle_get_slru_segment_request(&shard, &req, &ctx)
                             .instrument(span.clone())
                             .await
                             .map(|msg| (msg, timer))
@@ -1345,12 +1365,13 @@ impl PageServerHandler {
                 requests,
             } => {
                 fail::fail_point!("ps::handle-pagerequest-message::test");
+                let (shard, ctx) = upgrade_handle_and_set_context!(shard);
                 (
                     {
                         let npages = requests.len();
                         trace!(npages, "handling getpage request");
                         let res = self
-                            .handle_test_request_batch(&*shard.upgrade()?, requests, ctx)
+                            .handle_test_request_batch(&shard, requests, &ctx)
                             .instrument(span.clone())
                             .await;
                         assert_eq!(res.len(), npages);
@@ -2126,6 +2147,7 @@ impl PageServerHandler {
             .get(tenant_id, timeline_id, ShardSelector::Zero)
             .await?;
         set_tracing_field_shard_id(&timeline);
+        let ctx = ctx.with_scope_timeline(&timeline);
 
         if timeline.is_archived() == Some(true) {
             tracing::info!(
@@ -2143,7 +2165,7 @@ impl PageServerHandler {
                     lsn,
                     crate::tenant::timeline::WaitLsnWaiter::PageService,
                     crate::tenant::timeline::WaitLsnTimeout::Default,
-                    ctx,
+                    &ctx,
                 )
                 .await?;
             timeline
@@ -2169,7 +2191,7 @@ impl PageServerHandler {
                 prev_lsn,
                 full_backup,
                 replica,
-                ctx,
+                &ctx,
             )
             .await
             .map_err(map_basebackup_error)?;
@@ -2192,7 +2214,7 @@ impl PageServerHandler {
                     prev_lsn,
                     full_backup,
                     replica,
-                    ctx,
+                    &ctx,
                 )
                 .await
                 .map_err(map_basebackup_error)?;
@@ -2209,7 +2231,7 @@ impl PageServerHandler {
                     prev_lsn,
                     full_backup,
                     replica,
-                    ctx,
+                    &ctx,
                 )
                 .await
                 .map_err(map_basebackup_error)?;

pageserver/src/pgdatadir_mapping.rs

@@ -1964,14 +1964,12 @@ impl DatadirModification<'_> {
                 .context("deserialize db")?
         };
 
-        // Add the new relation to the rel directory entry, and write it back
-        if !rel_dir.rels.insert((rel.relnode, rel.forknum)) {
-            return Err(RelationError::AlreadyExists);
-        }
-
         let v2_enabled = self.maybe_enable_rel_size_v2()?;
 
         if v2_enabled {
+            if rel_dir.rels.contains(&(rel.relnode, rel.forknum)) {
+                return Err(RelationError::AlreadyExists);
+            }
             let sparse_rel_dir_key =
                 rel_tag_sparse_key(rel.spcnode, rel.dbnode, rel.relnode, rel.forknum);
             // check if the rel_dir_key exists in v2
@@ -2006,6 +2004,10 @@ impl DatadirModification<'_> {
             self.pending_directory_entries
                 .push((DirectoryKind::RelV2, MetricsUpdate::Add(1)));
         } else {
+            // Add the new relation to the rel directory entry, and write it back
+            if !rel_dir.rels.insert((rel.relnode, rel.forknum)) {
+                return Err(RelationError::AlreadyExists);
+            }
             if !dbdir_exists {
                 self.pending_directory_entries
                     .push((DirectoryKind::Rel, MetricsUpdate::Set(0)))
@@ -2019,6 +2021,7 @@ impl DatadirModification<'_> {
                 )),
             );
         }
+
         // Put size
         let size_key = rel_size_to_key(rel);
         let buf = nblocks.to_le_bytes();
@@ -2141,7 +2144,7 @@ impl DatadirModification<'_> {
                     // Remove entry from relation size cache
                     self.tline.remove_cached_rel_size(&rel_tag);
 
-                    // Delete size entry, as well as all blocks
+                    // Delete size entry, as well as all blocks; this is currently a no-op because we haven't implemented tombstones in storage.
                     self.delete(rel_key_range(rel_tag));
                 }
             }
@@ -2755,7 +2758,7 @@ mod tests {
             TimelineId::from_array(hex!("11223344556677881122334455667788"));
 
         let (tenant, ctx) = harness.load().await;
-        let tline = tenant
+        let (tline, ctx) = tenant
             .create_empty_timeline(TIMELINE_ID, Lsn(0x10), DEFAULT_PG_VERSION, &ctx)
             .await?;
         let tline = tline.raw_timeline().unwrap();

pageserver/src/tenant.rs

@@ -77,6 +77,8 @@ use self::timeline::{
     EvictionTaskTenantState, GcCutoffs, TimelineDeleteProgress, TimelineResources, WaitLsnError,
 };
 use crate::config::PageServerConf;
+use crate::context;
+use crate::context::RequestContextBuilder;
 use crate::context::{DownloadBehavior, RequestContext};
 use crate::deletion_queue::{DeletionQueueClient, DeletionQueueError};
 use crate::l0_flush::L0FlushGlobalState;
@@ -1114,7 +1116,7 @@ impl Tenant {
             }
         };
 
-        let timeline = self.create_timeline_struct(
+        let (timeline, timeline_ctx) = self.create_timeline_struct(
             timeline_id,
             &metadata,
             previous_heatmap,
@@ -1124,6 +1126,7 @@ impl Tenant {
             idempotency.clone(),
             index_part.gc_compaction.clone(),
             index_part.rel_size_migration.clone(),
+            ctx,
         )?;
         let disk_consistent_lsn = timeline.get_disk_consistent_lsn();
         anyhow::ensure!(
@@ -1150,7 +1153,7 @@ impl Tenant {
         // a previous heatmap which contains all visible layers in the layer map.
         // This previous heatmap will be used whenever a fresh heatmap is generated
         // for the timeline.
-        if matches!(cause, LoadTimelineCause::Unoffload) {
+        if self.conf.generate_unarchival_heatmap && matches!(cause, LoadTimelineCause::Unoffload) {
             let mut tline_ending_at = Some((&timeline, timeline.get_last_record_lsn()));
             while let Some((tline, end_lsn)) = tline_ending_at {
                 let unarchival_heatmap = tline.generate_unarchival_heatmap(end_lsn).await;
@@ -1257,7 +1260,7 @@ impl Tenant {
                         match activate {
                             ActivateTimelineArgs::Yes { broker_client } => {
                                 info!("activating timeline after reload from pgdata import task");
-                                timeline.activate(self.clone(), broker_client, None, ctx);
+                                timeline.activate(self.clone(), broker_client, None, &timeline_ctx);
                             }
                             ActivateTimelineArgs::No => (),
                         }
@@ -1582,6 +1585,10 @@ impl Tenant {
     }
 
     async fn read_on_disk_heatmap(&self) -> Option<(HeatMapTenant, std::time::Instant)> {
+        if !self.conf.load_previous_heatmap {
+            return None;
+        }
+
         let on_disk_heatmap_path = self.conf.tenant_heatmap_path(&self.tenant_shard_id);
         match tokio::fs::read_to_string(on_disk_heatmap_path).await {
             Ok(heatmap) => match serde_json::from_str::<HeatMapTenant>(&heatmap) {
@@ -1761,6 +1768,7 @@ impl Tenant {
                         import_pgdata,
                         ActivateTimelineArgs::No,
                         guard,
+                        ctx.detached_child(TaskKind::ImportPgdata, DownloadBehavior::Warn),
                     ));
                 }
             }
@@ -1778,6 +1786,7 @@ impl Tenant {
                 timeline_id,
                 &index_part.metadata,
                 remote_timeline_client,
+                ctx,
             )
             .instrument(tracing::info_span!("timeline_delete", %timeline_id))
             .await
@@ -2215,7 +2224,7 @@ impl Tenant {
                 self.clone(),
                 broker_client.clone(),
                 background_jobs_can_start,
-                &ctx,
+                &ctx.with_scope_timeline(&timeline),
             );
         }
 
@@ -2412,8 +2421,8 @@ impl Tenant {
         new_timeline_id: TimelineId,
         initdb_lsn: Lsn,
         pg_version: u32,
-        _ctx: &RequestContext,
-    ) -> anyhow::Result<UninitializedTimeline> {
+        ctx: &RequestContext,
+    ) -> anyhow::Result<(UninitializedTimeline, RequestContext)> {
         anyhow::ensure!(
             self.is_active(),
             "Cannot create empty timelines on inactive tenant"
@@ -2447,6 +2456,8 @@ impl Tenant {
             create_guard,
             initdb_lsn,
             None,
+            None,
+            ctx,
         )
         .await
     }
@@ -2464,7 +2475,7 @@ impl Tenant {
         pg_version: u32,
         ctx: &RequestContext,
     ) -> anyhow::Result<Arc<Timeline>> {
-        let uninit_tl = self
+        let (uninit_tl, ctx) = self
             .create_empty_timeline(new_timeline_id, initdb_lsn, pg_version, ctx)
             .await?;
         let tline = uninit_tl.raw_timeline().expect("we just created it");
@@ -2476,7 +2487,7 @@ impl Tenant {
             .init_empty_test_timeline()
             .context("init_empty_test_timeline")?;
         modification
-            .commit(ctx)
+            .commit(&ctx)
             .await
             .context("commit init_empty_test_timeline modification")?;
 
@@ -2694,7 +2705,12 @@ impl Tenant {
         // doing stuff before the IndexPart is durable in S3, which is done by the previous section.
         let activated_timeline = match result {
             CreateTimelineResult::Created(timeline) => {
-                timeline.activate(self.clone(), broker_client, None, ctx);
+                timeline.activate(
+                    self.clone(),
+                    broker_client,
+                    None,
+                    &ctx.with_scope_timeline(&timeline),
+                );
                 timeline
             }
             CreateTimelineResult::Idempotent(timeline) => {
@@ -2756,10 +2772,9 @@ impl Tenant {
             }
         };
 
-        let mut uninit_timeline = {
+        let (mut uninit_timeline, timeline_ctx) = {
             let this = &self;
             let initdb_lsn = Lsn(0);
-            let _ctx = ctx;
             async move {
                 let new_metadata = TimelineMetadata::new(
                     // Initialize disk_consistent LSN to 0, The caller must import some data to
@@ -2778,6 +2793,8 @@ impl Tenant {
                     timeline_create_guard,
                     initdb_lsn,
                     None,
+                    None,
+                    ctx,
                 )
                 .await
             }
@@ -2807,6 +2824,7 @@ impl Tenant {
             index_part,
             activate,
             timeline_create_guard,
+            timeline_ctx.detached_child(TaskKind::ImportPgdata, DownloadBehavior::Warn),
         ));
 
         // NB: the timeline doesn't exist in self.timelines at this point
@@ -2820,6 +2838,7 @@ impl Tenant {
         index_part: import_pgdata::index_part_format::Root,
         activate: ActivateTimelineArgs,
         timeline_create_guard: TimelineCreateGuard,
+        ctx: RequestContext,
     ) {
         debug_assert_current_span_has_tenant_and_timeline_id();
         info!("starting");
@@ -2831,6 +2850,7 @@ impl Tenant {
                 index_part,
                 activate,
                 timeline_create_guard,
+                ctx,
             )
             .await;
         if let Err(err) = &res {
@@ -2846,9 +2866,8 @@ impl Tenant {
         index_part: import_pgdata::index_part_format::Root,
         activate: ActivateTimelineArgs,
         timeline_create_guard: TimelineCreateGuard,
+        ctx: RequestContext,
     ) -> Result<(), anyhow::Error> {
-        let ctx = RequestContext::new(TaskKind::ImportPgdata, DownloadBehavior::Warn);
-
         info!("importing pgdata");
         import_pgdata::doit(&timeline, index_part, &ctx, self.cancel.clone())
             .await
@@ -3057,6 +3076,7 @@ impl Tenant {
 
             let mut has_pending_l0 = false;
             for timeline in compact_l0 {
+                let ctx = &ctx.with_scope_timeline(&timeline);
                 let outcome = timeline
                     .compact(cancel, CompactFlags::OnlyL0Compaction.into(), ctx)
                     .instrument(info_span!("compact_timeline", timeline_id = %timeline.timeline_id))
@@ -3090,6 +3110,7 @@ impl Tenant {
             if !timeline.is_active() {
                 continue;
             }
+            let ctx = &ctx.with_scope_timeline(&timeline);
 
             let mut outcome = timeline
                 .compact(cancel, EnumSet::default(), ctx)
@@ -3315,7 +3336,7 @@ impl Tenant {
                     self.clone(),
                     broker_client.clone(),
                     background_jobs_can_start,
-                    ctx,
+                    &ctx.with_scope_timeline(timeline),
                 );
                 activated_timelines += 1;
             }
@@ -3821,6 +3842,7 @@ impl Tenant {
             resident_size: 0,
             physical_size: 0,
             max_logical_size: 0,
+            max_logical_size_per_shard: 0,
         };
 
         for timeline in self.timelines.lock().unwrap().values() {
@@ -3837,6 +3859,10 @@ impl Tenant {
             );
         }
 
+        result.max_logical_size_per_shard = result
+            .max_logical_size
+            .div_ceil(self.tenant_shard_id.shard_count.count() as u64);
+
         result
     }
 }
@@ -4130,7 +4156,8 @@ impl Tenant {
         create_idempotency: CreateTimelineIdempotency,
         gc_compaction_state: Option<GcCompactionState>,
         rel_size_v2_status: Option<RelSizeMigration>,
-    ) -> anyhow::Result<Arc<Timeline>> {
+        ctx: &RequestContext,
+    ) -> anyhow::Result<(Arc<Timeline>, RequestContext)> {
         let state = match cause {
             CreateTimelineCause::Load => {
                 let ancestor_id = new_metadata.ancestor_timeline();
@@ -4166,7 +4193,11 @@ impl Tenant {
             self.cancel.child_token(),
         );
 
-        Ok(timeline)
+        let timeline_ctx = RequestContextBuilder::extend(ctx)
+            .scope(context::Scope::new_timeline(&timeline))
+            .build();
+
+        Ok((timeline, timeline_ctx))
     }
 
     /// [`Tenant::shutdown`] must be called before dropping the returned [`Tenant`] object
@@ -4582,6 +4613,7 @@ impl Tenant {
         // Ensures all timelines use the same start time when computing the time cutoff.
         let now_ts_for_pitr_calc = SystemTime::now();
         for timeline in timelines.iter() {
+            let ctx = &ctx.with_scope_timeline(timeline);
             let cutoff = timeline
                 .get_last_record_lsn()
                 .checked_sub(horizon)
@@ -4755,7 +4787,7 @@ impl Tenant {
         src_timeline: &Arc<Timeline>,
         dst_id: TimelineId,
         start_lsn: Option<Lsn>,
-        _ctx: &RequestContext,
+        ctx: &RequestContext,
     ) -> Result<CreateTimelineResult, CreateTimelineError> {
         let src_id = src_timeline.timeline_id;
 
@@ -4858,13 +4890,15 @@ impl Tenant {
             src_timeline.pg_version,
         );
 
-        let uninitialized_timeline = self
+        let (uninitialized_timeline, _timeline_ctx) = self
             .prepare_new_timeline(
                 dst_id,
                 &metadata,
                 timeline_create_guard,
                 start_lsn + 1,
                 Some(Arc::clone(src_timeline)),
+                Some(src_timeline.get_rel_size_v2_status()),
+                ctx,
             )
             .await?;
 
@@ -5131,13 +5165,15 @@ impl Tenant {
             pgdata_lsn,
             pg_version,
         );
-        let mut raw_timeline = self
+        let (mut raw_timeline, timeline_ctx) = self
             .prepare_new_timeline(
                 timeline_id,
                 &new_metadata,
                 timeline_create_guard,
                 pgdata_lsn,
                 None,
+                None,
+                ctx,
             )
             .await?;
 
@@ -5148,7 +5184,7 @@ impl Tenant {
                     &unfinished_timeline,
                     &pgdata_path,
                     pgdata_lsn,
-                    ctx,
+                    &timeline_ctx,
                 )
                 .await
                 .with_context(|| {
@@ -5209,6 +5245,7 @@ impl Tenant {
     /// An empty layer map is initialized, and new data and WAL can be imported starting
     /// at 'disk_consistent_lsn'. After any initial data has been imported, call
     /// `finish_creation` to insert the Timeline into the timelines map.
+    #[allow(clippy::too_many_arguments)]
     async fn prepare_new_timeline<'a>(
         &'a self,
         new_timeline_id: TimelineId,
@@ -5216,15 +5253,17 @@ impl Tenant {
         create_guard: TimelineCreateGuard,
         start_lsn: Lsn,
         ancestor: Option<Arc<Timeline>>,
-    ) -> anyhow::Result<UninitializedTimeline<'a>> {
+        rel_size_v2_status: Option<RelSizeMigration>,
+        ctx: &RequestContext,
+    ) -> anyhow::Result<(UninitializedTimeline<'a>, RequestContext)> {
         let tenant_shard_id = self.tenant_shard_id;
 
         let resources = self.build_timeline_resources(new_timeline_id);
         resources
             .remote_client
-            .init_upload_queue_for_empty_remote(new_metadata)?;
+            .init_upload_queue_for_empty_remote(new_metadata, rel_size_v2_status.clone())?;
 
-        let timeline_struct = self
+        let (timeline_struct, timeline_ctx) = self
             .create_timeline_struct(
                 new_timeline_id,
                 new_metadata,
@@ -5234,7 +5273,8 @@ impl Tenant {
                 CreateTimelineCause::Load,
                 create_guard.idempotency.clone(),
                 None,
-                None,
+                rel_size_v2_status,
+                ctx,
             )
             .context("Failed to create timeline data structure")?;
 
@@ -5255,10 +5295,13 @@ impl Tenant {
             "Successfully created initial files for timeline {tenant_shard_id}/{new_timeline_id}"
         );
 
-        Ok(UninitializedTimeline::new(
-            self,
-            new_timeline_id,
-            Some((timeline_struct, create_guard)),
+        Ok((
+            UninitializedTimeline::new(
+                self,
+                new_timeline_id,
+                Some((timeline_struct, create_guard)),
+            ),
+            timeline_ctx,
         ))
     }
 
@@ -5793,7 +5836,8 @@ pub(crate) mod harness {
         }
 
         pub(crate) async fn load(&self) -> (Arc<Tenant>, RequestContext) {
-            let ctx = RequestContext::new(TaskKind::UnitTest, DownloadBehavior::Error);
+            let ctx = RequestContext::new(TaskKind::UnitTest, DownloadBehavior::Error)
+                .with_scope_unit_test();
             (
                 self.do_try_load(&ctx)
                     .await
@@ -6816,7 +6860,7 @@ mod tests {
 
         let (tenant, ctx) = harness.load().await;
         let io_concurrency = IoConcurrency::spawn_for_test();
-        let tline = tenant
+        let (tline, ctx) = tenant
             .create_empty_timeline(TIMELINE_ID, Lsn(0), DEFAULT_PG_VERSION, &ctx)
             .await?;
         let tline = tline.raw_timeline().unwrap();
@@ -7438,7 +7482,7 @@ mod tests {
             .await;
 
         let initdb_lsn = Lsn(0x20);
-        let utline = tenant
+        let (utline, ctx) = tenant
             .create_empty_timeline(TIMELINE_ID, initdb_lsn, DEFAULT_PG_VERSION, &ctx)
             .await?;
         let tline = utline.raw_timeline().unwrap();
@@ -7505,7 +7549,7 @@ mod tests {
         let harness = TenantHarness::create(name).await?;
         {
             let (tenant, ctx) = harness.load().await;
-            let tline = tenant
+            let (tline, _ctx) = tenant
                 .create_empty_timeline(TIMELINE_ID, Lsn(0), DEFAULT_PG_VERSION, &ctx)
                 .await?;
             // Leave the timeline ID in [`Tenant::timelines_creating`] to exclude attempting to create it again

pageserver/src/tenant/blob_io.rs

@@ -471,7 +471,8 @@ pub(crate) mod tests {
         blobs: &[Vec<u8>],
         compression: bool,
     ) -> Result<(), Error> {
-        let ctx = RequestContext::new(TaskKind::UnitTest, DownloadBehavior::Error);
+        let ctx =
+            RequestContext::new(TaskKind::UnitTest, DownloadBehavior::Error).with_scope_unit_test();
         let (_temp_dir, pathbuf, offsets) =
             write_maybe_compressed::<BUFFERED>(blobs, compression, &ctx).await?;
 

pageserver/src/tenant/disk_btree.rs

@@ -32,8 +32,7 @@ use hex;
 use thiserror::Error;
 use tracing::error;
 
-use crate::context::{DownloadBehavior, RequestContext};
-use crate::task_mgr::TaskKind;
+use crate::context::RequestContext;
 use crate::tenant::block_io::{BlockReader, BlockWriter};
 
 // The maximum size of a value stored in the B-tree. 5 bytes is enough currently.
@@ -477,16 +476,15 @@ where
     }
 
     #[allow(dead_code)]
-    pub async fn dump(&self) -> Result<()> {
+    pub async fn dump(&self, ctx: &RequestContext) -> Result<()> {
         let mut stack = Vec::new();
-        let ctx = RequestContext::new(TaskKind::DebugTool, DownloadBehavior::Error);
 
         stack.push((self.root_blk, String::new(), 0, 0, 0));
 
         let block_cursor = self.reader.block_cursor();
 
         while let Some((blknum, path, depth, child_idx, key_off)) = stack.pop() {
-            let blk = block_cursor.read_blk(self.start_blk + blknum, &ctx).await?;
+            let blk = block_cursor.read_blk(self.start_blk + blknum, ctx).await?;
             let buf: &[u8] = blk.as_ref();
             let node = OnDiskNode::<L>::deparse(buf)?;
 
@@ -835,6 +833,8 @@ pub(crate) mod tests {
     use rand::Rng;
 
     use super::*;
+    use crate::context::DownloadBehavior;
+    use crate::task_mgr::TaskKind;
     use crate::tenant::block_io::{BlockCursor, BlockLease, BlockReaderRef};
 
     #[derive(Clone, Default)]
@@ -869,7 +869,8 @@ pub(crate) mod tests {
         let mut disk = TestDisk::new();
         let mut writer = DiskBtreeBuilder::<_, 6>::new(&mut disk);
 
-        let ctx = RequestContext::new(TaskKind::UnitTest, DownloadBehavior::Error);
+        let ctx =
+            RequestContext::new(TaskKind::UnitTest, DownloadBehavior::Error).with_scope_unit_test();
 
         let all_keys: Vec<&[u8; 6]> = vec![
             b"xaaaaa", b"xaaaba", b"xaaaca", b"xabaaa", b"xababa", b"xabaca", b"xabada", b"xabadb",
@@ -887,7 +888,7 @@ pub(crate) mod tests {
 
         let reader = DiskBtreeReader::new(0, root_offset, disk);
 
-        reader.dump().await?;
+        reader.dump(&ctx).await?;
 
         // Test the `get` function on all the keys.
         for (key, val) in all_data.iter() {
@@ -979,7 +980,8 @@ pub(crate) mod tests {
     async fn lots_of_keys() -> Result<()> {
         let mut disk = TestDisk::new();
         let mut writer = DiskBtreeBuilder::<_, 8>::new(&mut disk);
-        let ctx = RequestContext::new(TaskKind::UnitTest, DownloadBehavior::Error);
+        let ctx =
+            RequestContext::new(TaskKind::UnitTest, DownloadBehavior::Error).with_scope_unit_test();
 
         const NUM_KEYS: u64 = 1000;
 
@@ -997,7 +999,7 @@ pub(crate) mod tests {
 
         let reader = DiskBtreeReader::new(0, root_offset, disk);
 
-        reader.dump().await?;
+        reader.dump(&ctx).await?;
 
         use std::sync::Mutex;
 
@@ -1167,7 +1169,8 @@ pub(crate) mod tests {
         // Build a tree from it
         let mut disk = TestDisk::new();
         let mut writer = DiskBtreeBuilder::<_, 26>::new(&mut disk);
-        let ctx = RequestContext::new(TaskKind::UnitTest, DownloadBehavior::Error);
+        let ctx =
+            RequestContext::new(TaskKind::UnitTest, DownloadBehavior::Error).with_scope_unit_test();
 
         for (key, val) in disk_btree_test_data::TEST_DATA {
             writer.append(&key, val)?;
@@ -1198,7 +1201,7 @@ pub(crate) mod tests {
             .await?;
         assert_eq!(count, disk_btree_test_data::TEST_DATA.len());
 
-        reader.dump().await?;
+        reader.dump(&ctx).await?;
 
         Ok(())
     }

pageserver/src/tenant/ephemeral_file.rs

@@ -351,7 +351,8 @@ mod tests {
         let timeline_id = TimelineId::from_str("22000000000000000000000000000000").unwrap();
         fs::create_dir_all(conf.timeline_path(&tenant_shard_id, &timeline_id))?;
 
-        let ctx = RequestContext::new(TaskKind::UnitTest, DownloadBehavior::Error);
+        let ctx =
+            RequestContext::new(TaskKind::UnitTest, DownloadBehavior::Error).with_scope_unit_test();
 
         Ok((conf, tenant_shard_id, timeline_id, ctx))
     }

pageserver/src/tenant/remote_timeline_client.rs

@@ -437,9 +437,13 @@ impl RemoteTimelineClient {
 
     /// Initialize the upload queue for the case where the remote storage is empty,
     /// i.e., it doesn't have an `IndexPart`.
+    ///
+    /// `rel_size_v2_status` needs to be carried over during branching, and that's why
+    /// it's passed in here.
     pub fn init_upload_queue_for_empty_remote(
         &self,
         local_metadata: &TimelineMetadata,
+        rel_size_v2_status: Option<RelSizeMigration>,
     ) -> anyhow::Result<()> {
         // Set the maximum number of inprogress tasks to the remote storage concurrency. There's
         // certainly no point in starting more upload tasks than this.
@@ -449,7 +453,9 @@ impl RemoteTimelineClient {
             .as_ref()
             .map_or(0, |r| r.concurrency_limit());
         let mut upload_queue = self.upload_queue.lock().unwrap();
-        upload_queue.initialize_empty_remote(local_metadata, inprogress_limit)?;
+        let initialized_queue =
+            upload_queue.initialize_empty_remote(local_metadata, inprogress_limit)?;
+        initialized_queue.dirty.rel_size_migration = rel_size_v2_status;
         self.update_remote_physical_size_gauge(None);
         info!("initialized upload queue as empty");
         Ok(())
@@ -948,6 +954,14 @@ impl RemoteTimelineClient {
         Ok(())
     }
 
+    /// Only used in the `patch_index_part` HTTP API to force trigger an index upload.
+    pub fn force_schedule_index_upload(self: &Arc<Self>) -> Result<(), NotInitialized> {
+        let mut guard = self.upload_queue.lock().unwrap();
+        let upload_queue = guard.initialized_mut()?;
+        self.schedule_index_upload(upload_queue);
+        Ok(())
+    }
+
     /// Launch an index-file upload operation in the background (internal function)
     fn schedule_index_upload(self: &Arc<Self>, upload_queue: &mut UploadQueueInitialized) {
         let disk_consistent_lsn = upload_queue.dirty.metadata.disk_consistent_lsn();

pageserver/src/tenant/secondary/downloader.rs

@@ -491,7 +491,10 @@ impl JobGenerator<PendingDownload, RunningDownload, CompleteDownload, DownloadCo
         let remote_storage = self.remote_storage.clone();
         let conf = self.tenant_manager.get_conf();
         let tenant_shard_id = *secondary_state.get_tenant_shard_id();
-        let download_ctx = self.root_ctx.attached_child();
+        let download_ctx = self
+            .root_ctx
+            .attached_child()
+            .with_scope_secondary_tenant(&tenant_shard_id);
         (RunningDownload { barrier }, Box::pin(async move {
             let _completion = completion;
 
@@ -771,6 +774,7 @@ impl<'a> TenantDownloader<'a> {
 
         // Download the layers in the heatmap
         for timeline in heatmap.timelines {
+            let ctx = &ctx.with_scope_secondary_timeline(tenant_shard_id, &timeline.timeline_id);
             let timeline_state = timeline_states
                 .remove(&timeline.timeline_id)
                 .expect("Just populated above");
@@ -869,8 +873,7 @@ impl<'a> TenantDownloader<'a> {
                 let heatmap_timeline = heatmap.timelines.get(heatmap_timeline_index).unwrap();
 
                 let layers_in_heatmap = heatmap_timeline
-                    .layers
-                    .iter()
+                    .hot_layers()
                     .map(|l| (&l.name, l.metadata.generation))
                     .collect::<HashSet<_>>();
                 let layers_on_disk = timeline_state
@@ -1015,7 +1018,8 @@ impl<'a> TenantDownloader<'a> {
         // Accumulate updates to the state
         let mut touched = Vec::new();
 
-        for layer in timeline.layers {
+        let timeline_id = timeline.timeline_id;
+        for layer in timeline.into_hot_layers() {
             if self.secondary_state.cancel.is_cancelled() {
                 tracing::debug!("Cancelled -- dropping out of layer loop");
                 return (Err(UpdateError::Cancelled), touched);
@@ -1040,7 +1044,7 @@ impl<'a> TenantDownloader<'a> {
             }
 
             match self
-                .download_layer(tenant_shard_id, &timeline.timeline_id, layer, ctx)
+                .download_layer(tenant_shard_id, &timeline_id, layer, ctx)
                 .await
             {
                 Ok(Some(layer)) => touched.push(layer),
@@ -1148,7 +1152,7 @@ impl<'a> TenantDownloader<'a> {
         let tenant_shard_id = self.secondary_state.get_tenant_shard_id();
         let timeline_id = timeline.timeline_id;
 
-        tracing::debug!(timeline_id=%timeline_id, "Downloading layers, {} in heatmap", timeline.layers.len());
+        tracing::debug!(timeline_id=%timeline_id, "Downloading layers, {} in heatmap", timeline.hot_layers().count());
 
         let (result, touched) = self
             .download_timeline_layers(tenant_shard_id, timeline, timeline_state, deadline, ctx)
@@ -1316,11 +1320,11 @@ async fn init_timeline_state(
     // As we iterate through layers found on disk, we will look up their metadata from this map.
     // Layers not present in metadata will be discarded.
     let heatmap_metadata: HashMap<&LayerName, &HeatMapLayer> =
-        heatmap.layers.iter().map(|l| (&l.name, l)).collect();
+        heatmap.hot_layers().map(|l| (&l.name, l)).collect();
 
     let last_heatmap_metadata: HashMap<&LayerName, &HeatMapLayer> =
         if let Some(last_heatmap) = last_heatmap {
-            last_heatmap.layers.iter().map(|l| (&l.name, l)).collect()
+            last_heatmap.hot_layers().map(|l| (&l.name, l)).collect()
         } else {
             HashMap::new()
         };

pageserver/src/tenant/secondary/heatmap.rs

@@ -42,7 +42,7 @@ pub(crate) struct HeatMapTimeline {
     #[serde_as(as = "DisplayFromStr")]
     pub(crate) timeline_id: TimelineId,
 
-    pub(crate) layers: Vec<HeatMapLayer>,
+    layers: Vec<HeatMapLayer>,
 }
 
 #[serde_as]
@@ -53,8 +53,10 @@ pub(crate) struct HeatMapLayer {
 
     #[serde_as(as = "TimestampSeconds<i64>")]
     pub(crate) access_time: SystemTime,
-    // TODO: an actual 'heat' score that would let secondary locations prioritize downloading
-    // the hottest layers, rather than trying to simply mirror whatever layers are on-disk on the primary.
+
+    #[serde(default)]
+    pub(crate) cold: bool, // TODO: an actual 'heat' score that would let secondary locations prioritize downloading
+                           // the hottest layers, rather than trying to simply mirror whatever layers are on-disk on the primary.
 }
 
 impl HeatMapLayer {
@@ -62,11 +64,13 @@ impl HeatMapLayer {
         name: LayerName,
         metadata: LayerFileMetadata,
         access_time: SystemTime,
+        cold: bool,
     ) -> Self {
         Self {
             name,
             metadata,
             access_time,
+            cold,
         }
     }
 }
@@ -78,6 +82,18 @@ impl HeatMapTimeline {
             layers,
         }
     }
+
+    pub(crate) fn into_hot_layers(self) -> impl Iterator<Item = HeatMapLayer> {
+        self.layers.into_iter().filter(|l| !l.cold)
+    }
+
+    pub(crate) fn hot_layers(&self) -> impl Iterator<Item = &HeatMapLayer> {
+        self.layers.iter().filter(|l| !l.cold)
+    }
+
+    pub(crate) fn all_layers(&self) -> impl Iterator<Item = &HeatMapLayer> {
+        self.layers.iter()
+    }
 }
 
 pub(crate) struct HeatMapStats {
@@ -92,7 +108,7 @@ impl HeatMapTenant {
             layers: 0,
         };
         for timeline in &self.timelines {
-            for layer in &timeline.layers {
+            for layer in timeline.hot_layers() {
                 stats.layers += 1;
                 stats.bytes += layer.metadata.file_size;
             }

pageserver/src/tenant/size.rs

@@ -474,7 +474,7 @@ async fn fill_logical_sizes(
             if cached_size.is_none() {
                 let timeline = Arc::clone(timeline_hash.get(&timeline_id).unwrap());
                 let parallel_size_calcs = Arc::clone(limit);
-                let ctx = ctx.attached_child();
+                let ctx = ctx.attached_child().with_scope_timeline(&timeline);
                 joinset.spawn(
                     calculate_logical_size(parallel_size_calcs, timeline, lsn, cause, ctx)
                         .in_current_span(),

pageserver/src/tenant/storage_layer/delta_layer.rs

@@ -1334,7 +1334,7 @@ impl DeltaLayerInner {
             block_reader,
         );
 
-        tree_reader.dump().await?;
+        tree_reader.dump(ctx).await?;
 
         let keys = self.index_entries(ctx).await?;
 
@@ -1972,6 +1972,7 @@ pub(crate) mod test {
             .create_test_timeline(TimelineId::generate(), Lsn(0x10), 14, ctx)
             .await
             .unwrap();
+        let ctx = &ctx.with_scope_timeline(&timeline);
 
         let initdb_layer = timeline
             .layers

pageserver/src/tenant/storage_layer/image_layer.rs

@@ -199,7 +199,7 @@ impl ImageLayerInner {
             block_reader,
         );
 
-        tree_reader.dump().await?;
+        tree_reader.dump(ctx).await?;
 
         tree_reader
             .visit(

pageserver/src/tenant/storage_layer/layer.rs

@@ -1563,10 +1563,10 @@ impl LayerInner {
 
         self.access_stats.record_residence_event();
 
-        self.status.as_ref().unwrap().send_replace(Status::Evicted);
-
         *self.last_evicted_at.lock().unwrap() = Some(std::time::Instant::now());
 
+        self.status.as_ref().unwrap().send_replace(Status::Evicted);
+
         Ok(())
     }
 

pageserver/src/tenant/storage_layer/layer/tests.rs

@@ -8,7 +8,6 @@ use utils::id::TimelineId;
 use super::failpoints::{Failpoint, FailpointKind};
 use super::*;
 use crate::context::DownloadBehavior;
-use crate::task_mgr::TaskKind;
 use crate::tenant::harness::{TenantHarness, test_img};
 use crate::tenant::storage_layer::{IoConcurrency, LayerVisibilityHint};
 
@@ -27,11 +26,9 @@ async fn smoke_test() {
     let h = TenantHarness::create("smoke_test").await.unwrap();
     let span = h.span();
     let download_span = span.in_scope(|| tracing::info_span!("downloading", timeline_id = 1));
-    let (tenant, _) = h.load().await;
+    let (tenant, ctx) = h.load().await;
     let io_concurrency = IoConcurrency::spawn_for_test();
 
-    let ctx = RequestContext::new(TaskKind::UnitTest, DownloadBehavior::Download);
-
     let image_layers = vec![(
         Lsn(0x40),
         vec![(
@@ -56,6 +53,7 @@ async fn smoke_test() {
         )
         .await
         .unwrap();
+    let ctx = &ctx.with_scope_timeline(&timeline);
 
     // Grab one of the timeline's layers to exercise in the test, and the other layer that is just
     // there to avoid the timeline being illegally empty
@@ -94,7 +92,7 @@ async fn smoke_test() {
                 controlfile_keyspace.clone(),
                 Lsn(0x10)..Lsn(0x11),
                 &mut data,
-                &ctx,
+                ctx,
             )
             .await
             .unwrap();
@@ -129,7 +127,7 @@ async fn smoke_test() {
                 controlfile_keyspace.clone(),
                 Lsn(0x10)..Lsn(0x11),
                 &mut data,
-                &ctx,
+                ctx,
             )
             .instrument(download_span.clone())
             .await
@@ -179,7 +177,7 @@ async fn smoke_test() {
 
     // plain downloading is rarely needed
     layer
-        .download_and_keep_resident(&ctx)
+        .download_and_keep_resident(ctx)
         .instrument(download_span)
         .await
         .unwrap();
@@ -341,6 +339,7 @@ fn read_wins_pending_eviction() {
             .create_test_timeline(TimelineId::generate(), Lsn(0x10), 14, &ctx)
             .await
             .unwrap();
+        let ctx = ctx.with_scope_timeline(&timeline);
 
         let layer = {
             let mut layers = {
@@ -473,6 +472,7 @@ fn multiple_pending_evictions_scenario(name: &'static str, in_order: bool) {
             .create_test_timeline(TimelineId::generate(), Lsn(0x10), 14, &ctx)
             .await
             .unwrap();
+        let ctx = ctx.with_scope_timeline(&timeline);
 
         let layer = {
             let mut layers = {
@@ -642,12 +642,12 @@ async fn cancelled_get_or_maybe_download_does_not_cancel_eviction() {
         .create_test_timeline(TimelineId::generate(), Lsn(0x10), 14, &ctx)
         .await
         .unwrap();
+    let ctx = ctx.with_scope_timeline(&timeline);
 
     // This test does downloads
     let ctx = RequestContextBuilder::extend(&ctx)
         .download_behavior(DownloadBehavior::Download)
         .build();
-
     let layer = {
         let mut layers = {
             let layers = timeline.layers.read().await;
@@ -727,6 +727,7 @@ async fn evict_and_wait_does_not_wait_for_download() {
         .create_test_timeline(TimelineId::generate(), Lsn(0x10), 14, &ctx)
         .await
         .unwrap();
+    let ctx = ctx.with_scope_timeline(&timeline);
 
     // This test does downloads
     let ctx = RequestContextBuilder::extend(&ctx)

pageserver/src/tenant/timeline.rs

@@ -99,7 +99,8 @@ use crate::disk_usage_eviction_task::{DiskUsageEvictionInfo, EvictionCandidate,
 use crate::keyspace::{KeyPartitioning, KeySpace};
 use crate::l0_flush::{self, L0FlushGlobalState};
 use crate::metrics::{
-    DELTAS_PER_READ_GLOBAL, LAYERS_PER_READ_GLOBAL, ScanLatencyOngoingRecording, TimelineMetrics,
+    DELTAS_PER_READ_GLOBAL, LAYERS_PER_READ_AMORTIZED_GLOBAL, LAYERS_PER_READ_BATCH_GLOBAL,
+    LAYERS_PER_READ_GLOBAL, ScanLatencyOngoingRecording, TimelineMetrics,
 };
 use crate::page_service::TenantManagerTypes;
 use crate::pgdatadir_mapping::{
@@ -286,7 +287,7 @@ pub struct Timeline {
     // The LSN of gc-compaction that was last applied to this timeline.
     gc_compaction_state: ArcSwap<Option<GcCompactionState>>,
 
-    pub(super) metrics: TimelineMetrics,
+    pub(crate) metrics: Arc<TimelineMetrics>,
 
     // `Timeline` doesn't write these metrics itself, but it manages the lifetime.  Code
     // in `crate::page_service` writes these metrics.
@@ -1330,10 +1331,6 @@ impl Timeline {
         // (this is a requirement, not a bug). Skip updating the metric in these cases
         // to avoid infinite results.
         if !results.is_empty() {
-            // Record the total number of layers visited towards each key in the batch. While some
-            // layers may not intersect with a given read, and the cost of layer visits are
-            // amortized across the batch, each visited layer contributes directly to the observed
-            // latency for every read in the batch, which is what we care about.
             if layers_visited >= Self::LAYERS_VISITED_WARN_THRESHOLD {
                 static LOG_PACER: Lazy<Mutex<RateLimit>> =
                     Lazy::new(|| Mutex::new(RateLimit::new(Duration::from_secs(60))));
@@ -1348,9 +1345,23 @@ impl Timeline {
                 });
             }
 
+            // Records the number of layers visited in a few different ways:
+            //
+            // * LAYERS_PER_READ: all layers count towards every read in the batch, because each
+            //   layer directly affects its observed latency.
+            //
+            // * LAYERS_PER_READ_BATCH: all layers count towards each batch, to get the per-batch
+            //   layer visits and access cost.
+            //
+            // * LAYERS_PER_READ_AMORTIZED: the average layer count per read, to get the amortized
+            //   read amplification after batching.
+            let layers_visited = layers_visited as f64;
+            let avg_layers_visited = layers_visited / results.len() as f64;
+            LAYERS_PER_READ_BATCH_GLOBAL.observe(layers_visited);
             for _ in &results {
-                self.metrics.layers_per_read.observe(layers_visited as f64);
-                LAYERS_PER_READ_GLOBAL.observe(layers_visited as f64);
+                self.metrics.layers_per_read.observe(layers_visited);
+                LAYERS_PER_READ_GLOBAL.observe(layers_visited);
+                LAYERS_PER_READ_AMORTIZED_GLOBAL.observe(avg_layers_visited);
             }
         }
 
@@ -2674,14 +2685,14 @@ impl Timeline {
         }
 
         Arc::new_cyclic(|myself| {
-            let metrics = TimelineMetrics::new(
+            let metrics = Arc::new(TimelineMetrics::new(
                 &tenant_shard_id,
                 &timeline_id,
                 crate::metrics::EvictionsWithLowResidenceDurationBuilder::new(
                     "mtime",
                     evictions_low_residence_duration_metric_threshold,
                 ),
-            );
+            ));
             let aux_file_metrics = metrics.aux_file_size_gauge.clone();
 
             let mut result = Timeline {
@@ -2865,7 +2876,7 @@ impl Timeline {
             "layer flush task",
             async move {
                 let _guard = guard;
-                let background_ctx = RequestContext::todo_child(TaskKind::LayerFlushTask, DownloadBehavior::Error);
+                let background_ctx = RequestContext::todo_child(TaskKind::LayerFlushTask, DownloadBehavior::Error).with_scope_timeline(&self_clone);
                 self_clone.flush_loop(layer_flush_start_rx, &background_ctx).await;
                 let mut flush_loop_state = self_clone.flush_loop_state.lock().unwrap();
                 assert!(matches!(*flush_loop_state, FlushLoopState::Running{..}));
@@ -3637,7 +3648,7 @@ impl Timeline {
         let visible_non_resident = match previous_heatmap.as_deref() {
             Some(PreviousHeatmap::Active {
                 heatmap, read_at, ..
-            }) => Some(heatmap.layers.iter().filter_map(|hl| {
+            }) => Some(heatmap.all_layers().filter_map(|hl| {
                 let desc: PersistentLayerDesc = hl.name.clone().into();
                 let layer = guard.try_get_from_key(&desc.key())?;
 
@@ -3653,7 +3664,7 @@ impl Timeline {
                     return None;
                 }
 
-                Some((desc, hl.metadata.clone(), hl.access_time))
+                Some((desc, hl.metadata.clone(), hl.access_time, hl.cold))
             })),
             Some(PreviousHeatmap::Obsolete) => None,
             None => None,
@@ -3669,6 +3680,7 @@ impl Timeline {
                         layer.layer_desc().clone(),
                         layer.metadata(),
                         last_activity_ts,
+                        false, // these layers are not cold
                     ))
                 }
                 LayerVisibilityHint::Covered => {
@@ -3695,12 +3707,14 @@ impl Timeline {
         // Sort layers in order of which to download first.  For a large set of layers to download, we
         // want to prioritize those layers which are most likely to still be in the resident many minutes
         // or hours later:
+        // - Cold layers go last for convenience when a human inspects the heatmap.
         // - Download L0s last, because they churn the fastest: L0s on a fast-writing tenant might
         //   only exist for a few minutes before being compacted into L1s.
         // - For L1 & image layers, download most recent LSNs first: the older the LSN, the sooner
         //   the layer is likely to be covered by an image layer during compaction.
-        layers.sort_by_key(|(desc, _meta, _atime)| {
+        layers.sort_by_key(|(desc, _meta, _atime, cold)| {
             std::cmp::Reverse((
+                *cold,
                 !LayerMap::is_l0(&desc.key_range, desc.is_delta),
                 desc.lsn_range.end,
             ))
@@ -3708,7 +3722,9 @@ impl Timeline {
 
         let layers = layers
             .into_iter()
-            .map(|(desc, meta, atime)| HeatMapLayer::new(desc.layer_name(), meta, atime))
+            .map(|(desc, meta, atime, cold)| {
+                HeatMapLayer::new(desc.layer_name(), meta, atime, cold)
+            })
             .collect();
 
         Some(HeatMapTimeline::new(self.timeline_id, layers))
@@ -3728,6 +3744,7 @@ impl Timeline {
                 name: vl.layer_desc().layer_name(),
                 metadata: vl.metadata(),
                 access_time: now,
+                cold: true,
             };
             heatmap_layers.push(hl);
         }
@@ -7029,6 +7046,7 @@ mod tests {
 
     use pageserver_api::key::Key;
     use pageserver_api::value::Value;
+    use std::iter::Iterator;
     use tracing::Instrument;
     use utils::id::TimelineId;
     use utils::lsn::Lsn;
@@ -7042,8 +7060,8 @@ mod tests {
     use crate::tenant::{PreviousHeatmap, Timeline};
 
     fn assert_heatmaps_have_same_layers(lhs: &HeatMapTimeline, rhs: &HeatMapTimeline) {
-        assert_eq!(lhs.layers.len(), rhs.layers.len());
-        let lhs_rhs = lhs.layers.iter().zip(rhs.layers.iter());
+        assert_eq!(lhs.all_layers().count(), rhs.all_layers().count());
+        let lhs_rhs = lhs.all_layers().zip(rhs.all_layers());
         for (l, r) in lhs_rhs {
             assert_eq!(l.name, r.name);
             assert_eq!(l.metadata, r.metadata);
@@ -7109,6 +7127,7 @@ mod tests {
             )
             .await
             .unwrap();
+        let ctx = &ctx.with_scope_timeline(&timeline);
 
         // Layer visibility is an input to heatmap generation, so refresh it first
         timeline.update_layer_visibility().await.unwrap();
@@ -7121,10 +7140,11 @@ mod tests {
         assert_eq!(heatmap.timeline_id, timeline.timeline_id);
 
         // L0 should come last
-        assert_eq!(heatmap.layers.last().unwrap().name, l0_delta.layer_name());
+        let heatmap_layers = heatmap.all_layers().collect::<Vec<_>>();
+        assert_eq!(heatmap_layers.last().unwrap().name, l0_delta.layer_name());
 
         let mut last_lsn = Lsn::MAX;
-        for layer in &heatmap.layers {
+        for layer in heatmap_layers {
             // Covered layer should be omitted
             assert!(layer.name != covered_delta.layer_name());
 
@@ -7173,7 +7193,7 @@ mod tests {
 
             eprintln!("Downloading {layer} and re-generating heatmap");
 
-            let ctx = &RequestContextBuilder::extend(&ctx)
+            let ctx = &RequestContextBuilder::extend(ctx)
                 .download_behavior(crate::context::DownloadBehavior::Download)
                 .build();
 
@@ -7253,7 +7273,7 @@ mod tests {
             .expect("Infallible while timeline is not shut down");
 
         // Both layers should be in the heatmap
-        assert!(!heatmap.layers.is_empty());
+        assert!(heatmap.all_layers().count() > 0);
 
         // Now simulate a migration.
         timeline
@@ -7279,7 +7299,7 @@ mod tests {
             .await
             .expect("Infallible while timeline is not shut down");
 
-        assert!(post_eviction_heatmap.layers.is_empty());
+        assert_eq!(post_eviction_heatmap.all_layers().count(), 0);
         assert!(matches!(
             timeline.previous_heatmap.load().as_deref(),
             Some(PreviousHeatmap::Obsolete)

pageserver/src/tenant/timeline/compaction.rs

@@ -7,6 +7,7 @@
 use std::collections::{BinaryHeap, HashMap, HashSet, VecDeque};
 use std::ops::{Deref, Range};
 use std::sync::Arc;
+use std::time::Instant;
 
 use super::layer_manager::LayerManager;
 use super::{
@@ -15,10 +16,11 @@ use super::{
     Timeline,
 };
 
-use anyhow::{Context, anyhow, bail};
+use anyhow::{Context, anyhow};
 use bytes::Bytes;
 use enumset::EnumSet;
 use fail::fail_point;
+use futures::FutureExt;
 use itertools::Itertools;
 use once_cell::sync::Lazy;
 use pageserver_api::config::tenant_conf_defaults::DEFAULT_CHECKPOINT_DISTANCE;
@@ -234,6 +236,12 @@ impl GcCompactionQueue {
             // it enough in staging yet.
             return Ok(());
         }
+        if timeline.get_gc_compaction_watermark() == Lsn::INVALID {
+            // If the gc watermark is not set, we don't need to trigger auto compaction.
+            // This check is the same as in `gc_compaction_split_jobs` but we don't log
+            // here and we can also skip the computation of the trigger condition earlier.
+            return Ok(());
+        }
 
         let Ok(permit) = CONCURRENT_GC_COMPACTION_TASKS.clone().try_acquire_owned() else {
             // Only allow one compaction run at a time. TODO: As we do `try_acquire_owned`, we cannot ensure
@@ -357,8 +365,7 @@ impl GcCompactionQueue {
                 GcCompactJob::from_compact_options(options.clone()),
                 options.sub_compaction_max_job_size_mb,
             )
-            .await
-            .map_err(CompactionError::Other)?;
+            .await?;
         if jobs.is_empty() {
             info!("no jobs to run, skipping scheduled compaction task");
             self.notify_and_unblock(id);
@@ -386,6 +393,9 @@ impl GcCompactionQueue {
                 if job.dry_run {
                     flags |= CompactFlags::DryRun;
                 }
+                if options.flags.contains(CompactFlags::NoYield) {
+                    flags |= CompactFlags::NoYield;
+                }
                 let options = CompactOptions {
                     flags,
                     sub_compaction: false,
@@ -437,6 +447,7 @@ impl GcCompactionQueue {
             ));
         };
         let has_pending_tasks;
+        let mut yield_for_l0 = false;
         let Some((id, item)) = ({
             let mut guard = self.inner.lock().unwrap();
             if let Some((id, item)) = guard.queued.pop_front() {
@@ -486,13 +497,23 @@ impl GcCompactionQueue {
                         let mut guard = self.inner.lock().unwrap();
                         guard.guards.entry(id).or_default().gc_guard = Some(gc_guard);
                     }
-                    let _ = timeline.compact_with_options(cancel, options, ctx).await?;
+                    let compaction_result =
+                        timeline.compact_with_options(cancel, options, ctx).await?;
                     self.notify_and_unblock(id);
+                    if compaction_result == CompactionOutcome::YieldForL0 {
+                        yield_for_l0 = true;
+                    }
                 }
             }
             GcCompactionQueueItem::SubCompactionJob(options) => {
                 // TODO: error handling, clear the queue if any task fails?
-                let _ = timeline.compact_with_options(cancel, options, ctx).await?;
+                let compaction_result = timeline.compact_with_options(cancel, options, ctx).await?;
+                if compaction_result == CompactionOutcome::YieldForL0 {
+                    // We will permenantly give up a task if we yield for L0 compaction: the preempted subcompaction job won't be running
+                    // again. This ensures that we don't keep doing duplicated work within gc-compaction. Not directly returning here because
+                    // we need to clean things up before returning from the function.
+                    yield_for_l0 = true;
+                }
             }
             GcCompactionQueueItem::Notify(id, l2_lsn) => {
                 self.notify_and_unblock(id);
@@ -521,7 +542,10 @@ impl GcCompactionQueue {
             let mut guard = self.inner.lock().unwrap();
             guard.running = None;
         }
-        Ok(if has_pending_tasks {
+        Ok(if yield_for_l0 {
+            tracing::info!("give up gc-compaction: yield for L0 compaction");
+            CompactionOutcome::YieldForL0
+        } else if has_pending_tasks {
             CompactionOutcome::Pending
         } else {
             CompactionOutcome::Done
@@ -719,17 +743,41 @@ struct CompactionStatisticsNumSize {
 
 #[derive(Debug, Serialize, Default)]
 pub struct CompactionStatistics {
+    /// Delta layer visited (maybe compressed, physical size)
     delta_layer_visited: CompactionStatisticsNumSize,
+    /// Image layer visited (maybe compressed, physical size)
     image_layer_visited: CompactionStatisticsNumSize,
+    /// Delta layer produced (maybe compressed, physical size)
     delta_layer_produced: CompactionStatisticsNumSize,
+    /// Image layer produced (maybe compressed, physical size)
     image_layer_produced: CompactionStatisticsNumSize,
-    num_delta_layer_discarded: usize,
-    num_image_layer_discarded: usize,
+    /// Delta layer discarded (maybe compressed, physical size of the layer being discarded instead of the original layer)
+    delta_layer_discarded: CompactionStatisticsNumSize,
+    /// Image layer discarded (maybe compressed, physical size of the layer being discarded instead of the original layer)
+    image_layer_discarded: CompactionStatisticsNumSize,
     num_unique_keys_visited: usize,
+    /// Delta visited (uncompressed, original size)
     wal_keys_visited: CompactionStatisticsNumSize,
+    /// Image visited (uncompressed, original size)
     image_keys_visited: CompactionStatisticsNumSize,
+    /// Delta produced (uncompressed, original size)
     wal_produced: CompactionStatisticsNumSize,
+    /// Image produced (uncompressed, original size)
     image_produced: CompactionStatisticsNumSize,
+
+    // Time spent in each phase
+    time_acquire_lock_secs: f64,
+    time_analyze_secs: f64,
+    time_download_layer_secs: f64,
+    time_main_loop_secs: f64,
+    time_final_phase_secs: f64,
+    time_total_secs: f64,
+
+    // Summary
+    /// Ratio of the key-value size before/after gc-compaction.
+    uncompressed_size_ratio: f64,
+    /// Ratio of the physical size before/after gc-compaction.
+    physical_size_ratio: f64,
 }
 
 impl CompactionStatistics {
@@ -779,11 +827,13 @@ impl CompactionStatistics {
         self.image_produced.num += 1;
         self.image_produced.size += val.len() as u64 + Self::estimated_size_of_key() as u64;
     }
-    fn discard_delta_layer(&mut self) {
-        self.num_delta_layer_discarded += 1;
+    fn discard_delta_layer(&mut self, original_size: u64) {
+        self.delta_layer_discarded.num += 1;
+        self.delta_layer_discarded.size += original_size;
     }
-    fn discard_image_layer(&mut self) {
-        self.num_image_layer_discarded += 1;
+    fn discard_image_layer(&mut self, original_size: u64) {
+        self.image_layer_discarded.num += 1;
+        self.image_layer_discarded.size += original_size;
     }
     fn produce_delta_layer(&mut self, size: u64) {
         self.delta_layer_produced.num += 1;
@@ -793,6 +843,19 @@ impl CompactionStatistics {
         self.image_layer_produced.num += 1;
         self.image_layer_produced.size += size;
     }
+    fn finalize(&mut self) {
+        let original_key_value_size = self.image_keys_visited.size + self.wal_keys_visited.size;
+        let produced_key_value_size = self.image_produced.size + self.wal_produced.size;
+        self.uncompressed_size_ratio =
+            original_key_value_size as f64 / (produced_key_value_size as f64 + 1.0); // avoid div by 0
+        let original_physical_size = self.image_layer_visited.size + self.delta_layer_visited.size;
+        let produced_physical_size = self.image_layer_produced.size
+            + self.delta_layer_produced.size
+            + self.image_layer_discarded.size
+            + self.delta_layer_discarded.size; // Also include the discarded layers to make the ratio accurate
+        self.physical_size_ratio =
+            original_physical_size as f64 / (produced_physical_size as f64 + 1.0); // avoid div by 0
+    }
 }
 
 #[derive(Default, Debug, Clone, Copy, PartialEq, Eq)]
@@ -825,9 +888,7 @@ impl Timeline {
             .flags
             .contains(CompactFlags::EnhancedGcBottomMostCompaction)
         {
-            self.compact_with_gc(cancel, options, ctx)
-                .await
-                .map_err(CompactionError::Other)?;
+            self.compact_with_gc(cancel, options, ctx).await?;
             return Ok(CompactionOutcome::Done);
         }
 
@@ -1033,7 +1094,7 @@ impl Timeline {
         let latest_gc_cutoff = self.get_applied_gc_cutoff_lsn();
 
         tracing::info!(
-            "latest_gc_cutoff: {}, pitr cutoff {}",
+            "starting shard ancestor compaction, latest_gc_cutoff: {}, pitr cutoff {}",
             *latest_gc_cutoff,
             self.gc_info.read().unwrap().cutoffs.time
         );
@@ -1062,6 +1123,7 @@ impl Timeline {
                     // Expensive, exhaustive check of keys in this layer: this guards against ShardedRange's calculations being
                     // wrong.  If ShardedRange claims the local page count is zero, then no keys in this layer
                     // should be !is_key_disposable()
+                    // TODO: exclude sparse keyspace from this check, otherwise it will infinitely loop.
                     let range = layer_desc.get_key_range();
                     let mut key = range.start;
                     while key < range.end {
@@ -2345,12 +2407,19 @@ impl Timeline {
     async fn check_compaction_space(
         self: &Arc<Self>,
         layer_selection: &[Layer],
-    ) -> anyhow::Result<()> {
-        let available_space = self.check_available_space().await?;
+    ) -> Result<(), CompactionError> {
+        let available_space = self
+            .check_available_space()
+            .await
+            .map_err(CompactionError::Other)?;
         let mut remote_layer_size = 0;
         let mut all_layer_size = 0;
         for layer in layer_selection {
-            let needs_download = layer.needs_download().await?;
+            let needs_download = layer
+                .needs_download()
+                .await
+                .context("failed to check if layer needs download")
+                .map_err(CompactionError::Other)?;
             if needs_download.is_some() {
                 remote_layer_size += layer.layer_desc().file_size;
             }
@@ -2359,14 +2428,14 @@ impl Timeline {
         let allocated_space = (available_space as f64 * 0.8) as u64; /* reserve 20% space for other tasks */
         if all_layer_size /* space needed for newly-generated file */ + remote_layer_size /* space for downloading layers */ > allocated_space
         {
-            return Err(anyhow!(
+            return Err(CompactionError::Other(anyhow!(
                 "not enough space for compaction: available_space={}, allocated_space={}, all_layer_size={}, remote_layer_size={}, required_space={}",
                 available_space,
                 allocated_space,
                 all_layer_size,
                 remote_layer_size,
                 all_layer_size + remote_layer_size
-            ));
+            )));
         }
         Ok(())
     }
@@ -2397,7 +2466,7 @@ impl Timeline {
         self: &Arc<Self>,
         job: GcCompactJob,
         sub_compaction_max_job_size_mb: Option<u64>,
-    ) -> anyhow::Result<Vec<GcCompactJob>> {
+    ) -> Result<Vec<GcCompactJob>, CompactionError> {
         let compact_below_lsn = if job.compact_lsn_range.end != Lsn::MAX {
             job.compact_lsn_range.end
         } else {
@@ -2548,9 +2617,10 @@ impl Timeline {
         cancel: &CancellationToken,
         options: CompactOptions,
         ctx: &RequestContext,
-    ) -> anyhow::Result<()> {
+    ) -> Result<CompactionOutcome, CompactionError> {
         let sub_compaction = options.sub_compaction;
         let job = GcCompactJob::from_compact_options(options.clone());
+        let no_yield = options.flags.contains(CompactFlags::NoYield);
         if sub_compaction {
             info!(
                 "running enhanced gc bottom-most compaction with sub-compaction, splitting compaction jobs"
@@ -2565,14 +2635,15 @@ impl Timeline {
                     idx + 1,
                     jobs_len
                 );
-                self.compact_with_gc_inner(cancel, job, ctx).await?;
+                self.compact_with_gc_inner(cancel, job, ctx, no_yield)
+                    .await?;
             }
             if jobs_len == 0 {
                 info!("no jobs to run, skipping gc bottom-most compaction");
             }
-            return Ok(());
+            return Ok(CompactionOutcome::Done);
         }
-        self.compact_with_gc_inner(cancel, job, ctx).await
+        self.compact_with_gc_inner(cancel, job, ctx, no_yield).await
     }
 
     async fn compact_with_gc_inner(
@@ -2580,19 +2651,25 @@ impl Timeline {
         cancel: &CancellationToken,
         job: GcCompactJob,
         ctx: &RequestContext,
-    ) -> anyhow::Result<()> {
+        no_yield: bool,
+    ) -> Result<CompactionOutcome, CompactionError> {
         // Block other compaction/GC tasks from running for now. GC-compaction could run along
         // with legacy compaction tasks in the future. Always ensure the lock order is compaction -> gc.
         // Note that we already acquired the compaction lock when the outer `compact` function gets called.
 
+        let timer = Instant::now();
+        let begin_timer = timer;
+
         let gc_lock = async {
             tokio::select! {
                 guard = self.gc_lock.lock() => Ok(guard),
-                // TODO: refactor to CompactionError to correctly pass cancelled error
-                _ = cancel.cancelled() => Err(anyhow!("cancelled")),
+                _ = cancel.cancelled() => Err(CompactionError::ShuttingDown),
             }
         };
 
+        let time_acquire_lock = timer.elapsed();
+        let timer = Instant::now();
+
         let gc_lock = crate::timed(
             gc_lock,
             "acquires gc lock",
@@ -2644,7 +2721,7 @@ impl Timeline {
                         tracing::warn!(
                             "no layers to compact with gc: gc_cutoff not generated yet, skipping gc bottom-most compaction"
                         );
-                        return Ok(());
+                        return Ok(CompactionOutcome::Skipped);
                     }
                     real_gc_cutoff
                 } else {
@@ -2682,7 +2759,7 @@ impl Timeline {
                     "no layers to compact with gc: no historic layers below gc_cutoff, gc_cutoff={}",
                     gc_cutoff
                 );
-                return Ok(());
+                return Ok(CompactionOutcome::Done);
             };
             // Next, if the user specifies compact_lsn_range.start, we need to filter some layers out. All the layers (strictly) below
             // the min_layer_lsn computed as below will be filtered out and the data will be accessed using the normal read path, as if
@@ -2703,7 +2780,7 @@ impl Timeline {
                     "no layers to compact with gc: no historic layers above compact_above_lsn, compact_above_lsn={}",
                     compact_lsn_range.end
                 );
-                return Ok(());
+                return Ok(CompactionOutcome::Done);
             };
             // Then, pick all the layers that are below the max_layer_lsn. This is to ensure we can pick all single-key
             // layers to compact.
@@ -2729,7 +2806,7 @@ impl Timeline {
                     "no layers to compact with gc: no layers within the key range, gc_cutoff={}, key_range={}..{}",
                     gc_cutoff, compact_key_range.start, compact_key_range.end
                 );
-                return Ok(());
+                return Ok(CompactionOutcome::Done);
             }
             retain_lsns_below_horizon.sort();
             GcCompactionJobDescription {
@@ -2782,6 +2859,9 @@ impl Timeline {
             has_data_below,
         );
 
+        let time_analyze = timer.elapsed();
+        let timer = Instant::now();
+
         for layer in &job_desc.selected_layers {
             debug!("read layer: {}", layer.layer_desc().key());
         }
@@ -2810,10 +2890,10 @@ impl Timeline {
             .map(|layer| layer.layer_desc().layer_name())
             .collect_vec();
         if let Some(err) = check_valid_layermap(&layer_names) {
-            bail!(
+            return Err(CompactionError::Other(anyhow!(
                 "gc-compaction layer map check failed because {}, cannot proceed with compaction due to potential data loss",
                 err
-            );
+            )));
         }
         // The maximum LSN we are processing in this compaction loop
         let end_lsn = job_desc
@@ -2828,11 +2908,37 @@ impl Timeline {
         let mut total_downloaded_size = 0;
         let mut total_layer_size = 0;
         for layer in &job_desc.selected_layers {
-            if layer.needs_download().await?.is_some() {
+            if layer
+                .needs_download()
+                .await
+                .context("failed to check if layer needs download")
+                .map_err(CompactionError::Other)?
+                .is_some()
+            {
                 total_downloaded_size += layer.layer_desc().file_size;
             }
             total_layer_size += layer.layer_desc().file_size;
-            let resident_layer = layer.download_and_keep_resident(ctx).await?;
+            if cancel.is_cancelled() {
+                return Err(CompactionError::ShuttingDown);
+            }
+            if !no_yield {
+                let should_yield = self
+                    .l0_compaction_trigger
+                    .notified()
+                    .now_or_never()
+                    .is_some();
+                if should_yield {
+                    tracing::info!(
+                        "preempt gc-compaction when downloading layers: too many L0 layers"
+                    );
+                    return Ok(CompactionOutcome::YieldForL0);
+                }
+            }
+            let resident_layer = layer
+                .download_and_keep_resident(ctx)
+                .await
+                .context("failed to download and keep resident layer")
+                .map_err(CompactionError::Other)?;
             downloaded_layers.push(resident_layer);
         }
         info!(
@@ -2843,19 +2949,36 @@ impl Timeline {
         );
         for resident_layer in &downloaded_layers {
             if resident_layer.layer_desc().is_delta() {
-                let layer = resident_layer.get_as_delta(ctx).await?;
+                let layer = resident_layer
+                    .get_as_delta(ctx)
+                    .await
+                    .context("failed to get delta layer")
+                    .map_err(CompactionError::Other)?;
                 delta_layers.push(layer);
             } else {
-                let layer = resident_layer.get_as_image(ctx).await?;
+                let layer = resident_layer
+                    .get_as_image(ctx)
+                    .await
+                    .context("failed to get image layer")
+                    .map_err(CompactionError::Other)?;
                 image_layers.push(layer);
             }
         }
-        let (dense_ks, sparse_ks) = self.collect_gc_compaction_keyspace().await?;
+        let (dense_ks, sparse_ks) = self
+            .collect_gc_compaction_keyspace()
+            .await
+            .context("failed to collect gc compaction keyspace")
+            .map_err(CompactionError::Other)?;
         let mut merge_iter = FilterIterator::create(
             MergeIterator::create(&delta_layers, &image_layers, ctx),
             dense_ks,
             sparse_ks,
-        )?;
+        )
+        .context("failed to create filter iterator")
+        .map_err(CompactionError::Other)?;
+
+        let time_download_layer = timer.elapsed();
+        let timer = Instant::now();
 
         // Step 2: Produce images+deltas.
         let mut accumulated_values = Vec::new();
@@ -2874,7 +2997,9 @@ impl Timeline {
                     self.get_compaction_target_size(),
                     ctx,
                 )
-                .await?,
+                .await
+                .context("failed to create image layer writer")
+                .map_err(CompactionError::Other)?,
             )
         } else {
             None
@@ -2887,7 +3012,9 @@ impl Timeline {
             lowest_retain_lsn..end_lsn,
             self.get_compaction_target_size(),
         )
-        .await?;
+        .await
+        .context("failed to create delta layer writer")
+        .map_err(CompactionError::Other)?;
 
         #[derive(Default)]
         struct RewritingLayers {
@@ -2927,9 +3054,33 @@ impl Timeline {
         // the key and LSN range are determined. However, to keep things simple here, we still
         // create this writer, and discard the writer in the end.
 
-        while let Some(((key, lsn, val), desc)) = merge_iter.next_with_trace().await? {
+        let mut keys_processed = 0;
+
+        while let Some(((key, lsn, val), desc)) = merge_iter
+            .next_with_trace()
+            .await
+            .context("failed to get next key-value pair")
+            .map_err(CompactionError::Other)?
+        {
             if cancel.is_cancelled() {
-                return Err(anyhow!("cancelled")); // TODO: refactor to CompactionError and pass cancel error
+                return Err(CompactionError::ShuttingDown);
+            }
+
+            if !no_yield {
+                keys_processed += 1;
+                if keys_processed % 1000 == 0 {
+                    let should_yield = self
+                        .l0_compaction_trigger
+                        .notified()
+                        .now_or_never()
+                        .is_some();
+                    if should_yield {
+                        tracing::info!(
+                            "preempt gc-compaction in the main loop: too many L0 layers"
+                        );
+                        return Ok(CompactionOutcome::YieldForL0);
+                    }
+                }
             }
             if self.shard_identity.is_key_disposable(&key) {
                 // If this shard does not need to store this key, simply skip it.
@@ -2960,7 +3111,9 @@ impl Timeline {
                                 desc.lsn_range.clone(),
                                 ctx,
                             )
-                            .await?,
+                            .await
+                            .context("failed to create delta layer writer")
+                            .map_err(CompactionError::Other)?,
                         );
                     }
                     rewriter.before.as_mut().unwrap()
@@ -2975,14 +3128,20 @@ impl Timeline {
                                 desc.lsn_range.clone(),
                                 ctx,
                             )
-                            .await?,
+                            .await
+                            .context("failed to create delta layer writer")
+                            .map_err(CompactionError::Other)?,
                         );
                     }
                     rewriter.after.as_mut().unwrap()
                 } else {
                     unreachable!()
                 };
-                rewriter.put_value(key, lsn, val, ctx).await?;
+                rewriter
+                    .put_value(key, lsn, val, ctx)
+                    .await
+                    .context("failed to put value")
+                    .map_err(CompactionError::Other)?;
                 continue;
             }
             match val {
@@ -3005,9 +3164,13 @@ impl Timeline {
                         &job_desc.retain_lsns_below_horizon,
                         COMPACTION_DELTA_THRESHOLD,
                         get_ancestor_image(self, *last_key, ctx, has_data_below, lowest_retain_lsn)
-                            .await?,
+                            .await
+                            .context("failed to get ancestor image")
+                            .map_err(CompactionError::Other)?,
                     )
-                    .await?;
+                    .await
+                    .context("failed to generate key retention")
+                    .map_err(CompactionError::Other)?;
                 retention
                     .pipe_to(
                         *last_key,
@@ -3016,7 +3179,9 @@ impl Timeline {
                         &mut stat,
                         ctx,
                     )
-                    .await?;
+                    .await
+                    .context("failed to pipe to delta layer writer")
+                    .map_err(CompactionError::Other)?;
                 accumulated_values.clear();
                 *last_key = key;
                 accumulated_values.push((key, lsn, val));
@@ -3034,9 +3199,14 @@ impl Timeline {
                 job_desc.gc_cutoff,
                 &job_desc.retain_lsns_below_horizon,
                 COMPACTION_DELTA_THRESHOLD,
-                get_ancestor_image(self, last_key, ctx, has_data_below, lowest_retain_lsn).await?,
+                get_ancestor_image(self, last_key, ctx, has_data_below, lowest_retain_lsn)
+                    .await
+                    .context("failed to get ancestor image")
+                    .map_err(CompactionError::Other)?,
             )
-            .await?;
+            .await
+            .context("failed to generate key retention")
+            .map_err(CompactionError::Other)?;
         retention
             .pipe_to(
                 last_key,
@@ -3045,21 +3215,36 @@ impl Timeline {
                 &mut stat,
                 ctx,
             )
-            .await?;
+            .await
+            .context("failed to pipe to delta layer writer")
+            .map_err(CompactionError::Other)?;
         // end: move the above part to the loop body
 
+        let time_main_loop = timer.elapsed();
+        let timer = Instant::now();
+
         let mut rewrote_delta_layers = Vec::new();
         for (key, writers) in delta_layer_rewriters {
             if let Some(delta_writer_before) = writers.before {
                 let (desc, path) = delta_writer_before
                     .finish(job_desc.compaction_key_range.start, ctx)
-                    .await?;
-                let layer = Layer::finish_creating(self.conf, self, desc, &path)?;
+                    .await
+                    .context("failed to finish delta layer writer")
+                    .map_err(CompactionError::Other)?;
+                let layer = Layer::finish_creating(self.conf, self, desc, &path)
+                    .context("failed to finish creating delta layer")
+                    .map_err(CompactionError::Other)?;
                 rewrote_delta_layers.push(layer);
             }
             if let Some(delta_writer_after) = writers.after {
-                let (desc, path) = delta_writer_after.finish(key.key_range.end, ctx).await?;
-                let layer = Layer::finish_creating(self.conf, self, desc, &path)?;
+                let (desc, path) = delta_writer_after
+                    .finish(key.key_range.end, ctx)
+                    .await
+                    .context("failed to finish delta layer writer")
+                    .map_err(CompactionError::Other)?;
+                let layer = Layer::finish_creating(self.conf, self, desc, &path)
+                    .context("failed to finish creating delta layer")
+                    .map_err(CompactionError::Other)?;
                 rewrote_delta_layers.push(layer);
             }
         }
@@ -3074,7 +3259,9 @@ impl Timeline {
                 let end_key = job_desc.compaction_key_range.end;
                 writer
                     .finish_with_discard_fn(self, ctx, end_key, discard)
-                    .await?
+                    .await
+                    .context("failed to finish image layer writer")
+                    .map_err(CompactionError::Other)?
             } else {
                 drop(writer);
                 Vec::new()
@@ -3086,7 +3273,9 @@ impl Timeline {
         let produced_delta_layers = if !dry_run {
             delta_layer_writer
                 .finish_with_discard_fn(self, ctx, discard)
-                .await?
+                .await
+                .context("failed to finish delta layer writer")
+                .map_err(CompactionError::Other)?
         } else {
             drop(delta_layer_writer);
             Vec::new()
@@ -3098,6 +3287,13 @@ impl Timeline {
         let mut keep_layers = HashSet::new();
         let produced_delta_layers_len = produced_delta_layers.len();
         let produced_image_layers_len = produced_image_layers.len();
+
+        let layer_selection_by_key = job_desc
+            .selected_layers
+            .iter()
+            .map(|l| (l.layer_desc().key(), l.layer_desc().clone()))
+            .collect::<HashMap<_, _>>();
+
         for action in produced_delta_layers {
             match action {
                 BatchWriterResult::Produced(layer) => {
@@ -3111,8 +3307,16 @@ impl Timeline {
                     if cfg!(debug_assertions) {
                         info!("discarded delta layer: {}", l);
                     }
+                    if let Some(layer_desc) = layer_selection_by_key.get(&l) {
+                        stat.discard_delta_layer(layer_desc.file_size());
+                    } else {
+                        tracing::warn!(
+                            "discarded delta layer not in layer_selection: {}, produced a layer outside of the compaction key range?",
+                            l
+                        );
+                        stat.discard_delta_layer(0);
+                    }
                     keep_layers.insert(l);
-                    stat.discard_delta_layer();
                 }
             }
         }
@@ -3121,6 +3325,9 @@ impl Timeline {
                 "produced rewritten delta layer: {}",
                 layer.layer_desc().key()
             );
+            // For now, we include rewritten delta layer size in the "produce_delta_layer". We could
+            // make it a separate statistics in the future.
+            stat.produce_delta_layer(layer.layer_desc().file_size());
         }
         compact_to.extend(rewrote_delta_layers);
         for action in produced_image_layers {
@@ -3132,8 +3339,16 @@ impl Timeline {
                 }
                 BatchWriterResult::Discarded(l) => {
                     debug!("discarded image layer: {}", l);
+                    if let Some(layer_desc) = layer_selection_by_key.get(&l) {
+                        stat.discard_image_layer(layer_desc.file_size());
+                    } else {
+                        tracing::warn!(
+                            "discarded image layer not in layer_selection: {}, produced a layer outside of the compaction key range?",
+                            l
+                        );
+                        stat.discard_image_layer(0);
+                    }
                     keep_layers.insert(l);
-                    stat.discard_image_layer();
                 }
             }
         }
@@ -3166,7 +3381,9 @@ impl Timeline {
                     &layer.layer_desc().key_range,
                     &job_desc.compaction_key_range,
                 ) {
-                    bail!("violated constraint: image layer outside of compaction key range");
+                    return Err(CompactionError::Other(anyhow!(
+                        "violated constraint: image layer outside of compaction key range"
+                    )));
                 }
                 if !fully_contains(
                     &job_desc.compaction_key_range,
@@ -3179,13 +3396,25 @@ impl Timeline {
 
         layer_selection.retain(|x| !keep_layers.contains(&x.layer_desc().key()));
 
+        let time_final_phase = timer.elapsed();
+
+        stat.time_final_phase_secs = time_final_phase.as_secs_f64();
+        stat.time_main_loop_secs = time_main_loop.as_secs_f64();
+        stat.time_acquire_lock_secs = time_acquire_lock.as_secs_f64();
+        stat.time_download_layer_secs = time_download_layer.as_secs_f64();
+        stat.time_analyze_secs = time_analyze.as_secs_f64();
+        stat.time_total_secs = begin_timer.elapsed().as_secs_f64();
+        stat.finalize();
+
         info!(
             "gc-compaction statistics: {}",
-            serde_json::to_string(&stat)?
+            serde_json::to_string(&stat)
+                .context("failed to serialize gc-compaction statistics")
+                .map_err(CompactionError::Other)?
         );
 
         if dry_run {
-            return Ok(());
+            return Ok(CompactionOutcome::Done);
         }
 
         info!(
@@ -3220,10 +3449,10 @@ impl Timeline {
         // the writer, so potentially, we will need a function like `ImageLayerBatchWriter::get_all_pending_layer_keys` to get all the keys that are
         // in the writer before finalizing the persistent layers. Now we would leave some dangling layers on the disk if the check fails.
         if let Some(err) = check_valid_layermap(&final_layers) {
-            bail!(
+            return Err(CompactionError::Other(anyhow!(
                 "gc-compaction layer map check failed after compaction because {}, compaction result not applied to the layer map due to potential data loss",
                 err
-            );
+            )));
         }
 
         // Between the sanity check and this compaction update, there could be new layers being flushed, but it should be fine because we only
@@ -3275,7 +3504,9 @@ impl Timeline {
         // find_gc_cutoffs will try accessing things below the cutoff. TODO: ideally, this should
         // be batched into `schedule_compaction_update`.
         let disk_consistent_lsn = self.disk_consistent_lsn.load();
-        self.schedule_uploads(disk_consistent_lsn, None)?;
+        self.schedule_uploads(disk_consistent_lsn, None)
+            .context("failed to schedule uploads")
+            .map_err(CompactionError::Other)?;
         // If a layer gets rewritten throughout gc-compaction, we need to keep that layer only in `compact_to` instead
         // of `compact_from`.
         let compact_from = {
@@ -3302,7 +3533,7 @@ impl Timeline {
 
         drop(gc_lock);
 
-        Ok(())
+        Ok(CompactionOutcome::Done)
     }
 }
 

pageserver/src/tenant/timeline/delete.rs

@@ -11,6 +11,7 @@ use utils::id::TimelineId;
 use utils::{crashsafe, fs_ext, pausable_failpoint};
 
 use crate::config::PageServerConf;
+use crate::context::RequestContext;
 use crate::task_mgr::{self, TaskKind};
 use crate::tenant::metadata::TimelineMetadata;
 use crate::tenant::remote_timeline_client::{
@@ -291,10 +292,11 @@ impl DeleteTimelineFlow {
         timeline_id: TimelineId,
         local_metadata: &TimelineMetadata,
         remote_client: RemoteTimelineClient,
+        ctx: &RequestContext,
     ) -> anyhow::Result<()> {
         // Note: here we even skip populating layer map. Timeline is essentially uninitialized.
         // RemoteTimelineClient is the only functioning part.
-        let timeline = tenant
+        let (timeline, _timeline_ctx) = tenant
             .create_timeline_struct(
                 timeline_id,
                 local_metadata,
@@ -307,6 +309,7 @@ impl DeleteTimelineFlow {
                 crate::tenant::CreateTimelineIdempotency::FailWithConflict, // doesn't matter what we put here
                 None, // doesn't matter what we put here
                 None, // doesn't matter what we put here
+                ctx,
             )
             .context("create_timeline_struct")?;
 

pageserver/src/tenant/timeline/eviction_task.rs

@@ -93,7 +93,8 @@ impl Timeline {
             }
         }
 
-        let ctx = RequestContext::new(TaskKind::Eviction, DownloadBehavior::Warn);
+        let ctx = RequestContext::new(TaskKind::Eviction, DownloadBehavior::Warn)
+            .with_scope_timeline(&self);
         loop {
             let policy = self.get_eviction_policy();
             let cf = self

pageserver/src/tenant/timeline/heatmap_layers_downloader.rs

@@ -61,11 +61,11 @@ impl HeatmapLayersDownloader {
 
                 tracing::info!(
                     resident_size=%timeline.resident_physical_size(),
-                    heatmap_layers=%heatmap.layers.len(),
+                    heatmap_layers=%heatmap.all_layers().count(),
                     "Starting heatmap layers download"
                 );
 
-                let stream = futures::stream::iter(heatmap.layers.into_iter().filter_map(
+                let stream = futures::stream::iter(heatmap.all_layers().cloned().filter_map(
                     |layer| {
                         let ctx = ctx.attached_child();
                         let tl = timeline.clone();

pageserver/src/tenant/vectored_blob_io.rs

@@ -961,7 +961,8 @@ mod tests {
     }
 
     async fn round_trip_test_compressed(blobs: &[Vec<u8>], compression: bool) -> Result<(), Error> {
-        let ctx = RequestContext::new(TaskKind::UnitTest, DownloadBehavior::Error);
+        let ctx =
+            RequestContext::new(TaskKind::UnitTest, DownloadBehavior::Error).with_scope_unit_test();
         let (_temp_dir, pathbuf, offsets) =
             write_maybe_compressed::<true>(blobs, compression, &ctx).await?;
 

pageserver/src/virtual_file.rs

@@ -26,15 +26,14 @@ use owned_buffers_io::io_buf_aligned::{IoBufAligned, IoBufAlignedMut};
 use owned_buffers_io::io_buf_ext::FullSlice;
 use pageserver_api::config::defaults::DEFAULT_IO_BUFFER_ALIGNMENT;
 pub use pageserver_api::models::virtual_file as api;
-use pageserver_api::shard::TenantShardId;
 use tokio::sync::{RwLock, RwLockReadGuard, RwLockWriteGuard};
 use tokio::time::Instant;
 use tokio_epoll_uring::{BoundedBuf, IoBuf, IoBufMut, Slice};
 
+use crate::assert_u64_eq_usize::UsizeIsU64;
 use crate::context::RequestContext;
-use crate::metrics::{STORAGE_IO_SIZE, STORAGE_IO_TIME_METRIC, StorageIoOperation};
+use crate::metrics::{STORAGE_IO_TIME_METRIC, StorageIoOperation};
 use crate::page_cache::{PAGE_SZ, PageWriteGuard};
-use crate::tenant::TENANTS_SEGMENT_NAME;
 pub(crate) mod io_engine;
 pub use io_engine::{
     FeatureTestResult as IoEngineFeatureTestResult, feature_test as io_engine_feature_test,
@@ -121,7 +120,7 @@ impl VirtualFile {
     pub async fn open_with_options<P: AsRef<Utf8Path>>(
         path: P,
         open_options: &OpenOptions,
-        ctx: &RequestContext, /* TODO: carry a pointer to the metrics in the RequestContext instead of the parsing https://github.com/neondatabase/neon/issues/6107 */
+        ctx: &RequestContext,
     ) -> Result<Self, std::io::Error> {
         let inner = VirtualFileInner::open_with_options(path, open_options, ctx).await?;
         Ok(VirtualFile {
@@ -133,7 +132,7 @@ impl VirtualFile {
     pub async fn open_with_options_v2<P: AsRef<Utf8Path>>(
         path: P,
         open_options: &OpenOptions,
-        ctx: &RequestContext, /* TODO: carry a pointer to the metrics in the RequestContext instead of the parsing https://github.com/neondatabase/neon/issues/6107 */
+        ctx: &RequestContext,
     ) -> Result<Self, std::io::Error> {
         let file = match get_io_mode() {
             IoMode::Buffered => {
@@ -304,13 +303,6 @@ pub struct VirtualFileInner {
     /// storing it here.
     pub path: Utf8PathBuf,
     open_options: OpenOptions,
-
-    // These are strings becase we only use them for metrics, and those expect strings.
-    // It makes no sense for us to constantly turn the `TimelineId` and `TenantId` into
-    // strings.
-    tenant_id: String,
-    shard_id: String,
-    timeline_id: String,
 }
 
 #[derive(Debug, PartialEq, Clone, Copy)]
@@ -592,36 +584,16 @@ impl VirtualFileInner {
     pub async fn open_with_options<P: AsRef<Utf8Path>>(
         path: P,
         open_options: &OpenOptions,
-        _ctx: &RequestContext, /* TODO: carry a pointer to the metrics in the RequestContext instead of the parsing https://github.com/neondatabase/neon/issues/6107 */
+        _ctx: &RequestContext,
     ) -> Result<VirtualFileInner, std::io::Error> {
-        let path_ref = path.as_ref();
-        let path_str = path_ref.to_string();
-        let parts = path_str.split('/').collect::<Vec<&str>>();
-        let (tenant_id, shard_id, timeline_id) =
-            if parts.len() > 5 && parts[parts.len() - 5] == TENANTS_SEGMENT_NAME {
-                let tenant_shard_part = parts[parts.len() - 4];
-                let (tenant_id, shard_id) = match tenant_shard_part.parse::<TenantShardId>() {
-                    Ok(tenant_shard_id) => (
-                        tenant_shard_id.tenant_id.to_string(),
-                        format!("{}", tenant_shard_id.shard_slug()),
-                    ),
-                    Err(_) => {
-                        // Malformed path: this ID is just for observability, so tolerate it
-                        // and pass through
-                        (tenant_shard_part.to_string(), "*".to_string())
-                    }
-                };
-                (tenant_id, shard_id, parts[parts.len() - 2].to_string())
-            } else {
-                ("*".to_string(), "*".to_string(), "*".to_string())
-            };
+        let path = path.as_ref();
         let (handle, mut slot_guard) = get_open_files().find_victim_slot().await;
 
         // NB: there is also StorageIoOperation::OpenAfterReplace which is for the case
         // where our caller doesn't get to use the returned VirtualFile before its
         // slot gets re-used by someone else.
         let file = observe_duration!(StorageIoOperation::Open, {
-            open_options.open(path_ref.as_std_path()).await?
+            open_options.open(path.as_std_path()).await?
         });
 
         // Strip all options other than read and write.
@@ -637,11 +609,8 @@ impl VirtualFileInner {
         let vfile = VirtualFileInner {
             handle: RwLock::new(handle),
             pos: 0,
-            path: path_ref.to_path_buf(),
+            path: path.to_owned(),
             open_options: reopen_options,
-            tenant_id,
-            shard_id,
-            timeline_id,
         };
 
         // TODO: Under pressure, it's likely the slot will get re-used and
@@ -944,7 +913,7 @@ impl VirtualFileInner {
         &self,
         buf: tokio_epoll_uring::Slice<Buf>,
         offset: u64,
-        _ctx: &RequestContext, /* TODO: use for metrics: https://github.com/neondatabase/neon/issues/6107 */
+        ctx: &RequestContext,
     ) -> (tokio_epoll_uring::Slice<Buf>, Result<usize, Error>)
     where
         Buf: tokio_epoll_uring::IoBufMut + Send,
@@ -962,14 +931,7 @@ impl VirtualFileInner {
             let ((_file_guard, buf), res) = io_engine::get().read_at(file_guard, offset, buf).await;
             let res = res.maybe_fatal_err("io_engine read_at inside VirtualFileInner::read_at");
             if let Ok(size) = res {
-                STORAGE_IO_SIZE
-                    .with_label_values(&[
-                        "read",
-                        &self.tenant_id,
-                        &self.shard_id,
-                        &self.timeline_id,
-                    ])
-                    .add(size as i64);
+                ctx.io_size_metrics().read.add(size.into_u64());
             }
             (buf, res)
         })
@@ -980,9 +942,9 @@ impl VirtualFileInner {
         &self,
         buf: FullSlice<B>,
         offset: u64,
-        _ctx: &RequestContext, /* TODO: use for metrics: https://github.com/neondatabase/neon/issues/6107 */
+        ctx: &RequestContext,
     ) -> (FullSlice<B>, Result<usize, Error>) {
-        let (slice, result) = self.write_at_inner(buf, offset, _ctx).await;
+        let (slice, result) = self.write_at_inner(buf, offset, ctx).await;
         let result = result.maybe_fatal_err("write_at");
         (slice, result)
     }
@@ -991,7 +953,7 @@ impl VirtualFileInner {
         &self,
         buf: FullSlice<B>,
         offset: u64,
-        _ctx: &RequestContext, /* TODO: use for metrics: https://github.com/neondatabase/neon/issues/6107 */
+        ctx: &RequestContext,
     ) -> (FullSlice<B>, Result<usize, Error>) {
         let file_guard = match self.lock_file().await {
             Ok(file_guard) => file_guard,
@@ -1001,14 +963,7 @@ impl VirtualFileInner {
             let ((_file_guard, buf), result) =
                 io_engine::get().write_at(file_guard, offset, buf).await;
             if let Ok(size) = result {
-                STORAGE_IO_SIZE
-                    .with_label_values(&[
-                        "write",
-                        &self.tenant_id,
-                        &self.shard_id,
-                        &self.timeline_id,
-                    ])
-                    .add(size as i64);
+                ctx.io_size_metrics().write.add(size.into_u64());
             }
             (buf, result)
         })
@@ -1593,7 +1548,8 @@ mod tests {
     where
         A: Adapter,
     {
-        let ctx = RequestContext::new(TaskKind::UnitTest, DownloadBehavior::Error);
+        let ctx =
+            RequestContext::new(TaskKind::UnitTest, DownloadBehavior::Error).with_scope_unit_test();
         let testdir = crate::config::PageServerConf::test_repo_dir(testname);
         std::fs::create_dir_all(&testdir)?;
 
@@ -1720,7 +1676,8 @@ mod tests {
         const THREADS: usize = 100;
         const SAMPLE: [u8; SIZE] = [0xADu8; SIZE];
 
-        let ctx = RequestContext::new(TaskKind::UnitTest, DownloadBehavior::Error);
+        let ctx =
+            RequestContext::new(TaskKind::UnitTest, DownloadBehavior::Error).with_scope_unit_test();
         let testdir = crate::config::PageServerConf::test_repo_dir("vfile_concurrency");
         std::fs::create_dir_all(&testdir)?;
 
@@ -1779,7 +1736,8 @@ mod tests {
 
     #[tokio::test]
     async fn test_atomic_overwrite_basic() {
-        let ctx = RequestContext::new(TaskKind::UnitTest, DownloadBehavior::Error);
+        let ctx =
+            RequestContext::new(TaskKind::UnitTest, DownloadBehavior::Error).with_scope_unit_test();
         let testdir = crate::config::PageServerConf::test_repo_dir("test_atomic_overwrite_basic");
         std::fs::create_dir_all(&testdir).unwrap();
 
@@ -1807,7 +1765,8 @@ mod tests {
 
     #[tokio::test]
     async fn test_atomic_overwrite_preexisting_tmp() {
-        let ctx = RequestContext::new(TaskKind::UnitTest, DownloadBehavior::Error);
+        let ctx =
+            RequestContext::new(TaskKind::UnitTest, DownloadBehavior::Error).with_scope_unit_test();
         let testdir =
             crate::config::PageServerConf::test_repo_dir("test_atomic_overwrite_preexisting_tmp");
         std::fs::create_dir_all(&testdir).unwrap();

pgxn/neon/file_cache.c

@@ -1195,9 +1195,11 @@ lfc_writev(NRelFileInfo rinfo, ForkNumber forkNum, BlockNumber blkno,
 				state = GET_STATE(entry, chunk_offs + i);
 				if (state == PENDING) {
 					SET_STATE(entry, chunk_offs + i, REQUESTED);
-				} else if (state != REQUESTED) {
+				} else if (state == UNAVAILABLE) {
 					SET_STATE(entry, chunk_offs + i, PENDING);
 					break;
+				} else if (state == AVAILABLE) {
+					break;
 				}
 				if (!sleeping)
 				{
@@ -1369,6 +1371,10 @@ neon_get_lfc_stats(PG_FUNCTION_ARGS)
 			if (lfc_ctl)
 				value = lfc_ctl->limit;
 			break;
+		case 8:
+			key = "file_cache_chunk_size_pages";
+			value = BLOCKS_PER_CHUNK;
+			break;
 		default:
 			SRF_RETURN_DONE(funcctx);
 	}

pgxn/neon/libpagestore.c

@@ -16,6 +16,8 @@
 
 #include <math.h>
 
+#include "libpq-int.h"
+
 #include "access/xlog.h"
 #include "common/hashfn.h"
 #include "fmgr.h"
@@ -48,6 +50,20 @@
 #define MIN_RECONNECT_INTERVAL_USEC 1000
 #define MAX_RECONNECT_INTERVAL_USEC 1000000
 
+
+enum NeonEndpointType {
+	EP_TYPE_PRIMARY = 1,
+	EP_TYPE_REPLICA,
+	EP_TYPE_STATIC
+};
+
+static const struct config_enum_entry neon_endpoint_types[] = {
+	{"primary", EP_TYPE_PRIMARY, false},
+	{"replica", EP_TYPE_REPLICA, false},
+	{"static", EP_TYPE_STATIC, false},
+	{NULL, 0, false}
+};
+
 /* GUCs */
 char	   *neon_timeline;
 char	   *neon_tenant;
@@ -60,6 +76,8 @@ int			flush_every_n_requests = 8;
 
 int         neon_protocol_version = 2;
 
+static int neon_endpoint_type = 0;
+
 static int	max_reconnect_attempts = 60;
 static int	stripe_size;
 
@@ -390,7 +408,7 @@ pageserver_connect(shardno_t shard_no, int elevel)
 	{
 		const char *keywords[4];
 		const char *values[4];
-		char pid_str[16];
+		char pid_str[24];
 		int			n_pgsql_params;
 		TimestampTz	now;
 		int64		us_since_last_attempt;
@@ -443,7 +461,22 @@ pageserver_connect(shardno_t shard_no, int elevel)
 		 */
 		keywords[n_pgsql_params] = "application_name";
 		{
-			int ret = snprintf(pid_str, sizeof(pid_str), "%d", MyProcPid);
+			int ret;
+			switch (neon_endpoint_type)
+			{
+				case EP_TYPE_PRIMARY:
+					ret = snprintf(pid_str, sizeof(pid_str), "%d-%s", MyProcPid, "primary");
+					break;
+				case EP_TYPE_REPLICA:
+					ret = snprintf(pid_str, sizeof(pid_str), "%d-%s", MyProcPid, "replica");
+					break;
+				case EP_TYPE_STATIC:
+					ret = snprintf(pid_str, sizeof(pid_str), "%d-%s", MyProcPid, "static");
+					break;
+				default:
+					ret = snprintf(pid_str, sizeof(pid_str), "%d", MyProcPid);
+					break;
+			}
 			if (ret < 0 || ret >= (int)(sizeof(pid_str)))
 				elog(FATAL, "stack-allocated buffer too small to hold pid");
 		}
@@ -815,9 +848,10 @@ retry:
 			get_socket_stats(PQsocket(pageserver_conn), &sndbuf, &recvbuf);
 
 			neon_shard_log(shard_no, LOG,
-						   "no response received from pageserver for %0.3f s, still waiting (sent " UINT64_FORMAT " requests, received " UINT64_FORMAT " responses) (socket sndbuf=%d recvbuf=%d)",
+						   "no response received from pageserver for %0.3f s, still waiting (sent " UINT64_FORMAT " requests, received " UINT64_FORMAT " responses) (socket sndbuf=%d recvbuf=%d) (conn start=%d end=%d)",
 						   INSTR_TIME_GET_DOUBLE(since_start),
-						   shard->nrequests_sent, shard->nresponses_received, sndbuf, recvbuf);
+						   shard->nrequests_sent, shard->nresponses_received, sndbuf, recvbuf,
+				           pageserver_conn->inStart, pageserver_conn->inEnd);
 			shard->receive_last_log_time = now;
 			shard->receive_logged = true;
 		}
@@ -1367,6 +1401,17 @@ pg_init_libpagestore(void)
 							GUC_UNIT_MS,
 							NULL, NULL, NULL);
 
+	DefineCustomEnumVariable(
+							"neon.endpoint_type",
+							"The compute endpoint node type",
+							NULL,
+							&neon_endpoint_type,
+							EP_TYPE_PRIMARY,
+							neon_endpoint_types,
+							PGC_POSTMASTER,
+							0,
+							NULL, NULL, NULL);
+
 	relsize_hash_init();
 
 	if (page_server != NULL)

pgxn/neon/neon.c

@@ -457,6 +457,15 @@ _PG_init(void)
 							PGC_SIGHUP,
 							0,
 							NULL, NULL, NULL);
+	DefineCustomBoolVariable(
+							"neon.disable_wal_prevlink_checks",
+							"Disable validation of prev link in WAL records",
+							NULL,
+							&disable_wal_prev_lsn_checks,
+							false,
+							PGC_SIGHUP,
+							0,
+							NULL, NULL, NULL);
 
 	DefineCustomBoolVariable(
 							"neon.allow_replica_misconfig",

pgxn/neon/neon.h

@@ -23,6 +23,7 @@ extern char *wal_acceptors_list;
 extern int	wal_acceptor_reconnect_timeout;
 extern int	wal_acceptor_connection_timeout;
 extern int	readahead_getpage_pull_timeout_ms;
+extern bool	disable_wal_prev_lsn_checks;
 
 #if PG_MAJORVERSION_NUM >= 17
 extern uint32		WAIT_EVENT_NEON_LFC_MAINTENANCE;

pgxn/neon/walproposer.c

@@ -57,10 +57,11 @@ static void SendProposerGreeting(Safekeeper *sk);
 static void RecvAcceptorGreeting(Safekeeper *sk);
 static void SendVoteRequest(Safekeeper *sk);
 static void RecvVoteResponse(Safekeeper *sk);
+static bool VotesCollected(WalProposer *wp);
 static void HandleElectedProposer(WalProposer *wp);
 static term_t GetHighestTerm(TermHistory *th);
-static term_t GetEpoch(Safekeeper *sk);
-static void DetermineEpochStartLsn(WalProposer *wp);
+static term_t GetLastLogTerm(Safekeeper *sk);
+static void ProcessPropStartPos(WalProposer *wp);
 static void SendProposerElected(Safekeeper *sk);
 static void StartStreaming(Safekeeper *sk);
 static void SendMessageToNode(Safekeeper *sk);
@@ -97,6 +98,7 @@ WalProposerCreate(WalProposerConfig *config, walproposer_api api)
 	wp = palloc0(sizeof(WalProposer));
 	wp->config = config;
 	wp->api = api;
+	wp->state = WPS_COLLECTING_TERMS;
 
 	wp_log(LOG, "neon.safekeepers=%s", wp->config->safekeepers_list);
 
@@ -518,7 +520,7 @@ AdvancePollState(Safekeeper *sk, uint32 events)
 			 * nodes are transferred from SS_VOTING to sending actual vote
 			 * requests.
 			 */
-		case SS_VOTING:
+		case SS_WAIT_VOTING:
 			wp_log(WARNING, "EOF from node %s:%s in %s state", sk->host,
 				   sk->port, FormatSafekeeperState(sk));
 			ResetConnection(sk);
@@ -547,7 +549,7 @@ AdvancePollState(Safekeeper *sk, uint32 events)
 			/*
 			 * Idle state for waiting votes from quorum.
 			 */
-		case SS_IDLE:
+		case SS_WAIT_ELECTED:
 			wp_log(WARNING, "EOF from node %s:%s in %s state", sk->host,
 				   sk->port, FormatSafekeeperState(sk));
 			ResetConnection(sk);
@@ -721,6 +723,15 @@ SendProposerGreeting(Safekeeper *sk)
 	BlockingWrite(sk, sk->outbuf.data, sk->outbuf.len, SS_HANDSHAKE_RECV);
 }
 
+/*
+ * Have we received greeting from enough (quorum) safekeepers to start voting?
+ */
+static bool
+TermsCollected(WalProposer *wp)
+{
+	return wp->n_connected >= wp->quorum;
+}
+
 static void
 RecvAcceptorGreeting(Safekeeper *sk)
 {
@@ -754,7 +765,7 @@ RecvAcceptorGreeting(Safekeeper *sk)
 	}
 
 	/* Protocol is all good, move to voting. */
-	sk->state = SS_VOTING;
+	sk->state = SS_WAIT_VOTING;
 
 	/*
 	 * Note: it would be better to track the counter on per safekeeper basis,
@@ -762,17 +773,18 @@ RecvAcceptorGreeting(Safekeeper *sk)
 	 * as is for now.
 	 */
 	++wp->n_connected;
-	if (wp->n_connected <= wp->quorum)
+	if (wp->state == WPS_COLLECTING_TERMS)
 	{
 		/* We're still collecting terms from the majority. */
 		wp->propTerm = Max(sk->greetResponse.term, wp->propTerm);
 
 		/* Quorum is acquried, prepare the vote request. */
-		if (wp->n_connected == wp->quorum)
+		if (TermsCollected(wp))
 		{
 			wp->propTerm++;
 			wp_log(LOG, "proposer connected to quorum (%d) safekeepers, propTerm=" INT64_FORMAT, wp->quorum, wp->propTerm);
 
+			wp->state = WPS_CAMPAIGN;
 			wp->voteRequest.pam.tag = 'v';
 			wp->voteRequest.generation = wp->mconf.generation;
 			wp->voteRequest.term = wp->propTerm;
@@ -787,12 +799,10 @@ RecvAcceptorGreeting(Safekeeper *sk)
 	}
 
 	/*
-	 * Check if we have quorum. If there aren't enough safekeepers, wait and
-	 * do nothing. We'll eventually get a task when the election starts.
-	 *
-	 * If we do have quorum, we can start an election.
+	 * If we have quorum, start (or just send vote request to newly connected
+	 * node) election, otherwise wait until we have more greetings.
 	 */
-	if (wp->n_connected < wp->quorum)
+	if (wp->state == WPS_COLLECTING_TERMS)
 	{
 		/*
 		 * SS_VOTING is an idle state; read-ready indicates the connection
@@ -807,11 +817,7 @@ RecvAcceptorGreeting(Safekeeper *sk)
 		 */
 		for (int j = 0; j < wp->n_safekeepers; j++)
 		{
-			/*
-			 * Remember: SS_VOTING indicates that the safekeeper is
-			 * participating in voting, but hasn't sent anything yet.
-			 */
-			if (wp->safekeeper[j].state == SS_VOTING)
+			if (wp->safekeeper[j].state == SS_WAIT_VOTING)
 				SendVoteRequest(&wp->safekeeper[j]);
 		}
 	}
@@ -838,6 +844,8 @@ RecvVoteResponse(Safekeeper *sk)
 {
 	WalProposer *wp = sk->wp;
 
+	Assert(wp->state >= WPS_CAMPAIGN);
+
 	sk->voteResponse.apm.tag = 'v';
 	if (!AsyncReadMessage(sk, (AcceptorProposerMessage *) &sk->voteResponse))
 		return;
@@ -856,7 +864,7 @@ RecvVoteResponse(Safekeeper *sk)
 	 * we are not elected yet and thus need the vote.
 	 */
 	if ((!sk->voteResponse.voteGiven) &&
-		(sk->voteResponse.term > wp->propTerm || wp->n_votes < wp->quorum))
+		(sk->voteResponse.term > wp->propTerm || wp->state == WPS_CAMPAIGN))
 	{
 		wp_log(FATAL, "WAL acceptor %s:%s with term " INT64_FORMAT " rejects our connection request with term " INT64_FORMAT "",
 			   sk->host, sk->port,
@@ -864,38 +872,83 @@ RecvVoteResponse(Safekeeper *sk)
 	}
 	Assert(sk->voteResponse.term == wp->propTerm);
 
-	/* Handshake completed, do we have quorum? */
+	/* ready for elected message */
+	sk->state = SS_WAIT_ELECTED;
+
 	wp->n_votes++;
-	if (wp->n_votes < wp->quorum)
+	/* Are we already elected? */
+	if (wp->state == WPS_CAMPAIGN)
 	{
-		sk->state = SS_IDLE;	/* can't do much yet, no quorum */
-	}
-	else if (wp->n_votes > wp->quorum)
-	{
-		/* already elected, start streaming */
-		SendProposerElected(sk);
+		/* no; check if this vote makes us elected */
+		if (VotesCollected(wp))
+		{
+			wp->state = WPS_ELECTED;
+			HandleElectedProposer(wp);
+		}
+		else
+		{
+			/* can't do much yet, no quorum */
+			return;
+		}
 	}
 	else
 	{
-		sk->state = SS_IDLE;
-		/* Idle state waits for read-ready events */
-		wp->api.update_event_set(sk, WL_SOCKET_READABLE);
-
-		HandleElectedProposer(sk->wp);
+		Assert(wp->state == WPS_ELECTED);
+		/* send elected only to this sk */
+		SendProposerElected(sk);
 	}
 }
 
+/*
+ * Checks if enough votes has been collected to get elected and if that's the
+ * case finds the highest vote, setting donor, donorLastLogTerm,
+ * propTermStartLsn fields. Also sets truncateLsn.
+ */
+static bool
+VotesCollected(WalProposer *wp)
+{
+	int			n_ready = 0;
+
+	/* assumed to be called only when not elected yet */
+	Assert(wp->state == WPS_CAMPAIGN);
+
+	wp->propTermStartLsn = InvalidXLogRecPtr;
+	wp->donorLastLogTerm = 0;
+	wp->truncateLsn = InvalidXLogRecPtr;
+
+	for (int i = 0; i < wp->n_safekeepers; i++)
+	{
+		if (wp->safekeeper[i].state == SS_WAIT_ELECTED)
+		{
+			n_ready++;
+
+			if (GetLastLogTerm(&wp->safekeeper[i]) > wp->donorLastLogTerm ||
+				(GetLastLogTerm(&wp->safekeeper[i]) == wp->donorLastLogTerm &&
+				 wp->safekeeper[i].voteResponse.flushLsn > wp->propTermStartLsn))
+			{
+				wp->donorLastLogTerm = GetLastLogTerm(&wp->safekeeper[i]);
+				wp->propTermStartLsn = wp->safekeeper[i].voteResponse.flushLsn;
+				wp->donor = i;
+			}
+			wp->truncateLsn = Max(wp->safekeeper[i].voteResponse.truncateLsn, wp->truncateLsn);
+		}
+	}
+
+	return n_ready >= wp->quorum;
+}
+
 /*
  * Called once a majority of acceptors have voted for us and current proposer
  * has been elected.
  *
- * Sends ProposerElected message to all acceptors in SS_IDLE state and starts
+ * Sends ProposerElected message to all acceptors in SS_WAIT_ELECTED state and starts
  * replication from walsender.
  */
 static void
 HandleElectedProposer(WalProposer *wp)
 {
-	DetermineEpochStartLsn(wp);
+	ProcessPropStartPos(wp);
+	Assert(wp->propTermStartLsn != InvalidXLogRecPtr);
 
 	/*
 	 * Synchronously download WAL from the most advanced safekeeper. We do
@@ -907,40 +960,24 @@ HandleElectedProposer(WalProposer *wp)
 		wp_log(FATAL, "failed to download WAL for logical replicaiton");
 	}
 
-	/*
-	 * Zero propEpochStartLsn means majority of safekeepers doesn't have any
-	 * WAL, timeline was just created. Compute bumps it to basebackup LSN,
-	 * otherwise we must be sync-safekeepers and we have nothing to do then.
-	 *
-	 * Proceeding is not only pointless but harmful, because we'd give
-	 * safekeepers term history starting with 0/0. These hacks will go away
-	 * once we disable implicit timeline creation on safekeepers and create it
-	 * with non zero LSN from the start.
-	 */
-	if (wp->propEpochStartLsn == InvalidXLogRecPtr)
-	{
-		Assert(wp->config->syncSafekeepers);
-		wp_log(LOG, "elected with zero propEpochStartLsn in sync-safekeepers, exiting");
-		wp->api.finish_sync_safekeepers(wp, wp->propEpochStartLsn);
-	}
-
-	if (wp->truncateLsn == wp->propEpochStartLsn && wp->config->syncSafekeepers)
+	if (wp->truncateLsn == wp->propTermStartLsn && wp->config->syncSafekeepers)
 	{
 		/* Sync is not needed: just exit */
-		wp->api.finish_sync_safekeepers(wp, wp->propEpochStartLsn);
+		wp->api.finish_sync_safekeepers(wp, wp->propTermStartLsn);
 		/* unreachable */
 	}
 
 	for (int i = 0; i < wp->n_safekeepers; i++)
 	{
-		if (wp->safekeeper[i].state == SS_IDLE)
+		if (wp->safekeeper[i].state == SS_WAIT_ELECTED)
 			SendProposerElected(&wp->safekeeper[i]);
 	}
 
 	/*
 	 * The proposer has been elected, and there will be no quorum waiting
-	 * after this point. There will be no safekeeper with state SS_IDLE also,
-	 * because that state is used only for quorum waiting.
+	 * after this point. There will be no safekeeper with state
+	 * SS_WAIT_ELECTED also, because that state is used only for quorum
+	 * waiting.
 	 */
 
 	if (wp->config->syncSafekeepers)
@@ -957,7 +994,7 @@ HandleElectedProposer(WalProposer *wp)
 		return;
 	}
 
-	wp->api.start_streaming(wp, wp->propEpochStartLsn);
+	wp->api.start_streaming(wp, wp->propTermStartLsn);
 	/* Should not return here */
 }
 
@@ -970,7 +1007,7 @@ GetHighestTerm(TermHistory *th)
 
 /* safekeeper's epoch is the term of the highest entry in the log */
 static term_t
-GetEpoch(Safekeeper *sk)
+GetLastLogTerm(Safekeeper *sk)
 {
 	return GetHighestTerm(&sk->voteResponse.termHistory);
 }
@@ -991,72 +1028,52 @@ SkipXLogPageHeader(WalProposer *wp, XLogRecPtr lsn)
 }
 
 /*
- * Called after majority of acceptors gave votes, it calculates the most
- * advanced safekeeper (who will be the donor) and epochStartLsn -- LSN since
- * which we'll write WAL in our term.
- *
- * Sets truncateLsn along the way (though it is not of much use at this point --
- * only for skipping recovery).
+ * Called after quorum gave votes and proposer starting position (highest vote
+ * term + flush LSN) -- is determined (VotesCollected true), this function
+ * adopts it: pushes LSN to shmem, sets wp term history, verifies that the
+ * basebackup matches.
  */
 static void
-DetermineEpochStartLsn(WalProposer *wp)
+ProcessPropStartPos(WalProposer *wp)
 {
 	TermHistory *dth;
-	int			n_ready = 0;
 	WalproposerShmemState *walprop_shared;
 
-	wp->propEpochStartLsn = InvalidXLogRecPtr;
-	wp->donorEpoch = 0;
-	wp->truncateLsn = InvalidXLogRecPtr;
-
-	for (int i = 0; i < wp->n_safekeepers; i++)
-	{
-		if (wp->safekeeper[i].state == SS_IDLE)
-		{
-			n_ready++;
-
-			if (GetEpoch(&wp->safekeeper[i]) > wp->donorEpoch ||
-				(GetEpoch(&wp->safekeeper[i]) == wp->donorEpoch &&
-				 wp->safekeeper[i].voteResponse.flushLsn > wp->propEpochStartLsn))
-			{
-				wp->donorEpoch = GetEpoch(&wp->safekeeper[i]);
-				wp->propEpochStartLsn = wp->safekeeper[i].voteResponse.flushLsn;
-				wp->donor = i;
-			}
-			wp->truncateLsn = Max(wp->safekeeper[i].voteResponse.truncateLsn, wp->truncateLsn);
-		}
-	}
-
-	if (n_ready < wp->quorum)
-	{
-		/*
-		 * This is a rare case that can be triggered if safekeeper has voted
-		 * and disconnected. In this case, its state will not be SS_IDLE and
-		 * its vote cannot be used, because we clean up `voteResponse` in
-		 * `ShutdownConnection`.
-		 */
-		wp_log(FATAL, "missing majority of votes, collected %d, expected %d, got %d", wp->n_votes, wp->quorum, n_ready);
-	}
+	/* must have collected votes */
+	Assert(wp->state == WPS_ELECTED);
 
 	/*
-	 * If propEpochStartLsn is 0, it means flushLsn is 0 everywhere, we are
-	 * bootstrapping and nothing was committed yet. Start streaming then from
-	 * the basebackup LSN.
+	 * If propTermStartLsn is 0, it means flushLsn is 0 everywhere, we are
+	 * bootstrapping and nothing was committed yet. Start streaming from the
+	 * basebackup LSN then.
+	 *
+	 * In case of sync-safekeepers just exit: proceeding is not only pointless
+	 * but harmful, because we'd give safekeepers term history starting with
+	 * 0/0. These hacks will go away once we disable implicit timeline
+	 * creation on safekeepers and create it with non zero LSN from the start.
 	 */
-	if (wp->propEpochStartLsn == InvalidXLogRecPtr && !wp->config->syncSafekeepers)
+	if (wp->propTermStartLsn == InvalidXLogRecPtr)
 	{
-		wp->propEpochStartLsn = wp->truncateLsn = wp->api.get_redo_start_lsn(wp);
-		wp_log(LOG, "bumped epochStartLsn to the first record %X/%X", LSN_FORMAT_ARGS(wp->propEpochStartLsn));
+		if (!wp->config->syncSafekeepers)
+		{
+			wp->propTermStartLsn = wp->truncateLsn = wp->api.get_redo_start_lsn(wp);
+			wp_log(LOG, "bumped epochStartLsn to the first record %X/%X", LSN_FORMAT_ARGS(wp->propTermStartLsn));
+		}
+		else
+		{
+			wp_log(LOG, "elected with zero propTermStartLsn in sync-safekeepers, exiting");
+			wp->api.finish_sync_safekeepers(wp, wp->propTermStartLsn);
+		}
 	}
-	pg_atomic_write_u64(&wp->api.get_shmem_state(wp)->propEpochStartLsn, wp->propEpochStartLsn);
+	pg_atomic_write_u64(&wp->api.get_shmem_state(wp)->propEpochStartLsn, wp->propTermStartLsn);
 
 	Assert(wp->truncateLsn != InvalidXLogRecPtr || wp->config->syncSafekeepers);
 
 	/*
-	 * We will be generating WAL since propEpochStartLsn, so we should set
+	 * We will be generating WAL since propTermStartLsn, so we should set
 	 * availableLsn to mark this LSN as the latest available position.
 	 */
-	wp->availableLsn = wp->propEpochStartLsn;
+	wp->availableLsn = wp->propTermStartLsn;
 
 	/*
 	 * Proposer's term history is the donor's + its own entry.
@@ -1067,12 +1084,12 @@ DetermineEpochStartLsn(WalProposer *wp)
 	if (dth->n_entries > 0)
 		memcpy(wp->propTermHistory.entries, dth->entries, sizeof(TermSwitchEntry) * dth->n_entries);
 	wp->propTermHistory.entries[wp->propTermHistory.n_entries - 1].term = wp->propTerm;
-	wp->propTermHistory.entries[wp->propTermHistory.n_entries - 1].lsn = wp->propEpochStartLsn;
+	wp->propTermHistory.entries[wp->propTermHistory.n_entries - 1].lsn = wp->propTermStartLsn;
 
 	wp_log(LOG, "got votes from majority (%d) of nodes, term " UINT64_FORMAT ", epochStartLsn %X/%X, donor %s:%s, truncate_lsn %X/%X",
 		   wp->quorum,
 		   wp->propTerm,
-		   LSN_FORMAT_ARGS(wp->propEpochStartLsn),
+		   LSN_FORMAT_ARGS(wp->propTermStartLsn),
 		   wp->safekeeper[wp->donor].host, wp->safekeeper[wp->donor].port,
 		   LSN_FORMAT_ARGS(wp->truncateLsn));
 
@@ -1090,7 +1107,7 @@ DetermineEpochStartLsn(WalProposer *wp)
 		 * Safekeepers don't skip header as they need continious stream of
 		 * data, so correct LSN for comparison.
 		 */
-		if (SkipXLogPageHeader(wp, wp->propEpochStartLsn) != wp->api.get_redo_start_lsn(wp))
+		if (SkipXLogPageHeader(wp, wp->propTermStartLsn) != wp->api.get_redo_start_lsn(wp))
 		{
 			/*
 			 * However, allow to proceed if last_log_term on the node which
@@ -1111,8 +1128,8 @@ DetermineEpochStartLsn(WalProposer *wp)
 				 */
 				disable_core_dump();
 				wp_log(PANIC,
-					   "collected propEpochStartLsn %X/%X, but basebackup LSN %X/%X",
-					   LSN_FORMAT_ARGS(wp->propEpochStartLsn),
+					   "collected propTermStartLsn %X/%X, but basebackup LSN %X/%X",
+					   LSN_FORMAT_ARGS(wp->propTermStartLsn),
 					   LSN_FORMAT_ARGS(wp->api.get_redo_start_lsn(wp)));
 			}
 		}
@@ -1623,7 +1640,7 @@ GetAcknowledgedByQuorumWALPosition(WalProposer *wp)
 		 * Like in Raft, we aren't allowed to commit entries from previous
 		 * terms, so ignore reported LSN until it gets to epochStartLsn.
 		 */
-		responses[i] = wp->safekeeper[i].appendResponse.flushLsn >= wp->propEpochStartLsn ? wp->safekeeper[i].appendResponse.flushLsn : 0;
+		responses[i] = wp->safekeeper[i].appendResponse.flushLsn >= wp->propTermStartLsn ? wp->safekeeper[i].appendResponse.flushLsn : 0;
 	}
 	qsort(responses, wp->n_safekeepers, sizeof(XLogRecPtr), CompareLsn);
 
@@ -1656,10 +1673,10 @@ UpdateDonorShmem(WalProposer *wp)
 	 * about its position immediately after election before any feedbacks are
 	 * sent.
 	 */
-	if (wp->safekeeper[wp->donor].state >= SS_IDLE)
+	if (wp->safekeeper[wp->donor].state >= SS_WAIT_ELECTED)
 	{
 		donor = &wp->safekeeper[wp->donor];
-		donor_lsn = wp->propEpochStartLsn;
+		donor_lsn = wp->propTermStartLsn;
 	}
 
 	/*
@@ -1748,7 +1765,7 @@ HandleSafekeeperResponse(WalProposer *wp, Safekeeper *fromsk)
 		for (int i = 0; i < wp->n_safekeepers; i++)
 		{
 			Safekeeper *sk = &wp->safekeeper[i];
-			bool		synced = sk->appendResponse.commitLsn >= wp->propEpochStartLsn;
+			bool		synced = sk->appendResponse.commitLsn >= wp->propTermStartLsn;
 
 			/* alive safekeeper which is not synced yet; wait for it */
 			if (sk->state != SS_OFFLINE && !synced)
@@ -1772,7 +1789,7 @@ HandleSafekeeperResponse(WalProposer *wp, Safekeeper *fromsk)
 			 */
 			BroadcastAppendRequest(wp);
 
-			wp->api.finish_sync_safekeepers(wp, wp->propEpochStartLsn);
+			wp->api.finish_sync_safekeepers(wp, wp->propTermStartLsn);
 			/* unreachable */
 		}
 	}
@@ -2378,7 +2395,7 @@ FormatSafekeeperState(Safekeeper *sk)
 		case SS_HANDSHAKE_RECV:
 			return_val = "handshake (receiving)";
 			break;
-		case SS_VOTING:
+		case SS_WAIT_VOTING:
 			return_val = "voting";
 			break;
 		case SS_WAIT_VERDICT:
@@ -2387,7 +2404,7 @@ FormatSafekeeperState(Safekeeper *sk)
 		case SS_SEND_ELECTED_FLUSH:
 			return_val = "send-announcement-flush";
 			break;
-		case SS_IDLE:
+		case SS_WAIT_ELECTED:
 			return_val = "idle";
 			break;
 		case SS_ACTIVE:
@@ -2476,8 +2493,8 @@ SafekeeperStateDesiredEvents(Safekeeper *sk, uint32 *sk_events, uint32 *nwr_even
 			 * Idle states use read-readiness as a sign that the connection
 			 * has been disconnected.
 			 */
-		case SS_VOTING:
-		case SS_IDLE:
+		case SS_WAIT_VOTING:
+		case SS_WAIT_ELECTED:
 			*sk_events = WL_SOCKET_READABLE;
 			return;
 

pgxn/neon/walproposer.h

@@ -73,12 +73,12 @@ typedef enum
 	 * Moved externally by execution of SS_HANDSHAKE_RECV, when we received a
 	 * quorum of handshakes.
 	 */
-	SS_VOTING,
+	SS_WAIT_VOTING,
 
 	/*
 	 * Already sent voting information, waiting to receive confirmation from
-	 * the node. After receiving, moves to SS_IDLE, if the quorum isn't
-	 * reached yet.
+	 * the node. After receiving, moves to SS_WAIT_ELECTED, if the quorum
+	 * isn't reached yet.
 	 */
 	SS_WAIT_VERDICT,
 
@@ -91,7 +91,7 @@ typedef enum
 	 *
 	 * Moves to SS_ACTIVE only by call to StartStreaming.
 	 */
-	SS_IDLE,
+	SS_WAIT_ELECTED,
 
 	/*
 	 * Active phase, when we acquired quorum and have WAL to send or feedback
@@ -751,6 +751,15 @@ typedef struct WalProposerConfig
 #endif
 } WalProposerConfig;
 
+typedef enum
+{
+	/* collecting greetings to determine term to campaign for */
+	WPS_COLLECTING_TERMS,
+	/* campaing started, waiting for votes */
+	WPS_CAMPAIGN,
+	/* successfully elected */
+	WPS_ELECTED,
+} WalProposerState;
 
 /*
  * WAL proposer state.
@@ -758,6 +767,7 @@ typedef struct WalProposerConfig
 typedef struct WalProposer
 {
 	WalProposerConfig *config;
+	WalProposerState state;
 	/* Current walproposer membership configuration */
 	MembershipConfiguration mconf;
 
@@ -813,10 +823,10 @@ typedef struct WalProposer
 	TermHistory propTermHistory;
 
 	/* epoch start lsn of the proposer */
-	XLogRecPtr	propEpochStartLsn;
+	XLogRecPtr	propTermStartLsn;
 
 	/* Most advanced acceptor epoch */
-	term_t		donorEpoch;
+	term_t		donorLastLogTerm;
 
 	/* Most advanced acceptor */
 	int			donor;

pgxn/neon/walproposer_pg.c

@@ -1496,7 +1496,7 @@ walprop_pg_wal_reader_allocate(Safekeeper *sk)
 
 	snprintf(log_prefix, sizeof(log_prefix), WP_LOG_PREFIX "sk %s:%s nwr: ", sk->host, sk->port);
 	Assert(!sk->xlogreader);
-	sk->xlogreader = NeonWALReaderAllocate(wal_segment_size, sk->wp->propEpochStartLsn, log_prefix);
+	sk->xlogreader = NeonWALReaderAllocate(wal_segment_size, sk->wp->propTermStartLsn, log_prefix);
 	if (sk->xlogreader == NULL)
 		wpg_log(FATAL, "failed to allocate xlog reader");
 }

pgxn/neon/walsender_hooks.c

@@ -32,6 +32,8 @@ extern XLogRecPtr WalSndWaitForWal(XLogRecPtr loc);
 extern bool GetDonorShmem(XLogRecPtr *donor_lsn);
 extern XLogRecPtr GetXLogReplayRecPtr(TimeLineID *replayTLI);
 
+bool disable_wal_prev_lsn_checks = false;
+
 static XLogRecPtr
 NeonWALReadWaitForWAL(XLogRecPtr loc)
 {
@@ -82,6 +84,8 @@ NeonWALPageRead(
 	if (flushptr < targetPagePtr + reqLen)
 		return -1;
 
+	xlogreader->skip_lsn_checks = disable_wal_prev_lsn_checks;
+
 	/* Read at most XLOG_BLCKSZ bytes */
 	if (targetPagePtr + XLOG_BLCKSZ <= flushptr)
 		count = XLOG_BLCKSZ;

poetry.lock

@@ -1414,14 +1414,14 @@ files = [
 
 [[package]]
 name = "jinja2"
-version = "3.1.5"
+version = "3.1.6"
 description = "A very fast and expressive template engine."
 optional = false
 python-versions = ">=3.7"
 groups = ["main"]
 files = [
-    {file = "jinja2-3.1.5-py3-none-any.whl", hash = "sha256:aba0f4dc9ed8013c424088f68a5c226f7d6097ed89b246d7749c2ec4175c6adb"},
-    {file = "jinja2-3.1.5.tar.gz", hash = "sha256:8fefff8dc3034e27bb80d67c671eb8a9bc424c0ef4c0826edbff304cceff43bb"},
+    {file = "jinja2-3.1.6-py3-none-any.whl", hash = "sha256:85ece4451f492d0c13c5dd7c13a64681a86afae63a5f347908daf103ce6d2f67"},
+    {file = "jinja2-3.1.6.tar.gz", hash = "sha256:0137fb05990d35f1275a587e9aee6d56da821fc83491a0fb838183be43f66d6d"},
 ]
 
 [package.dependencies]
@@ -3820,4 +3820,4 @@ cffi = ["cffi (>=1.11)"]
 [metadata]
 lock-version = "2.1"
 python-versions = "^3.11"
-content-hash = "9711c5479c867fa614ce3d352f1bbc63dba1cb2376d347f96fbeda6f512ee308"
+content-hash = "010ffce959bb256880ab5a267048c182e4612b3151f9a94e3bf5d3a7807962fe"

proxy/src/compute.rs

@@ -287,9 +287,10 @@ impl ConnCfg {
         // TODO: lots of useful info but maybe we can move it elsewhere (eg traces?)
         info!(
             cold_start_info = ctx.cold_start_info().as_str(),
-            "connected to compute node at {host} ({socket_addr}) sslmode={:?}, latency={}",
+            "connected to compute node at {host} ({socket_addr}) sslmode={:?}, latency={}, query_id={}",
             self.0.get_ssl_mode(),
             ctx.get_proxy_latency(),
+            ctx.get_testodrome_id().unwrap_or_default(),
         );
 
         // NB: CancelToken is supposed to hold socket_addr, but we use connect_raw.

proxy/src/context/mod.rs

@@ -63,6 +63,7 @@ struct RequestContextInner {
     success: bool,
     pub(crate) cold_start_info: ColdStartInfo,
     pg_options: Option<StartupMessageParams>,
+    testodrome_query_id: Option<String>,
 
     // extra
     // This sender is here to keep the request monitoring channel open while requests are taking place.
@@ -110,6 +111,7 @@ impl Clone for RequestContext {
             rejected: inner.rejected,
             cold_start_info: inner.cold_start_info,
             pg_options: inner.pg_options.clone(),
+            testodrome_query_id: inner.testodrome_query_id.clone(),
 
             sender: None,
             disconnect_sender: None,
@@ -160,6 +162,7 @@ impl RequestContext {
             rejected: None,
             cold_start_info: ColdStartInfo::Unknown,
             pg_options: None,
+            testodrome_query_id: None,
 
             sender: LOG_CHAN.get().and_then(|tx| tx.upgrade()),
             disconnect_sender: LOG_CHAN_DISCONNECT.get().and_then(|tx| tx.upgrade()),
@@ -210,6 +213,19 @@ impl RequestContext {
             this.set_dbname(dbname.into());
         }
 
+        // Try to get testodrome_query_id directly from parameters
+        if let Some(options_str) = options.get("options") {
+            // If not found directly, try to extract it from the options string
+            for option in options_str.split_whitespace() {
+                if option.starts_with("neon_query_id:") {
+                    if let Some(value) = option.strip_prefix("neon_query_id:") {
+                        this.set_testodrome_id(value.to_string());
+                        break;
+                    }
+                }
+            }
+        }
+
         this.pg_options = Some(options);
     }
 
@@ -256,6 +272,13 @@ impl RequestContext {
             .set_user_agent(user_agent);
     }
 
+    pub(crate) fn set_testodrome_id(&self, query_id: String) {
+        self.0
+            .try_lock()
+            .expect("should not deadlock")
+            .set_testodrome_id(query_id);
+    }
+
     pub(crate) fn set_auth_method(&self, auth_method: AuthMethod) {
         let mut this = self.0.try_lock().expect("should not deadlock");
         this.auth_method = Some(auth_method);
@@ -355,6 +378,14 @@ impl RequestContext {
             .accumulated()
     }
 
+    pub(crate) fn get_testodrome_id(&self) -> Option<String> {
+        self.0
+            .try_lock()
+            .expect("should not deadlock")
+            .testodrome_query_id
+            .clone()
+    }
+
     pub(crate) fn success(&self) {
         self.0
             .try_lock()
@@ -416,6 +447,10 @@ impl RequestContextInner {
         self.user = Some(user);
     }
 
+    fn set_testodrome_id(&mut self, query_id: String) {
+        self.testodrome_query_id = Some(query_id);
+    }
+
     fn has_private_peer_addr(&self) -> bool {
         match self.conn_info.addr.ip() {
             IpAddr::V4(ip) => ip.is_private(),

proxy/src/logging.rs

@@ -286,17 +286,16 @@ where
 
     /// Registers a SpanFields instance as span extension.
     fn on_new_span(&self, attrs: &span::Attributes<'_>, id: &span::Id, ctx: Context<'_, S>) {
-        let csid = self.callsite_id(attrs.metadata().callsite());
         let span = ctx.span(id).expect("span must exist");
         let fields = SpanFields::default();
         fields.record_fields(attrs);
+
         // This could deadlock when there's a panic somewhere in the tracing
         // event handling and a read or write guard is still held. This includes
         // the OTel subscriber.
         let mut exts = span.extensions_mut();
 
         exts.insert(fields);
-        exts.insert(csid);
     }
 
     fn on_record(&self, id: &span::Id, values: &span::Record<'_>, ctx: Context<'_, S>) {
@@ -565,6 +564,13 @@ impl EventFormatter {
                 )?;
             }
 
+            let spans = SerializableSpans {
+                ctx,
+                callsite_ids,
+                extract: ExtractedSpanFields::<'_, F>::new(extract_fields),
+            };
+            serializer.serialize_entry("spans", &spans)?;
+
             // TODO: thread-local cache?
             let pid = std::process::id();
             // Skip adding pid 1 to reduce noise for services running in containers.
@@ -614,15 +620,9 @@ impl EventFormatter {
                 }
             }
 
-            let stack = SerializableSpans {
-                ctx,
-                callsite_ids,
-                fields: ExtractedSpanFields::<'_, F>::new(extract_fields),
-            };
-            serializer.serialize_entry("spans", &stack)?;
-
-            if stack.fields.has_values() {
-                serializer.serialize_entry("extract", &stack.fields)?;
+            if spans.extract.has_values() {
+                // TODO: add fields from event, too?
+                serializer.serialize_entry("extract", &spans.extract)?;
             }
 
             serializer.end()
@@ -911,7 +911,7 @@ where
 {
     ctx: &'a Context<'ctx, Span>,
     callsite_ids: &'a papaya::HashMap<callsite::Identifier, CallsiteId>,
-    fields: ExtractedSpanFields<'a, F>,
+    extract: ExtractedSpanFields<'a, F>,
 }
 
 impl<Span, const F: usize> serde::ser::Serialize for SerializableSpans<'_, '_, Span, F>
@@ -940,7 +940,7 @@ where
 
                 serializer.serialize_value(&SerializableSpanFields {
                     span: &span,
-                    fields: &self.fields,
+                    extract: &self.extract,
                 })?;
             }
         }
@@ -955,7 +955,7 @@ where
     Span: for<'lookup> LookupSpan<'lookup>,
 {
     span: &'a SpanRef<'span, Span>,
-    fields: &'a ExtractedSpanFields<'a, F>,
+    extract: &'a ExtractedSpanFields<'a, F>,
 }
 
 impl<Span, const F: usize> serde::ser::Serialize for SerializableSpanFields<'_, '_, Span, F>
@@ -973,7 +973,7 @@ where
             for (name, value) in &data.fields.pin() {
                 serializer.serialize_entry(name, value)?;
                 // TODO: replace clone with reference, if possible.
-                self.fields.set(name, value.clone());
+                self.extract.set(name, value.clone());
             }
         }
 

proxy/src/metrics.rs

@@ -406,8 +406,11 @@ impl std::fmt::Display for LatencyAccumulated {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         write!(
             f,
-            "client: {:?}, cplane: {:?}, compute: {:?}, retry: {:?}",
-            self.client, self.cplane, self.compute, self.retry
+            "client: {}, cplane: {}, compute: {}, retry: {}",
+            self.client.as_micros(),
+            self.cplane.as_micros(),
+            self.compute.as_micros(),
+            self.retry.as_micros()
         )
     }
 }