refactor(pageserver): make image layer creation atomic

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

.github/workflows/benchmarking.yml

@@ -683,7 +683,7 @@ jobs:
       with:
         aws-region: eu-central-1
         role-to-assume: ${{ vars.DEV_AWS_OIDC_ROLE_ARN }}
-        role-duration-seconds: 18000 # 5 hours
+        role-duration-seconds: 43200 # 12 hours
 
     - name: Download Neon artifact
       uses: ./.github/actions/download

build-tools.Dockerfile

@@ -57,6 +57,18 @@ RUN set -e \
         zstd \
     && rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/*
 
+# sql_exporter
+
+# Keep the version the same as in compute/compute-node.Dockerfile and
+# test_runner/regress/test_compute_metrics.py.
+ENV SQL_EXPORTER_VERSION=0.13.1
+RUN curl -fsSL \
+    "https://github.com/burningalchemist/sql_exporter/releases/download/${SQL_EXPORTER_VERSION}/sql_exporter-${SQL_EXPORTER_VERSION}.linux-$(case "$(uname -m)" in x86_64) echo amd64;; aarch64) echo arm64;; esac).tar.gz" \
+    --output sql_exporter.tar.gz \
+    && mkdir /tmp/sql_exporter \
+    && tar xzvf sql_exporter.tar.gz -C /tmp/sql_exporter --strip-components=1 \
+    && mv /tmp/sql_exporter/sql_exporter /usr/local/bin/sql_exporter
+
 # protobuf-compiler (protoc)
 ENV PROTOC_VERSION=25.1
 RUN curl -fsSL "https://github.com/protocolbuffers/protobuf/releases/download/v${PROTOC_VERSION}/protoc-${PROTOC_VERSION}-linux-$(uname -m | sed 's/aarch64/aarch_64/g').zip" -o "protoc.zip" \

compute/Makefile

@@ -22,6 +22,7 @@ sql_exporter.yml: $(jsonnet_files)
 		--output-file etc/$@ \
 		--tla-str collector_name=neon_collector \
 		--tla-str collector_file=neon_collector.yml \
+		--tla-str 'connection_string=postgresql://cloud_admin@127.0.0.1:5432/postgres?sslmode=disable&application_name=sql_exporter' \
 		etc/sql_exporter.jsonnet
 
 sql_exporter_autoscaling.yml: $(jsonnet_files)
@@ -29,7 +30,7 @@ sql_exporter_autoscaling.yml: $(jsonnet_files)
 		--output-file etc/$@ \
 		--tla-str collector_name=neon_collector_autoscaling \
 		--tla-str collector_file=neon_collector_autoscaling.yml \
-		--tla-str application_name=sql_exporter_autoscaling \
+		--tla-str 'connection_string=postgresql://cloud_admin@127.0.0.1:5432/postgres?sslmode=disable&application_name=sql_exporter_autoscaling' \
 		etc/sql_exporter.jsonnet
 
 .PHONY: clean

compute/compute-node.Dockerfile

@@ -431,14 +431,11 @@ COPY --from=pg-build /usr/local/pgsql/ /usr/local/pgsql/
 
 COPY compute/patches/rum.patch /rum.patch
 
-# maybe version-specific
-# support for v17 is unknown
-# last release 1.3.13 - Sep 19, 2022
-RUN case "${PG_VERSION}" in "v17") \
-    echo "v17 extensions are not supported yet. Quit" && exit 0;; \
-    esac && \
-    wget https://github.com/postgrespro/rum/archive/refs/tags/1.3.13.tar.gz -O rum.tar.gz && \
-    echo "6ab370532c965568df6210bd844ac6ba649f53055e48243525b0b7e5c4d69a7d rum.tar.gz" | sha256sum --check && \
+# supports v17 since https://github.com/postgrespro/rum/commit/cb1edffc57736cd2a4455f8d0feab0d69928da25
+# doesn't use releases since 1.3.13 - Sep 19, 2022
+# use latest commit from the master branch
+RUN wget https://github.com/postgrespro/rum/archive/cb1edffc57736cd2a4455f8d0feab0d69928da25.tar.gz -O rum.tar.gz && \
+    echo "65e0a752e99f4c3226400c9b899f997049e93503db8bf5c8072efa136d32fd83 rum.tar.gz" | sha256sum --check && \
     mkdir rum-src && cd rum-src && tar xzf ../rum.tar.gz --strip-components=1 -C . && \
     patch -p1 < /rum.patch && \
     make -j $(getconf _NPROCESSORS_ONLN) PG_CONFIG=/usr/local/pgsql/bin/pg_config USE_PGXS=1 && \
@@ -959,21 +956,31 @@ RUN apt-get install -y protobuf-compiler && \
 #
 #########################################################################################
 
-FROM rust-extensions-build AS pg-jsonschema-pg-build
+FROM rust-extensions-build-pgrx12 AS pg-jsonschema-pg-build
 ARG PG_VERSION
-
-RUN case "${PG_VERSION}" in "v17") \
-    echo "pg_jsonschema does not yet have a release that supports pg17" && exit 0;; \
+# version 0.3.3 supports v17
+# last release v0.3.3 - Oct 16, 2024
+#
+# there were no breaking changes
+# so we can use the same version for all postgres versions
+RUN case "${PG_VERSION}" in \
+    "v14" | "v15" | "v16" | "v17") \
+        export PG_JSONSCHEMA_VERSION=0.3.3 \
+        export PG_JSONSCHEMA_CHECKSUM=40c2cffab4187e0233cb8c3bde013be92218c282f95f4469c5282f6b30d64eac \
+    ;; \
+    *) \
+        echo "unexpected PostgreSQL version" && exit 1 \
+    ;; \
     esac && \
-    wget https://github.com/supabase/pg_jsonschema/archive/refs/tags/v0.3.1.tar.gz -O pg_jsonschema.tar.gz && \
-    echo "61df3db1ed83cf24f6aa39c826f8818bfa4f0bd33b587fd6b2b1747985642297 pg_jsonschema.tar.gz" | sha256sum --check && \
+    wget https://github.com/supabase/pg_jsonschema/archive/refs/tags/v${PG_JSONSCHEMA_VERSION}.tar.gz -O pg_jsonschema.tar.gz && \
+    echo "${PG_JSONSCHEMA_CHECKSUM} pg_jsonschema.tar.gz" | sha256sum --check && \
     mkdir pg_jsonschema-src && cd pg_jsonschema-src && tar xzf ../pg_jsonschema.tar.gz --strip-components=1 -C . && \
     # see commit 252b3685a27a0f4c31a0f91e983c6314838e89e8
     # `unsafe-postgres` feature allows to build pgx extensions
     # against postgres forks that decided to change their ABI name (like us).
     # With that we can build extensions without forking them and using stock
     # pgx. As this feature is new few manual version bumps were required.
-    sed -i 's/pgrx = "0.11.3"/pgrx = { version = "0.11.3", features = [ "unsafe-postgres" ] }/g' Cargo.toml && \
+    sed -i 's/pgrx = "0.12.6"/pgrx = { version = "0.12.6", features = [ "unsafe-postgres" ] }/g' Cargo.toml && \
     cargo pgrx install --release && \
     echo "trusted = true" >> /usr/local/pgsql/share/extension/pg_jsonschema.control
 
@@ -984,16 +991,27 @@ RUN case "${PG_VERSION}" in "v17") \
 #
 #########################################################################################
 
-FROM rust-extensions-build AS pg-graphql-pg-build
+FROM rust-extensions-build-pgrx12 AS pg-graphql-pg-build
 ARG PG_VERSION
 
-RUN case "${PG_VERSION}" in "v17") \
-    echo "pg_graphql does not yet have a release that supports pg17 as of now" && exit 0;; \
+# version 1.5.9 supports v17
+# last release v1.5.9 - Oct 16, 2024
+#
+# there were no breaking changes
+# so we can use the same version for all postgres versions
+RUN case "${PG_VERSION}" in \
+    "v14" | "v15" | "v16" | "v17") \
+        export PG_GRAPHQL_VERSION=1.5.9 \
+        export PG_GRAPHQL_CHECKSUM=cf768385a41278be1333472204fc0328118644ae443182cf52f7b9b23277e497 \
+    ;; \
+    *) \
+        echo "unexpected PostgreSQL version" && exit 1 \
+    ;; \
     esac && \
-    wget https://github.com/supabase/pg_graphql/archive/refs/tags/v1.5.7.tar.gz -O pg_graphql.tar.gz && \
-    echo "2b3e567a5b31019cb97ae0e33263c1bcc28580be5a444ac4c8ece5c4be2aea41 pg_graphql.tar.gz" | sha256sum --check && \
+    wget https://github.com/supabase/pg_graphql/archive/refs/tags/v${PG_GRAPHQL_VERSION}.tar.gz -O pg_graphql.tar.gz && \
+    echo "${PG_GRAPHQL_CHECKSUM} pg_graphql.tar.gz" | sha256sum --check && \
     mkdir pg_graphql-src && cd pg_graphql-src && tar xzf ../pg_graphql.tar.gz --strip-components=1 -C . && \
-    sed -i 's/pgrx = "=0.11.3"/pgrx = { version = "0.11.3", features = [ "unsafe-postgres" ] }/g' Cargo.toml && \
+    sed -i 's/pgrx = "=0.12.6"/pgrx = { version = "0.12.6", features = [ "unsafe-postgres" ] }/g' Cargo.toml && \
     cargo pgrx install --release && \
     # it's needed to enable extension because it uses untrusted C language
     sed -i 's/superuser = false/superuser = true/g' /usr/local/pgsql/share/extension/pg_graphql.control && \
@@ -1006,15 +1024,13 @@ RUN case "${PG_VERSION}" in "v17") \
 #
 #########################################################################################
 
-FROM rust-extensions-build AS pg-tiktoken-pg-build
+FROM rust-extensions-build-pgrx12 AS pg-tiktoken-pg-build
 ARG PG_VERSION
 
-# 26806147b17b60763039c6a6878884c41a262318 made on 26/09/2023
-RUN case "${PG_VERSION}" in "v17") \
-    echo "pg_tiktoken does not have versions, nor support for pg17" && exit 0;; \
-    esac && \
-    wget https://github.com/kelvich/pg_tiktoken/archive/26806147b17b60763039c6a6878884c41a262318.tar.gz -O pg_tiktoken.tar.gz && \
-    echo "e64e55aaa38c259512d3e27c572da22c4637418cf124caba904cd50944e5004e pg_tiktoken.tar.gz" | sha256sum --check && \
+# doesn't use releases
+# 9118dd4549b7d8c0bbc98e04322499f7bf2fa6f7 - on Oct 29, 2024
+RUN wget https://github.com/kelvich/pg_tiktoken/archive/9118dd4549b7d8c0bbc98e04322499f7bf2fa6f7.tar.gz -O pg_tiktoken.tar.gz && \
+    echo "a5bc447e7920ee149d3c064b8b9f0086c0e83939499753178f7d35788416f628 pg_tiktoken.tar.gz" | sha256sum --check && \
     mkdir pg_tiktoken-src && cd pg_tiktoken-src && tar xzf ../pg_tiktoken.tar.gz --strip-components=1 -C . && \
     # TODO update pgrx version in the pg_tiktoken repo and remove this line
     sed -i 's/pgrx = { version = "=0.10.2",/pgrx = { version = "0.11.3",/g' Cargo.toml && \
@@ -1032,6 +1048,8 @@ RUN case "${PG_VERSION}" in "v17") \
 FROM rust-extensions-build AS pg-pgx-ulid-build
 ARG PG_VERSION
 
+# doesn't support v17 yet
+# https://github.com/pksunkara/pgx_ulid/pull/52
 RUN case "${PG_VERSION}" in "v17") \
     echo "pgx_ulid does not support pg17 as of the latest version (0.1.5)" && exit 0;; \
     esac && \
@@ -1049,16 +1067,16 @@ RUN case "${PG_VERSION}" in "v17") \
 #
 #########################################################################################
 
-FROM rust-extensions-build AS pg-session-jwt-build
+FROM rust-extensions-build-pgrx12 AS pg-session-jwt-build
 ARG PG_VERSION
 
-RUN case "${PG_VERSION}" in "v17") \
-    echo "pg_session_jwt does not yet have a release that supports pg17" && exit 0;; \
-    esac && \
-    wget https://github.com/neondatabase/pg_session_jwt/archive/e1310b08ba51377a19e0559e4d1194883b9b2ba2.tar.gz -O pg_session_jwt.tar.gz && \
-    echo "837932a077888d5545fd54b0abcc79e5f8e37017c2769a930afc2f5c94df6f4e pg_session_jwt.tar.gz" | sha256sum --check && \
+# NOTE: local_proxy depends on the version of pg_session_jwt
+# Do not update without approve from proxy team
+# Make sure the version is reflected in proxy/src/serverless/local_conn_pool.rs
+RUN wget https://github.com/neondatabase/pg_session_jwt/archive/refs/tags/v0.1.2-v17.tar.gz -O pg_session_jwt.tar.gz && \
+    echo "c8ecbed9cb8c6441bce5134a176002b043018adf9d05a08e457dda233090a86e pg_session_jwt.tar.gz" | sha256sum --check && \
     mkdir pg_session_jwt-src && cd pg_session_jwt-src && tar xzf ../pg_session_jwt.tar.gz --strip-components=1 -C . && \
-    sed -i 's/pgrx = "=0.11.3"/pgrx = { version = "=0.11.3", features = [ "unsafe-postgres" ] }/g' Cargo.toml && \
+    sed -i 's/pgrx = "0.12.6"/pgrx = { version = "=0.12.6", features = [ "unsafe-postgres" ] }/g' Cargo.toml && \
     cargo pgrx install --release
 
 #########################################################################################
@@ -1131,7 +1149,10 @@ FROM rust-extensions-build AS pg-mooncake-build
 ARG PG_VERSION
 COPY --from=pg-build /usr/local/pgsql/ /usr/local/pgsql/
 
-ENV PG_MOONCAKE_VERSION=0a7de4c0b5c7b1a5e2175e1c5f4625b97b7346f1
+# The topmost commit in the `neon` branch at the time of writing this
+# https://github.com/Mooncake-Labs/pg_mooncake/commits/neon/
+# https://github.com/Mooncake-Labs/pg_mooncake/commit/568b5a82b5fc16136bdf4ca5aac3e0cc261ab48d
+ENV PG_MOONCAKE_VERSION=568b5a82b5fc16136bdf4ca5aac3e0cc261ab48d
 ENV PATH="/usr/local/pgsql/bin/:$PATH"
 
 RUN case "${PG_VERSION}" in \
@@ -1298,7 +1319,10 @@ RUN mold -run cargo build --locked --profile release-line-debug-size-lto --bin l
 #########################################################################################
 
 FROM quay.io/prometheuscommunity/postgres-exporter:v0.12.1 AS postgres-exporter
-FROM burningalchemist/sql_exporter:0.13 AS sql-exporter
+
+# Keep the version the same as in build-tools.Dockerfile and
+# test_runner/regress/test_compute_metrics.py.
+FROM burningalchemist/sql_exporter:0.13.1 AS sql-exporter
 
 #########################################################################################
 #

compute/etc/neon_collector.jsonnet

@@ -3,6 +3,7 @@
   metrics: [
     import 'sql_exporter/checkpoints_req.libsonnet',
     import 'sql_exporter/checkpoints_timed.libsonnet',
+    import 'sql_exporter/compute_backpressure_throttling_ms.libsonnet',
     import 'sql_exporter/compute_current_lsn.libsonnet',
     import 'sql_exporter/compute_logical_snapshot_files.libsonnet',
     import 'sql_exporter/compute_receive_lsn.libsonnet',

compute/etc/sql_exporter.jsonnet

@@ -1,4 +1,4 @@
-function(collector_name, collector_file, application_name='sql_exporter') {
+function(collector_name, collector_file, connection_string) {
   // Configuration for sql_exporter for autoscaling-agent
   // Global defaults.
   global: {
@@ -23,7 +23,7 @@ function(collector_name, collector_file, application_name='sql_exporter') {
   target: {
     // Data source name always has a URI schema that matches the driver name. In some cases (e.g. MySQL)
     // the schema gets dropped or replaced to match the driver expected DSN format.
-    data_source_name: std.format('postgresql://cloud_admin@127.0.0.1:5432/postgres?sslmode=disable&application_name=%s', [application_name]),
+    data_source_name: connection_string,
 
     // Collectors (referenced by name) to execute on the target.
     // Glob patterns are supported (see <https://pkg.go.dev/path/filepath#Match> for syntax).

compute/etc/sql_exporter/compute_backpressure_throttling_ms.libsonnet

@@ -0,0 +1,10 @@
+{
+  metric_name: 'compute_backpressure_throttling_ms',
+  type: 'gauge',
+  help: 'Time compute has spent throttled',
+  key_labels: null,
+  values: [
+    'throttled',
+  ],
+  query: importstr 'sql_exporter/compute_backpressure_throttling_ms.sql',
+}

compute/etc/sql_exporter/compute_backpressure_throttling_ms.sql

@@ -0,0 +1 @@
+SELECT neon.backpressure_throttling_time() AS throttled;

compute_tools/src/compute.rs

@@ -1,7 +1,6 @@
 use std::collections::HashMap;
 use std::env;
 use std::fs;
-use std::io::BufRead;
 use std::os::unix::fs::{symlink, PermissionsExt};
 use std::path::Path;
 use std::process::{Command, Stdio};
@@ -365,8 +364,7 @@ impl ComputeNode {
         let pageserver_connect_micros = start_time.elapsed().as_micros() as u64;
 
         let basebackup_cmd = match lsn {
-            // HACK We don't use compression on first start (Lsn(0)) because there's no API for it
-            Lsn(0) => format!("basebackup {} {}", spec.tenant_id, spec.timeline_id),
+            Lsn(0) => format!("basebackup {} {} --gzip", spec.tenant_id, spec.timeline_id),
             _ => format!(
                 "basebackup {} {} {} --gzip",
                 spec.tenant_id, spec.timeline_id, lsn
@@ -375,38 +373,16 @@ impl ComputeNode {
 
         let copyreader = client.copy_out(basebackup_cmd.as_str())?;
         let mut measured_reader = MeasuredReader::new(copyreader);
-
-        // Check the magic number to see if it's a gzip or not. Even though
-        // we might explicitly ask for gzip, an old pageserver with no implementation
-        // of gzip compression might send us uncompressed data. After some time
-        // passes we can assume all pageservers know how to compress and we can
-        // delete this check.
-        //
-        // If the data is not gzip, it will be tar. It will not be mistakenly
-        // recognized as gzip because tar starts with an ascii encoding of a filename,
-        // and 0x1f and 0x8b are unlikely first characters for any filename. Moreover,
-        // we send the "global" directory first from the pageserver, so it definitely
-        // won't be recognized as gzip.
         let mut bufreader = std::io::BufReader::new(&mut measured_reader);
-        let gzip = {
-            let peek = bufreader.fill_buf().unwrap();
-            peek[0] == 0x1f && peek[1] == 0x8b
-        };
 
         // Read the archive directly from the `CopyOutReader`
         //
         // Set `ignore_zeros` so that unpack() reads all the Copy data and
         // doesn't stop at the end-of-archive marker. Otherwise, if the server
         // sends an Error after finishing the tarball, we will not notice it.
-        if gzip {
-            let mut ar = tar::Archive::new(flate2::read::GzDecoder::new(&mut bufreader));
-            ar.set_ignore_zeros(true);
-            ar.unpack(&self.pgdata)?;
-        } else {
-            let mut ar = tar::Archive::new(&mut bufreader);
-            ar.set_ignore_zeros(true);
-            ar.unpack(&self.pgdata)?;
-        };
+        let mut ar = tar::Archive::new(flate2::read::GzDecoder::new(&mut bufreader));
+        ar.set_ignore_zeros(true);
+        ar.unpack(&self.pgdata)?;
 
         // Report metrics
         let mut state = self.state.lock().unwrap();

libs/postgres_ffi/src/lib.rs

@@ -36,6 +36,7 @@ macro_rules! postgres_ffi {
             pub mod controlfile_utils;
             pub mod nonrelfile_utils;
             pub mod wal_craft_test_export;
+            pub mod wal_generator;
             pub mod waldecoder_handler;
             pub mod xlog_utils;
 

libs/postgres_ffi/src/wal_generator.rs

@@ -0,0 +1,203 @@
+use std::ffi::CStr;
+
+use bytes::{Bytes, BytesMut};
+use crc32c::crc32c_append;
+use utils::lsn::Lsn;
+
+use super::bindings::{XLogLongPageHeaderData, XLogPageHeaderData, XLOG_PAGE_MAGIC};
+use super::xlog_utils::{
+    XlLogicalMessage, XLOG_RECORD_CRC_OFFS, XLOG_SIZE_OF_XLOG_RECORD, XLP_BKP_REMOVABLE,
+    XLP_FIRST_IS_CONTRECORD,
+};
+use super::XLogRecord;
+use crate::pg_constants::{
+    RM_LOGICALMSG_ID, XLOG_LOGICAL_MESSAGE, XLP_LONG_HEADER, XLR_BLOCK_ID_DATA_LONG,
+    XLR_BLOCK_ID_DATA_SHORT,
+};
+use crate::{WAL_SEGMENT_SIZE, XLOG_BLCKSZ};
+
+/// Generates binary WAL records for use in tests and benchmarks. Currently only generates logical
+/// messages (effectively noops) with a fixed payload. It is used as an iterator which yields
+/// encoded bytes for a single WAL record, including internal page headers if it spans pages.
+/// Concatenating the bytes will yield a complete, well-formed WAL, which can be chunked at segment
+/// boundaries if desired. Not optimized for performance.
+///
+/// The WAL format is version-dependant (see e.g. `XLOG_PAGE_MAGIC`), so make sure to import this
+/// for the appropriate Postgres version (e.g. `postgres_ffi::v17::wal_generator::WalGenerator`).
+///
+/// A WAL is split into 16 MB segments. Each segment is split into 8 KB pages, with headers.
+/// Records are arbitrary length, 8-byte aligned, and may span pages. The layout is e.g.:
+///
+/// |        Segment 1         |        Segment 2         |        Segment 3         |
+/// | Page 1 | Page 2 | Page 3 | Page 4 | Page 5 | Page 6 | Page 7 | Page 8 | Page 9 |
+/// | R1 |   R2  |R3|  R4  | R5  |  R6  |                 R7            | R8  |
+///
+/// TODO: support generating actual tables and rows.
+#[derive(Default)]
+pub struct WalGenerator {
+    /// Current LSN to append the next record at.
+    ///
+    /// Callers can modify this (and prev_lsn) to restart generation at a different LSN, but should
+    /// ensure that the LSN is on a valid record boundary (i.e. we can't start appending in the
+    /// middle on an existing record or header, or beyond the end of the existing WAL).
+    pub lsn: Lsn,
+    /// The starting LSN of the previous record. Used in WAL record headers. The Safekeeper doesn't
+    /// care about this, unlike Postgres, but we include it for completeness.
+    pub prev_lsn: Lsn,
+}
+
+impl WalGenerator {
+    // For now, hardcode the message payload.
+    // TODO: support specifying the payload size.
+    const PREFIX: &CStr = c"prefix";
+    const MESSAGE: &[u8] = b"message";
+
+    // Hardcode the sys, timeline, and DB IDs. We can make them configurable if we care about them.
+    const SYS_ID: u64 = 0;
+    const TIMELINE_ID: u32 = 1;
+    const DB_ID: u32 = 0;
+
+    /// Creates a new WAL generator, which emits logical message records (noops).
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    /// Encodes a logical message (basically a noop), with the given prefix and message.
+    pub(crate) fn encode_logical_message(prefix: &CStr, message: &[u8]) -> Bytes {
+        let prefix = prefix.to_bytes_with_nul();
+        let header = XlLogicalMessage {
+            db_id: Self::DB_ID,
+            transactional: 0,
+            prefix_size: prefix.len() as u64,
+            message_size: message.len() as u64,
+        };
+        [&header.encode(), prefix, message].concat().into()
+    }
+
+    /// Encode a WAL record with the given payload data (e.g. a logical message).
+    pub(crate) fn encode_record(data: Bytes, rmid: u8, info: u8, prev_lsn: Lsn) -> Bytes {
+        // Prefix data with block ID and length.
+        let data_header = Bytes::from(match data.len() {
+            0 => vec![],
+            1..=255 => vec![XLR_BLOCK_ID_DATA_SHORT, data.len() as u8],
+            256.. => {
+                let len_bytes = (data.len() as u32).to_le_bytes();
+                [&[XLR_BLOCK_ID_DATA_LONG], len_bytes.as_slice()].concat()
+            }
+        });
+
+        // Construct the WAL record header.
+        let mut header = XLogRecord {
+            xl_tot_len: (XLOG_SIZE_OF_XLOG_RECORD + data_header.len() + data.len()) as u32,
+            xl_xid: 0,
+            xl_prev: prev_lsn.into(),
+            xl_info: info,
+            xl_rmid: rmid,
+            __bindgen_padding_0: [0; 2],
+            xl_crc: 0, // see below
+        };
+
+        // Compute the CRC checksum for the data, and the header up to the CRC field.
+        let mut crc = 0;
+        crc = crc32c_append(crc, &data_header);
+        crc = crc32c_append(crc, &data);
+        crc = crc32c_append(crc, &header.encode().unwrap()[0..XLOG_RECORD_CRC_OFFS]);
+        header.xl_crc = crc;
+
+        // Encode the final header and record.
+        let header = header.encode().unwrap();
+
+        [header, data_header, data].concat().into()
+    }
+
+    /// Injects page headers on 8KB page boundaries. Takes the current LSN position where the record
+    /// is to be appended.
+    fn encode_pages(record: Bytes, mut lsn: Lsn) -> Bytes {
+        // Fast path: record fits in current page, and the page already has a header.
+        if lsn.remaining_in_block() as usize >= record.len() && lsn.block_offset() > 0 {
+            return record;
+        }
+
+        let mut pages = BytesMut::new();
+        let mut remaining = record.clone(); // Bytes::clone() is cheap
+        while !remaining.is_empty() {
+            // At new page boundary, inject page header.
+            if lsn.block_offset() == 0 {
+                let mut page_header = XLogPageHeaderData {
+                    xlp_magic: XLOG_PAGE_MAGIC as u16,
+                    xlp_info: XLP_BKP_REMOVABLE,
+                    xlp_tli: Self::TIMELINE_ID,
+                    xlp_pageaddr: lsn.0,
+                    xlp_rem_len: 0,
+                    __bindgen_padding_0: [0; 4],
+                };
+                // If the record was split across page boundaries, mark as continuation.
+                if remaining.len() < record.len() {
+                    page_header.xlp_rem_len = remaining.len() as u32;
+                    page_header.xlp_info |= XLP_FIRST_IS_CONTRECORD;
+                }
+                // At start of segment, use a long page header.
+                let page_header = if lsn.segment_offset(WAL_SEGMENT_SIZE) == 0 {
+                    page_header.xlp_info |= XLP_LONG_HEADER;
+                    XLogLongPageHeaderData {
+                        std: page_header,
+                        xlp_sysid: Self::SYS_ID,
+                        xlp_seg_size: WAL_SEGMENT_SIZE as u32,
+                        xlp_xlog_blcksz: XLOG_BLCKSZ as u32,
+                    }
+                    .encode()
+                    .unwrap()
+                } else {
+                    page_header.encode().unwrap()
+                };
+                pages.extend_from_slice(&page_header);
+                lsn += page_header.len() as u64;
+            }
+
+            // Append the record up to the next page boundary, if any.
+            let page_free = lsn.remaining_in_block() as usize;
+            let chunk = remaining.split_to(std::cmp::min(page_free, remaining.len()));
+            pages.extend_from_slice(&chunk);
+            lsn += chunk.len() as u64;
+        }
+        pages.freeze()
+    }
+
+    /// Records must be 8-byte aligned. Take an encoded record (including any injected page
+    /// boundaries), starting at the given LSN, and add any necessary padding at the end.
+    fn pad_record(record: Bytes, mut lsn: Lsn) -> Bytes {
+        lsn += record.len() as u64;
+        let padding = lsn.calc_padding(8u64) as usize;
+        if padding == 0 {
+            return record;
+        }
+        [record, Bytes::from(vec![0; padding])].concat().into()
+    }
+
+    /// Generates a record with an arbitrary payload at the current LSN, then increments the LSN.
+    pub fn generate_record(&mut self, data: Bytes, rmid: u8, info: u8) -> Bytes {
+        let record = Self::encode_record(data, rmid, info, self.prev_lsn);
+        let record = Self::encode_pages(record, self.lsn);
+        let record = Self::pad_record(record, self.lsn);
+        self.prev_lsn = self.lsn;
+        self.lsn += record.len() as u64;
+        record
+    }
+
+    /// Generates a logical message at the current LSN. Can be used to construct arbitrary messages.
+    pub fn generate_logical_message(&mut self, prefix: &CStr, message: &[u8]) -> Bytes {
+        let data = Self::encode_logical_message(prefix, message);
+        self.generate_record(data, RM_LOGICALMSG_ID, XLOG_LOGICAL_MESSAGE)
+    }
+}
+
+/// Generate WAL records as an iterator.
+impl Iterator for WalGenerator {
+    type Item = (Lsn, Bytes);
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let lsn = self.lsn;
+        let record = self.generate_logical_message(Self::PREFIX, Self::MESSAGE);
+        Some((lsn, record))
+    }
+}

libs/postgres_ffi/src/xlog_utils.rs

@@ -7,15 +7,14 @@
 // have been named the same as the corresponding PostgreSQL functions instead.
 //
 
-use crc32c::crc32c_append;
-
 use super::super::waldecoder::WalStreamDecoder;
 use super::bindings::{
     CheckPoint, ControlFileData, DBState_DB_SHUTDOWNED, FullTransactionId, TimeLineID, TimestampTz,
     XLogLongPageHeaderData, XLogPageHeaderData, XLogRecPtr, XLogRecord, XLogSegNo, XLOG_PAGE_MAGIC,
 };
+use super::wal_generator::WalGenerator;
 use super::PG_MAJORVERSION;
-use crate::pg_constants;
+use crate::pg_constants::{self, RM_LOGICALMSG_ID, XLOG_LOGICAL_MESSAGE};
 use crate::PG_TLI;
 use crate::{uint32, uint64, Oid};
 use crate::{WAL_SEGMENT_SIZE, XLOG_BLCKSZ};
@@ -26,7 +25,7 @@ use bytes::{Buf, Bytes};
 use log::*;
 
 use serde::Serialize;
-use std::ffi::OsStr;
+use std::ffi::{CString, OsStr};
 use std::fs::File;
 use std::io::prelude::*;
 use std::io::ErrorKind;
@@ -39,6 +38,7 @@ use utils::bin_ser::SerializeError;
 use utils::lsn::Lsn;
 
 pub const XLOG_FNAME_LEN: usize = 24;
+pub const XLP_BKP_REMOVABLE: u16 = 0x0004;
 pub const XLP_FIRST_IS_CONTRECORD: u16 = 0x0001;
 pub const XLP_REM_LEN_OFFS: usize = 2 + 2 + 4 + 8;
 pub const XLOG_RECORD_CRC_OFFS: usize = 4 + 4 + 8 + 1 + 1 + 2;
@@ -489,64 +489,16 @@ impl XlLogicalMessage {
 /// Create new WAL record for non-transactional logical message.
 /// Used for creating artificial WAL for tests, as LogicalMessage
 /// record is basically no-op.
-///
-/// NOTE: This leaves the xl_prev field zero. The safekeeper and
-/// pageserver tolerate that, but PostgreSQL does not.
-pub fn encode_logical_message(prefix: &str, message: &str) -> Vec<u8> {
-    let mut prefix_bytes: Vec<u8> = Vec::with_capacity(prefix.len() + 1);
-    prefix_bytes.write_all(prefix.as_bytes()).unwrap();
-    prefix_bytes.push(0);
-
-    let message_bytes = message.as_bytes();
-
-    let logical_message = XlLogicalMessage {
-        db_id: 0,
-        transactional: 0,
-        prefix_size: prefix_bytes.len() as u64,
-        message_size: message_bytes.len() as u64,
-    };
-
-    let mainrdata = logical_message.encode();
-    let mainrdata_len: usize = mainrdata.len() + prefix_bytes.len() + message_bytes.len();
-    // only short mainrdata is supported for now
-    assert!(mainrdata_len <= 255);
-    let mainrdata_len = mainrdata_len as u8;
-
-    let mut data: Vec<u8> = vec![pg_constants::XLR_BLOCK_ID_DATA_SHORT, mainrdata_len];
-    data.extend_from_slice(&mainrdata);
-    data.extend_from_slice(&prefix_bytes);
-    data.extend_from_slice(message_bytes);
-
-    let total_len = XLOG_SIZE_OF_XLOG_RECORD + data.len();
-
-    let mut header = XLogRecord {
-        xl_tot_len: total_len as u32,
-        xl_xid: 0,
-        xl_prev: 0,
-        xl_info: 0,
-        xl_rmid: 21,
-        __bindgen_padding_0: [0u8; 2usize],
-        xl_crc: 0, // crc will be calculated later
-    };
-
-    let header_bytes = header.encode().expect("failed to encode header");
-    let crc = crc32c_append(0, &data);
-    let crc = crc32c_append(crc, &header_bytes[0..XLOG_RECORD_CRC_OFFS]);
-    header.xl_crc = crc;
-
-    let mut wal: Vec<u8> = Vec::new();
-    wal.extend_from_slice(&header.encode().expect("failed to encode header"));
-    wal.extend_from_slice(&data);
-
-    // WAL start position must be aligned at 8 bytes,
-    // this will add padding for the next WAL record.
-    const PADDING: usize = 8;
-    let padding_rem = wal.len() % PADDING;
-    if padding_rem != 0 {
-        wal.resize(wal.len() + PADDING - padding_rem, 0);
-    }
-
-    wal
+pub fn encode_logical_message(prefix: &str, message: &str) -> Bytes {
+    // This function can take untrusted input, so discard any NUL bytes in the prefix string.
+    let prefix = CString::new(prefix.replace('\0', "")).expect("no NULs");
+    let message = message.as_bytes();
+    WalGenerator::encode_record(
+        WalGenerator::encode_logical_message(&prefix, message),
+        RM_LOGICALMSG_ID,
+        XLOG_LOGICAL_MESSAGE,
+        Lsn(0),
+    )
 }
 
 #[cfg(test)]

libs/utils/src/lsn.rs

@@ -12,7 +12,7 @@ use crate::seqwait::MonotonicCounter;
 pub const XLOG_BLCKSZ: u32 = 8192;
 
 /// A Postgres LSN (Log Sequence Number), also known as an XLogRecPtr
-#[derive(Clone, Copy, Eq, Ord, PartialEq, PartialOrd, Hash)]
+#[derive(Clone, Copy, Default, Eq, Ord, PartialEq, PartialOrd, Hash)]
 pub struct Lsn(pub u64);
 
 impl Serialize for Lsn {

libs/wal_decoder/src/decoder.rs

@@ -1 +1,970 @@
+//! This module contains logic for decoding and interpreting
+//! raw bytes which represent a raw Postgres WAL record.
 
+use crate::models::*;
+use bytes::{Buf, Bytes, BytesMut};
+use pageserver_api::key::rel_block_to_key;
+use pageserver_api::record::NeonWalRecord;
+use pageserver_api::reltag::{RelTag, SlruKind};
+use pageserver_api::shard::ShardIdentity;
+use pageserver_api::value::Value;
+use postgres_ffi::relfile_utils::VISIBILITYMAP_FORKNUM;
+use postgres_ffi::walrecord::*;
+use postgres_ffi::{page_is_new, page_set_lsn, pg_constants, BLCKSZ};
+use utils::lsn::Lsn;
+
+impl InterpretedWalRecord {
+    /// Decode and interpreted raw bytes which represent one Postgres WAL record.
+    /// Data blocks which do not match the provided shard identity are filtered out.
+    /// Shard 0 is a special case since it tracks all relation sizes. We only give it
+    /// the keys that are being written as that is enough for updating relation sizes.
+    pub fn from_bytes_filtered(
+        buf: Bytes,
+        shard: &ShardIdentity,
+        lsn: Lsn,
+        pg_version: u32,
+    ) -> anyhow::Result<InterpretedWalRecord> {
+        let mut decoded = DecodedWALRecord::default();
+        decode_wal_record(buf, &mut decoded, pg_version)?;
+
+        let flush_uncommitted = if decoded.is_dbase_create_copy(pg_version) {
+            FlushUncommittedRecords::Yes
+        } else {
+            FlushUncommittedRecords::No
+        };
+
+        let metadata_record = MetadataRecord::from_decoded(&decoded, lsn, pg_version)?;
+
+        let mut blocks = Vec::default();
+        for blk in decoded.blocks.iter() {
+            let rel = RelTag {
+                spcnode: blk.rnode_spcnode,
+                dbnode: blk.rnode_dbnode,
+                relnode: blk.rnode_relnode,
+                forknum: blk.forknum,
+            };
+
+            let key = rel_block_to_key(rel, blk.blkno);
+
+            if !key.is_valid_key_on_write_path() {
+                anyhow::bail!("Unsupported key decoded at LSN {}: {}", lsn, key);
+            }
+
+            let key_is_local = shard.is_key_local(&key);
+
+            tracing::debug!(
+                lsn=%lsn,
+                key=%key,
+                "ingest: shard decision {}",
+                if !key_is_local { "drop" } else { "keep" },
+            );
+
+            if !key_is_local {
+                if shard.is_shard_zero() {
+                    // Shard 0 tracks relation sizes.  Although we will not store this block, we will observe
+                    // its blkno in case it implicitly extends a relation.
+                    blocks.push((key.to_compact(), None));
+                }
+
+                continue;
+            }
+
+            // Instead of storing full-page-image WAL record,
+            // it is better to store extracted image: we can skip wal-redo
+            // in this case. Also some FPI records may contain multiple (up to 32) pages,
+            // so them have to be copied multiple times.
+            //
+            let value = if blk.apply_image
+                && blk.has_image
+                && decoded.xl_rmid == pg_constants::RM_XLOG_ID
+                && (decoded.xl_info == pg_constants::XLOG_FPI
+                || decoded.xl_info == pg_constants::XLOG_FPI_FOR_HINT)
+                // compression of WAL is not yet supported: fall back to storing the original WAL record
+                && !postgres_ffi::bkpimage_is_compressed(blk.bimg_info, pg_version)
+                // do not materialize null pages because them most likely be soon replaced with real data
+                && blk.bimg_len != 0
+            {
+                // Extract page image from FPI record
+                let img_len = blk.bimg_len as usize;
+                let img_offs = blk.bimg_offset as usize;
+                let mut image = BytesMut::with_capacity(BLCKSZ as usize);
+                // TODO(vlad): skip the copy
+                image.extend_from_slice(&decoded.record[img_offs..img_offs + img_len]);
+
+                if blk.hole_length != 0 {
+                    let tail = image.split_off(blk.hole_offset as usize);
+                    image.resize(image.len() + blk.hole_length as usize, 0u8);
+                    image.unsplit(tail);
+                }
+                //
+                // Match the logic of XLogReadBufferForRedoExtended:
+                // The page may be uninitialized. If so, we can't set the LSN because
+                // that would corrupt the page.
+                //
+                if !page_is_new(&image) {
+                    page_set_lsn(&mut image, lsn)
+                }
+                assert_eq!(image.len(), BLCKSZ as usize);
+
+                Value::Image(image.freeze())
+            } else {
+                Value::WalRecord(NeonWalRecord::Postgres {
+                    will_init: blk.will_init || blk.apply_image,
+                    rec: decoded.record.clone(),
+                })
+            };
+
+            blocks.push((key.to_compact(), Some(value)));
+        }
+
+        Ok(InterpretedWalRecord {
+            metadata_record,
+            blocks,
+            lsn,
+            flush_uncommitted,
+            xid: decoded.xl_xid,
+        })
+    }
+}
+
+impl MetadataRecord {
+    fn from_decoded(
+        decoded: &DecodedWALRecord,
+        lsn: Lsn,
+        pg_version: u32,
+    ) -> anyhow::Result<Option<MetadataRecord>> {
+        // Note: this doesn't actually copy the bytes since
+        // the [`Bytes`] type implements it via a level of indirection.
+        let mut buf = decoded.record.clone();
+        buf.advance(decoded.main_data_offset);
+
+        match decoded.xl_rmid {
+            pg_constants::RM_HEAP_ID | pg_constants::RM_HEAP2_ID => {
+                Self::decode_heapam_record(&mut buf, decoded, pg_version)
+            }
+            pg_constants::RM_NEON_ID => Self::decode_neonmgr_record(&mut buf, decoded, pg_version),
+            // Handle other special record types
+            pg_constants::RM_SMGR_ID => Self::decode_smgr_record(&mut buf, decoded),
+            pg_constants::RM_DBASE_ID => Self::decode_dbase_record(&mut buf, decoded, pg_version),
+            pg_constants::RM_TBLSPC_ID => {
+                tracing::trace!("XLOG_TBLSPC_CREATE/DROP is not handled yet");
+                Ok(None)
+            }
+            pg_constants::RM_CLOG_ID => Self::decode_clog_record(&mut buf, decoded, pg_version),
+            pg_constants::RM_XACT_ID => Self::decode_xact_record(&mut buf, decoded, lsn),
+            pg_constants::RM_MULTIXACT_ID => {
+                Self::decode_multixact_record(&mut buf, decoded, pg_version)
+            }
+            pg_constants::RM_RELMAP_ID => Self::decode_relmap_record(&mut buf, decoded),
+            // This is an odd duck. It needs to go to all shards.
+            // Since it uses the checkpoint image (that's initialized from CHECKPOINT_KEY
+            // in WalIngest::new), we have to send the whole DecodedWalRecord::record to
+            // the pageserver and decode it there.
+            //
+            // Alternatively, one can make the checkpoint part of the subscription protocol
+            // to the pageserver. This should work fine, but can be done at a later point.
+            pg_constants::RM_XLOG_ID => Self::decode_xlog_record(&mut buf, decoded, lsn),
+            pg_constants::RM_LOGICALMSG_ID => {
+                Self::decode_logical_message_record(&mut buf, decoded)
+            }
+            pg_constants::RM_STANDBY_ID => Self::decode_standby_record(&mut buf, decoded),
+            pg_constants::RM_REPLORIGIN_ID => Self::decode_replorigin_record(&mut buf, decoded),
+            _unexpected => {
+                // TODO: consider failing here instead of blindly doing something without
+                // understanding the protocol
+                Ok(None)
+            }
+        }
+    }
+
+    fn decode_heapam_record(
+        buf: &mut Bytes,
+        decoded: &DecodedWALRecord,
+        pg_version: u32,
+    ) -> anyhow::Result<Option<MetadataRecord>> {
+        // Handle VM bit updates that are implicitly part of heap records.
+
+        // First, look at the record to determine which VM bits need
+        // to be cleared. If either of these variables is set, we
+        // need to clear the corresponding bits in the visibility map.
+        let mut new_heap_blkno: Option<u32> = None;
+        let mut old_heap_blkno: Option<u32> = None;
+        let mut flags = pg_constants::VISIBILITYMAP_VALID_BITS;
+
+        match pg_version {
+            14 => {
+                if decoded.xl_rmid == pg_constants::RM_HEAP_ID {
+                    let info = decoded.xl_info & pg_constants::XLOG_HEAP_OPMASK;
+
+                    if info == pg_constants::XLOG_HEAP_INSERT {
+                        let xlrec = v14::XlHeapInsert::decode(buf);
+                        assert_eq!(0, buf.remaining());
+                        if (xlrec.flags & pg_constants::XLH_INSERT_ALL_VISIBLE_CLEARED) != 0 {
+                            new_heap_blkno = Some(decoded.blocks[0].blkno);
+                        }
+                    } else if info == pg_constants::XLOG_HEAP_DELETE {
+                        let xlrec = v14::XlHeapDelete::decode(buf);
+                        if (xlrec.flags & pg_constants::XLH_DELETE_ALL_VISIBLE_CLEARED) != 0 {
+                            new_heap_blkno = Some(decoded.blocks[0].blkno);
+                        }
+                    } else if info == pg_constants::XLOG_HEAP_UPDATE
+                        || info == pg_constants::XLOG_HEAP_HOT_UPDATE
+                    {
+                        let xlrec = v14::XlHeapUpdate::decode(buf);
+                        // the size of tuple data is inferred from the size of the record.
+                        // we can't validate the remaining number of bytes without parsing
+                        // the tuple data.
+                        if (xlrec.flags & pg_constants::XLH_UPDATE_OLD_ALL_VISIBLE_CLEARED) != 0 {
+                            old_heap_blkno = Some(decoded.blocks.last().unwrap().blkno);
+                        }
+                        if (xlrec.flags & pg_constants::XLH_UPDATE_NEW_ALL_VISIBLE_CLEARED) != 0 {
+                            // PostgreSQL only uses XLH_UPDATE_NEW_ALL_VISIBLE_CLEARED on a
+                            // non-HOT update where the new tuple goes to different page than
+                            // the old one. Otherwise, only XLH_UPDATE_OLD_ALL_VISIBLE_CLEARED is
+                            // set.
+                            new_heap_blkno = Some(decoded.blocks[0].blkno);
+                        }
+                    } else if info == pg_constants::XLOG_HEAP_LOCK {
+                        let xlrec = v14::XlHeapLock::decode(buf);
+                        if (xlrec.flags & pg_constants::XLH_LOCK_ALL_FROZEN_CLEARED) != 0 {
+                            old_heap_blkno = Some(decoded.blocks[0].blkno);
+                            flags = pg_constants::VISIBILITYMAP_ALL_FROZEN;
+                        }
+                    }
+                } else if decoded.xl_rmid == pg_constants::RM_HEAP2_ID {
+                    let info = decoded.xl_info & pg_constants::XLOG_HEAP_OPMASK;
+                    if info == pg_constants::XLOG_HEAP2_MULTI_INSERT {
+                        let xlrec = v14::XlHeapMultiInsert::decode(buf);
+
+                        let offset_array_len =
+                            if decoded.xl_info & pg_constants::XLOG_HEAP_INIT_PAGE > 0 {
+                                // the offsets array is omitted if XLOG_HEAP_INIT_PAGE is set
+                                0
+                            } else {
+                                size_of::<u16>() * xlrec.ntuples as usize
+                            };
+                        assert_eq!(offset_array_len, buf.remaining());
+
+                        if (xlrec.flags & pg_constants::XLH_INSERT_ALL_VISIBLE_CLEARED) != 0 {
+                            new_heap_blkno = Some(decoded.blocks[0].blkno);
+                        }
+                    } else if info == pg_constants::XLOG_HEAP2_LOCK_UPDATED {
+                        let xlrec = v14::XlHeapLockUpdated::decode(buf);
+                        if (xlrec.flags & pg_constants::XLH_LOCK_ALL_FROZEN_CLEARED) != 0 {
+                            old_heap_blkno = Some(decoded.blocks[0].blkno);
+                            flags = pg_constants::VISIBILITYMAP_ALL_FROZEN;
+                        }
+                    }
+                } else {
+                    anyhow::bail!("Unknown RMGR {} for Heap decoding", decoded.xl_rmid);
+                }
+            }
+            15 => {
+                if decoded.xl_rmid == pg_constants::RM_HEAP_ID {
+                    let info = decoded.xl_info & pg_constants::XLOG_HEAP_OPMASK;
+
+                    if info == pg_constants::XLOG_HEAP_INSERT {
+                        let xlrec = v15::XlHeapInsert::decode(buf);
+                        assert_eq!(0, buf.remaining());
+                        if (xlrec.flags & pg_constants::XLH_INSERT_ALL_VISIBLE_CLEARED) != 0 {
+                            new_heap_blkno = Some(decoded.blocks[0].blkno);
+                        }
+                    } else if info == pg_constants::XLOG_HEAP_DELETE {
+                        let xlrec = v15::XlHeapDelete::decode(buf);
+                        if (xlrec.flags & pg_constants::XLH_DELETE_ALL_VISIBLE_CLEARED) != 0 {
+                            new_heap_blkno = Some(decoded.blocks[0].blkno);
+                        }
+                    } else if info == pg_constants::XLOG_HEAP_UPDATE
+                        || info == pg_constants::XLOG_HEAP_HOT_UPDATE
+                    {
+                        let xlrec = v15::XlHeapUpdate::decode(buf);
+                        // the size of tuple data is inferred from the size of the record.
+                        // we can't validate the remaining number of bytes without parsing
+                        // the tuple data.
+                        if (xlrec.flags & pg_constants::XLH_UPDATE_OLD_ALL_VISIBLE_CLEARED) != 0 {
+                            old_heap_blkno = Some(decoded.blocks.last().unwrap().blkno);
+                        }
+                        if (xlrec.flags & pg_constants::XLH_UPDATE_NEW_ALL_VISIBLE_CLEARED) != 0 {
+                            // PostgreSQL only uses XLH_UPDATE_NEW_ALL_VISIBLE_CLEARED on a
+                            // non-HOT update where the new tuple goes to different page than
+                            // the old one. Otherwise, only XLH_UPDATE_OLD_ALL_VISIBLE_CLEARED is
+                            // set.
+                            new_heap_blkno = Some(decoded.blocks[0].blkno);
+                        }
+                    } else if info == pg_constants::XLOG_HEAP_LOCK {
+                        let xlrec = v15::XlHeapLock::decode(buf);
+                        if (xlrec.flags & pg_constants::XLH_LOCK_ALL_FROZEN_CLEARED) != 0 {
+                            old_heap_blkno = Some(decoded.blocks[0].blkno);
+                            flags = pg_constants::VISIBILITYMAP_ALL_FROZEN;
+                        }
+                    }
+                } else if decoded.xl_rmid == pg_constants::RM_HEAP2_ID {
+                    let info = decoded.xl_info & pg_constants::XLOG_HEAP_OPMASK;
+                    if info == pg_constants::XLOG_HEAP2_MULTI_INSERT {
+                        let xlrec = v15::XlHeapMultiInsert::decode(buf);
+
+                        let offset_array_len =
+                            if decoded.xl_info & pg_constants::XLOG_HEAP_INIT_PAGE > 0 {
+                                // the offsets array is omitted if XLOG_HEAP_INIT_PAGE is set
+                                0
+                            } else {
+                                size_of::<u16>() * xlrec.ntuples as usize
+                            };
+                        assert_eq!(offset_array_len, buf.remaining());
+
+                        if (xlrec.flags & pg_constants::XLH_INSERT_ALL_VISIBLE_CLEARED) != 0 {
+                            new_heap_blkno = Some(decoded.blocks[0].blkno);
+                        }
+                    } else if info == pg_constants::XLOG_HEAP2_LOCK_UPDATED {
+                        let xlrec = v15::XlHeapLockUpdated::decode(buf);
+                        if (xlrec.flags & pg_constants::XLH_LOCK_ALL_FROZEN_CLEARED) != 0 {
+                            old_heap_blkno = Some(decoded.blocks[0].blkno);
+                            flags = pg_constants::VISIBILITYMAP_ALL_FROZEN;
+                        }
+                    }
+                } else {
+                    anyhow::bail!("Unknown RMGR {} for Heap decoding", decoded.xl_rmid);
+                }
+            }
+            16 => {
+                if decoded.xl_rmid == pg_constants::RM_HEAP_ID {
+                    let info = decoded.xl_info & pg_constants::XLOG_HEAP_OPMASK;
+
+                    if info == pg_constants::XLOG_HEAP_INSERT {
+                        let xlrec = v16::XlHeapInsert::decode(buf);
+                        assert_eq!(0, buf.remaining());
+                        if (xlrec.flags & pg_constants::XLH_INSERT_ALL_VISIBLE_CLEARED) != 0 {
+                            new_heap_blkno = Some(decoded.blocks[0].blkno);
+                        }
+                    } else if info == pg_constants::XLOG_HEAP_DELETE {
+                        let xlrec = v16::XlHeapDelete::decode(buf);
+                        if (xlrec.flags & pg_constants::XLH_DELETE_ALL_VISIBLE_CLEARED) != 0 {
+                            new_heap_blkno = Some(decoded.blocks[0].blkno);
+                        }
+                    } else if info == pg_constants::XLOG_HEAP_UPDATE
+                        || info == pg_constants::XLOG_HEAP_HOT_UPDATE
+                    {
+                        let xlrec = v16::XlHeapUpdate::decode(buf);
+                        // the size of tuple data is inferred from the size of the record.
+                        // we can't validate the remaining number of bytes without parsing
+                        // the tuple data.
+                        if (xlrec.flags & pg_constants::XLH_UPDATE_OLD_ALL_VISIBLE_CLEARED) != 0 {
+                            old_heap_blkno = Some(decoded.blocks.last().unwrap().blkno);
+                        }
+                        if (xlrec.flags & pg_constants::XLH_UPDATE_NEW_ALL_VISIBLE_CLEARED) != 0 {
+                            // PostgreSQL only uses XLH_UPDATE_NEW_ALL_VISIBLE_CLEARED on a
+                            // non-HOT update where the new tuple goes to different page than
+                            // the old one. Otherwise, only XLH_UPDATE_OLD_ALL_VISIBLE_CLEARED is
+                            // set.
+                            new_heap_blkno = Some(decoded.blocks[0].blkno);
+                        }
+                    } else if info == pg_constants::XLOG_HEAP_LOCK {
+                        let xlrec = v16::XlHeapLock::decode(buf);
+                        if (xlrec.flags & pg_constants::XLH_LOCK_ALL_FROZEN_CLEARED) != 0 {
+                            old_heap_blkno = Some(decoded.blocks[0].blkno);
+                            flags = pg_constants::VISIBILITYMAP_ALL_FROZEN;
+                        }
+                    }
+                } else if decoded.xl_rmid == pg_constants::RM_HEAP2_ID {
+                    let info = decoded.xl_info & pg_constants::XLOG_HEAP_OPMASK;
+                    if info == pg_constants::XLOG_HEAP2_MULTI_INSERT {
+                        let xlrec = v16::XlHeapMultiInsert::decode(buf);
+
+                        let offset_array_len =
+                            if decoded.xl_info & pg_constants::XLOG_HEAP_INIT_PAGE > 0 {
+                                // the offsets array is omitted if XLOG_HEAP_INIT_PAGE is set
+                                0
+                            } else {
+                                size_of::<u16>() * xlrec.ntuples as usize
+                            };
+                        assert_eq!(offset_array_len, buf.remaining());
+
+                        if (xlrec.flags & pg_constants::XLH_INSERT_ALL_VISIBLE_CLEARED) != 0 {
+                            new_heap_blkno = Some(decoded.blocks[0].blkno);
+                        }
+                    } else if info == pg_constants::XLOG_HEAP2_LOCK_UPDATED {
+                        let xlrec = v16::XlHeapLockUpdated::decode(buf);
+                        if (xlrec.flags & pg_constants::XLH_LOCK_ALL_FROZEN_CLEARED) != 0 {
+                            old_heap_blkno = Some(decoded.blocks[0].blkno);
+                            flags = pg_constants::VISIBILITYMAP_ALL_FROZEN;
+                        }
+                    }
+                } else {
+                    anyhow::bail!("Unknown RMGR {} for Heap decoding", decoded.xl_rmid);
+                }
+            }
+            17 => {
+                if decoded.xl_rmid == pg_constants::RM_HEAP_ID {
+                    let info = decoded.xl_info & pg_constants::XLOG_HEAP_OPMASK;
+
+                    if info == pg_constants::XLOG_HEAP_INSERT {
+                        let xlrec = v17::XlHeapInsert::decode(buf);
+                        assert_eq!(0, buf.remaining());
+                        if (xlrec.flags & pg_constants::XLH_INSERT_ALL_VISIBLE_CLEARED) != 0 {
+                            new_heap_blkno = Some(decoded.blocks[0].blkno);
+                        }
+                    } else if info == pg_constants::XLOG_HEAP_DELETE {
+                        let xlrec = v17::XlHeapDelete::decode(buf);
+                        if (xlrec.flags & pg_constants::XLH_DELETE_ALL_VISIBLE_CLEARED) != 0 {
+                            new_heap_blkno = Some(decoded.blocks[0].blkno);
+                        }
+                    } else if info == pg_constants::XLOG_HEAP_UPDATE
+                        || info == pg_constants::XLOG_HEAP_HOT_UPDATE
+                    {
+                        let xlrec = v17::XlHeapUpdate::decode(buf);
+                        // the size of tuple data is inferred from the size of the record.
+                        // we can't validate the remaining number of bytes without parsing
+                        // the tuple data.
+                        if (xlrec.flags & pg_constants::XLH_UPDATE_OLD_ALL_VISIBLE_CLEARED) != 0 {
+                            old_heap_blkno = Some(decoded.blocks.last().unwrap().blkno);
+                        }
+                        if (xlrec.flags & pg_constants::XLH_UPDATE_NEW_ALL_VISIBLE_CLEARED) != 0 {
+                            // PostgreSQL only uses XLH_UPDATE_NEW_ALL_VISIBLE_CLEARED on a
+                            // non-HOT update where the new tuple goes to different page than
+                            // the old one. Otherwise, only XLH_UPDATE_OLD_ALL_VISIBLE_CLEARED is
+                            // set.
+                            new_heap_blkno = Some(decoded.blocks[0].blkno);
+                        }
+                    } else if info == pg_constants::XLOG_HEAP_LOCK {
+                        let xlrec = v17::XlHeapLock::decode(buf);
+                        if (xlrec.flags & pg_constants::XLH_LOCK_ALL_FROZEN_CLEARED) != 0 {
+                            old_heap_blkno = Some(decoded.blocks[0].blkno);
+                            flags = pg_constants::VISIBILITYMAP_ALL_FROZEN;
+                        }
+                    }
+                } else if decoded.xl_rmid == pg_constants::RM_HEAP2_ID {
+                    let info = decoded.xl_info & pg_constants::XLOG_HEAP_OPMASK;
+                    if info == pg_constants::XLOG_HEAP2_MULTI_INSERT {
+                        let xlrec = v17::XlHeapMultiInsert::decode(buf);
+
+                        let offset_array_len =
+                            if decoded.xl_info & pg_constants::XLOG_HEAP_INIT_PAGE > 0 {
+                                // the offsets array is omitted if XLOG_HEAP_INIT_PAGE is set
+                                0
+                            } else {
+                                size_of::<u16>() * xlrec.ntuples as usize
+                            };
+                        assert_eq!(offset_array_len, buf.remaining());
+
+                        if (xlrec.flags & pg_constants::XLH_INSERT_ALL_VISIBLE_CLEARED) != 0 {
+                            new_heap_blkno = Some(decoded.blocks[0].blkno);
+                        }
+                    } else if info == pg_constants::XLOG_HEAP2_LOCK_UPDATED {
+                        let xlrec = v17::XlHeapLockUpdated::decode(buf);
+                        if (xlrec.flags & pg_constants::XLH_LOCK_ALL_FROZEN_CLEARED) != 0 {
+                            old_heap_blkno = Some(decoded.blocks[0].blkno);
+                            flags = pg_constants::VISIBILITYMAP_ALL_FROZEN;
+                        }
+                    }
+                } else {
+                    anyhow::bail!("Unknown RMGR {} for Heap decoding", decoded.xl_rmid);
+                }
+            }
+            _ => {}
+        }
+
+        if new_heap_blkno.is_some() || old_heap_blkno.is_some() {
+            let vm_rel = RelTag {
+                forknum: VISIBILITYMAP_FORKNUM,
+                spcnode: decoded.blocks[0].rnode_spcnode,
+                dbnode: decoded.blocks[0].rnode_dbnode,
+                relnode: decoded.blocks[0].rnode_relnode,
+            };
+
+            Ok(Some(MetadataRecord::Heapam(HeapamRecord::ClearVmBits(
+                ClearVmBits {
+                    new_heap_blkno,
+                    old_heap_blkno,
+                    vm_rel,
+                    flags,
+                },
+            ))))
+        } else {
+            Ok(None)
+        }
+    }
+
+    fn decode_neonmgr_record(
+        buf: &mut Bytes,
+        decoded: &DecodedWALRecord,
+        pg_version: u32,
+    ) -> anyhow::Result<Option<MetadataRecord>> {
+        // Handle VM bit updates that are implicitly part of heap records.
+
+        // First, look at the record to determine which VM bits need
+        // to be cleared. If either of these variables is set, we
+        // need to clear the corresponding bits in the visibility map.
+        let mut new_heap_blkno: Option<u32> = None;
+        let mut old_heap_blkno: Option<u32> = None;
+        let mut flags = pg_constants::VISIBILITYMAP_VALID_BITS;
+
+        assert_eq!(decoded.xl_rmid, pg_constants::RM_NEON_ID);
+
+        match pg_version {
+            16 | 17 => {
+                let info = decoded.xl_info & pg_constants::XLOG_HEAP_OPMASK;
+
+                match info {
+                    pg_constants::XLOG_NEON_HEAP_INSERT => {
+                        let xlrec = v17::rm_neon::XlNeonHeapInsert::decode(buf);
+                        assert_eq!(0, buf.remaining());
+                        if (xlrec.flags & pg_constants::XLH_INSERT_ALL_VISIBLE_CLEARED) != 0 {
+                            new_heap_blkno = Some(decoded.blocks[0].blkno);
+                        }
+                    }
+                    pg_constants::XLOG_NEON_HEAP_DELETE => {
+                        let xlrec = v17::rm_neon::XlNeonHeapDelete::decode(buf);
+                        if (xlrec.flags & pg_constants::XLH_DELETE_ALL_VISIBLE_CLEARED) != 0 {
+                            new_heap_blkno = Some(decoded.blocks[0].blkno);
+                        }
+                    }
+                    pg_constants::XLOG_NEON_HEAP_UPDATE
+                    | pg_constants::XLOG_NEON_HEAP_HOT_UPDATE => {
+                        let xlrec = v17::rm_neon::XlNeonHeapUpdate::decode(buf);
+                        // the size of tuple data is inferred from the size of the record.
+                        // we can't validate the remaining number of bytes without parsing
+                        // the tuple data.
+                        if (xlrec.flags & pg_constants::XLH_UPDATE_OLD_ALL_VISIBLE_CLEARED) != 0 {
+                            old_heap_blkno = Some(decoded.blocks.last().unwrap().blkno);
+                        }
+                        if (xlrec.flags & pg_constants::XLH_UPDATE_NEW_ALL_VISIBLE_CLEARED) != 0 {
+                            // PostgreSQL only uses XLH_UPDATE_NEW_ALL_VISIBLE_CLEARED on a
+                            // non-HOT update where the new tuple goes to different page than
+                            // the old one. Otherwise, only XLH_UPDATE_OLD_ALL_VISIBLE_CLEARED is
+                            // set.
+                            new_heap_blkno = Some(decoded.blocks[0].blkno);
+                        }
+                    }
+                    pg_constants::XLOG_NEON_HEAP_MULTI_INSERT => {
+                        let xlrec = v17::rm_neon::XlNeonHeapMultiInsert::decode(buf);
+
+                        let offset_array_len =
+                            if decoded.xl_info & pg_constants::XLOG_HEAP_INIT_PAGE > 0 {
+                                // the offsets array is omitted if XLOG_HEAP_INIT_PAGE is set
+                                0
+                            } else {
+                                size_of::<u16>() * xlrec.ntuples as usize
+                            };
+                        assert_eq!(offset_array_len, buf.remaining());
+
+                        if (xlrec.flags & pg_constants::XLH_INSERT_ALL_VISIBLE_CLEARED) != 0 {
+                            new_heap_blkno = Some(decoded.blocks[0].blkno);
+                        }
+                    }
+                    pg_constants::XLOG_NEON_HEAP_LOCK => {
+                        let xlrec = v17::rm_neon::XlNeonHeapLock::decode(buf);
+                        if (xlrec.flags & pg_constants::XLH_LOCK_ALL_FROZEN_CLEARED) != 0 {
+                            old_heap_blkno = Some(decoded.blocks[0].blkno);
+                            flags = pg_constants::VISIBILITYMAP_ALL_FROZEN;
+                        }
+                    }
+                    info => anyhow::bail!("Unknown WAL record type for Neon RMGR: {}", info),
+                }
+            }
+            _ => anyhow::bail!(
+                "Neon RMGR has no known compatibility with PostgreSQL version {}",
+                pg_version
+            ),
+        }
+
+        if new_heap_blkno.is_some() || old_heap_blkno.is_some() {
+            let vm_rel = RelTag {
+                forknum: VISIBILITYMAP_FORKNUM,
+                spcnode: decoded.blocks[0].rnode_spcnode,
+                dbnode: decoded.blocks[0].rnode_dbnode,
+                relnode: decoded.blocks[0].rnode_relnode,
+            };
+
+            Ok(Some(MetadataRecord::Neonrmgr(NeonrmgrRecord::ClearVmBits(
+                ClearVmBits {
+                    new_heap_blkno,
+                    old_heap_blkno,
+                    vm_rel,
+                    flags,
+                },
+            ))))
+        } else {
+            Ok(None)
+        }
+    }
+
+    fn decode_smgr_record(
+        buf: &mut Bytes,
+        decoded: &DecodedWALRecord,
+    ) -> anyhow::Result<Option<MetadataRecord>> {
+        let info = decoded.xl_info & pg_constants::XLR_RMGR_INFO_MASK;
+        if info == pg_constants::XLOG_SMGR_CREATE {
+            let create = XlSmgrCreate::decode(buf);
+            let rel = RelTag {
+                spcnode: create.rnode.spcnode,
+                dbnode: create.rnode.dbnode,
+                relnode: create.rnode.relnode,
+                forknum: create.forknum,
+            };
+
+            return Ok(Some(MetadataRecord::Smgr(SmgrRecord::Create(SmgrCreate {
+                rel,
+            }))));
+        } else if info == pg_constants::XLOG_SMGR_TRUNCATE {
+            let truncate = XlSmgrTruncate::decode(buf);
+            return Ok(Some(MetadataRecord::Smgr(SmgrRecord::Truncate(truncate))));
+        }
+
+        Ok(None)
+    }
+
+    fn decode_dbase_record(
+        buf: &mut Bytes,
+        decoded: &DecodedWALRecord,
+        pg_version: u32,
+    ) -> anyhow::Result<Option<MetadataRecord>> {
+        // TODO: Refactor this to avoid the duplication between postgres versions.
+
+        let info = decoded.xl_info & pg_constants::XLR_RMGR_INFO_MASK;
+        tracing::debug!(%info, %pg_version, "handle RM_DBASE_ID");
+
+        if pg_version == 14 {
+            if info == postgres_ffi::v14::bindings::XLOG_DBASE_CREATE {
+                let createdb = XlCreateDatabase::decode(buf);
+                tracing::debug!("XLOG_DBASE_CREATE v14");
+
+                let record = MetadataRecord::Dbase(DbaseRecord::Create(DbaseCreate {
+                    db_id: createdb.db_id,
+                    tablespace_id: createdb.tablespace_id,
+                    src_db_id: createdb.src_db_id,
+                    src_tablespace_id: createdb.src_tablespace_id,
+                }));
+
+                return Ok(Some(record));
+            } else if info == postgres_ffi::v14::bindings::XLOG_DBASE_DROP {
+                let dropdb = XlDropDatabase::decode(buf);
+
+                let record = MetadataRecord::Dbase(DbaseRecord::Drop(DbaseDrop {
+                    db_id: dropdb.db_id,
+                    tablespace_ids: dropdb.tablespace_ids,
+                }));
+
+                return Ok(Some(record));
+            }
+        } else if pg_version == 15 {
+            if info == postgres_ffi::v15::bindings::XLOG_DBASE_CREATE_WAL_LOG {
+                tracing::debug!("XLOG_DBASE_CREATE_WAL_LOG: noop");
+            } else if info == postgres_ffi::v15::bindings::XLOG_DBASE_CREATE_FILE_COPY {
+                // The XLOG record was renamed between v14 and v15,
+                // but the record format is the same.
+                // So we can reuse XlCreateDatabase here.
+                tracing::debug!("XLOG_DBASE_CREATE_FILE_COPY");
+
+                let createdb = XlCreateDatabase::decode(buf);
+                let record = MetadataRecord::Dbase(DbaseRecord::Create(DbaseCreate {
+                    db_id: createdb.db_id,
+                    tablespace_id: createdb.tablespace_id,
+                    src_db_id: createdb.src_db_id,
+                    src_tablespace_id: createdb.src_tablespace_id,
+                }));
+
+                return Ok(Some(record));
+            } else if info == postgres_ffi::v15::bindings::XLOG_DBASE_DROP {
+                let dropdb = XlDropDatabase::decode(buf);
+                let record = MetadataRecord::Dbase(DbaseRecord::Drop(DbaseDrop {
+                    db_id: dropdb.db_id,
+                    tablespace_ids: dropdb.tablespace_ids,
+                }));
+
+                return Ok(Some(record));
+            }
+        } else if pg_version == 16 {
+            if info == postgres_ffi::v16::bindings::XLOG_DBASE_CREATE_WAL_LOG {
+                tracing::debug!("XLOG_DBASE_CREATE_WAL_LOG: noop");
+            } else if info == postgres_ffi::v16::bindings::XLOG_DBASE_CREATE_FILE_COPY {
+                // The XLOG record was renamed between v14 and v15,
+                // but the record format is the same.
+                // So we can reuse XlCreateDatabase here.
+                tracing::debug!("XLOG_DBASE_CREATE_FILE_COPY");
+
+                let createdb = XlCreateDatabase::decode(buf);
+                let record = MetadataRecord::Dbase(DbaseRecord::Create(DbaseCreate {
+                    db_id: createdb.db_id,
+                    tablespace_id: createdb.tablespace_id,
+                    src_db_id: createdb.src_db_id,
+                    src_tablespace_id: createdb.src_tablespace_id,
+                }));
+
+                return Ok(Some(record));
+            } else if info == postgres_ffi::v16::bindings::XLOG_DBASE_DROP {
+                let dropdb = XlDropDatabase::decode(buf);
+                let record = MetadataRecord::Dbase(DbaseRecord::Drop(DbaseDrop {
+                    db_id: dropdb.db_id,
+                    tablespace_ids: dropdb.tablespace_ids,
+                }));
+
+                return Ok(Some(record));
+            }
+        } else if pg_version == 17 {
+            if info == postgres_ffi::v17::bindings::XLOG_DBASE_CREATE_WAL_LOG {
+                tracing::debug!("XLOG_DBASE_CREATE_WAL_LOG: noop");
+            } else if info == postgres_ffi::v17::bindings::XLOG_DBASE_CREATE_FILE_COPY {
+                // The XLOG record was renamed between v14 and v15,
+                // but the record format is the same.
+                // So we can reuse XlCreateDatabase here.
+                tracing::debug!("XLOG_DBASE_CREATE_FILE_COPY");
+
+                let createdb = XlCreateDatabase::decode(buf);
+                let record = MetadataRecord::Dbase(DbaseRecord::Create(DbaseCreate {
+                    db_id: createdb.db_id,
+                    tablespace_id: createdb.tablespace_id,
+                    src_db_id: createdb.src_db_id,
+                    src_tablespace_id: createdb.src_tablespace_id,
+                }));
+
+                return Ok(Some(record));
+            } else if info == postgres_ffi::v17::bindings::XLOG_DBASE_DROP {
+                let dropdb = XlDropDatabase::decode(buf);
+                let record = MetadataRecord::Dbase(DbaseRecord::Drop(DbaseDrop {
+                    db_id: dropdb.db_id,
+                    tablespace_ids: dropdb.tablespace_ids,
+                }));
+
+                return Ok(Some(record));
+            }
+        }
+
+        Ok(None)
+    }
+
+    fn decode_clog_record(
+        buf: &mut Bytes,
+        decoded: &DecodedWALRecord,
+        pg_version: u32,
+    ) -> anyhow::Result<Option<MetadataRecord>> {
+        let info = decoded.xl_info & !pg_constants::XLR_INFO_MASK;
+
+        if info == pg_constants::CLOG_ZEROPAGE {
+            let pageno = if pg_version < 17 {
+                buf.get_u32_le()
+            } else {
+                buf.get_u64_le() as u32
+            };
+            let segno = pageno / pg_constants::SLRU_PAGES_PER_SEGMENT;
+            let rpageno = pageno % pg_constants::SLRU_PAGES_PER_SEGMENT;
+
+            Ok(Some(MetadataRecord::Clog(ClogRecord::ZeroPage(
+                ClogZeroPage { segno, rpageno },
+            ))))
+        } else {
+            assert!(info == pg_constants::CLOG_TRUNCATE);
+            let xlrec = XlClogTruncate::decode(buf, pg_version);
+
+            Ok(Some(MetadataRecord::Clog(ClogRecord::Truncate(
+                ClogTruncate {
+                    pageno: xlrec.pageno,
+                    oldest_xid: xlrec.oldest_xid,
+                    oldest_xid_db: xlrec.oldest_xid_db,
+                },
+            ))))
+        }
+    }
+
+    fn decode_xact_record(
+        buf: &mut Bytes,
+        decoded: &DecodedWALRecord,
+        lsn: Lsn,
+    ) -> anyhow::Result<Option<MetadataRecord>> {
+        let info = decoded.xl_info & pg_constants::XLOG_XACT_OPMASK;
+        let origin_id = decoded.origin_id;
+        let xl_xid = decoded.xl_xid;
+
+        if info == pg_constants::XLOG_XACT_COMMIT {
+            let parsed = XlXactParsedRecord::decode(buf, decoded.xl_xid, decoded.xl_info);
+            return Ok(Some(MetadataRecord::Xact(XactRecord::Commit(XactCommon {
+                parsed,
+                origin_id,
+                xl_xid,
+                lsn,
+            }))));
+        } else if info == pg_constants::XLOG_XACT_ABORT {
+            let parsed = XlXactParsedRecord::decode(buf, decoded.xl_xid, decoded.xl_info);
+            return Ok(Some(MetadataRecord::Xact(XactRecord::Abort(XactCommon {
+                parsed,
+                origin_id,
+                xl_xid,
+                lsn,
+            }))));
+        } else if info == pg_constants::XLOG_XACT_COMMIT_PREPARED {
+            let parsed = XlXactParsedRecord::decode(buf, decoded.xl_xid, decoded.xl_info);
+            return Ok(Some(MetadataRecord::Xact(XactRecord::CommitPrepared(
+                XactCommon {
+                    parsed,
+                    origin_id,
+                    xl_xid,
+                    lsn,
+                },
+            ))));
+        } else if info == pg_constants::XLOG_XACT_ABORT_PREPARED {
+            let parsed = XlXactParsedRecord::decode(buf, decoded.xl_xid, decoded.xl_info);
+            return Ok(Some(MetadataRecord::Xact(XactRecord::AbortPrepared(
+                XactCommon {
+                    parsed,
+                    origin_id,
+                    xl_xid,
+                    lsn,
+                },
+            ))));
+        } else if info == pg_constants::XLOG_XACT_PREPARE {
+            return Ok(Some(MetadataRecord::Xact(XactRecord::Prepare(
+                XactPrepare {
+                    xl_xid: decoded.xl_xid,
+                    data: Bytes::copy_from_slice(&buf[..]),
+                },
+            ))));
+        }
+
+        Ok(None)
+    }
+
+    fn decode_multixact_record(
+        buf: &mut Bytes,
+        decoded: &DecodedWALRecord,
+        pg_version: u32,
+    ) -> anyhow::Result<Option<MetadataRecord>> {
+        let info = decoded.xl_info & pg_constants::XLR_RMGR_INFO_MASK;
+
+        if info == pg_constants::XLOG_MULTIXACT_ZERO_OFF_PAGE
+            || info == pg_constants::XLOG_MULTIXACT_ZERO_MEM_PAGE
+        {
+            let pageno = if pg_version < 17 {
+                buf.get_u32_le()
+            } else {
+                buf.get_u64_le() as u32
+            };
+            let segno = pageno / pg_constants::SLRU_PAGES_PER_SEGMENT;
+            let rpageno = pageno % pg_constants::SLRU_PAGES_PER_SEGMENT;
+
+            let slru_kind = match info {
+                pg_constants::XLOG_MULTIXACT_ZERO_OFF_PAGE => SlruKind::MultiXactOffsets,
+                pg_constants::XLOG_MULTIXACT_ZERO_MEM_PAGE => SlruKind::MultiXactMembers,
+                _ => unreachable!(),
+            };
+
+            return Ok(Some(MetadataRecord::MultiXact(MultiXactRecord::ZeroPage(
+                MultiXactZeroPage {
+                    slru_kind,
+                    segno,
+                    rpageno,
+                },
+            ))));
+        } else if info == pg_constants::XLOG_MULTIXACT_CREATE_ID {
+            let xlrec = XlMultiXactCreate::decode(buf);
+            return Ok(Some(MetadataRecord::MultiXact(MultiXactRecord::Create(
+                xlrec,
+            ))));
+        } else if info == pg_constants::XLOG_MULTIXACT_TRUNCATE_ID {
+            let xlrec = XlMultiXactTruncate::decode(buf);
+            return Ok(Some(MetadataRecord::MultiXact(MultiXactRecord::Truncate(
+                xlrec,
+            ))));
+        }
+
+        Ok(None)
+    }
+
+    fn decode_relmap_record(
+        buf: &mut Bytes,
+        decoded: &DecodedWALRecord,
+    ) -> anyhow::Result<Option<MetadataRecord>> {
+        let update = XlRelmapUpdate::decode(buf);
+
+        let mut buf = decoded.record.clone();
+        buf.advance(decoded.main_data_offset);
+        // skip xl_relmap_update
+        buf.advance(12);
+
+        Ok(Some(MetadataRecord::Relmap(RelmapRecord::Update(
+            RelmapUpdate {
+                update,
+                buf: Bytes::copy_from_slice(&buf[..]),
+            },
+        ))))
+    }
+
+    fn decode_xlog_record(
+        buf: &mut Bytes,
+        decoded: &DecodedWALRecord,
+        lsn: Lsn,
+    ) -> anyhow::Result<Option<MetadataRecord>> {
+        let info = decoded.xl_info & pg_constants::XLR_RMGR_INFO_MASK;
+        Ok(Some(MetadataRecord::Xlog(XlogRecord::Raw(RawXlogRecord {
+            info,
+            lsn,
+            buf: buf.clone(),
+        }))))
+    }
+
+    fn decode_logical_message_record(
+        buf: &mut Bytes,
+        decoded: &DecodedWALRecord,
+    ) -> anyhow::Result<Option<MetadataRecord>> {
+        let info = decoded.xl_info & pg_constants::XLR_RMGR_INFO_MASK;
+        if info == pg_constants::XLOG_LOGICAL_MESSAGE {
+            let xlrec = XlLogicalMessage::decode(buf);
+            let prefix = std::str::from_utf8(&buf[0..xlrec.prefix_size - 1])?;
+
+            #[cfg(feature = "testing")]
+            if prefix == "neon-test" {
+                return Ok(Some(MetadataRecord::LogicalMessage(
+                    LogicalMessageRecord::Failpoint,
+                )));
+            }
+
+            if let Some(path) = prefix.strip_prefix("neon-file:") {
+                let buf_size = xlrec.prefix_size + xlrec.message_size;
+                let buf = Bytes::copy_from_slice(&buf[xlrec.prefix_size..buf_size]);
+                return Ok(Some(MetadataRecord::LogicalMessage(
+                    LogicalMessageRecord::Put(PutLogicalMessage {
+                        path: path.to_string(),
+                        buf,
+                    }),
+                )));
+            }
+        }
+
+        Ok(None)
+    }
+
+    fn decode_standby_record(
+        buf: &mut Bytes,
+        decoded: &DecodedWALRecord,
+    ) -> anyhow::Result<Option<MetadataRecord>> {
+        let info = decoded.xl_info & pg_constants::XLR_RMGR_INFO_MASK;
+        if info == pg_constants::XLOG_RUNNING_XACTS {
+            let xlrec = XlRunningXacts::decode(buf);
+            return Ok(Some(MetadataRecord::Standby(StandbyRecord::RunningXacts(
+                StandbyRunningXacts {
+                    oldest_running_xid: xlrec.oldest_running_xid,
+                },
+            ))));
+        }
+
+        Ok(None)
+    }
+
+    fn decode_replorigin_record(
+        buf: &mut Bytes,
+        decoded: &DecodedWALRecord,
+    ) -> anyhow::Result<Option<MetadataRecord>> {
+        let info = decoded.xl_info & pg_constants::XLR_RMGR_INFO_MASK;
+        if info == pg_constants::XLOG_REPLORIGIN_SET {
+            let xlrec = XlReploriginSet::decode(buf);
+            return Ok(Some(MetadataRecord::Replorigin(ReploriginRecord::Set(
+                xlrec,
+            ))));
+        } else if info == pg_constants::XLOG_REPLORIGIN_DROP {
+            let xlrec = XlReploriginDrop::decode(buf);
+            return Ok(Some(MetadataRecord::Replorigin(ReploriginRecord::Drop(
+                xlrec,
+            ))));
+        }
+
+        Ok(None)
+    }
+}

libs/wal_decoder/src/models.rs

@@ -25,7 +25,9 @@
 //!                     |--> write to KV store within the pageserver
 
 use bytes::Bytes;
+use pageserver_api::key::CompactKey;
 use pageserver_api::reltag::{RelTag, SlruKind};
+use pageserver_api::value::Value;
 use postgres_ffi::walrecord::{
     XlMultiXactCreate, XlMultiXactTruncate, XlRelmapUpdate, XlReploriginDrop, XlReploriginSet,
     XlSmgrTruncate, XlXactParsedRecord,
@@ -33,6 +35,48 @@ use postgres_ffi::walrecord::{
 use postgres_ffi::{Oid, TransactionId};
 use utils::lsn::Lsn;
 
+pub enum FlushUncommittedRecords {
+    Yes,
+    No,
+}
+
+/// An interpreted Postgres WAL record, ready to be handled by the pageserver
+pub struct InterpretedWalRecord {
+    /// Optional metadata record - may cause writes to metadata keys
+    /// in the storage engine
+    pub metadata_record: Option<MetadataRecord>,
+    /// Images or deltas for blocks modified in the original WAL record.
+    /// The [`Value`] is optional to avoid sending superfluous data to
+    /// shard 0 for relation size tracking.
+    pub blocks: Vec<(CompactKey, Option<Value>)>,
+    /// Byte offset within WAL for the end of the original PG WAL record
+    pub lsn: Lsn,
+    /// Whether to flush all uncommitted modifications to the storage engine
+    /// before ingesting this record. This is currently only used for legacy PG
+    /// database creations which read pages from a template database. Such WAL
+    /// records require reading data blocks while ingesting, hence the need to flush.
+    pub flush_uncommitted: FlushUncommittedRecords,
+    /// Transaction id of the original PG WAL record
+    pub xid: TransactionId,
+}
+
+/// The interpreted part of the Postgres WAL record which requires metadata
+/// writes to the underlying storage engine.
+pub enum MetadataRecord {
+    Heapam(HeapamRecord),
+    Neonrmgr(NeonrmgrRecord),
+    Smgr(SmgrRecord),
+    Dbase(DbaseRecord),
+    Clog(ClogRecord),
+    Xact(XactRecord),
+    MultiXact(MultiXactRecord),
+    Relmap(RelmapRecord),
+    Xlog(XlogRecord),
+    LogicalMessage(LogicalMessageRecord),
+    Standby(StandbyRecord),
+    Replorigin(ReploriginRecord),
+}
+
 pub enum HeapamRecord {
     ClearVmBits(ClearVmBits),
 }

pageserver/ctl/src/layer_map_analyzer.rs

@@ -2,7 +2,7 @@
 //!
 //! Currently it only analyzes holes, which are regions within the layer range that the layer contains no updates for. In the future it might do more analysis (maybe key quantiles?) but it should never return sensitive data.
 
-use anyhow::Result;
+use anyhow::{anyhow, Result};
 use camino::{Utf8Path, Utf8PathBuf};
 use pageserver::context::{DownloadBehavior, RequestContext};
 use pageserver::task_mgr::TaskKind;
@@ -11,13 +11,14 @@ use pageserver::virtual_file::api::IoMode;
 use std::cmp::Ordering;
 use std::collections::BinaryHeap;
 use std::ops::Range;
+use std::str::FromStr;
 use std::{fs, str};
 
 use pageserver::page_cache::{self, PAGE_SZ};
 use pageserver::tenant::block_io::FileBlockReader;
 use pageserver::tenant::disk_btree::{DiskBtreeReader, VisitDirection};
 use pageserver::tenant::storage_layer::delta_layer::{Summary, DELTA_KEY_SIZE};
-use pageserver::tenant::storage_layer::range_overlaps;
+use pageserver::tenant::storage_layer::{range_overlaps, LayerName};
 use pageserver::virtual_file::{self, VirtualFile};
 use pageserver_api::key::{Key, KEY_SIZE};
 
@@ -74,35 +75,15 @@ impl LayerFile {
     }
 }
 
-pub(crate) fn parse_filename(name: &str) -> Option<LayerFile> {
-    let split: Vec<&str> = name.split("__").collect();
-    if split.len() != 2 {
-        return None;
-    }
-    let keys: Vec<&str> = split[0].split('-').collect();
-    let lsn_and_opt_generation: Vec<&str> = split[1].split('v').collect();
-    let lsns: Vec<&str> = lsn_and_opt_generation[0].split('-').collect();
-    let the_lsns: [&str; 2];
+pub(crate) fn parse_filename(name: &str) -> anyhow::Result<LayerFile> {
+    let layer_name =
+        LayerName::from_str(name).map_err(|e| anyhow!("failed to parse layer name: {e}"))?;
 
-    /*
-     * Generations add a -vX-XXXXXX postfix, which causes issues when we try to
-     * parse 'vX' as an LSN.
-     */
-    let is_delta = if lsns.len() == 1 || lsns[1].is_empty() {
-        the_lsns = [lsns[0], lsns[0]];
-        false
-    } else {
-        the_lsns = [lsns[0], lsns[1]];
-        true
-    };
-
-    let key_range = Key::from_hex(keys[0]).unwrap()..Key::from_hex(keys[1]).unwrap();
-    let lsn_range = Lsn::from_hex(the_lsns[0]).unwrap()..Lsn::from_hex(the_lsns[1]).unwrap();
     let holes = Vec::new();
-    Some(LayerFile {
-        key_range,
-        lsn_range,
-        is_delta,
+    Ok(LayerFile {
+        key_range: layer_name.key_range().clone(),
+        lsn_range: layer_name.lsn_as_range(),
+        is_delta: layer_name.is_delta(),
         holes,
     })
 }
@@ -179,7 +160,7 @@ pub(crate) async fn main(cmd: &AnalyzeLayerMapCmd) -> Result<()> {
 
             for layer in fs::read_dir(timeline.path())? {
                 let layer = layer?;
-                if let Some(mut layer_file) =
+                if let Ok(mut layer_file) =
                     parse_filename(&layer.file_name().into_string().unwrap())
                 {
                     if layer_file.is_delta {

pageserver/ctl/src/layers.rs

@@ -5,24 +5,12 @@ use camino::{Utf8Path, Utf8PathBuf};
 use clap::Subcommand;
 use pageserver::context::{DownloadBehavior, RequestContext};
 use pageserver::task_mgr::TaskKind;
-use pageserver::tenant::block_io::BlockCursor;
-use pageserver::tenant::disk_btree::DiskBtreeReader;
-use pageserver::tenant::storage_layer::delta_layer::{BlobRef, Summary};
 use pageserver::tenant::storage_layer::{delta_layer, image_layer};
 use pageserver::tenant::storage_layer::{DeltaLayer, ImageLayer};
 use pageserver::tenant::{TENANTS_SEGMENT_NAME, TIMELINES_SEGMENT_NAME};
 use pageserver::virtual_file::api::IoMode;
 use pageserver::{page_cache, virtual_file};
-use pageserver::{
-    tenant::{
-        block_io::FileBlockReader, disk_btree::VisitDirection,
-        storage_layer::delta_layer::DELTA_KEY_SIZE,
-    },
-    virtual_file::VirtualFile,
-};
-use pageserver_api::key::{Key, KEY_SIZE};
-use std::fs;
-use utils::bin_ser::BeSer;
+use std::fs::{self, File};
 use utils::id::{TenantId, TimelineId};
 
 use crate::layer_map_analyzer::parse_filename;
@@ -59,44 +47,30 @@ pub(crate) enum LayerCmd {
 }
 
 async fn read_delta_file(path: impl AsRef<Path>, ctx: &RequestContext) -> Result<()> {
-    let path = Utf8Path::from_path(path.as_ref()).expect("non-Unicode path");
     virtual_file::init(
         10,
         virtual_file::api::IoEngineKind::StdFs,
         IoMode::preferred(),
     );
     page_cache::init(100);
-    let file = VirtualFile::open(path, ctx).await?;
-    let file_id = page_cache::next_file_id();
-    let block_reader = FileBlockReader::new(&file, file_id);
-    let summary_blk = block_reader.read_blk(0, ctx).await?;
-    let actual_summary = Summary::des_prefix(summary_blk.as_ref())?;
-    let tree_reader = DiskBtreeReader::<_, DELTA_KEY_SIZE>::new(
-        actual_summary.index_start_blk,
-        actual_summary.index_root_blk,
-        &block_reader,
+    let path = Utf8Path::from_path(path.as_ref()).expect("non-Unicode path");
+    let file = File::open(path)?;
+    let delta_layer = DeltaLayer::new_for_path(path, file)?;
+    delta_layer.dump(true, ctx).await?;
+    Ok(())
+}
+
+async fn read_image_file(path: impl AsRef<Path>, ctx: &RequestContext) -> Result<()> {
+    virtual_file::init(
+        10,
+        virtual_file::api::IoEngineKind::StdFs,
+        IoMode::preferred(),
     );
-    // TODO(chi): dedup w/ `delta_layer.rs` by exposing the API.
-    let mut all = vec![];
-    tree_reader
-        .visit(
-            &[0u8; DELTA_KEY_SIZE],
-            VisitDirection::Forwards,
-            |key, value_offset| {
-                let curr = Key::from_slice(&key[..KEY_SIZE]);
-                all.push((curr, BlobRef(value_offset)));
-                true
-            },
-            ctx,
-        )
-        .await?;
-    let cursor = BlockCursor::new_fileblockreader(&block_reader);
-    for (k, v) in all {
-        let value = cursor.read_blob(v.pos(), ctx).await?;
-        println!("key:{} value_len:{}", k, value.len());
-        assert!(k.is_i128_representable(), "invalid key: ");
-    }
-    // TODO(chi): special handling for last key?
+    page_cache::init(100);
+    let path = Utf8Path::from_path(path.as_ref()).expect("non-Unicode path");
+    let file = File::open(path)?;
+    let image_layer = ImageLayer::new_for_path(path, file)?;
+    image_layer.dump(true, ctx).await?;
     Ok(())
 }
 
@@ -133,8 +107,7 @@ pub(crate) async fn main(cmd: &LayerCmd) -> Result<()> {
             let mut idx = 0;
             for layer in fs::read_dir(timeline_path)? {
                 let layer = layer?;
-                if let Some(layer_file) = parse_filename(&layer.file_name().into_string().unwrap())
-                {
+                if let Ok(layer_file) = parse_filename(&layer.file_name().into_string().unwrap()) {
                     println!(
                         "[{:3}]  key:{}-{}\n       lsn:{}-{}\n       delta:{}",
                         idx,
@@ -163,8 +136,7 @@ pub(crate) async fn main(cmd: &LayerCmd) -> Result<()> {
             let mut idx = 0;
             for layer in fs::read_dir(timeline_path)? {
                 let layer = layer?;
-                if let Some(layer_file) = parse_filename(&layer.file_name().into_string().unwrap())
-                {
+                if let Ok(layer_file) = parse_filename(&layer.file_name().into_string().unwrap()) {
                     if *id == idx {
                         // TODO(chi): dedup code
                         println!(
@@ -180,7 +152,7 @@ pub(crate) async fn main(cmd: &LayerCmd) -> Result<()> {
                         if layer_file.is_delta {
                             read_delta_file(layer.path(), &ctx).await?;
                         } else {
-                            anyhow::bail!("not supported yet :(");
+                            read_image_file(layer.path(), &ctx).await?;
                         }
 
                         break;

pageserver/src/basebackup.rs

@@ -59,6 +59,7 @@ pub async fn send_basebackup_tarball<'a, W>(
     req_lsn: Option<Lsn>,
     prev_lsn: Option<Lsn>,
     full_backup: bool,
+    replica: bool,
     ctx: &'a RequestContext,
 ) -> Result<(), BasebackupError>
 where
@@ -110,8 +111,8 @@ where
     };
 
     info!(
-        "taking basebackup lsn={}, prev_lsn={} (full_backup={})",
-        backup_lsn, prev_lsn, full_backup
+        "taking basebackup lsn={}, prev_lsn={} (full_backup={}, replica={})",
+        backup_lsn, prev_lsn, full_backup, replica
     );
 
     let basebackup = Basebackup {
@@ -120,6 +121,7 @@ where
         lsn: backup_lsn,
         prev_record_lsn: prev_lsn,
         full_backup,
+        replica,
         ctx,
     };
     basebackup
@@ -140,6 +142,7 @@ where
     lsn: Lsn,
     prev_record_lsn: Lsn,
     full_backup: bool,
+    replica: bool,
     ctx: &'a RequestContext,
 }
 
@@ -372,6 +375,10 @@ where
 
         for (path, content) in aux_files {
             if path.starts_with("pg_replslot") {
+                // Do not create LR slots at standby because they are not used but prevent WAL truncation
+                if self.replica {
+                    continue;
+                }
                 let offs = pg_constants::REPL_SLOT_ON_DISK_OFFSETOF_RESTART_LSN;
                 let restart_lsn = Lsn(u64::from_le_bytes(
                     content[offs..offs + 8].try_into().unwrap(),

pageserver/src/bin/pageserver.rs

@@ -398,9 +398,7 @@ fn start_pageserver(
         ControllerUpcallClient::new(conf, &shutdown_pageserver),
         conf,
     );
-    if let Some(deletion_workers) = deletion_workers {
-        deletion_workers.spawn_with(BACKGROUND_RUNTIME.handle());
-    }
+    deletion_workers.spawn_with(BACKGROUND_RUNTIME.handle());
 
     // Up to this point no significant I/O has been done: this should have been fast.  Record
     // duration prior to starting I/O intensive phase of startup.

pageserver/src/consumption_metrics.rs

@@ -14,6 +14,7 @@ use itertools::Itertools as _;
 use pageserver_api::models::TenantState;
 use remote_storage::{GenericRemoteStorage, RemoteStorageConfig};
 use reqwest::Url;
+use serde::{Deserialize, Serialize};
 use std::collections::HashMap;
 use std::sync::Arc;
 use std::time::{Duration, SystemTime};
@@ -35,12 +36,62 @@ const DEFAULT_HTTP_REPORTING_TIMEOUT: Duration = Duration::from_secs(60);
 /// upload attempts.
 type RawMetric = (MetricsKey, (EventType, u64));
 
+/// The new serializable metrics format
+#[derive(Serialize, Deserialize)]
+struct NewMetricsRoot {
+    version: usize,
+    metrics: Vec<NewRawMetric>,
+}
+
+impl NewMetricsRoot {
+    pub fn is_v2_metrics(json_value: &serde_json::Value) -> bool {
+        if let Some(ver) = json_value.get("version") {
+            if let Some(2) = ver.as_u64() {
+                return true;
+            }
+        }
+        false
+    }
+}
+
+/// The new serializable metrics format
+#[derive(Serialize)]
+struct NewMetricsRefRoot<'a> {
+    version: usize,
+    metrics: &'a [NewRawMetric],
+}
+
+impl<'a> NewMetricsRefRoot<'a> {
+    fn new(metrics: &'a [NewRawMetric]) -> Self {
+        Self {
+            version: 2,
+            metrics,
+        }
+    }
+}
+
+/// The new serializable metrics format
+#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq)]
+struct NewRawMetric {
+    key: MetricsKey,
+    kind: EventType,
+    value: u64,
+    // TODO: add generation field and check against generations
+}
+
+impl NewRawMetric {
+    #[cfg(test)]
+    fn to_kv_pair(&self) -> (MetricsKey, NewRawMetric) {
+        (self.key, self.clone())
+    }
+}
+
 /// Caches the [`RawMetric`]s
 ///
 /// In practice, during startup, last sent values are stored here to be used in calculating new
 /// ones. After successful uploading, the cached values are updated to cache. This used to be used
 /// for deduplication, but that is no longer needed.
-type Cache = HashMap<MetricsKey, (EventType, u64)>;
+type Cache = HashMap<MetricsKey, NewRawMetric>;
 
 pub async fn run(
     conf: &'static PageServerConf,
@@ -231,11 +282,14 @@ async fn restore_and_reschedule(
             // collect_all_metrics
             let earlier_metric_at = found_some
                 .iter()
-                .map(|(_, (et, _))| et.recorded_at())
+                .map(|item| item.kind.recorded_at())
                 .copied()
                 .next();
 
-            let cached = found_some.into_iter().collect::<Cache>();
+            let cached = found_some
+                .into_iter()
+                .map(|item| (item.key, item))
+                .collect::<Cache>();
 
             (cached, earlier_metric_at)
         }

pageserver/src/consumption_metrics/disk_cache.rs

@@ -2,11 +2,33 @@ use anyhow::Context;
 use camino::{Utf8Path, Utf8PathBuf};
 use std::sync::Arc;
 
-use super::RawMetric;
+use crate::consumption_metrics::NewMetricsRefRoot;
+
+use super::{NewMetricsRoot, NewRawMetric, RawMetric};
+
+pub(super) fn read_metrics_from_serde_value(
+    json_value: serde_json::Value,
+) -> anyhow::Result<Vec<NewRawMetric>> {
+    if NewMetricsRoot::is_v2_metrics(&json_value) {
+        let root = serde_json::from_value::<NewMetricsRoot>(json_value)?;
+        Ok(root.metrics)
+    } else {
+        let all_metrics = serde_json::from_value::<Vec<RawMetric>>(json_value)?;
+        let all_metrics = all_metrics
+            .into_iter()
+            .map(|(key, (event_type, value))| NewRawMetric {
+                key,
+                kind: event_type,
+                value,
+            })
+            .collect();
+        Ok(all_metrics)
+    }
+}
 
 pub(super) async fn read_metrics_from_disk(
     path: Arc<Utf8PathBuf>,
-) -> anyhow::Result<Vec<RawMetric>> {
+) -> anyhow::Result<Vec<NewRawMetric>> {
     // do not add context to each error, callsite will log with full path
     let span = tracing::Span::current();
     tokio::task::spawn_blocking(move || {
@@ -20,7 +42,8 @@ pub(super) async fn read_metrics_from_disk(
 
         let mut file = std::fs::File::open(&*path)?;
         let reader = std::io::BufReader::new(&mut file);
-        anyhow::Ok(serde_json::from_reader::<_, Vec<RawMetric>>(reader)?)
+        let json_value = serde_json::from_reader::<_, serde_json::Value>(reader)?;
+        read_metrics_from_serde_value(json_value)
     })
     .await
     .context("read metrics join error")
@@ -63,7 +86,7 @@ fn scan_and_delete_with_same_prefix(path: &Utf8Path) -> std::io::Result<()> {
 }
 
 pub(super) async fn flush_metrics_to_disk(
-    current_metrics: &Arc<Vec<RawMetric>>,
+    current_metrics: &Arc<Vec<NewRawMetric>>,
     path: &Arc<Utf8PathBuf>,
 ) -> anyhow::Result<()> {
     use std::io::Write;
@@ -93,8 +116,11 @@ pub(super) async fn flush_metrics_to_disk(
             // write out all of the raw metrics, to be read out later on restart as cached values
             {
                 let mut writer = std::io::BufWriter::new(&mut tempfile);
-                serde_json::to_writer(&mut writer, &*current_metrics)
-                    .context("serialize metrics")?;
+                serde_json::to_writer(
+                    &mut writer,
+                    &NewMetricsRefRoot::new(current_metrics.as_ref()),
+                )
+                .context("serialize metrics")?;
                 writer
                     .into_inner()
                     .map_err(|_| anyhow::anyhow!("flushing metrics failed"))?;

pageserver/src/consumption_metrics/metrics.rs

@@ -9,7 +9,7 @@ use utils::{
     lsn::Lsn,
 };
 
-use super::{Cache, RawMetric};
+use super::{Cache, NewRawMetric};
 
 /// Name of the metric, used by `MetricsKey` factory methods and `deserialize_cached_events`
 /// instead of static str.
@@ -64,11 +64,21 @@ impl MetricsKey {
 struct AbsoluteValueFactory(MetricsKey);
 
 impl AbsoluteValueFactory {
-    const fn at(self, time: DateTime<Utc>, val: u64) -> RawMetric {
+    #[cfg(test)]
+    const fn at_old_format(self, time: DateTime<Utc>, val: u64) -> super::RawMetric {
         let key = self.0;
         (key, (EventType::Absolute { time }, val))
     }
 
+    const fn at(self, time: DateTime<Utc>, val: u64) -> NewRawMetric {
+        let key = self.0;
+        NewRawMetric {
+            key,
+            kind: EventType::Absolute { time },
+            value: val,
+        }
+    }
+
     fn key(&self) -> &MetricsKey {
         &self.0
     }
@@ -84,7 +94,28 @@ impl IncrementalValueFactory {
         prev_end: DateTime<Utc>,
         up_to: DateTime<Utc>,
         val: u64,
-    ) -> RawMetric {
+    ) -> NewRawMetric {
+        let key = self.0;
+        // cannot assert prev_end < up_to because these are realtime clock based
+        let when = EventType::Incremental {
+            start_time: prev_end,
+            stop_time: up_to,
+        };
+        NewRawMetric {
+            key,
+            kind: when,
+            value: val,
+        }
+    }
+
+    #[allow(clippy::wrong_self_convention)]
+    #[cfg(test)]
+    const fn from_until_old_format(
+        self,
+        prev_end: DateTime<Utc>,
+        up_to: DateTime<Utc>,
+        val: u64,
+    ) -> super::RawMetric {
         let key = self.0;
         // cannot assert prev_end < up_to because these are realtime clock based
         let when = EventType::Incremental {
@@ -185,7 +216,7 @@ pub(super) async fn collect_all_metrics(
     tenant_manager: &Arc<TenantManager>,
     cached_metrics: &Cache,
     ctx: &RequestContext,
-) -> Vec<RawMetric> {
+) -> Vec<NewRawMetric> {
     use pageserver_api::models::TenantState;
 
     let started_at = std::time::Instant::now();
@@ -220,11 +251,11 @@ pub(super) async fn collect_all_metrics(
     res
 }
 
-async fn collect<S>(tenants: S, cache: &Cache, ctx: &RequestContext) -> Vec<RawMetric>
+async fn collect<S>(tenants: S, cache: &Cache, ctx: &RequestContext) -> Vec<NewRawMetric>
 where
     S: futures::stream::Stream<Item = (TenantId, Arc<crate::tenant::Tenant>)>,
 {
-    let mut current_metrics: Vec<RawMetric> = Vec::new();
+    let mut current_metrics: Vec<NewRawMetric> = Vec::new();
 
     let mut tenants = std::pin::pin!(tenants);
 
@@ -291,7 +322,7 @@ impl TenantSnapshot {
         tenant_id: TenantId,
         now: DateTime<Utc>,
         cached: &Cache,
-        metrics: &mut Vec<RawMetric>,
+        metrics: &mut Vec<NewRawMetric>,
     ) {
         let remote_size = MetricsKey::remote_storage_size(tenant_id).at(now, self.remote_size);
 
@@ -302,9 +333,9 @@ impl TenantSnapshot {
             let mut synthetic_size = self.synthetic_size;
 
             if synthetic_size == 0 {
-                if let Some((_, value)) = cached.get(factory.key()) {
-                    // use the latest value from previous session
-                    synthetic_size = *value;
+                if let Some(item) = cached.get(factory.key()) {
+                    // use the latest value from previous session, TODO: check generation number
+                    synthetic_size = item.value;
                 }
             }
 
@@ -381,37 +412,36 @@ impl TimelineSnapshot {
         tenant_id: TenantId,
         timeline_id: TimelineId,
         now: DateTime<Utc>,
-        metrics: &mut Vec<RawMetric>,
+        metrics: &mut Vec<NewRawMetric>,
         cache: &Cache,
     ) {
         let timeline_written_size = u64::from(self.last_record_lsn);
 
         let written_size_delta_key = MetricsKey::written_size_delta(tenant_id, timeline_id);
 
-        let last_stop_time = cache
-            .get(written_size_delta_key.key())
-            .map(|(until, _val)| {
-                until
-                    .incremental_timerange()
-                    .expect("never create EventType::Absolute for written_size_delta")
-                    .end
-            });
+        let last_stop_time = cache.get(written_size_delta_key.key()).map(|item| {
+            item.kind
+                .incremental_timerange()
+                .expect("never create EventType::Absolute for written_size_delta")
+                .end
+        });
 
-        let (key, written_size_now) =
+        let written_size_now =
             MetricsKey::written_size(tenant_id, timeline_id).at(now, timeline_written_size);
 
         // by default, use the last sent written_size as the basis for
         // calculating the delta. if we don't yet have one, use the load time value.
-        let prev = cache
-            .get(&key)
-            .map(|(prev_at, prev)| {
+        let prev: (DateTime<Utc>, u64) = cache
+            .get(&written_size_now.key)
+            .map(|item| {
                 // use the prev time from our last incremental update, or default to latest
                 // absolute update on the first round.
-                let prev_at = prev_at
+                let prev_at = item
+                    .kind
                     .absolute_time()
                     .expect("never create EventType::Incremental for written_size");
                 let prev_at = last_stop_time.unwrap_or(prev_at);
-                (*prev_at, *prev)
+                (*prev_at, item.value)
             })
             .unwrap_or_else(|| {
                 // if we don't have a previous point of comparison, compare to the load time
@@ -422,24 +452,28 @@ impl TimelineSnapshot {
 
         let up_to = now;
 
-        if let Some(delta) = written_size_now.1.checked_sub(prev.1) {
+        if let Some(delta) = written_size_now.value.checked_sub(prev.1) {
             let key_value = written_size_delta_key.from_until(prev.0, up_to, delta);
             // written_size_delta
             metrics.push(key_value);
             // written_size
-            metrics.push((key, written_size_now));
+            metrics.push(written_size_now);
         } else {
             // the cached value was ahead of us, report zero until we've caught up
             metrics.push(written_size_delta_key.from_until(prev.0, up_to, 0));
             // the cached value was ahead of us, report the same until we've caught up
-            metrics.push((key, (written_size_now.0, prev.1)));
+            metrics.push(NewRawMetric {
+                key: written_size_now.key,
+                kind: written_size_now.kind,
+                value: prev.1,
+            });
         }
 
         {
             let factory = MetricsKey::timeline_logical_size(tenant_id, timeline_id);
             let current_or_previous = self
                 .current_exact_logical_size
-                .or_else(|| cache.get(factory.key()).map(|(_, val)| *val));
+                .or_else(|| cache.get(factory.key()).map(|item| item.value));
 
             if let Some(size) = current_or_previous {
                 metrics.push(factory.at(now, size));
@@ -452,4 +486,4 @@ impl TimelineSnapshot {
 mod tests;
 
 #[cfg(test)]
-pub(crate) use tests::metric_examples;
+pub(crate) use tests::{metric_examples, metric_examples_old};

pageserver/src/consumption_metrics/metrics/tests.rs

@@ -1,3 +1,5 @@
+use crate::consumption_metrics::RawMetric;
+
 use super::*;
 use std::collections::HashMap;
 
@@ -50,9 +52,9 @@ fn startup_collected_timeline_metrics_second_round() {
     let disk_consistent_lsn = Lsn(initdb_lsn.0 * 2);
 
     let mut metrics = Vec::new();
-    let cache = HashMap::from([
-        MetricsKey::written_size(tenant_id, timeline_id).at(before, disk_consistent_lsn.0)
-    ]);
+    let cache = HashMap::from([MetricsKey::written_size(tenant_id, timeline_id)
+        .at(before, disk_consistent_lsn.0)
+        .to_kv_pair()]);
 
     let snap = TimelineSnapshot {
         loaded_at: (disk_consistent_lsn, init),
@@ -89,9 +91,13 @@ fn startup_collected_timeline_metrics_nth_round_at_same_lsn() {
     let mut metrics = Vec::new();
     let cache = HashMap::from([
         // at t=before was the last time the last_record_lsn changed
-        MetricsKey::written_size(tenant_id, timeline_id).at(before, disk_consistent_lsn.0),
+        MetricsKey::written_size(tenant_id, timeline_id)
+            .at(before, disk_consistent_lsn.0)
+            .to_kv_pair(),
         // end time of this event is used for the next ones
-        MetricsKey::written_size_delta(tenant_id, timeline_id).from_until(before, just_before, 0),
+        MetricsKey::written_size_delta(tenant_id, timeline_id)
+            .from_until(before, just_before, 0)
+            .to_kv_pair(),
     ]);
 
     let snap = TimelineSnapshot {
@@ -138,13 +144,17 @@ fn post_restart_written_sizes_with_rolled_back_last_record_lsn() {
     };
 
     let mut cache = HashMap::from([
-        MetricsKey::written_size(tenant_id, timeline_id).at(before_restart, 100),
-        MetricsKey::written_size_delta(tenant_id, timeline_id).from_until(
-            way_before,
-            before_restart,
-            // not taken into account, but the timestamps are important
-            999_999_999,
-        ),
+        MetricsKey::written_size(tenant_id, timeline_id)
+            .at(before_restart, 100)
+            .to_kv_pair(),
+        MetricsKey::written_size_delta(tenant_id, timeline_id)
+            .from_until(
+                way_before,
+                before_restart,
+                // not taken into account, but the timestamps are important
+                999_999_999,
+            )
+            .to_kv_pair(),
     ]);
 
     let mut metrics = Vec::new();
@@ -163,7 +173,7 @@ fn post_restart_written_sizes_with_rolled_back_last_record_lsn() {
     );
 
     // now if we cache these metrics, and re-run while "still in recovery"
-    cache.extend(metrics.drain(..));
+    cache.extend(metrics.drain(..).map(|x| x.to_kv_pair()));
 
     // "still in recovery", because our snapshot did not change
     snap.to_metrics(tenant_id, timeline_id, later, &mut metrics, &cache);
@@ -194,14 +204,14 @@ fn post_restart_current_exact_logical_size_uses_cached() {
         current_exact_logical_size: None,
     };
 
-    let cache = HashMap::from([
-        MetricsKey::timeline_logical_size(tenant_id, timeline_id).at(before_restart, 100)
-    ]);
+    let cache = HashMap::from([MetricsKey::timeline_logical_size(tenant_id, timeline_id)
+        .at(before_restart, 100)
+        .to_kv_pair()]);
 
     let mut metrics = Vec::new();
     snap.to_metrics(tenant_id, timeline_id, now, &mut metrics, &cache);
 
-    metrics.retain(|(key, _)| key.metric == Name::LogicalSize);
+    metrics.retain(|item| item.key.metric == Name::LogicalSize);
 
     assert_eq!(
         metrics,
@@ -224,7 +234,9 @@ fn post_restart_synthetic_size_uses_cached_if_available() {
     let before_restart = DateTime::<Utc>::from(now - std::time::Duration::from_secs(5 * 60));
     let now = DateTime::<Utc>::from(now);
 
-    let cached = HashMap::from([MetricsKey::synthetic_size(tenant_id).at(before_restart, 1000)]);
+    let cached = HashMap::from([MetricsKey::synthetic_size(tenant_id)
+        .at(before_restart, 1000)
+        .to_kv_pair()]);
 
     let mut metrics = Vec::new();
     ts.to_metrics(tenant_id, now, &cached, &mut metrics);
@@ -278,12 +290,29 @@ fn time_backwards<const N: usize>() -> [std::time::SystemTime; N] {
     times
 }
 
-pub(crate) const fn metric_examples(
+pub(crate) const fn metric_examples_old(
     tenant_id: TenantId,
     timeline_id: TimelineId,
     now: DateTime<Utc>,
     before: DateTime<Utc>,
 ) -> [RawMetric; 6] {
+    [
+        MetricsKey::written_size(tenant_id, timeline_id).at_old_format(now, 0),
+        MetricsKey::written_size_delta(tenant_id, timeline_id)
+            .from_until_old_format(before, now, 0),
+        MetricsKey::timeline_logical_size(tenant_id, timeline_id).at_old_format(now, 0),
+        MetricsKey::remote_storage_size(tenant_id).at_old_format(now, 0),
+        MetricsKey::resident_size(tenant_id).at_old_format(now, 0),
+        MetricsKey::synthetic_size(tenant_id).at_old_format(now, 1),
+    ]
+}
+
+pub(crate) const fn metric_examples(
+    tenant_id: TenantId,
+    timeline_id: TimelineId,
+    now: DateTime<Utc>,
+    before: DateTime<Utc>,
+) -> [NewRawMetric; 6] {
     [
         MetricsKey::written_size(tenant_id, timeline_id).at(now, 0),
         MetricsKey::written_size_delta(tenant_id, timeline_id).from_until(before, now, 0),

pageserver/src/consumption_metrics/upload.rs

@@ -7,7 +7,7 @@ use tokio::io::AsyncWriteExt;
 use tokio_util::sync::CancellationToken;
 use tracing::Instrument;
 
-use super::{metrics::Name, Cache, MetricsKey, RawMetric};
+use super::{metrics::Name, Cache, MetricsKey, NewRawMetric, RawMetric};
 use utils::id::{TenantId, TimelineId};
 
 /// How the metrics from pageserver are identified.
@@ -24,7 +24,7 @@ pub(super) async fn upload_metrics_http(
     client: &reqwest::Client,
     metric_collection_endpoint: &reqwest::Url,
     cancel: &CancellationToken,
-    metrics: &[RawMetric],
+    metrics: &[NewRawMetric],
     cached_metrics: &mut Cache,
     idempotency_keys: &[IdempotencyKey<'_>],
 ) -> anyhow::Result<()> {
@@ -53,8 +53,8 @@ pub(super) async fn upload_metrics_http(
 
         match res {
             Ok(()) => {
-                for (curr_key, curr_val) in chunk {
-                    cached_metrics.insert(*curr_key, *curr_val);
+                for item in chunk {
+                    cached_metrics.insert(item.key, item.clone());
                 }
                 uploaded += chunk.len();
             }
@@ -86,7 +86,7 @@ pub(super) async fn upload_metrics_bucket(
     client: &GenericRemoteStorage,
     cancel: &CancellationToken,
     node_id: &str,
-    metrics: &[RawMetric],
+    metrics: &[NewRawMetric],
     idempotency_keys: &[IdempotencyKey<'_>],
 ) -> anyhow::Result<()> {
     if metrics.is_empty() {
@@ -140,16 +140,16 @@ pub(super) async fn upload_metrics_bucket(
 /// across different metrics sinks), and must have the same length as input.
 fn serialize_in_chunks<'a>(
     chunk_size: usize,
-    input: &'a [RawMetric],
+    input: &'a [NewRawMetric],
     idempotency_keys: &'a [IdempotencyKey<'a>],
-) -> impl ExactSizeIterator<Item = Result<(&'a [RawMetric], bytes::Bytes), serde_json::Error>> + 'a
+) -> impl ExactSizeIterator<Item = Result<(&'a [NewRawMetric], bytes::Bytes), serde_json::Error>> + 'a
 {
     use bytes::BufMut;
 
     assert_eq!(input.len(), idempotency_keys.len());
 
     struct Iter<'a> {
-        inner: std::slice::Chunks<'a, RawMetric>,
+        inner: std::slice::Chunks<'a, NewRawMetric>,
         idempotency_keys: std::slice::Iter<'a, IdempotencyKey<'a>>,
         chunk_size: usize,
 
@@ -160,7 +160,7 @@ fn serialize_in_chunks<'a>(
     }
 
     impl<'a> Iterator for Iter<'a> {
-        type Item = Result<(&'a [RawMetric], bytes::Bytes), serde_json::Error>;
+        type Item = Result<(&'a [NewRawMetric], bytes::Bytes), serde_json::Error>;
 
         fn next(&mut self) -> Option<Self::Item> {
             let chunk = self.inner.next()?;
@@ -269,6 +269,58 @@ impl RawMetricExt for RawMetric {
     }
 }
 
+impl RawMetricExt for NewRawMetric {
+    fn as_event(&self, key: &IdempotencyKey<'_>) -> Event<Ids, Name> {
+        let MetricsKey {
+            metric,
+            tenant_id,
+            timeline_id,
+        } = self.key;
+
+        let kind = self.kind;
+        let value = self.value;
+
+        Event {
+            kind,
+            metric,
+            idempotency_key: key.to_string(),
+            value,
+            extra: Ids {
+                tenant_id,
+                timeline_id,
+            },
+        }
+    }
+
+    fn update_in_place(&self, event: &mut Event<Ids, Name>, key: &IdempotencyKey<'_>) {
+        use std::fmt::Write;
+
+        let MetricsKey {
+            metric,
+            tenant_id,
+            timeline_id,
+        } = self.key;
+
+        let kind = self.kind;
+        let value = self.value;
+
+        *event = Event {
+            kind,
+            metric,
+            idempotency_key: {
+                event.idempotency_key.clear();
+                write!(event.idempotency_key, "{key}").unwrap();
+                std::mem::take(&mut event.idempotency_key)
+            },
+            value,
+            extra: Ids {
+                tenant_id,
+                timeline_id,
+            },
+        };
+    }
+}
+
 pub(crate) trait KeyGen<'a> {
     fn generate(&self) -> IdempotencyKey<'a>;
 }
@@ -381,6 +433,10 @@ async fn upload(
 
 #[cfg(test)]
 mod tests {
+    use crate::consumption_metrics::{
+        disk_cache::read_metrics_from_serde_value, NewMetricsRefRoot,
+    };
+
     use super::*;
     use chrono::{DateTime, Utc};
     use once_cell::sync::Lazy;
@@ -473,23 +529,49 @@ mod tests {
         let idempotency_key = consumption_metrics::IdempotencyKey::for_tests(*SAMPLES_NOW, "1", 0);
         let examples = examples.into_iter().zip(metric_samples());
 
-        for ((line, expected), (key, (kind, value))) in examples {
+        for ((line, expected), item) in examples {
             let e = consumption_metrics::Event {
-                kind,
-                metric: key.metric,
+                kind: item.kind,
+                metric: item.key.metric,
                 idempotency_key: idempotency_key.to_string(),
-                value,
+                value: item.value,
                 extra: Ids {
-                    tenant_id: key.tenant_id,
-                    timeline_id: key.timeline_id,
+                    tenant_id: item.key.tenant_id,
+                    timeline_id: item.key.timeline_id,
                 },
             };
             let actual = serde_json::to_string(&e).unwrap();
-            assert_eq!(expected, actual, "example for {kind:?} from line {line}");
+            assert_eq!(
+                expected, actual,
+                "example for {:?} from line {line}",
+                item.kind
+            );
         }
     }
 
-    fn metric_samples() -> [RawMetric; 6] {
+    #[test]
+    fn disk_format_upgrade() {
+        let old_samples_json = serde_json::to_value(metric_samples_old()).unwrap();
+        let new_samples =
+            serde_json::to_value(NewMetricsRefRoot::new(metric_samples().as_ref())).unwrap();
+        let upgraded_samples = read_metrics_from_serde_value(old_samples_json).unwrap();
+        let new_samples = read_metrics_from_serde_value(new_samples).unwrap();
+        assert_eq!(upgraded_samples, new_samples);
+    }
+
+    fn metric_samples_old() -> [RawMetric; 6] {
+        let tenant_id = TenantId::from_array([0; 16]);
+        let timeline_id = TimelineId::from_array([0xff; 16]);
+
+        let before = DateTime::parse_from_rfc3339("2023-09-14T00:00:00.123456789Z")
+            .unwrap()
+            .into();
+        let [now, before] = [*SAMPLES_NOW, before];
+
+        super::super::metrics::metric_examples_old(tenant_id, timeline_id, now, before)
+    }
+
+    fn metric_samples() -> [NewRawMetric; 6] {
         let tenant_id = TenantId::from_array([0; 16]);
         let timeline_id = TimelineId::from_array([0xff; 16]);
 

pageserver/src/deletion_queue.rs

@@ -618,13 +618,11 @@ impl DeletionQueue {
     /// Caller may use the returned object to construct clients with new_client.
     /// Caller should tokio::spawn the background() members of the two worker objects returned:
     /// we don't spawn those inside new() so that the caller can use their runtime/spans of choice.
-    ///
-    /// If remote_storage is None, then the returned workers will also be None.
     pub fn new<C>(
         remote_storage: GenericRemoteStorage,
         controller_upcall_client: Option<C>,
         conf: &'static PageServerConf,
-    ) -> (Self, Option<DeletionQueueWorkers<C>>)
+    ) -> (Self, DeletionQueueWorkers<C>)
     where
         C: ControlPlaneGenerationsApi + Send + Sync,
     {
@@ -656,7 +654,7 @@ impl DeletionQueue {
                 },
                 cancel: cancel.clone(),
             },
-            Some(DeletionQueueWorkers {
+            DeletionQueueWorkers {
                 frontend: ListWriter::new(conf, rx, backend_tx, cancel.clone()),
                 backend: Validator::new(
                     conf,
@@ -667,7 +665,7 @@ impl DeletionQueue {
                     cancel.clone(),
                 ),
                 executor: Deleter::new(remote_storage, executor_rx, cancel.clone()),
-            }),
+            },
         )
     }
 
@@ -742,9 +740,7 @@ mod test {
             );
 
             tracing::debug!("Spawning worker for new queue queue");
-            let worker_join = workers
-                .unwrap()
-                .spawn_with(&tokio::runtime::Handle::current());
+            let worker_join = workers.spawn_with(&tokio::runtime::Handle::current());
 
             let old_worker_join = std::mem::replace(&mut self.worker_join, worker_join);
             let old_deletion_queue = std::mem::replace(&mut self.deletion_queue, deletion_queue);
@@ -855,7 +851,6 @@ mod test {
             harness.conf,
         );
 
-        let worker = worker.unwrap();
         let worker_join = worker.spawn_with(&tokio::runtime::Handle::current());
 
         Ok(TestSetup {

pageserver/src/http/routes.rs

@@ -80,6 +80,7 @@ use crate::tenant::size::ModelInputs;
 use crate::tenant::storage_layer::LayerAccessStatsReset;
 use crate::tenant::storage_layer::LayerName;
 use crate::tenant::timeline::offload::offload_timeline;
+use crate::tenant::timeline::offload::OffloadError;
 use crate::tenant::timeline::CompactFlags;
 use crate::tenant::timeline::CompactionError;
 use crate::tenant::timeline::Timeline;
@@ -2004,7 +2005,12 @@ async fn timeline_offload_handler(
         }
         offload_timeline(&tenant, &timeline)
             .await
-            .map_err(ApiError::InternalServerError)?;
+            .map_err(|e| {
+                match e {
+                    OffloadError::Cancelled => ApiError::ResourceUnavailable("Timeline shutting down".into()),
+                    _ => ApiError::InternalServerError(anyhow!(e))
+                }
+            })?;
 
         json_response(StatusCode::OK, ())
     }
@@ -2060,6 +2066,7 @@ async fn timeline_checkpoint_handler(
                 .map_err(|e|
                     match e {
                         CompactionError::ShuttingDown => ApiError::ShuttingDown,
+                        CompactionError::Offload(e) => ApiError::InternalServerError(anyhow::anyhow!(e)),
                         CompactionError::Other(e) => ApiError::InternalServerError(e)
                     }
                 )?;

pageserver/src/import_datadir.rs

@@ -12,6 +12,7 @@ use pageserver_api::key::rel_block_to_key;
 use tokio::io::{AsyncRead, AsyncReadExt};
 use tokio_tar::Archive;
 use tracing::*;
+use wal_decoder::models::InterpretedWalRecord;
 use walkdir::WalkDir;
 
 use crate::context::RequestContext;
@@ -23,7 +24,6 @@ use pageserver_api::reltag::{RelTag, SlruKind};
 use postgres_ffi::pg_constants;
 use postgres_ffi::relfile_utils::*;
 use postgres_ffi::waldecoder::WalStreamDecoder;
-use postgres_ffi::walrecord::{decode_wal_record, DecodedWALRecord};
 use postgres_ffi::ControlFileData;
 use postgres_ffi::DBState_DB_SHUTDOWNED;
 use postgres_ffi::Oid;
@@ -312,11 +312,15 @@ async fn import_wal(
         let mut modification = tline.begin_modification(last_lsn);
         while last_lsn <= endpoint {
             if let Some((lsn, recdata)) = waldecoder.poll_decode()? {
-                let mut decoded = DecodedWALRecord::default();
-                decode_wal_record(recdata, &mut decoded, tline.pg_version)?;
+                let interpreted = InterpretedWalRecord::from_bytes_filtered(
+                    recdata,
+                    tline.get_shard_identity(),
+                    lsn,
+                    tline.pg_version,
+                )?;
 
                 walingest
-                    .ingest_record(decoded, lsn, &mut modification, ctx)
+                    .ingest_record(interpreted, &mut modification, ctx)
                     .await?;
                 WAL_INGEST.records_committed.inc();
 
@@ -453,10 +457,15 @@ pub async fn import_wal_from_tar(
         let mut modification = tline.begin_modification(last_lsn);
         while last_lsn <= end_lsn {
             if let Some((lsn, recdata)) = waldecoder.poll_decode()? {
-                let mut decoded = DecodedWALRecord::default();
-                decode_wal_record(recdata, &mut decoded, tline.pg_version)?;
+                let interpreted = InterpretedWalRecord::from_bytes_filtered(
+                    recdata,
+                    tline.get_shard_identity(),
+                    lsn,
+                    tline.pg_version,
+                )?;
+
                 walingest
-                    .ingest_record(decoded, lsn, &mut modification, ctx)
+                    .ingest_record(interpreted, &mut modification, ctx)
                     .await?;
                 modification.commit(ctx).await?;
                 last_lsn = lsn;

pageserver/src/page_service.rs

@@ -1080,6 +1080,7 @@ impl PageServerHandler {
         prev_lsn: Option<Lsn>,
         full_backup: bool,
         gzip: bool,
+        replica: bool,
         ctx: &RequestContext,
     ) -> Result<(), QueryError>
     where
@@ -1132,6 +1133,7 @@ impl PageServerHandler {
                 lsn,
                 prev_lsn,
                 full_backup,
+                replica,
                 ctx,
             )
             .await
@@ -1154,6 +1156,7 @@ impl PageServerHandler {
                     lsn,
                     prev_lsn,
                     full_backup,
+                    replica,
                     ctx,
                 )
                 .await
@@ -1170,6 +1173,7 @@ impl PageServerHandler {
                     lsn,
                     prev_lsn,
                     full_backup,
+                    replica,
                     ctx,
                 )
                 .await
@@ -1326,24 +1330,27 @@ where
                 .for_command(ComputeCommandKind::Basebackup)
                 .inc();
 
-            let (lsn, gzip) = match (params.get(2), params.get(3)) {
-                (None, _) => (None, false),
-                (Some(&"--gzip"), _) => (None, true),
-                (Some(lsn_str), gzip_str_opt) => {
-                    let lsn = Lsn::from_str(lsn_str)
-                        .with_context(|| format!("Failed to parse Lsn from {lsn_str}"))?;
-                    let gzip = match gzip_str_opt {
-                        Some(&"--gzip") => true,
-                        None => false,
-                        Some(third_param) => {
+            let mut lsn = None;
+            let mut replica = false;
+            let mut gzip = false;
+            for param in &params[2..] {
+                if param.starts_with("--") {
+                    match *param {
+                        "--gzip" => gzip = true,
+                        "--replica" => replica = true,
+                        _ => {
                             return Err(QueryError::Other(anyhow::anyhow!(
-                                "Parameter in position 3 unknown {third_param}",
+                                "Unknown parameter {param}",
                             )))
                         }
-                    };
-                    (Some(lsn), gzip)
+                    }
+                } else {
+                    lsn = Some(
+                        Lsn::from_str(param)
+                            .with_context(|| format!("Failed to parse Lsn from {param}"))?,
+                    );
                 }
-            };
+            }
 
             let metric_recording = metrics::BASEBACKUP_QUERY_TIME.start_recording(&ctx);
             let res = async {
@@ -1355,6 +1362,7 @@ where
                     None,
                     false,
                     gzip,
+                    replica,
                     &ctx,
                 )
                 .await?;
@@ -1415,6 +1423,7 @@ where
                 prev_lsn,
                 true,
                 false,
+                false,
                 &ctx,
             )
             .await?;

pageserver/src/tenant.rs

@@ -302,6 +302,13 @@ pub struct Tenant {
     /// **Lock order**: if acquiring all (or a subset), acquire them in order `timelines`, `timelines_offloaded`, `timelines_creating`
     timelines_offloaded: Mutex<HashMap<TimelineId, Arc<OffloadedTimeline>>>,
 
+    /// Serialize writes of the tenant manifest to remote storage.  If there are concurrent operations
+    /// affecting the manifest, such as timeline deletion and timeline offload, they must wait for
+    /// each other (this could be optimized to coalesce writes if necessary).
+    ///
+    /// The contents of the Mutex are the last manifest we successfully uploaded
+    tenant_manifest_upload: tokio::sync::Mutex<Option<TenantManifest>>,
+
     // This mutex prevents creation of new timelines during GC.
     // Adding yet another mutex (in addition to `timelines`) is needed because holding
     // `timelines` mutex during all GC iteration
@@ -514,13 +521,6 @@ pub struct OffloadedTimeline {
     /// Present for future flattening deliberations.
     pub archived_at: NaiveDateTime,
 
-    /// Lazily constructed remote client for the timeline
-    ///
-    /// If we offload a timeline, we keep around the remote client
-    /// for the duration of the process. If we find it through the
-    /// manifest, we don't construct it up until it's needed (deletion).
-    pub remote_client: Option<Arc<RemoteTimelineClient>>,
-
     /// Prevent two tasks from deleting the timeline at the same time. If held, the
     /// timeline is being deleted. If 'true', the timeline has already been deleted.
     pub delete_progress: TimelineDeleteProgress,
@@ -547,7 +547,6 @@ impl OffloadedTimeline {
             ancestor_retain_lsn,
             archived_at,
 
-            remote_client: Some(timeline.remote_client.clone()),
             delete_progress: timeline.delete_progress.clone(),
         })
     }
@@ -564,7 +563,6 @@ impl OffloadedTimeline {
             ancestor_timeline_id,
             ancestor_retain_lsn,
             archived_at,
-            remote_client: None,
             delete_progress: TimelineDeleteProgress::default(),
         }
     }
@@ -629,7 +627,7 @@ impl TimelineOrOffloaded {
     fn maybe_remote_client(&self) -> Option<Arc<RemoteTimelineClient>> {
         match self {
             TimelineOrOffloaded::Timeline(timeline) => Some(timeline.remote_client.clone()),
-            TimelineOrOffloaded::Offloaded(offloaded) => offloaded.remote_client.clone(),
+            TimelineOrOffloaded::Offloaded(_offloaded) => None,
         }
     }
 }
@@ -741,6 +739,24 @@ pub enum TimelineArchivalError {
     Other(anyhow::Error),
 }
 
+#[derive(thiserror::Error, Debug)]
+pub(crate) enum TenantManifestError {
+    #[error("Remote storage error: {0}")]
+    RemoteStorage(anyhow::Error),
+
+    #[error("Cancelled")]
+    Cancelled,
+}
+
+impl From<TenantManifestError> for TimelineArchivalError {
+    fn from(e: TenantManifestError) -> Self {
+        match e {
+            TenantManifestError::RemoteStorage(e) => Self::Other(e),
+            TenantManifestError::Cancelled => Self::Cancelled,
+        }
+    }
+}
+
 impl Debug for TimelineArchivalError {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {
@@ -1327,14 +1343,15 @@ impl Tenant {
         )
         .await?;
         let (offloaded_add, tenant_manifest) =
-            match remote_timeline_client::do_download_tenant_manifest(
+            match remote_timeline_client::download_tenant_manifest(
                 remote_storage,
                 &self.tenant_shard_id,
+                self.generation,
                 &cancel,
             )
             .await
             {
-                Ok((tenant_manifest, _generation)) => (
+                Ok((tenant_manifest, _generation, _manifest_mtime)) => (
                     format!("{} offloaded", tenant_manifest.offloaded_timelines.len()),
                     tenant_manifest,
                 ),
@@ -1526,18 +1543,7 @@ impl Tenant {
             offloaded_timelines_accessor.extend(offloaded_timelines_list.into_iter());
         }
         if !offloaded_timeline_ids.is_empty() {
-            let manifest = self.tenant_manifest();
-            // TODO: generation support
-            let generation = remote_timeline_client::TENANT_MANIFEST_GENERATION;
-            upload_tenant_manifest(
-                &self.remote_storage,
-                &self.tenant_shard_id,
-                generation,
-                &manifest,
-                &self.cancel,
-            )
-            .await
-            .map_err(TimelineArchivalError::Other)?;
+            self.store_tenant_manifest().await?;
         }
 
         // The local filesystem contents are a cache of what's in the remote IndexPart;
@@ -1918,18 +1924,7 @@ impl Tenant {
         };
 
         // Upload new list of offloaded timelines to S3
-        let manifest = self.tenant_manifest();
-        // TODO: generation support
-        let generation = remote_timeline_client::TENANT_MANIFEST_GENERATION;
-        upload_tenant_manifest(
-            &self.remote_storage,
-            &self.tenant_shard_id,
-            generation,
-            &manifest,
-            &cancel,
-        )
-        .await
-        .map_err(TimelineArchivalError::Other)?;
+        self.store_tenant_manifest().await?;
 
         // Activate the timeline (if it makes sense)
         if !(timeline.is_broken() || timeline.is_stopping()) {
@@ -2534,6 +2529,11 @@ impl Tenant {
                         .await
                         .inspect_err(|e| match e {
                             timeline::CompactionError::ShuttingDown => (),
+                            timeline::CompactionError::Offload(_) => {
+                                // Failures to offload timelines do not trip the circuit breaker, because
+                                // they do not do lots of writes the way compaction itself does: it is cheap
+                                // to retry, and it would be bad to stop all compaction because of an issue with offloading.
+                            }
                             timeline::CompactionError::Other(e) => {
                                 self.compaction_circuit_breaker
                                     .lock()
@@ -2549,8 +2549,7 @@ impl Tenant {
             if pending_task_left == Some(false) && *can_offload {
                 offload_timeline(self, timeline)
                     .instrument(info_span!("offload_timeline", %timeline_id))
-                    .await
-                    .map_err(timeline::CompactionError::Other)?;
+                    .await?;
             }
         }
 
@@ -3126,9 +3125,7 @@ impl Tenant {
             }
         }
 
-        let tenant_manifest = self.tenant_manifest();
-        // TODO: generation support
-        let generation = remote_timeline_client::TENANT_MANIFEST_GENERATION;
+        let tenant_manifest = self.build_tenant_manifest();
         for child_shard in child_shards {
             tracing::info!(
                 "Uploading tenant manifest for child {}",
@@ -3137,7 +3134,7 @@ impl Tenant {
             upload_tenant_manifest(
                 &self.remote_storage,
                 child_shard,
-                generation,
+                self.generation,
                 &tenant_manifest,
                 &self.cancel,
             )
@@ -3321,7 +3318,8 @@ impl Tenant {
             .unwrap_or(self.conf.default_tenant_conf.lsn_lease_length)
     }
 
-    pub(crate) fn tenant_manifest(&self) -> TenantManifest {
+    /// Generate an up-to-date TenantManifest based on the state of this Tenant.
+    fn build_tenant_manifest(&self) -> TenantManifest {
         let timelines_offloaded = self.timelines_offloaded.lock().unwrap();
 
         let mut timeline_manifests = timelines_offloaded
@@ -3529,6 +3527,7 @@ impl Tenant {
             timelines: Mutex::new(HashMap::new()),
             timelines_creating: Mutex::new(HashSet::new()),
             timelines_offloaded: Mutex::new(HashMap::new()),
+            tenant_manifest_upload: Default::default(),
             gc_cs: tokio::sync::Mutex::new(()),
             walredo_mgr,
             remote_storage,
@@ -4708,6 +4707,49 @@ impl Tenant {
             .max()
             .unwrap_or(0)
     }
+
+    /// Serialize and write the latest TenantManifest to remote storage.
+    pub(crate) async fn store_tenant_manifest(&self) -> Result<(), TenantManifestError> {
+        // Only one manifest write may be done at at time, and the contents of the manifest
+        // must be loaded while holding this lock. This makes it safe to call this function
+        // from anywhere without worrying about colliding updates.
+        let mut guard = tokio::select! {
+            g = self.tenant_manifest_upload.lock() => {
+                g
+            },
+            _ = self.cancel.cancelled() => {
+                return Err(TenantManifestError::Cancelled);
+            }
+        };
+
+        let manifest = self.build_tenant_manifest();
+        if Some(&manifest) == (*guard).as_ref() {
+            // Optimisation: skip uploads that don't change anything.
+            return Ok(());
+        }
+
+        upload_tenant_manifest(
+            &self.remote_storage,
+            &self.tenant_shard_id,
+            self.generation,
+            &manifest,
+            &self.cancel,
+        )
+        .await
+        .map_err(|e| {
+            if self.cancel.is_cancelled() {
+                TenantManifestError::Cancelled
+            } else {
+                TenantManifestError::RemoteStorage(e)
+            }
+        })?;
+
+        // Store the successfully uploaded manifest, so that future callers can avoid
+        // re-uploading the same thing.
+        *guard = Some(manifest);
+
+        Ok(())
+    }
 }
 
 /// Create the cluster temporarily in 'initdbpath' directory inside the repository
@@ -5090,6 +5132,7 @@ mod tests {
     use pageserver_api::keyspace::KeySpace;
     use pageserver_api::models::{CompactionAlgorithm, CompactionAlgorithmSettings};
     use pageserver_api::value::Value;
+    use pageserver_compaction::helpers::overlaps_with;
     use rand::{thread_rng, Rng};
     use storage_layer::PersistentLayerKey;
     use tests::storage_layer::ValuesReconstructState;
@@ -7612,23 +7655,7 @@ mod tests {
         }
 
         // Check if old layers are removed / new layers have the expected LSN
-        let mut all_layers = tline.inspect_historic_layers().await.unwrap();
-        all_layers.sort_by(|k1, k2| {
-            (
-                k1.is_delta,
-                k1.key_range.start,
-                k1.key_range.end,
-                k1.lsn_range.start,
-                k1.lsn_range.end,
-            )
-                .cmp(&(
-                    k2.is_delta,
-                    k2.key_range.start,
-                    k2.key_range.end,
-                    k2.lsn_range.start,
-                    k2.lsn_range.end,
-                ))
-        });
+        let all_layers = inspect_and_sort(&tline, None).await;
         assert_eq!(
             all_layers,
             vec![
@@ -9172,4 +9199,249 @@ mod tests {
 
         Ok(())
     }
+
+    async fn inspect_and_sort(
+        tline: &Arc<Timeline>,
+        filter: Option<std::ops::Range<Key>>,
+    ) -> Vec<PersistentLayerKey> {
+        let mut all_layers = tline.inspect_historic_layers().await.unwrap();
+        if let Some(filter) = filter {
+            all_layers.retain(|layer| overlaps_with(&layer.key_range, &filter));
+        }
+        all_layers.sort_by(|k1, k2| {
+            (
+                k1.is_delta,
+                k1.key_range.start,
+                k1.key_range.end,
+                k1.lsn_range.start,
+                k1.lsn_range.end,
+            )
+                .cmp(&(
+                    k2.is_delta,
+                    k2.key_range.start,
+                    k2.key_range.end,
+                    k2.lsn_range.start,
+                    k2.lsn_range.end,
+                ))
+        });
+        all_layers
+    }
+
+    #[cfg(feature = "testing")]
+    #[tokio::test]
+    async fn test_simple_partial_bottom_most_compaction() -> anyhow::Result<()> {
+        let harness = TenantHarness::create("test_simple_partial_bottom_most_compaction").await?;
+        let (tenant, ctx) = harness.load().await;
+
+        fn get_key(id: u32) -> Key {
+            // using aux key here b/c they are guaranteed to be inside `collect_keyspace`.
+            let mut key = Key::from_hex("620000000033333333444444445500000000").unwrap();
+            key.field6 = id;
+            key
+        }
+
+        // img layer at 0x10
+        let img_layer = (0..10)
+            .map(|id| (get_key(id), Bytes::from(format!("value {id}@0x10"))))
+            .collect_vec();
+
+        let delta1 = vec![
+            (
+                get_key(1),
+                Lsn(0x20),
+                Value::Image(Bytes::from("value 1@0x20")),
+            ),
+            (
+                get_key(2),
+                Lsn(0x30),
+                Value::Image(Bytes::from("value 2@0x30")),
+            ),
+            (
+                get_key(3),
+                Lsn(0x40),
+                Value::Image(Bytes::from("value 3@0x40")),
+            ),
+        ];
+        let delta2 = vec![
+            (
+                get_key(5),
+                Lsn(0x20),
+                Value::Image(Bytes::from("value 5@0x20")),
+            ),
+            (
+                get_key(6),
+                Lsn(0x20),
+                Value::Image(Bytes::from("value 6@0x20")),
+            ),
+        ];
+        let delta3 = vec![
+            (
+                get_key(8),
+                Lsn(0x48),
+                Value::Image(Bytes::from("value 8@0x48")),
+            ),
+            (
+                get_key(9),
+                Lsn(0x48),
+                Value::Image(Bytes::from("value 9@0x48")),
+            ),
+        ];
+
+        let tline = tenant
+            .create_test_timeline_with_layers(
+                TIMELINE_ID,
+                Lsn(0x10),
+                DEFAULT_PG_VERSION,
+                &ctx,
+                vec![
+                    DeltaLayerTestDesc::new_with_inferred_key_range(Lsn(0x20)..Lsn(0x48), delta1),
+                    DeltaLayerTestDesc::new_with_inferred_key_range(Lsn(0x20)..Lsn(0x48), delta2),
+                    DeltaLayerTestDesc::new_with_inferred_key_range(Lsn(0x48)..Lsn(0x50), delta3),
+                ], // delta layers
+                vec![(Lsn(0x10), img_layer)], // image layers
+                Lsn(0x50),
+            )
+            .await?;
+
+        {
+            // Update GC info
+            let mut guard = tline.gc_info.write().unwrap();
+            *guard = GcInfo {
+                retain_lsns: vec![(Lsn(0x20), tline.timeline_id, MaybeOffloaded::No)],
+                cutoffs: GcCutoffs {
+                    time: Lsn(0x30),
+                    space: Lsn(0x30),
+                },
+                leases: Default::default(),
+                within_ancestor_pitr: false,
+            };
+        }
+
+        let cancel = CancellationToken::new();
+
+        // Do a partial compaction on key range 0..4, we should generate a image layer; no other layers
+        // can be removed because they might be used for other key ranges.
+        tline
+            .partial_compact_with_gc(Some(get_key(0)..get_key(4)), &cancel, EnumSet::new(), &ctx)
+            .await
+            .unwrap();
+        let all_layers = inspect_and_sort(&tline, Some(get_key(0)..get_key(10))).await;
+        assert_eq!(
+            all_layers,
+            vec![
+                PersistentLayerKey {
+                    key_range: get_key(0)..get_key(4),
+                    lsn_range: Lsn(0x20)..Lsn(0x21),
+                    is_delta: false
+                },
+                PersistentLayerKey {
+                    key_range: get_key(0)..get_key(10),
+                    lsn_range: Lsn(0x10)..Lsn(0x11),
+                    is_delta: false
+                },
+                PersistentLayerKey {
+                    key_range: get_key(1)..get_key(4),
+                    lsn_range: Lsn(0x20)..Lsn(0x48),
+                    is_delta: true
+                },
+                PersistentLayerKey {
+                    key_range: get_key(5)..get_key(7),
+                    lsn_range: Lsn(0x20)..Lsn(0x48),
+                    is_delta: true
+                },
+                PersistentLayerKey {
+                    key_range: get_key(8)..get_key(10),
+                    lsn_range: Lsn(0x48)..Lsn(0x50),
+                    is_delta: true
+                }
+            ]
+        );
+
+        // Do a partial compaction on key range 4..10
+        tline
+            .partial_compact_with_gc(Some(get_key(4)..get_key(10)), &cancel, EnumSet::new(), &ctx)
+            .await
+            .unwrap();
+        let all_layers = inspect_and_sort(&tline, Some(get_key(0)..get_key(10))).await;
+        assert_eq!(
+            all_layers,
+            vec![
+                PersistentLayerKey {
+                    key_range: get_key(0)..get_key(4),
+                    lsn_range: Lsn(0x20)..Lsn(0x21),
+                    is_delta: false
+                },
+                PersistentLayerKey {
+                    // if (in the future) GC kicks in, this layer will be removed
+                    key_range: get_key(0)..get_key(10),
+                    lsn_range: Lsn(0x10)..Lsn(0x11),
+                    is_delta: false
+                },
+                PersistentLayerKey {
+                    key_range: get_key(4)..get_key(10),
+                    lsn_range: Lsn(0x20)..Lsn(0x21),
+                    is_delta: false
+                },
+                PersistentLayerKey {
+                    key_range: get_key(1)..get_key(4),
+                    lsn_range: Lsn(0x20)..Lsn(0x48),
+                    is_delta: true
+                },
+                PersistentLayerKey {
+                    key_range: get_key(5)..get_key(7),
+                    lsn_range: Lsn(0x20)..Lsn(0x48),
+                    is_delta: true
+                },
+                PersistentLayerKey {
+                    key_range: get_key(8)..get_key(10),
+                    lsn_range: Lsn(0x48)..Lsn(0x50),
+                    is_delta: true
+                }
+            ]
+        );
+
+        // Do a partial compaction on key range 0..10, all image layers below LSN 20 can be replaced with new ones.
+        tline
+            .partial_compact_with_gc(Some(get_key(0)..get_key(10)), &cancel, EnumSet::new(), &ctx)
+            .await
+            .unwrap();
+        let all_layers = inspect_and_sort(&tline, Some(get_key(0)..get_key(10))).await;
+        assert_eq!(
+            all_layers,
+            vec![
+                PersistentLayerKey {
+                    key_range: get_key(0)..get_key(4),
+                    lsn_range: Lsn(0x20)..Lsn(0x21),
+                    is_delta: false
+                },
+                PersistentLayerKey {
+                    key_range: get_key(0)..get_key(10),
+                    lsn_range: Lsn(0x20)..Lsn(0x21),
+                    is_delta: false
+                },
+                PersistentLayerKey {
+                    key_range: get_key(4)..get_key(10),
+                    lsn_range: Lsn(0x20)..Lsn(0x21),
+                    is_delta: false
+                },
+                PersistentLayerKey {
+                    key_range: get_key(1)..get_key(4),
+                    lsn_range: Lsn(0x20)..Lsn(0x48),
+                    is_delta: true
+                },
+                PersistentLayerKey {
+                    key_range: get_key(5)..get_key(7),
+                    lsn_range: Lsn(0x20)..Lsn(0x48),
+                    is_delta: true
+                },
+                PersistentLayerKey {
+                    key_range: get_key(8)..get_key(10),
+                    lsn_range: Lsn(0x48)..Lsn(0x50),
+                    is_delta: true
+                }
+            ]
+        );
+
+        Ok(())
+    }
 }

pageserver/src/tenant/remote_timeline_client.rs

@@ -190,6 +190,7 @@ use chrono::{NaiveDateTime, Utc};
 pub(crate) use download::download_initdb_tar_zst;
 use pageserver_api::models::TimelineArchivalState;
 use pageserver_api::shard::{ShardIndex, TenantShardId};
+use regex::Regex;
 use scopeguard::ScopeGuard;
 use tokio_util::sync::CancellationToken;
 use utils::backoff::{
@@ -199,7 +200,7 @@ use utils::pausable_failpoint;
 
 use std::collections::{HashMap, VecDeque};
 use std::sync::atomic::{AtomicU32, Ordering};
-use std::sync::{Arc, Mutex};
+use std::sync::{Arc, Mutex, OnceLock};
 use std::time::Duration;
 
 use remote_storage::{
@@ -245,11 +246,11 @@ use super::upload_queue::{NotInitialized, SetDeletedFlagProgress};
 use super::Generation;
 
 pub(crate) use download::{
-    do_download_tenant_manifest, download_index_part, is_temp_download_file,
+    download_index_part, download_tenant_manifest, is_temp_download_file,
     list_remote_tenant_shards, list_remote_timelines,
 };
 pub(crate) use index::LayerFileMetadata;
-pub(crate) use upload::{upload_initdb_dir, upload_tenant_manifest};
+pub(crate) use upload::upload_initdb_dir;
 
 // Occasional network issues and such can cause remote operations to fail, and
 // that's expected. If a download fails, we log it at info-level, and retry.
@@ -274,12 +275,6 @@ pub(crate) const BUFFER_SIZE: usize = 32 * 1024;
 /// which we warn and skip.
 const DELETION_QUEUE_FLUSH_TIMEOUT: Duration = Duration::from_secs(10);
 
-/// Hardcode a generation for the tenant manifest for now so that we don't
-/// need to deal with generation-less manifests in the future.
-///
-/// TODO: add proper generation support to all the places that use this.
-pub(crate) const TENANT_MANIFEST_GENERATION: Generation = Generation::new(1);
-
 pub enum MaybeDeletedIndexPart {
     IndexPart(IndexPart),
     Deleted(IndexPart),
@@ -2239,6 +2234,12 @@ pub fn remote_tenant_manifest_path(
     RemotePath::from_string(&path).expect("Failed to construct path")
 }
 
+/// Prefix to all generations' manifest objects in a tenant shard
+pub fn remote_tenant_manifest_prefix(tenant_shard_id: &TenantShardId) -> RemotePath {
+    let path = format!("tenants/{tenant_shard_id}/tenant-manifest",);
+    RemotePath::from_string(&path).expect("Failed to construct path")
+}
+
 pub fn remote_timelines_path(tenant_shard_id: &TenantShardId) -> RemotePath {
     let path = format!("tenants/{tenant_shard_id}/{TIMELINES_SEGMENT_NAME}");
     RemotePath::from_string(&path).expect("Failed to construct path")
@@ -2333,6 +2334,15 @@ pub fn parse_remote_index_path(path: RemotePath) -> Option<Generation> {
     }
 }
 
+/// Given the key of a tenant manifest, parse out the generation number
+pub(crate) fn parse_remote_tenant_manifest_path(path: RemotePath) -> Option<Generation> {
+    static RE: OnceLock<Regex> = OnceLock::new();
+    let re = RE.get_or_init(|| Regex::new(r".+tenant-manifest-([0-9a-f]{8}).json").unwrap());
+    re.captures(path.get_path().as_str())
+        .and_then(|c| c.get(1))
+        .and_then(|m| Generation::parse_suffix(m.as_str()))
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

pageserver/src/tenant/remote_timeline_client/download.rs

@@ -20,7 +20,9 @@ use utils::backoff;
 
 use crate::config::PageServerConf;
 use crate::context::RequestContext;
-use crate::span::debug_assert_current_span_has_tenant_and_timeline_id;
+use crate::span::{
+    debug_assert_current_span_has_tenant_and_timeline_id, debug_assert_current_span_has_tenant_id,
+};
 use crate::tenant::remote_timeline_client::{remote_layer_path, remote_timelines_path};
 use crate::tenant::storage_layer::LayerName;
 use crate::tenant::Generation;
@@ -36,9 +38,10 @@ use utils::pausable_failpoint;
 use super::index::{IndexPart, LayerFileMetadata};
 use super::manifest::TenantManifest;
 use super::{
-    parse_remote_index_path, remote_index_path, remote_initdb_archive_path,
-    remote_initdb_preserved_archive_path, remote_tenant_manifest_path, remote_tenant_path,
-    FAILED_DOWNLOAD_WARN_THRESHOLD, FAILED_REMOTE_OP_RETRIES, INITDB_PATH,
+    parse_remote_index_path, parse_remote_tenant_manifest_path, remote_index_path,
+    remote_initdb_archive_path, remote_initdb_preserved_archive_path, remote_tenant_manifest_path,
+    remote_tenant_manifest_prefix, remote_tenant_path, FAILED_DOWNLOAD_WARN_THRESHOLD,
+    FAILED_REMOTE_OP_RETRIES, INITDB_PATH,
 };
 
 ///
@@ -365,32 +368,34 @@ async fn do_download_remote_path_retry_forever(
     .await
 }
 
-pub async fn do_download_tenant_manifest(
+async fn do_download_tenant_manifest(
     storage: &GenericRemoteStorage,
     tenant_shard_id: &TenantShardId,
+    _timeline_id: Option<&TimelineId>,
+    generation: Generation,
     cancel: &CancellationToken,
-) -> Result<(TenantManifest, Generation), DownloadError> {
-    // TODO: generation support
-    let generation = super::TENANT_MANIFEST_GENERATION;
+) -> Result<(TenantManifest, Generation, SystemTime), DownloadError> {
     let remote_path = remote_tenant_manifest_path(tenant_shard_id, generation);
 
-    let (manifest_bytes, _manifest_bytes_mtime) =
+    let (manifest_bytes, manifest_bytes_mtime) =
         do_download_remote_path_retry_forever(storage, &remote_path, cancel).await?;
 
     let tenant_manifest = TenantManifest::from_json_bytes(&manifest_bytes)
         .with_context(|| format!("deserialize tenant manifest file at {remote_path:?}"))
         .map_err(DownloadError::Other)?;
 
-    Ok((tenant_manifest, generation))
+    Ok((tenant_manifest, generation, manifest_bytes_mtime))
 }
 
 async fn do_download_index_part(
     storage: &GenericRemoteStorage,
     tenant_shard_id: &TenantShardId,
-    timeline_id: &TimelineId,
+    timeline_id: Option<&TimelineId>,
     index_generation: Generation,
     cancel: &CancellationToken,
 ) -> Result<(IndexPart, Generation, SystemTime), DownloadError> {
+    let timeline_id =
+        timeline_id.expect("A timeline ID is always provided when downloading an index");
     let remote_path = remote_index_path(tenant_shard_id, timeline_id, index_generation);
 
     let (index_part_bytes, index_part_mtime) =
@@ -403,59 +408,79 @@ async fn do_download_index_part(
     Ok((index_part, index_generation, index_part_mtime))
 }
 
-/// index_part.json objects are suffixed with a generation number, so we cannot
-/// directly GET the latest index part without doing some probing.
+/// Metadata objects are "generationed", meaning that they include a generation suffix.  This
+/// function downloads the object with the highest generation <= `my_generation`.
 ///
-/// In this function we probe for the most recent index in a generation <= our current generation.
-/// See "Finding the remote indices for timelines" in docs/rfcs/025-generation-numbers.md
+/// Data objects (layer files) also include a generation in their path, but there is no equivalent
+/// search process, because their reference from an index includes the generation.
+///
+/// An expensive object listing operation is only done if necessary: the typical fast path is to issue two
+/// GET operations, one to our own generation (stale attachment case), and one to the immediately preceding
+/// generation (normal case when migrating/restarting).  Only if both of these return 404 do we fall back
+/// to listing objects.
+///
+/// * `my_generation`: the value of `[crate::tenant::Tenant::generation]`
+/// * `what`: for logging, what object are we downloading
+/// * `prefix`: when listing objects, use this prefix (i.e. the part of the object path before the generation)
+/// * `do_download`: a GET of the object in a particular generation, which should **retry indefinitely** unless
+///                  `cancel`` has fired.  This function does not do its own retries of GET operations, and relies
+///                  on the function passed in to do so.
+/// * `parse_path`: parse a fully qualified remote storage path to get the generation of the object.
+#[allow(clippy::too_many_arguments)]
 #[tracing::instrument(skip_all, fields(generation=?my_generation))]
-pub(crate) async fn download_index_part(
-    storage: &GenericRemoteStorage,
-    tenant_shard_id: &TenantShardId,
-    timeline_id: &TimelineId,
+pub(crate) async fn download_generation_object<'a, T, DF, DFF, PF>(
+    storage: &'a GenericRemoteStorage,
+    tenant_shard_id: &'a TenantShardId,
+    timeline_id: Option<&'a TimelineId>,
     my_generation: Generation,
-    cancel: &CancellationToken,
-) -> Result<(IndexPart, Generation, SystemTime), DownloadError> {
-    debug_assert_current_span_has_tenant_and_timeline_id();
+    what: &str,
+    prefix: RemotePath,
+    do_download: DF,
+    parse_path: PF,
+    cancel: &'a CancellationToken,
+) -> Result<(T, Generation, SystemTime), DownloadError>
+where
+    DF: Fn(
+        &'a GenericRemoteStorage,
+        &'a TenantShardId,
+        Option<&'a TimelineId>,
+        Generation,
+        &'a CancellationToken,
+    ) -> DFF,
+    DFF: Future<Output = Result<(T, Generation, SystemTime), DownloadError>>,
+    PF: Fn(RemotePath) -> Option<Generation>,
+    T: 'static,
+{
+    debug_assert_current_span_has_tenant_id();
 
     if my_generation.is_none() {
         // Operating without generations: just fetch the generation-less path
-        return do_download_index_part(
-            storage,
-            tenant_shard_id,
-            timeline_id,
-            my_generation,
-            cancel,
-        )
-        .await;
+        return do_download(storage, tenant_shard_id, timeline_id, my_generation, cancel).await;
     }
 
-    // Stale case: If we were intentionally attached in a stale generation, there may already be a remote
-    // index in our generation.
+    // Stale case: If we were intentionally attached in a stale generation, the remote object may already
+    // exist in our generation.
     //
     // This is an optimization to avoid doing the listing for the general case below.
-    let res =
-        do_download_index_part(storage, tenant_shard_id, timeline_id, my_generation, cancel).await;
+    let res = do_download(storage, tenant_shard_id, timeline_id, my_generation, cancel).await;
     match res {
-        Ok(index_part) => {
-            tracing::debug!(
-                "Found index_part from current generation (this is a stale attachment)"
-            );
-            return Ok(index_part);
+        Ok(decoded) => {
+            tracing::debug!("Found {what} from current generation (this is a stale attachment)");
+            return Ok(decoded);
         }
         Err(DownloadError::NotFound) => {}
         Err(e) => return Err(e),
     };
 
-    // Typical case: the previous generation of this tenant was running healthily, and had uploaded
-    // and index part.  We may safely start from this index without doing a listing, because:
+    // Typical case: the previous generation of this tenant was running healthily, and had uploaded the object
+    // we are seeking in that generation.  We may safely start from this index without doing a listing, because:
     //  - We checked for current generation case above
     //  - generations > my_generation are to be ignored
-    //  - any other indices that exist would have an older generation than `previous_gen`, and
-    //    we want to find the most recent index from a previous generation.
+    //  - any other objects that exist would have an older generation than `previous_gen`, and
+    //    we want to find the most recent object from a previous generation.
     //
     // This is an optimization to avoid doing the listing for the general case below.
-    let res = do_download_index_part(
+    let res = do_download(
         storage,
         tenant_shard_id,
         timeline_id,
@@ -464,14 +489,12 @@ pub(crate) async fn download_index_part(
     )
     .await;
     match res {
-        Ok(index_part) => {
-            tracing::debug!("Found index_part from previous generation");
-            return Ok(index_part);
+        Ok(decoded) => {
+            tracing::debug!("Found {what} from previous generation");
+            return Ok(decoded);
         }
         Err(DownloadError::NotFound) => {
-            tracing::debug!(
-                "No index_part found from previous generation, falling back to listing"
-            );
+            tracing::debug!("No {what} found from previous generation, falling back to listing");
         }
         Err(e) => {
             return Err(e);
@@ -481,12 +504,10 @@ pub(crate) async fn download_index_part(
     // General case/fallback: if there is no index at my_generation or prev_generation, then list all index_part.json
     // objects, and select the highest one with a generation <= my_generation.  Constructing the prefix is equivalent
     // to constructing a full index path with no generation, because the generation is a suffix.
-    let index_prefix = remote_index_path(tenant_shard_id, timeline_id, Generation::none());
-
-    let indices = download_retry(
+    let paths = download_retry(
         || async {
             storage
-                .list(Some(&index_prefix), ListingMode::NoDelimiter, None, cancel)
+                .list(Some(&prefix), ListingMode::NoDelimiter, None, cancel)
                 .await
         },
         "list index_part files",
@@ -497,22 +518,22 @@ pub(crate) async fn download_index_part(
 
     // General case logic for which index to use: the latest index whose generation
     // is <= our own.  See "Finding the remote indices for timelines" in docs/rfcs/025-generation-numbers.md
-    let max_previous_generation = indices
+    let max_previous_generation = paths
         .into_iter()
-        .filter_map(|o| parse_remote_index_path(o.key))
+        .filter_map(|o| parse_path(o.key))
         .filter(|g| g <= &my_generation)
         .max();
 
     match max_previous_generation {
         Some(g) => {
-            tracing::debug!("Found index_part in generation {g:?}");
-            do_download_index_part(storage, tenant_shard_id, timeline_id, g, cancel).await
+            tracing::debug!("Found {what} in generation {g:?}");
+            do_download(storage, tenant_shard_id, timeline_id, g, cancel).await
         }
         None => {
             // Migration from legacy pre-generation state: we have a generation but no prior
             // attached pageservers did.  Try to load from a no-generation path.
-            tracing::debug!("No index_part.json* found");
-            do_download_index_part(
+            tracing::debug!("No {what}* found");
+            do_download(
                 storage,
                 tenant_shard_id,
                 timeline_id,
@@ -524,6 +545,57 @@ pub(crate) async fn download_index_part(
     }
 }
 
+/// index_part.json objects are suffixed with a generation number, so we cannot
+/// directly GET the latest index part without doing some probing.
+///
+/// In this function we probe for the most recent index in a generation <= our current generation.
+/// See "Finding the remote indices for timelines" in docs/rfcs/025-generation-numbers.md
+pub(crate) async fn download_index_part(
+    storage: &GenericRemoteStorage,
+    tenant_shard_id: &TenantShardId,
+    timeline_id: &TimelineId,
+    my_generation: Generation,
+    cancel: &CancellationToken,
+) -> Result<(IndexPart, Generation, SystemTime), DownloadError> {
+    debug_assert_current_span_has_tenant_and_timeline_id();
+
+    let index_prefix = remote_index_path(tenant_shard_id, timeline_id, Generation::none());
+    download_generation_object(
+        storage,
+        tenant_shard_id,
+        Some(timeline_id),
+        my_generation,
+        "index_part",
+        index_prefix,
+        do_download_index_part,
+        parse_remote_index_path,
+        cancel,
+    )
+    .await
+}
+
+pub(crate) async fn download_tenant_manifest(
+    storage: &GenericRemoteStorage,
+    tenant_shard_id: &TenantShardId,
+    my_generation: Generation,
+    cancel: &CancellationToken,
+) -> Result<(TenantManifest, Generation, SystemTime), DownloadError> {
+    let manifest_prefix = remote_tenant_manifest_prefix(tenant_shard_id);
+
+    download_generation_object(
+        storage,
+        tenant_shard_id,
+        None,
+        my_generation,
+        "tenant-manifest",
+        manifest_prefix,
+        do_download_tenant_manifest,
+        parse_remote_tenant_manifest_path,
+        cancel,
+    )
+    .await
+}
+
 pub(crate) async fn download_initdb_tar_zst(
     conf: &'static PageServerConf,
     storage: &GenericRemoteStorage,

pageserver/src/tenant/remote_timeline_client/manifest.rs

@@ -3,7 +3,7 @@ use serde::{Deserialize, Serialize};
 use utils::{id::TimelineId, lsn::Lsn};
 
 /// Tenant-shard scoped manifest
-#[derive(Clone, Serialize, Deserialize)]
+#[derive(Clone, Serialize, Deserialize, PartialEq, Eq)]
 pub struct TenantManifest {
     /// Debugging aid describing the version of this manifest.
     /// Can also be used for distinguishing breaking changes later on.
@@ -23,7 +23,7 @@ pub struct TenantManifest {
 /// Very similar to [`pageserver_api::models::OffloadedTimelineInfo`],
 /// but the two datastructures serve different needs, this is for a persistent disk format
 /// that must be backwards compatible, while the other is only for informative purposes.
-#[derive(Clone, Serialize, Deserialize, Copy)]
+#[derive(Clone, Serialize, Deserialize, Copy, PartialEq, Eq)]
 pub struct OffloadedTimelineManifest {
     pub timeline_id: TimelineId,
     /// Whether the timeline has a parent it has been branched off from or not

pageserver/src/tenant/storage_layer/batch_split_writer.rs

@@ -120,6 +120,8 @@ impl BatchLayerWriter {
                         writer.finish(layer_key.key_range.end, ctx).await
                     }
                     LayerWriterWrapper::Image(writer) => {
+                        assert_eq!(writer.key_range().start, layer_key.key_range.start);
+                        assert_eq!(writer.lsn(), layer_key.lsn_range.start);
                         writer
                             .finish_with_end_key(layer_key.key_range.end, ctx)
                             .await

pageserver/src/tenant/storage_layer/delta_layer.rs

@@ -270,7 +270,7 @@ impl AsLayerDesc for DeltaLayer {
 }
 
 impl DeltaLayer {
-    pub(crate) async fn dump(&self, verbose: bool, ctx: &RequestContext) -> Result<()> {
+    pub async fn dump(&self, verbose: bool, ctx: &RequestContext) -> Result<()> {
         self.desc.dump();
 
         if !verbose {
@@ -1438,7 +1438,7 @@ impl DeltaLayerInner {
         offset
     }
 
-    pub(crate) fn iter<'a>(&'a self, ctx: &'a RequestContext) -> DeltaLayerIterator<'a> {
+    pub fn iter<'a>(&'a self, ctx: &'a RequestContext) -> DeltaLayerIterator<'a> {
         let block_reader = FileBlockReader::new(&self.file, self.file_id);
         let tree_reader =
             DiskBtreeReader::new(self.index_start_blk, self.index_root_blk, block_reader);

pageserver/src/tenant/storage_layer/image_layer.rs

@@ -231,7 +231,7 @@ impl AsLayerDesc for ImageLayer {
 }
 
 impl ImageLayer {
-    pub(crate) async fn dump(&self, verbose: bool, ctx: &RequestContext) -> Result<()> {
+    pub async fn dump(&self, verbose: bool, ctx: &RequestContext) -> Result<()> {
         self.desc.dump();
 
         if !verbose {
@@ -885,6 +885,7 @@ impl ImageLayerWriterInner {
         }
 
         let final_key_range = if let Some(end_key) = end_key {
+            assert!(end_key <= self.key_range.end);
             self.key_range.start..end_key
         } else {
             self.key_range.clone()
@@ -1033,6 +1034,14 @@ impl ImageLayerWriter {
     ) -> anyhow::Result<(PersistentLayerDesc, Utf8PathBuf)> {
         self.inner.take().unwrap().finish(ctx, Some(end_key)).await
     }
+
+    pub(crate) fn key_range(&self) -> Range<Key> {
+        self.inner.as_ref().unwrap().key_range.clone()
+    }
+
+    pub(crate) fn lsn(&self) -> Lsn {
+        self.inner.as_ref().unwrap().lsn
+    }
 }
 
 impl Drop for ImageLayerWriter {

pageserver/src/tenant/storage_layer/layer_name.rs

@@ -2,13 +2,11 @@
 //! Helper functions for dealing with filenames of the image and delta layer files.
 //!
 use pageserver_api::key::Key;
-use std::borrow::Cow;
 use std::cmp::Ordering;
 use std::fmt;
 use std::ops::Range;
 use std::str::FromStr;
 
-use regex::Regex;
 use utils::lsn::Lsn;
 
 use super::PersistentLayerDesc;
@@ -60,32 +58,31 @@ impl Ord for DeltaLayerName {
 /// Represents the region of the LSN-Key space covered by a DeltaLayer
 ///
 /// ```text
-///    <key start>-<key end>__<LSN start>-<LSN end>
+///    <key start>-<key end>__<LSN start>-<LSN end>-<generation>
 /// ```
 impl DeltaLayerName {
     /// Parse the part of a delta layer's file name that represents the LayerName. Returns None
     /// if the filename does not match the expected pattern.
     pub fn parse_str(fname: &str) -> Option<Self> {
-        let mut parts = fname.split("__");
-        let mut key_parts = parts.next()?.split('-');
-        let mut lsn_parts = parts.next()?.split('-');
-
-        let key_start_str = key_parts.next()?;
-        let key_end_str = key_parts.next()?;
-        let lsn_start_str = lsn_parts.next()?;
-        let lsn_end_str = lsn_parts.next()?;
-
-        if parts.next().is_some() || key_parts.next().is_some() || key_parts.next().is_some() {
-            return None;
-        }
-
-        if key_start_str.len() != 36
-            || key_end_str.len() != 36
-            || lsn_start_str.len() != 16
-            || lsn_end_str.len() != 16
+        let (key_parts, lsn_generation_parts) = fname.split_once("__")?;
+        let (key_start_str, key_end_str) = key_parts.split_once('-')?;
+        let (lsn_start_str, lsn_end_generation_parts) = lsn_generation_parts.split_once('-')?;
+        let lsn_end_str = if let Some((lsn_end_str, maybe_generation)) =
+            lsn_end_generation_parts.split_once('-')
         {
-            return None;
-        }
+            if maybe_generation.starts_with("v") {
+                // vY-XXXXXXXX
+                lsn_end_str
+            } else if maybe_generation.len() == 8 {
+                // XXXXXXXX
+                lsn_end_str
+            } else {
+                // no idea what this is
+                return None;
+            }
+        } else {
+            lsn_end_generation_parts
+        };
 
         let key_start = Key::from_hex(key_start_str).ok()?;
         let key_end = Key::from_hex(key_end_str).ok()?;
@@ -173,25 +170,29 @@ impl ImageLayerName {
 /// Represents the part of the Key-LSN space covered by an ImageLayer
 ///
 /// ```text
-///    <key start>-<key end>__<LSN>
+///    <key start>-<key end>__<LSN>-<generation>
 /// ```
 impl ImageLayerName {
     /// Parse a string as then LayerName part of an image layer file name. Returns None if the
     /// filename does not match the expected pattern.
     pub fn parse_str(fname: &str) -> Option<Self> {
-        let mut parts = fname.split("__");
-        let mut key_parts = parts.next()?.split('-');
-
-        let key_start_str = key_parts.next()?;
-        let key_end_str = key_parts.next()?;
-        let lsn_str = parts.next()?;
-        if parts.next().is_some() || key_parts.next().is_some() {
-            return None;
-        }
-
-        if key_start_str.len() != 36 || key_end_str.len() != 36 || lsn_str.len() != 16 {
-            return None;
-        }
+        let (key_parts, lsn_generation_parts) = fname.split_once("__")?;
+        let (key_start_str, key_end_str) = key_parts.split_once('-')?;
+        let lsn_str =
+            if let Some((lsn_str, maybe_generation)) = lsn_generation_parts.split_once('-') {
+                if maybe_generation.starts_with("v") {
+                    // vY-XXXXXXXX
+                    lsn_str
+                } else if maybe_generation.len() == 8 {
+                    // XXXXXXXX
+                    lsn_str
+                } else {
+                    // likely a delta layer
+                    return None;
+                }
+            } else {
+                lsn_generation_parts
+            };
 
         let key_start = Key::from_hex(key_start_str).ok()?;
         let key_end = Key::from_hex(key_end_str).ok()?;
@@ -258,6 +259,14 @@ impl LayerName {
         }
     }
 
+    /// Gets the LSN range encoded in the layer name.
+    pub fn lsn_as_range(&self) -> Range<Lsn> {
+        match &self {
+            LayerName::Image(layer) => layer.lsn_as_range(),
+            LayerName::Delta(layer) => layer.lsn_range.clone(),
+        }
+    }
+
     pub fn is_delta(&self) -> bool {
         matches!(self, LayerName::Delta(_))
     }
@@ -290,18 +299,8 @@ impl FromStr for LayerName {
     /// Self. When loading a physical layer filename, we drop any extra information
     /// not needed to build Self.
     fn from_str(value: &str) -> Result<Self, Self::Err> {
-        let gen_suffix_regex = Regex::new("^(?<base>.+)(?<gen>-v1-[0-9a-f]{8})$").unwrap();
-        let file_name: Cow<str> = match gen_suffix_regex.captures(value) {
-            Some(captures) => captures
-                .name("base")
-                .expect("Non-optional group")
-                .as_str()
-                .into(),
-            None => value.into(),
-        };
-
-        let delta = DeltaLayerName::parse_str(&file_name);
-        let image = ImageLayerName::parse_str(&file_name);
+        let delta = DeltaLayerName::parse_str(value);
+        let image = ImageLayerName::parse_str(value);
         let ok = match (delta, image) {
             (None, None) => {
                 return Err(format!(
@@ -367,11 +366,14 @@ mod test {
             lsn: Lsn::from_hex("00000000014FED58").unwrap(),
         });
         let parsed = LayerName::from_str("000000000000000000000000000000000000-000000067F00000001000004DF0000000006__00000000014FED58-v1-00000001").unwrap();
-        assert_eq!(parsed, expected,);
+        assert_eq!(parsed, expected);
+
+        let parsed = LayerName::from_str("000000000000000000000000000000000000-000000067F00000001000004DF0000000006__00000000014FED58-00000001").unwrap();
+        assert_eq!(parsed, expected);
 
         // Omitting generation suffix is valid
         let parsed = LayerName::from_str("000000000000000000000000000000000000-000000067F00000001000004DF0000000006__00000000014FED58").unwrap();
-        assert_eq!(parsed, expected,);
+        assert_eq!(parsed, expected);
     }
 
     #[test]
@@ -385,6 +387,9 @@ mod test {
         let parsed = LayerName::from_str("000000000000000000000000000000000000-000000067F00000001000004DF0000000006__00000000014FED58-000000000154C481-v1-00000001").unwrap();
         assert_eq!(parsed, expected);
 
+        let parsed = LayerName::from_str("000000000000000000000000000000000000-000000067F00000001000004DF0000000006__00000000014FED58-000000000154C481-00000001").unwrap();
+        assert_eq!(parsed, expected);
+
         // Omitting generation suffix is valid
         let parsed = LayerName::from_str("000000000000000000000000000000000000-000000067F00000001000004DF0000000006__00000000014FED58-000000000154C481").unwrap();
         assert_eq!(parsed, expected);

pageserver/src/tenant/tasks.rs

@@ -279,6 +279,7 @@ fn log_compaction_error(
 
     let decision = match e {
         ShuttingDown => None,
+        Offload(_) => Some(LooksLike::Error),
         _ if task_cancelled => Some(LooksLike::Info),
         Other(e) => {
             let root_cause = e.root_cause();

pageserver/src/tenant/timeline.rs

@@ -20,6 +20,8 @@ use chrono::{DateTime, Utc};
 use enumset::EnumSet;
 use fail::fail_point;
 use handle::ShardTimelineId;
+use itertools::Itertools;
+use offload::OffloadError;
 use once_cell::sync::Lazy;
 use pageserver_api::{
     key::{
@@ -64,13 +66,14 @@ use std::{
 };
 use std::{pin::pin, sync::OnceLock};
 
+use crate::pgdatadir_mapping::MAX_AUX_FILE_V2_DELTAS;
 use crate::{
     aux_file::AuxFileSizeEstimator,
     tenant::{
         config::AttachmentMode,
         layer_map::{LayerMap, SearchResult},
         metadata::TimelineMetadata,
-        storage_layer::{inmemory_layer::IndexEntry, PersistentLayerDesc},
+        storage_layer::inmemory_layer::IndexEntry,
     },
     walingest::WalLagCooldown,
     walredo,
@@ -103,7 +106,6 @@ use crate::{
     virtual_file::{MaybeFatalIo, VirtualFile},
 };
 use crate::{pgdatadir_mapping::LsnForTimestamp, tenant::tasks::BackgroundLoopKind};
-use crate::{pgdatadir_mapping::MAX_AUX_FILE_V2_DELTAS, tenant::storage_layer::PersistentLayerKey};
 use pageserver_api::config::tenant_conf_defaults::DEFAULT_PITR_INTERVAL;
 
 use crate::config::PageServerConf;
@@ -141,7 +143,10 @@ use self::walreceiver::{WalReceiver, WalReceiverConf};
 
 use super::{
     config::TenantConf,
-    storage_layer::{inmemory_layer, LayerVisibilityHint},
+    storage_layer::{
+        batch_split_writer::{BatchLayerWriter, BatchWriterResult},
+        inmemory_layer, LayerVisibilityHint,
+    },
     upload_queue::NotInitialized,
     MaybeOffloaded,
 };
@@ -855,7 +860,7 @@ pub(crate) enum ShutdownMode {
 }
 
 struct ImageLayerCreationOutcome {
-    image: Option<ResidentLayer>,
+    image: Option<ImageLayerWriter>,
     next_start_key: Key,
 }
 
@@ -4052,11 +4057,14 @@ impl Timeline {
         if wrote_keys {
             // Normal path: we have written some data into the new image layer for this
             // partition, so flush it to disk.
-            let (desc, path) = image_layer_writer.finish(ctx).await?;
-            let image_layer = Layer::finish_creating(self.conf, self, desc, &path)?;
-            info!("created image layer for rel {}", image_layer.local_path());
+            info!(
+                "flushed image layer for rel key_start={} key_end={} lsn={}",
+                image_layer_writer.key_range().start,
+                image_layer_writer.key_range().end,
+                image_layer_writer.lsn()
+            );
             Ok(ImageLayerCreationOutcome {
-                image: Some(image_layer),
+                image: Some(image_layer_writer),
                 next_start_key: img_range.end,
             })
         } else {
@@ -4142,14 +4150,14 @@ impl Timeline {
         if wrote_any_image {
             // Normal path: we have written some data into the new image layer for this
             // partition, so flush it to disk.
-            let (desc, path) = image_layer_writer.finish(ctx).await?;
-            let image_layer = Layer::finish_creating(self.conf, self, desc, &path)?;
             info!(
-                "created image layer for metadata {}",
-                image_layer.local_path()
+                "flushed image layer for metadata key_start={} key_end={} lsn={}",
+                image_layer_writer.key_range().start,
+                image_layer_writer.key_range().end,
+                image_layer_writer.lsn()
             );
             Ok(ImageLayerCreationOutcome {
-                image: Some(image_layer),
+                image: Some(image_layer_writer),
                 next_start_key: img_range.end,
             })
         } else {
@@ -4226,7 +4234,6 @@ impl Timeline {
         ctx: &RequestContext,
     ) -> Result<Vec<ResidentLayer>, CreateImageLayersError> {
         let timer = self.metrics.create_images_time_histo.start_timer();
-        let mut image_layers = Vec::new();
 
         // We need to avoid holes between generated image layers.
         // Otherwise LayerMap::image_layer_exists will return false if key range of some layer is covered by more than one
@@ -4241,6 +4248,8 @@ impl Timeline {
 
         let check_for_image_layers = self.should_check_if_image_layers_required(lsn);
 
+        let mut batch_image_writer = BatchLayerWriter::new(self.conf).await?;
+
         for partition in partitioning.parts.iter() {
             if self.cancel.is_cancelled() {
                 return Err(CreateImageLayersError::Cancelled);
@@ -4271,24 +4280,6 @@ impl Timeline {
                     continue;
                 }
             }
-            if let ImageLayerCreationMode::Force = mode {
-                // When forced to create image layers, we might try and create them where they already
-                // exist.  This mode is only used in tests/debug.
-                let layers = self.layers.read().await;
-                if layers.contains_key(&PersistentLayerKey {
-                    key_range: img_range.clone(),
-                    lsn_range: PersistentLayerDesc::image_layer_lsn_range(lsn),
-                    is_delta: false,
-                }) {
-                    tracing::info!(
-                        "Skipping image layer at {lsn} {}..{}, already exists",
-                        img_range.start,
-                        img_range.end
-                    );
-                    start = img_range.end;
-                    continue;
-                }
-            }
 
             let image_layer_writer = ImageLayerWriter::new(
                 self.conf,
@@ -4322,7 +4313,11 @@ impl Timeline {
                     .await?;
 
                 start = next_start_key;
-                image_layers.extend(image);
+                if let Some(image) = image {
+                    let key_range = image.key_range();
+                    let lsn = image.lsn();
+                    batch_image_writer.add_unfinished_image_writer(image, key_range, lsn);
+                }
             } else {
                 let ImageLayerCreationOutcome {
                     image,
@@ -4339,10 +4334,48 @@ impl Timeline {
                     )
                     .await?;
                 start = next_start_key;
-                image_layers.extend(image);
+                if let Some(image) = image {
+                    let key_range = image.key_range();
+                    let lsn = image.lsn();
+                    batch_image_writer.add_unfinished_image_writer(image, key_range, lsn);
+                }
             }
         }
 
+        let image_layers = batch_image_writer
+            .finish_with_discard_fn(self, ctx, |key| {
+                // TODO: remove this clone when Rust Edition 2024 is available, closure should capture this
+                // lifetime.
+                let key = key.clone();
+                async move {
+                    // When forced to create image layers, we might try and create them where they already
+                    // exist.  The force mode is only used in tests/debug.
+                    let layers = self.layers.read().await;
+                    if layers.contains_key(&key) {
+                        tracing::info!(
+                            "Skipping image layer at {} {}..{}, already exists",
+                            key.lsn_range.start,
+                            key.key_range.start,
+                            key.key_range.end
+                        );
+                        true
+                    } else {
+                        false
+                    }
+                }
+            })
+            .await?;
+        let image_layers = image_layers
+            .into_iter()
+            .filter_map(|x| {
+                if let BatchWriterResult::Produced(x) = x {
+                    Some(x)
+                } else {
+                    None
+                }
+            })
+            .collect_vec();
+
         let mut guard = self.layers.write().await;
 
         // FIXME: we could add the images to be uploaded *before* returning from here, but right
@@ -4475,11 +4508,23 @@ impl Drop for Timeline {
 pub(crate) enum CompactionError {
     #[error("The timeline or pageserver is shutting down")]
     ShuttingDown,
+    /// Compaction tried to offload a timeline and failed
+    #[error("Failed to offload timeline: {0}")]
+    Offload(OffloadError),
     /// Compaction cannot be done right now; page reconstruction and so on.
     #[error(transparent)]
     Other(anyhow::Error),
 }
 
+impl From<OffloadError> for CompactionError {
+    fn from(e: OffloadError) -> Self {
+        match e {
+            OffloadError::Cancelled => Self::ShuttingDown,
+            _ => Self::Offload(e),
+        }
+    }
+}
+
 impl CompactionError {
     pub fn is_cancelled(&self) -> bool {
         matches!(self, CompactionError::ShuttingDown)

pageserver/src/tenant/timeline/compaction.rs

@@ -1716,20 +1716,32 @@ impl Timeline {
         Ok(())
     }
 
-    /// An experimental compaction building block that combines compaction with garbage collection.
-    ///
-    /// The current implementation picks all delta + image layers that are below or intersecting with
-    /// the GC horizon without considering retain_lsns. Then, it does a full compaction over all these delta
-    /// layers and image layers, which generates image layers on the gc horizon, drop deltas below gc horizon,
-    /// and create delta layers with all deltas >= gc horizon.
     pub(crate) async fn compact_with_gc(
         self: &Arc<Self>,
         cancel: &CancellationToken,
         flags: EnumSet<CompactFlags>,
         ctx: &RequestContext,
     ) -> anyhow::Result<()> {
-        use std::collections::BTreeSet;
+        self.partial_compact_with_gc(None, cancel, flags, ctx).await
+    }
 
+    /// An experimental compaction building block that combines compaction with garbage collection.
+    ///
+    /// The current implementation picks all delta + image layers that are below or intersecting with
+    /// the GC horizon without considering retain_lsns. Then, it does a full compaction over all these delta
+    /// layers and image layers, which generates image layers on the gc horizon, drop deltas below gc horizon,
+    /// and create delta layers with all deltas >= gc horizon.
+    ///
+    /// If `key_range`, it will only compact the keys within the range, aka partial compaction. This functionality
+    /// is not complete yet, and if it is set, only image layers will be generated.
+    ///
+    pub(crate) async fn partial_compact_with_gc(
+        self: &Arc<Self>,
+        compaction_key_range: Option<Range<Key>>,
+        cancel: &CancellationToken,
+        flags: EnumSet<CompactFlags>,
+        ctx: &RequestContext,
+    ) -> anyhow::Result<()> {
         // 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.
@@ -1750,8 +1762,13 @@ impl Timeline {
         .await?;
 
         let dry_run = flags.contains(CompactFlags::DryRun);
+        let partial_compaction = compaction_key_range.is_some();
 
-        info!("running enhanced gc bottom-most compaction, dry_run={dry_run}");
+        if let Some(ref compaction_key_range) = compaction_key_range {
+            info!("running enhanced gc bottom-most compaction, dry_run={dry_run}, compaction_key_range={}..{}", compaction_key_range.start, compaction_key_range.end);
+        } else {
+            info!("running enhanced gc bottom-most compaction, dry_run={dry_run}");
+        }
 
         scopeguard::defer! {
             info!("done enhanced gc bottom-most compaction");
@@ -1763,7 +1780,7 @@ impl Timeline {
         // The layer selection has the following properties:
         // 1. If a layer is in the selection, all layers below it are in the selection.
         // 2. Inferred from (1), for each key in the layer selection, the value can be reconstructed only with the layers in the layer selection.
-        let (layer_selection, gc_cutoff, retain_lsns_below_horizon) = {
+        let (layer_selection, gc_cutoff, retain_lsns_below_horizon) = if !partial_compaction {
             let guard = self.layers.read().await;
             let layers = guard.layer_map()?;
             let gc_info = self.gc_info.read().unwrap();
@@ -1779,7 +1796,7 @@ impl Timeline {
                     retain_lsns_below_horizon.push(*lsn);
                 }
             }
-            let mut selected_layers = Vec::new();
+            let mut selected_layers: Vec<Layer> = Vec::new();
             drop(gc_info);
             // Pick all the layers intersect or below the gc_cutoff, get the largest LSN in the selected layers.
             let Some(max_layer_lsn) = layers
@@ -1804,8 +1821,52 @@ impl Timeline {
             }
             retain_lsns_below_horizon.sort();
             (selected_layers, gc_cutoff, retain_lsns_below_horizon)
+        } else {
+            // In case of partial compaction, we currently only support generating image layers, and therefore,
+            // we pick all layers that are below the lowest retain_lsn and does not intersect with any of the layers.
+            let guard = self.layers.read().await;
+            let layers = guard.layer_map()?;
+            let gc_info = self.gc_info.read().unwrap();
+            let mut min_lsn = gc_info.cutoffs.select_min();
+            for (lsn, _, _) in &gc_info.retain_lsns {
+                if lsn < &min_lsn {
+                    min_lsn = *lsn;
+                }
+            }
+            for lsn in gc_info.leases.keys() {
+                if lsn < &min_lsn {
+                    min_lsn = *lsn;
+                }
+            }
+            let mut selected_layers = Vec::new();
+            drop(gc_info);
+            // |-------| |-------| |-------|
+            // | Delta | | Delta | | Delta | -- min_lsn could be intersecting with the layers
+            // |-------| |-------| |-------| <- we want to pick all the layers below min_lsn, so that
+            // | Delta | | Delta | | Delta |    ...we can remove them after compaction
+            // |-------| |-------| |-------|
+            // Pick all the layers intersect or below the min_lsn, get the largest LSN in the selected layers.
+            let Some(compaction_key_range) = compaction_key_range.as_ref() else {
+                unreachable!()
+            };
+            for desc in layers.iter_historic_layers() {
+                if desc.get_lsn_range().end <= min_lsn
+                    && overlaps_with(&desc.key_range, compaction_key_range)
+                {
+                    selected_layers.push(guard.get_from_desc(&desc));
+                }
+            }
+            if selected_layers.is_empty() {
+                info!("no layers to compact with gc");
+                return Ok(());
+            }
+            (selected_layers, min_lsn, Vec::new())
         };
         let lowest_retain_lsn = if self.ancestor_timeline.is_some() {
+            if partial_compaction {
+                warn!("partial compaction cannot run on child branches (for now)");
+                return Ok(());
+            }
             Lsn(self.ancestor_lsn.0 + 1)
         } else {
             let res = retain_lsns_below_horizon
@@ -1833,23 +1894,18 @@ impl Timeline {
 
         self.check_compaction_space(&layer_selection).await?;
 
-        // Step 1: (In the future) construct a k-merge iterator over all layers. For now, simply collect all keys + LSNs.
-        // Also, verify if the layer map can be split by drawing a horizontal line at every LSN start/end split point.
-        let mut lsn_split_point = BTreeSet::new(); // TODO: use a better data structure (range tree / range set?)
+        // Generate statistics for the compaction
         for layer in &layer_selection {
             let desc = layer.layer_desc();
             if desc.is_delta() {
-                // ignore single-key layer files
-                if desc.key_range.start.next() != desc.key_range.end {
-                    let lsn_range = &desc.lsn_range;
-                    lsn_split_point.insert(lsn_range.start);
-                    lsn_split_point.insert(lsn_range.end);
-                }
                 stat.visit_delta_layer(desc.file_size());
             } else {
                 stat.visit_image_layer(desc.file_size());
             }
         }
+
+        // Step 1: construct a k-merge iterator over all layers.
+        // Also, verify if the layer map can be split by drawing a horizontal line at every LSN start/end split point.
         let layer_names: Vec<crate::tenant::storage_layer::LayerName> = layer_selection
             .iter()
             .map(|layer| layer.layer_desc().layer_name())
@@ -1900,7 +1956,10 @@ impl Timeline {
                     self.conf,
                     self.timeline_id,
                     self.tenant_shard_id,
-                    Key::MIN,
+                    compaction_key_range
+                        .as_ref()
+                        .map(|x| x.start)
+                        .unwrap_or(Key::MIN),
                     lowest_retain_lsn,
                     self.get_compaction_target_size(),
                     ctx,
@@ -1961,55 +2020,71 @@ impl Timeline {
             } else {
                 let last_key = last_key.as_mut().unwrap();
                 stat.on_unique_key_visited();
-                let retention = self
-                    .generate_key_retention(
-                        *last_key,
-                        &accumulated_values,
-                        gc_cutoff,
-                        &retain_lsns_below_horizon,
-                        COMPACTION_DELTA_THRESHOLD,
-                        get_ancestor_image(self, *last_key, ctx).await?,
-                    )
-                    .await?;
-                // Put the image into the image layer. Currently we have a single big layer for the compaction.
-                retention
-                    .pipe_to(
-                        *last_key,
-                        &mut delta_layer_writer,
-                        image_layer_writer.as_mut(),
-                        &mut stat,
-                        ctx,
-                    )
-                    .await?;
+                let skip_adding_key = if let Some(ref compaction_key_range) = compaction_key_range {
+                    !compaction_key_range.contains(last_key)
+                } else {
+                    false
+                };
+                if !skip_adding_key {
+                    let retention = self
+                        .generate_key_retention(
+                            *last_key,
+                            &accumulated_values,
+                            gc_cutoff,
+                            &retain_lsns_below_horizon,
+                            COMPACTION_DELTA_THRESHOLD,
+                            get_ancestor_image(self, *last_key, ctx).await?,
+                        )
+                        .await?;
+                    // Put the image into the image layer. Currently we have a single big layer for the compaction.
+                    retention
+                        .pipe_to(
+                            *last_key,
+                            &mut delta_layer_writer,
+                            image_layer_writer.as_mut(),
+                            &mut stat,
+                            ctx,
+                        )
+                        .await?;
+                }
                 accumulated_values.clear();
                 *last_key = key;
                 accumulated_values.push((key, lsn, val));
             }
         }
 
+        // TODO: move the below part to the loop body
         let last_key = last_key.expect("no keys produced during compaction");
-        // TODO: move this part to the loop body
         stat.on_unique_key_visited();
-        let retention = self
-            .generate_key_retention(
-                last_key,
-                &accumulated_values,
-                gc_cutoff,
-                &retain_lsns_below_horizon,
-                COMPACTION_DELTA_THRESHOLD,
-                get_ancestor_image(self, last_key, ctx).await?,
-            )
-            .await?;
-        // Put the image into the image layer. Currently we have a single big layer for the compaction.
-        retention
-            .pipe_to(
-                last_key,
-                &mut delta_layer_writer,
-                image_layer_writer.as_mut(),
-                &mut stat,
-                ctx,
-            )
-            .await?;
+
+        let skip_adding_key = if let Some(ref compaction_key_range) = compaction_key_range {
+            !compaction_key_range.contains(&last_key)
+        } else {
+            false
+        };
+        if !skip_adding_key {
+            let retention = self
+                .generate_key_retention(
+                    last_key,
+                    &accumulated_values,
+                    gc_cutoff,
+                    &retain_lsns_below_horizon,
+                    COMPACTION_DELTA_THRESHOLD,
+                    get_ancestor_image(self, last_key, ctx).await?,
+                )
+                .await?;
+            // Put the image into the image layer. Currently we have a single big layer for the compaction.
+            retention
+                .pipe_to(
+                    last_key,
+                    &mut delta_layer_writer,
+                    image_layer_writer.as_mut(),
+                    &mut stat,
+                    ctx,
+                )
+                .await?;
+        }
+        // end: move the above part to the loop body
 
         let discard = |key: &PersistentLayerKey| {
             let key = key.clone();
@@ -2018,8 +2093,12 @@ impl Timeline {
 
         let produced_image_layers = if let Some(writer) = image_layer_writer {
             if !dry_run {
+                let end_key = compaction_key_range
+                    .as_ref()
+                    .map(|x| x.end)
+                    .unwrap_or(Key::MAX);
                 writer
-                    .finish_with_discard_fn(self, ctx, Key::MAX, discard)
+                    .finish_with_discard_fn(self, ctx, end_key, discard)
                     .await?
             } else {
                 drop(writer);
@@ -2038,6 +2117,10 @@ impl Timeline {
             Vec::new()
         };
 
+        if partial_compaction && !produced_delta_layers.is_empty() {
+            bail!("implementation error: partial compaction should not be producing delta layers (for now)");
+        }
+
         let mut compact_to = Vec::new();
         let mut keep_layers = HashSet::new();
         let produced_delta_layers_len = produced_delta_layers.len();
@@ -2068,6 +2151,28 @@ impl Timeline {
         }
         let mut layer_selection = layer_selection;
         layer_selection.retain(|x| !keep_layers.contains(&x.layer_desc().key()));
+        if let Some(ref compaction_key_range) = compaction_key_range {
+            // Partial compaction might select more data than it processes, e.g., if
+            // the compaction_key_range only partially overlaps:
+            //
+            //         [---compaction_key_range---]
+            //   [---A----][----B----][----C----][----D----]
+            //
+            // A,B,C,D are all in the `layer_selection`. The created image layers contain
+            // whatever is needed from B, C, and from `----]` of A, and from  `[--` of D.
+            //
+            // In contrast, `[--A-` and `--D----]` have not been processed, so, we must
+            // keep that data.
+            //
+            // The solution for now is to keep A and D completely.
+            // (layer_selection is what we'll remove from the layer map, so,
+            //  retain what is _not_ fully covered by compaction_key_range).
+            layer_selection.retain(|x| {
+                let key_range = &x.layer_desc().key_range;
+                key_range.start >= compaction_key_range.start
+                    && key_range.end <= compaction_key_range.end
+            });
+        }
 
         info!(
             "gc-compaction statistics: {}",
@@ -2356,7 +2461,10 @@ impl TimelineAdaptor {
             )
             .await?;
 
-        if let Some(image_layer) = image {
+        if let Some(image_layer_writer) = image {
+            let (desc, path) = image_layer_writer.finish(ctx).await?;
+            let image_layer =
+                Layer::finish_creating(self.timeline.conf, &self.timeline, desc, &path)?;
             self.new_images.push(image_layer);
         }
 

pageserver/src/tenant/timeline/delete.rs

@@ -14,10 +14,9 @@ use crate::{
     task_mgr::{self, TaskKind},
     tenant::{
         metadata::TimelineMetadata,
-        remote_timeline_client::{
-            self, MaybeDeletedIndexPart, PersistIndexPartWithDeletedFlagError, RemoteTimelineClient,
-        },
-        CreateTimelineCause, DeleteTimelineError, Tenant, TimelineOrOffloaded,
+        remote_timeline_client::{PersistIndexPartWithDeletedFlagError, RemoteTimelineClient},
+        CreateTimelineCause, DeleteTimelineError, MaybeDeletedIndexPart, Tenant,
+        TimelineOrOffloaded,
     },
 };
 
@@ -176,32 +175,6 @@ async fn remove_maybe_offloaded_timeline_from_tenant(
     Ok(())
 }
 
-/// It is important that this gets called when DeletionGuard is being held.
-/// For more context see comments in [`DeleteTimelineFlow::prepare`]
-async fn upload_new_tenant_manifest(
-    tenant: &Tenant,
-    _: &DeletionGuard, // using it as a witness
-) -> anyhow::Result<()> {
-    // This is susceptible to race conditions, i.e. we won't continue deletions if there is a crash
-    // between the deletion of the index-part.json and reaching of this code.
-    // So indeed, the tenant manifest might refer to an offloaded timeline which has already been deleted.
-    // However, we handle this case in tenant loading code so the next time we attach, the issue is
-    // resolved.
-    let manifest = tenant.tenant_manifest();
-    // TODO: generation support
-    let generation = remote_timeline_client::TENANT_MANIFEST_GENERATION;
-    remote_timeline_client::upload_tenant_manifest(
-        &tenant.remote_storage,
-        &tenant.tenant_shard_id,
-        generation,
-        &manifest,
-        &tenant.cancel,
-    )
-    .await?;
-
-    Ok(())
-}
-
 /// Orchestrates timeline shut down of all timeline tasks, removes its in-memory structures,
 /// and deletes its data from both disk and s3.
 /// The sequence of steps:
@@ -241,7 +214,8 @@ impl DeleteTimelineFlow {
     ) -> Result<(), DeleteTimelineError> {
         super::debug_assert_current_span_has_tenant_and_timeline_id();
 
-        let (timeline, mut guard) = Self::prepare(tenant, timeline_id)?;
+        let allow_offloaded_children = false;
+        let (timeline, mut guard) = Self::prepare(tenant, timeline_id, allow_offloaded_children)?;
 
         guard.mark_in_progress()?;
 
@@ -367,6 +341,7 @@ impl DeleteTimelineFlow {
     pub(super) fn prepare(
         tenant: &Tenant,
         timeline_id: TimelineId,
+        allow_offloaded_children: bool,
     ) -> Result<(TimelineOrOffloaded, DeletionGuard), DeleteTimelineError> {
         // Note the interaction between this guard and deletion guard.
         // Here we attempt to lock deletion guard when we're holding a lock on timelines.
@@ -379,30 +354,27 @@ impl DeleteTimelineFlow {
         // T1: acquire deletion lock, do another `DeleteTimelineFlow::run`
         // For more context see this discussion: `https://github.com/neondatabase/neon/pull/4552#discussion_r1253437346`
         let timelines = tenant.timelines.lock().unwrap();
+        let timelines_offloaded = tenant.timelines_offloaded.lock().unwrap();
 
         let timeline = match timelines.get(&timeline_id) {
             Some(t) => TimelineOrOffloaded::Timeline(Arc::clone(t)),
-            None => {
-                let offloaded_timelines = tenant.timelines_offloaded.lock().unwrap();
-                match offloaded_timelines.get(&timeline_id) {
-                    Some(t) => TimelineOrOffloaded::Offloaded(Arc::clone(t)),
-                    None => return Err(DeleteTimelineError::NotFound),
-                }
-            }
+            None => match timelines_offloaded.get(&timeline_id) {
+                Some(t) => TimelineOrOffloaded::Offloaded(Arc::clone(t)),
+                None => return Err(DeleteTimelineError::NotFound),
+            },
         };
 
-        // Ensure that there are no child timelines **attached to that pageserver**,
-        // because detach removes files, which will break child branches
-        let children: Vec<TimelineId> = timelines
-            .iter()
-            .filter_map(|(id, entry)| {
-                if entry.get_ancestor_timeline_id() == Some(timeline_id) {
-                    Some(*id)
-                } else {
-                    None
-                }
-            })
-            .collect();
+        // Ensure that there are no child timelines, because we are about to remove files,
+        // which will break child branches
+        let mut children = Vec::new();
+        if !allow_offloaded_children {
+            children.extend(timelines_offloaded.iter().filter_map(|(id, entry)| {
+                (entry.ancestor_timeline_id == Some(timeline_id)).then_some(*id)
+            }));
+        }
+        children.extend(timelines.iter().filter_map(|(id, entry)| {
+            (entry.get_ancestor_timeline_id() == Some(timeline_id)).then_some(*id)
+        }));
 
         if !children.is_empty() {
             return Err(DeleteTimelineError::HasChildren(children));
@@ -480,7 +452,15 @@ impl DeleteTimelineFlow {
 
         remove_maybe_offloaded_timeline_from_tenant(tenant, timeline, &guard).await?;
 
-        upload_new_tenant_manifest(tenant, &guard).await?;
+        // This is susceptible to race conditions, i.e. we won't continue deletions if there is a crash
+        // between the deletion of the index-part.json and reaching of this code.
+        // So indeed, the tenant manifest might refer to an offloaded timeline which has already been deleted.
+        // However, we handle this case in tenant loading code so the next time we attach, the issue is
+        // resolved.
+        tenant
+            .store_tenant_manifest()
+            .await
+            .map_err(|e| DeleteTimelineError::Other(anyhow::anyhow!(e)))?;
 
         *guard = Self::Finished;
 

pageserver/src/tenant/timeline/detach_ancestor.rs

@@ -29,6 +29,9 @@ pub(crate) enum Error {
     #[error("shutting down, please retry later")]
     ShuttingDown,
 
+    #[error("archived: {}", .0)]
+    Archived(TimelineId),
+
     #[error(transparent)]
     NotFound(crate::tenant::GetTimelineError),
 
@@ -79,8 +82,9 @@ impl From<Error> for ApiError {
     fn from(value: Error) -> Self {
         match value {
             Error::NoAncestor => ApiError::Conflict(value.to_string()),
-            Error::TooManyAncestors => ApiError::BadRequest(anyhow::anyhow!("{}", value)),
+            Error::TooManyAncestors => ApiError::BadRequest(anyhow::anyhow!("{value}")),
             Error::ShuttingDown => ApiError::ShuttingDown,
+            Error::Archived(_) => ApiError::BadRequest(anyhow::anyhow!("{value}")),
             Error::OtherTimelineDetachOngoing(_) | Error::FailedToReparentAll => {
                 ApiError::ResourceUnavailable(value.to_string().into())
             }
@@ -201,12 +205,18 @@ pub(super) async fn prepare(
         }));
     };
 
+    if detached.is_archived() != Some(false) {
+        return Err(Archived(detached.timeline_id));
+    }
+
     if !ancestor_lsn.is_valid() {
         // rare case, probably wouldn't even load
         tracing::error!("ancestor is set, but ancestor_lsn is invalid, this timeline needs fixing");
         return Err(NoAncestor);
     }
 
+    check_no_archived_children_of_ancestor(tenant, detached, &ancestor, ancestor_lsn)?;
+
     if ancestor.ancestor_timeline.is_some() {
         // non-technical requirement; we could flatten N ancestors just as easily but we chose
         // not to, at least initially
@@ -950,3 +960,36 @@ where
         }
     })
 }
+
+fn check_no_archived_children_of_ancestor(
+    tenant: &Tenant,
+    detached: &Arc<Timeline>,
+    ancestor: &Arc<Timeline>,
+    ancestor_lsn: Lsn,
+) -> Result<(), Error> {
+    let timelines = tenant.timelines.lock().unwrap();
+    let timelines_offloaded = tenant.timelines_offloaded.lock().unwrap();
+    for timeline in reparentable_timelines(timelines.values(), detached, ancestor, ancestor_lsn) {
+        if timeline.is_archived() == Some(true) {
+            return Err(Error::Archived(timeline.timeline_id));
+        }
+    }
+    for timeline_offloaded in timelines_offloaded.values() {
+        if timeline_offloaded.ancestor_timeline_id != Some(ancestor.timeline_id) {
+            continue;
+        }
+        // This forbids the detach ancestor feature if flattened timelines are present,
+        // even if the ancestor_lsn is from after the branchpoint of the detached timeline.
+        // But as per current design, we don't record the ancestor_lsn of flattened timelines.
+        // This is a bit unfortunate, but as of writing this we don't support flattening
+        // anyway. Maybe we can evolve the data model in the future.
+        if let Some(retain_lsn) = timeline_offloaded.ancestor_retain_lsn {
+            let is_earlier = retain_lsn <= ancestor_lsn;
+            if !is_earlier {
+                continue;
+            }
+        }
+        return Err(Error::Archived(timeline_offloaded.timeline_id));
+    }
+    Ok(())
+}

pageserver/src/tenant/timeline/offload.rs

@@ -3,16 +3,40 @@ use std::sync::Arc;
 use super::delete::{delete_local_timeline_directory, DeleteTimelineFlow, DeletionGuard};
 use super::Timeline;
 use crate::span::debug_assert_current_span_has_tenant_and_timeline_id;
-use crate::tenant::{remote_timeline_client, OffloadedTimeline, Tenant, TimelineOrOffloaded};
+use crate::tenant::{OffloadedTimeline, Tenant, TenantManifestError, TimelineOrOffloaded};
+
+#[derive(thiserror::Error, Debug)]
+pub(crate) enum OffloadError {
+    #[error("Cancelled")]
+    Cancelled,
+    #[error("Timeline is not archived")]
+    NotArchived,
+    #[error(transparent)]
+    RemoteStorage(anyhow::Error),
+    #[error("Unexpected offload error: {0}")]
+    Other(anyhow::Error),
+}
+
+impl From<TenantManifestError> for OffloadError {
+    fn from(e: TenantManifestError) -> Self {
+        match e {
+            TenantManifestError::Cancelled => Self::Cancelled,
+            TenantManifestError::RemoteStorage(e) => Self::RemoteStorage(e),
+        }
+    }
+}
 
 pub(crate) async fn offload_timeline(
     tenant: &Tenant,
     timeline: &Arc<Timeline>,
-) -> anyhow::Result<()> {
+) -> Result<(), OffloadError> {
     debug_assert_current_span_has_tenant_and_timeline_id();
     tracing::info!("offloading archived timeline");
 
-    let (timeline, guard) = DeleteTimelineFlow::prepare(tenant, timeline.timeline_id)?;
+    let allow_offloaded_children = true;
+    let (timeline, guard) =
+        DeleteTimelineFlow::prepare(tenant, timeline.timeline_id, allow_offloaded_children)
+            .map_err(|e| OffloadError::Other(anyhow::anyhow!(e)))?;
 
     let TimelineOrOffloaded::Timeline(timeline) = timeline else {
         tracing::error!("timeline already offloaded, but given timeline object");
@@ -24,14 +48,15 @@ pub(crate) async fn offload_timeline(
         Some(true) => (),
         Some(false) => {
             tracing::warn!(?is_archived, "tried offloading a non-archived timeline");
-            anyhow::bail!("timeline isn't archived");
+            return Err(OffloadError::NotArchived);
         }
         None => {
-            tracing::warn!(
+            // This is legal: calls to this function can race with the timeline shutting down
+            tracing::info!(
                 ?is_archived,
-                "tried offloading a timeline where manifest is not yet available"
+                "tried offloading a timeline whose remote storage is not initialized"
             );
-            anyhow::bail!("timeline manifest hasn't been loaded yet");
+            return Err(OffloadError::Cancelled);
         }
     }
 
@@ -42,9 +67,11 @@ pub(crate) async fn offload_timeline(
     // to make deletions possible while offloading is in progress
 
     let conf = &tenant.conf;
-    delete_local_timeline_directory(conf, tenant.tenant_shard_id, &timeline).await?;
+    delete_local_timeline_directory(conf, tenant.tenant_shard_id, &timeline)
+        .await
+        .map_err(OffloadError::Other)?;
 
-    remove_timeline_from_tenant(tenant, &timeline, &guard).await?;
+    remove_timeline_from_tenant(tenant, &timeline, &guard);
 
     {
         let mut offloaded_timelines = tenant.timelines_offloaded.lock().unwrap();
@@ -63,28 +90,18 @@ pub(crate) async fn offload_timeline(
     // at the next restart attach it again.
     // For that to happen, we'd need to make the manifest reflect our *intended* state,
     // not our actual state of offloaded timelines.
-    let manifest = tenant.tenant_manifest();
-    // TODO: generation support
-    let generation = remote_timeline_client::TENANT_MANIFEST_GENERATION;
-    remote_timeline_client::upload_tenant_manifest(
-        &tenant.remote_storage,
-        &tenant.tenant_shard_id,
-        generation,
-        &manifest,
-        &tenant.cancel,
-    )
-    .await?;
+    tenant.store_tenant_manifest().await?;
 
     Ok(())
 }
 
 /// It is important that this gets called when DeletionGuard is being held.
 /// For more context see comments in [`DeleteTimelineFlow::prepare`]
-async fn remove_timeline_from_tenant(
+fn remove_timeline_from_tenant(
     tenant: &Tenant,
     timeline: &Timeline,
     _: &DeletionGuard, // using it as a witness
-) -> anyhow::Result<()> {
+) {
     // Remove the timeline from the map.
     let mut timelines = tenant.timelines.lock().unwrap();
     let children_exist = timelines
@@ -100,8 +117,4 @@ async fn remove_timeline_from_tenant(
     timelines
         .remove(&timeline.timeline_id)
         .expect("timeline that we were deleting was concurrently removed from 'timelines' map");
-
-    drop(timelines);
-
-    Ok(())
 }

pageserver/src/tenant/timeline/uninit.rs

@@ -141,7 +141,9 @@ impl Drop for UninitializedTimeline<'_> {
     fn drop(&mut self) {
         if let Some((_, create_guard)) = self.raw_timeline.take() {
             let _entered = info_span!("drop_uninitialized_timeline", tenant_id = %self.owning_tenant.tenant_shard_id.tenant_id, shard_id = %self.owning_tenant.tenant_shard_id.shard_slug(), timeline_id = %self.timeline_id).entered();
-            error!("Timeline got dropped without initializing, cleaning its files");
+            // This is unusual, but can happen harmlessly if the pageserver is stopped while
+            // creating a timeline.
+            info!("Timeline got dropped without initializing, cleaning its files");
             cleanup_timeline_directory(create_guard);
         }
     }

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

@@ -22,6 +22,7 @@ use tokio::{select, sync::watch, time};
 use tokio_postgres::{replication::ReplicationStream, Client};
 use tokio_util::sync::CancellationToken;
 use tracing::{debug, error, info, trace, warn, Instrument};
+use wal_decoder::models::{FlushUncommittedRecords, InterpretedWalRecord};
 
 use super::TaskStateUpdate;
 use crate::{
@@ -35,7 +36,6 @@ use crate::{
 use postgres_backend::is_expected_io_error;
 use postgres_connection::PgConnectionConfig;
 use postgres_ffi::waldecoder::WalStreamDecoder;
-use postgres_ffi::walrecord::{decode_wal_record, DecodedWALRecord};
 use utils::{id::NodeId, lsn::Lsn};
 use utils::{pageserver_feedback::PageserverFeedback, sync::gate::GateError};
 
@@ -339,11 +339,15 @@ pub(super) async fn handle_walreceiver_connection(
                             return Err(WalReceiverError::Other(anyhow!("LSN not aligned")));
                         }
 
-                        // Deserialize WAL record
-                        let mut decoded = DecodedWALRecord::default();
-                        decode_wal_record(recdata, &mut decoded, modification.tline.pg_version)?;
+                        // Deserialize and interpret WAL record
+                        let interpreted = InterpretedWalRecord::from_bytes_filtered(
+                            recdata,
+                            modification.tline.get_shard_identity(),
+                            lsn,
+                            modification.tline.pg_version,
+                        )?;
 
-                        if decoded.is_dbase_create_copy(timeline.pg_version)
+                        if matches!(interpreted.flush_uncommitted, FlushUncommittedRecords::Yes)
                             && uncommitted_records > 0
                         {
                             // Special case: legacy PG database creations operate by reading pages from a 'template' database:
@@ -360,7 +364,7 @@ pub(super) async fn handle_walreceiver_connection(
 
                         // Ingest the records without immediately committing them.
                         let ingested = walingest
-                            .ingest_record(decoded, lsn, &mut modification, &ctx)
+                            .ingest_record(interpreted, &mut modification, &ctx)
                             .await
                             .with_context(|| format!("could not ingest record at {lsn}"))?;
                         if !ingested {

pgxn/neon/unstable_extensions.c

@@ -65,8 +65,8 @@ CheckUnstableExtension(
 			{
 				ereport(ERROR,
 						(errcode(ERRCODE_INSUFFICIENT_PRIVILEGE),
-						 errmsg("installing %s is currently prohibited", stmt->extname),
-						 errhint("Set neon.allow_unstable_extensions to true")));
+						 errmsg("%s extension is in beta and may be unstable or introduce backward-incompatible changes.\nWe recommend testing it in a separate, dedicated Neon project.", stmt->extname),
+						 errhint("to proceed with installation, run SET neon.allow_unstable_extensions='true'")));
 			}
 			break;
 		}
@@ -110,13 +110,13 @@ InitUnstableExtensionsSupport(void)
 		NULL,
 		&allow_unstable_extensions,
 		false,
-		PGC_SUSET,
+		PGC_USERSET,
 		0,
 		NULL, NULL, NULL);
 
 	DefineCustomStringVariable(
 		"neon.unstable_extensions",
-		"Allow unstable extensions to be installed and used",
+		"List of unstable extensions",
 		NULL,
 		&unstable_extensions,
 		NULL,

poetry.lock

@@ -1,4 +1,4 @@
-# This file is automatically @generated by Poetry 1.8.3 and should not be changed by hand.
+# This file is automatically @generated by Poetry 1.8.4 and should not be changed by hand.
 
 [[package]]
 name = "aiohappyeyeballs"
@@ -1034,24 +1034,25 @@ test-randomorder = ["pytest-randomly"]
 
 [[package]]
 name = "docker"
-version = "4.2.2"
+version = "7.1.0"
 description = "A Python library for the Docker Engine API."
 optional = false
-python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*, !=3.4.*"
+python-versions = ">=3.8"
 files = [
-    {file = "docker-4.2.2-py2.py3-none-any.whl", hash = "sha256:03a46400c4080cb6f7aa997f881ddd84fef855499ece219d75fbdb53289c17ab"},
-    {file = "docker-4.2.2.tar.gz", hash = "sha256:26eebadce7e298f55b76a88c4f8802476c5eaddbdbe38dbc6cce8781c47c9b54"},
+    {file = "docker-7.1.0-py3-none-any.whl", hash = "sha256:c96b93b7f0a746f9e77d325bcfb87422a3d8bd4f03136ae8a85b37f1898d5fc0"},
+    {file = "docker-7.1.0.tar.gz", hash = "sha256:ad8c70e6e3f8926cb8a92619b832b4ea5299e2831c14284663184e200546fa6c"},
 ]
 
 [package.dependencies]
-pypiwin32 = {version = "223", markers = "sys_platform == \"win32\" and python_version >= \"3.6\""}
-requests = ">=2.14.2,<2.18.0 || >2.18.0"
-six = ">=1.4.0"
-websocket-client = ">=0.32.0"
+pywin32 = {version = ">=304", markers = "sys_platform == \"win32\""}
+requests = ">=2.26.0"
+urllib3 = ">=1.26.0"
 
 [package.extras]
-ssh = ["paramiko (>=2.4.2)"]
-tls = ["cryptography (>=1.3.4)", "idna (>=2.0.0)", "pyOpenSSL (>=17.5.0)"]
+dev = ["coverage (==7.2.7)", "pytest (==7.4.2)", "pytest-cov (==4.1.0)", "pytest-timeout (==2.1.0)", "ruff (==0.1.8)"]
+docs = ["myst-parser (==0.18.0)", "sphinx (==5.1.1)"]
+ssh = ["paramiko (>=2.4.3)"]
+websockets = ["websocket-client (>=1.3.0)"]
 
 [[package]]
 name = "exceptiongroup"
@@ -1416,6 +1417,16 @@ files = [
     {file = "jsondiff-2.0.0.tar.gz", hash = "sha256:2795844ef075ec8a2b8d385c4d59f5ea48b08e7180fce3cb2787be0db00b1fb4"},
 ]
 
+[[package]]
+name = "jsonnet"
+version = "0.20.0"
+description = "Python bindings for Jsonnet - The data templating language"
+optional = false
+python-versions = "*"
+files = [
+    {file = "jsonnet-0.20.0.tar.gz", hash = "sha256:7e770c7bf3a366b97b650a39430450f77612e74406731eb75c5bd59f3f104d4f"},
+]
+
 [[package]]
 name = "jsonpatch"
 version = "1.32"
@@ -1521,6 +1532,21 @@ files = [
 [package.dependencies]
 six = "*"
 
+[[package]]
+name = "jwcrypto"
+version = "1.5.6"
+description = "Implementation of JOSE Web standards"
+optional = false
+python-versions = ">= 3.8"
+files = [
+    {file = "jwcrypto-1.5.6-py3-none-any.whl", hash = "sha256:150d2b0ebbdb8f40b77f543fb44ffd2baeff48788be71f67f03566692fd55789"},
+    {file = "jwcrypto-1.5.6.tar.gz", hash = "sha256:771a87762a0c081ae6166958a954f80848820b2ab066937dc8b8379d65b1b039"},
+]
+
+[package.dependencies]
+cryptography = ">=3.4"
+typing-extensions = ">=4.5.0"
+
 [[package]]
 name = "kafka-python"
 version = "2.0.2"
@@ -2328,20 +2354,6 @@ files = [
 [package.extras]
 diagrams = ["jinja2", "railroad-diagrams"]
 
-[[package]]
-name = "pypiwin32"
-version = "223"
-description = ""
-optional = false
-python-versions = "*"
-files = [
-    {file = "pypiwin32-223-py3-none-any.whl", hash = "sha256:67adf399debc1d5d14dffc1ab5acacb800da569754fafdc576b2a039485aa775"},
-    {file = "pypiwin32-223.tar.gz", hash = "sha256:71be40c1fbd28594214ecaecb58e7aa8b708eabfa0125c8a109ebd51edbd776a"},
-]
-
-[package.dependencies]
-pywin32 = ">=223"
-
 [[package]]
 name = "pyrsistent"
 version = "0.18.1"
@@ -2561,81 +2573,91 @@ files = [
 
 [[package]]
 name = "pywin32"
-version = "301"
+version = "308"
 description = "Python for Window Extensions"
 optional = false
 python-versions = "*"
 files = [
-    {file = "pywin32-301-cp35-cp35m-win32.whl", hash = "sha256:93367c96e3a76dfe5003d8291ae16454ca7d84bb24d721e0b74a07610b7be4a7"},
-    {file = "pywin32-301-cp35-cp35m-win_amd64.whl", hash = "sha256:9635df6998a70282bd36e7ac2a5cef9ead1627b0a63b17c731312c7a0daebb72"},
-    {file = "pywin32-301-cp36-cp36m-win32.whl", hash = "sha256:c866f04a182a8cb9b7855de065113bbd2e40524f570db73ef1ee99ff0a5cc2f0"},
-    {file = "pywin32-301-cp36-cp36m-win_amd64.whl", hash = "sha256:dafa18e95bf2a92f298fe9c582b0e205aca45c55f989937c52c454ce65b93c78"},
-    {file = "pywin32-301-cp37-cp37m-win32.whl", hash = "sha256:98f62a3f60aa64894a290fb7494bfa0bfa0a199e9e052e1ac293b2ad3cd2818b"},
-    {file = "pywin32-301-cp37-cp37m-win_amd64.whl", hash = "sha256:fb3b4933e0382ba49305cc6cd3fb18525df7fd96aa434de19ce0878133bf8e4a"},
-    {file = "pywin32-301-cp38-cp38-win32.whl", hash = "sha256:88981dd3cfb07432625b180f49bf4e179fb8cbb5704cd512e38dd63636af7a17"},
-    {file = "pywin32-301-cp38-cp38-win_amd64.whl", hash = "sha256:8c9d33968aa7fcddf44e47750e18f3d034c3e443a707688a008a2e52bbef7e96"},
-    {file = "pywin32-301-cp39-cp39-win32.whl", hash = "sha256:595d397df65f1b2e0beaca63a883ae6d8b6df1cdea85c16ae85f6d2e648133fe"},
-    {file = "pywin32-301-cp39-cp39-win_amd64.whl", hash = "sha256:87604a4087434cd814ad8973bd47d6524bd1fa9e971ce428e76b62a5e0860fdf"},
+    {file = "pywin32-308-cp310-cp310-win32.whl", hash = "sha256:796ff4426437896550d2981b9c2ac0ffd75238ad9ea2d3bfa67a1abd546d262e"},
+    {file = "pywin32-308-cp310-cp310-win_amd64.whl", hash = "sha256:4fc888c59b3c0bef905ce7eb7e2106a07712015ea1c8234b703a088d46110e8e"},
+    {file = "pywin32-308-cp310-cp310-win_arm64.whl", hash = "sha256:a5ab5381813b40f264fa3495b98af850098f814a25a63589a8e9eb12560f450c"},
+    {file = "pywin32-308-cp311-cp311-win32.whl", hash = "sha256:5d8c8015b24a7d6855b1550d8e660d8daa09983c80e5daf89a273e5c6fb5095a"},
+    {file = "pywin32-308-cp311-cp311-win_amd64.whl", hash = "sha256:575621b90f0dc2695fec346b2d6302faebd4f0f45c05ea29404cefe35d89442b"},
+    {file = "pywin32-308-cp311-cp311-win_arm64.whl", hash = "sha256:100a5442b7332070983c4cd03f2e906a5648a5104b8a7f50175f7906efd16bb6"},
+    {file = "pywin32-308-cp312-cp312-win32.whl", hash = "sha256:587f3e19696f4bf96fde9d8a57cec74a57021ad5f204c9e627e15c33ff568897"},
+    {file = "pywin32-308-cp312-cp312-win_amd64.whl", hash = "sha256:00b3e11ef09ede56c6a43c71f2d31857cf7c54b0ab6e78ac659497abd2834f47"},
+    {file = "pywin32-308-cp312-cp312-win_arm64.whl", hash = "sha256:9b4de86c8d909aed15b7011182c8cab38c8850de36e6afb1f0db22b8959e3091"},
+    {file = "pywin32-308-cp313-cp313-win32.whl", hash = "sha256:1c44539a37a5b7b21d02ab34e6a4d314e0788f1690d65b48e9b0b89f31abbbed"},
+    {file = "pywin32-308-cp313-cp313-win_amd64.whl", hash = "sha256:fd380990e792eaf6827fcb7e187b2b4b1cede0585e3d0c9e84201ec27b9905e4"},
+    {file = "pywin32-308-cp313-cp313-win_arm64.whl", hash = "sha256:ef313c46d4c18dfb82a2431e3051ac8f112ccee1a34f29c263c583c568db63cd"},
+    {file = "pywin32-308-cp37-cp37m-win32.whl", hash = "sha256:1f696ab352a2ddd63bd07430080dd598e6369152ea13a25ebcdd2f503a38f1ff"},
+    {file = "pywin32-308-cp37-cp37m-win_amd64.whl", hash = "sha256:13dcb914ed4347019fbec6697a01a0aec61019c1046c2b905410d197856326a6"},
+    {file = "pywin32-308-cp38-cp38-win32.whl", hash = "sha256:5794e764ebcabf4ff08c555b31bd348c9025929371763b2183172ff4708152f0"},
+    {file = "pywin32-308-cp38-cp38-win_amd64.whl", hash = "sha256:3b92622e29d651c6b783e368ba7d6722b1634b8e70bd376fd7610fe1992e19de"},
+    {file = "pywin32-308-cp39-cp39-win32.whl", hash = "sha256:7873ca4dc60ab3287919881a7d4f88baee4a6e639aa6962de25a98ba6b193341"},
+    {file = "pywin32-308-cp39-cp39-win_amd64.whl", hash = "sha256:71b3322d949b4cc20776436a9c9ba0eeedcbc9c650daa536df63f0ff111bb920"},
 ]
 
 [[package]]
 name = "pyyaml"
-version = "6.0.1"
+version = "6.0.2"
 description = "YAML parser and emitter for Python"
 optional = false
-python-versions = ">=3.6"
+python-versions = ">=3.8"
 files = [
-    {file = "PyYAML-6.0.1-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:d858aa552c999bc8a8d57426ed01e40bef403cd8ccdd0fc5f6f04a00414cac2a"},
-    {file = "PyYAML-6.0.1-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:fd66fc5d0da6d9815ba2cebeb4205f95818ff4b79c3ebe268e75d961704af52f"},
-    {file = "PyYAML-6.0.1-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:69b023b2b4daa7548bcfbd4aa3da05b3a74b772db9e23b982788168117739938"},
-    {file = "PyYAML-6.0.1-cp310-cp310-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:81e0b275a9ecc9c0c0c07b4b90ba548307583c125f54d5b6946cfee6360c733d"},
-    {file = "PyYAML-6.0.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ba336e390cd8e4d1739f42dfe9bb83a3cc2e80f567d8805e11b46f4a943f5515"},
-    {file = "PyYAML-6.0.1-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:326c013efe8048858a6d312ddd31d56e468118ad4cdeda36c719bf5bb6192290"},
-    {file = "PyYAML-6.0.1-cp310-cp310-win32.whl", hash = "sha256:bd4af7373a854424dabd882decdc5579653d7868b8fb26dc7d0e99f823aa5924"},
-    {file = "PyYAML-6.0.1-cp310-cp310-win_amd64.whl", hash = "sha256:fd1592b3fdf65fff2ad0004b5e363300ef59ced41c2e6b3a99d4089fa8c5435d"},
-    {file = "PyYAML-6.0.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:6965a7bc3cf88e5a1c3bd2e0b5c22f8d677dc88a455344035f03399034eb3007"},
-    {file = "PyYAML-6.0.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:f003ed9ad21d6a4713f0a9b5a7a0a79e08dd0f221aff4525a2be4c346ee60aab"},
-    {file = "PyYAML-6.0.1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:42f8152b8dbc4fe7d96729ec2b99c7097d656dc1213a3229ca5383f973a5ed6d"},
-    {file = "PyYAML-6.0.1-cp311-cp311-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:062582fca9fabdd2c8b54a3ef1c978d786e0f6b3a1510e0ac93ef59e0ddae2bc"},
-    {file = "PyYAML-6.0.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:d2b04aac4d386b172d5b9692e2d2da8de7bfb6c387fa4f801fbf6fb2e6ba4673"},
-    {file = "PyYAML-6.0.1-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:e7d73685e87afe9f3b36c799222440d6cf362062f78be1013661b00c5c6f678b"},
-    {file = "PyYAML-6.0.1-cp311-cp311-win32.whl", hash = "sha256:1635fd110e8d85d55237ab316b5b011de701ea0f29d07611174a1b42f1444741"},
-    {file = "PyYAML-6.0.1-cp311-cp311-win_amd64.whl", hash = "sha256:bf07ee2fef7014951eeb99f56f39c9bb4af143d8aa3c21b1677805985307da34"},
-    {file = "PyYAML-6.0.1-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:855fb52b0dc35af121542a76b9a84f8d1cd886ea97c84703eaa6d88e37a2ad28"},
-    {file = "PyYAML-6.0.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:40df9b996c2b73138957fe23a16a4f0ba614f4c0efce1e9406a184b6d07fa3a9"},
-    {file = "PyYAML-6.0.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a08c6f0fe150303c1c6b71ebcd7213c2858041a7e01975da3a99aed1e7a378ef"},
-    {file = "PyYAML-6.0.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6c22bec3fbe2524cde73d7ada88f6566758a8f7227bfbf93a408a9d86bcc12a0"},
-    {file = "PyYAML-6.0.1-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:8d4e9c88387b0f5c7d5f281e55304de64cf7f9c0021a3525bd3b1c542da3b0e4"},
-    {file = "PyYAML-6.0.1-cp312-cp312-win32.whl", hash = "sha256:d483d2cdf104e7c9fa60c544d92981f12ad66a457afae824d146093b8c294c54"},
-    {file = "PyYAML-6.0.1-cp312-cp312-win_amd64.whl", hash = "sha256:0d3304d8c0adc42be59c5f8a4d9e3d7379e6955ad754aa9d6ab7a398b59dd1df"},
-    {file = "PyYAML-6.0.1-cp36-cp36m-macosx_10_9_x86_64.whl", hash = "sha256:50550eb667afee136e9a77d6dc71ae76a44df8b3e51e41b77f6de2932bfe0f47"},
-    {file = "PyYAML-6.0.1-cp36-cp36m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1fe35611261b29bd1de0070f0b2f47cb6ff71fa6595c077e42bd0c419fa27b98"},
-    {file = "PyYAML-6.0.1-cp36-cp36m-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:704219a11b772aea0d8ecd7058d0082713c3562b4e271b849ad7dc4a5c90c13c"},
-    {file = "PyYAML-6.0.1-cp36-cp36m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:afd7e57eddb1a54f0f1a974bc4391af8bcce0b444685d936840f125cf046d5bd"},
-    {file = "PyYAML-6.0.1-cp36-cp36m-win32.whl", hash = "sha256:fca0e3a251908a499833aa292323f32437106001d436eca0e6e7833256674585"},
-    {file = "PyYAML-6.0.1-cp36-cp36m-win_amd64.whl", hash = "sha256:f22ac1c3cac4dbc50079e965eba2c1058622631e526bd9afd45fedd49ba781fa"},
-    {file = "PyYAML-6.0.1-cp37-cp37m-macosx_10_9_x86_64.whl", hash = "sha256:b1275ad35a5d18c62a7220633c913e1b42d44b46ee12554e5fd39c70a243d6a3"},
-    {file = "PyYAML-6.0.1-cp37-cp37m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:18aeb1bf9a78867dc38b259769503436b7c72f7a1f1f4c93ff9a17de54319b27"},
-    {file = "PyYAML-6.0.1-cp37-cp37m-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:596106435fa6ad000c2991a98fa58eeb8656ef2325d7e158344fb33864ed87e3"},
-    {file = "PyYAML-6.0.1-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:baa90d3f661d43131ca170712d903e6295d1f7a0f595074f151c0aed377c9b9c"},
-    {file = "PyYAML-6.0.1-cp37-cp37m-win32.whl", hash = "sha256:9046c58c4395dff28dd494285c82ba00b546adfc7ef001486fbf0324bc174fba"},
-    {file = "PyYAML-6.0.1-cp37-cp37m-win_amd64.whl", hash = "sha256:4fb147e7a67ef577a588a0e2c17b6db51dda102c71de36f8549b6816a96e1867"},
-    {file = "PyYAML-6.0.1-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:1d4c7e777c441b20e32f52bd377e0c409713e8bb1386e1099c2415f26e479595"},
-    {file = "PyYAML-6.0.1-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a0cd17c15d3bb3fa06978b4e8958dcdc6e0174ccea823003a106c7d4d7899ac5"},
-    {file = "PyYAML-6.0.1-cp38-cp38-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:28c119d996beec18c05208a8bd78cbe4007878c6dd15091efb73a30e90539696"},
-    {file = "PyYAML-6.0.1-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7e07cbde391ba96ab58e532ff4803f79c4129397514e1413a7dc761ccd755735"},
-    {file = "PyYAML-6.0.1-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:49a183be227561de579b4a36efbb21b3eab9651dd81b1858589f796549873dd6"},
-    {file = "PyYAML-6.0.1-cp38-cp38-win32.whl", hash = "sha256:184c5108a2aca3c5b3d3bf9395d50893a7ab82a38004c8f61c258d4428e80206"},
-    {file = "PyYAML-6.0.1-cp38-cp38-win_amd64.whl", hash = "sha256:1e2722cc9fbb45d9b87631ac70924c11d3a401b2d7f410cc0e3bbf249f2dca62"},
-    {file = "PyYAML-6.0.1-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:9eb6caa9a297fc2c2fb8862bc5370d0303ddba53ba97e71f08023b6cd73d16a8"},
-    {file = "PyYAML-6.0.1-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:c8098ddcc2a85b61647b2590f825f3db38891662cfc2fc776415143f599bb859"},
-    {file = "PyYAML-6.0.1-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:5773183b6446b2c99bb77e77595dd486303b4faab2b086e7b17bc6bef28865f6"},
-    {file = "PyYAML-6.0.1-cp39-cp39-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:b786eecbdf8499b9ca1d697215862083bd6d2a99965554781d0d8d1ad31e13a0"},
-    {file = "PyYAML-6.0.1-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:bc1bf2925a1ecd43da378f4db9e4f799775d6367bdb94671027b73b393a7c42c"},
-    {file = "PyYAML-6.0.1-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:04ac92ad1925b2cff1db0cfebffb6ffc43457495c9b3c39d3fcae417d7125dc5"},
-    {file = "PyYAML-6.0.1-cp39-cp39-win32.whl", hash = "sha256:faca3bdcf85b2fc05d06ff3fbc1f83e1391b3e724afa3feba7d13eeab355484c"},
-    {file = "PyYAML-6.0.1-cp39-cp39-win_amd64.whl", hash = "sha256:510c9deebc5c0225e8c96813043e62b680ba2f9c50a08d3724c7f28a747d1486"},
-    {file = "PyYAML-6.0.1.tar.gz", hash = "sha256:bfdf460b1736c775f2ba9f6a92bca30bc2095067b8a9d77876d1fad6cc3b4a43"},
+    {file = "PyYAML-6.0.2-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:0a9a2848a5b7feac301353437eb7d5957887edbf81d56e903999a75a3d743086"},
+    {file = "PyYAML-6.0.2-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:29717114e51c84ddfba879543fb232a6ed60086602313ca38cce623c1d62cfbf"},
+    {file = "PyYAML-6.0.2-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:8824b5a04a04a047e72eea5cec3bc266db09e35de6bdfe34c9436ac5ee27d237"},
+    {file = "PyYAML-6.0.2-cp310-cp310-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:7c36280e6fb8385e520936c3cb3b8042851904eba0e58d277dca80a5cfed590b"},
+    {file = "PyYAML-6.0.2-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ec031d5d2feb36d1d1a24380e4db6d43695f3748343d99434e6f5f9156aaa2ed"},
+    {file = "PyYAML-6.0.2-cp310-cp310-musllinux_1_1_aarch64.whl", hash = "sha256:936d68689298c36b53b29f23c6dbb74de12b4ac12ca6cfe0e047bedceea56180"},
+    {file = "PyYAML-6.0.2-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:23502f431948090f597378482b4812b0caae32c22213aecf3b55325e049a6c68"},
+    {file = "PyYAML-6.0.2-cp310-cp310-win32.whl", hash = "sha256:2e99c6826ffa974fe6e27cdb5ed0021786b03fc98e5ee3c5bfe1fd5015f42b99"},
+    {file = "PyYAML-6.0.2-cp310-cp310-win_amd64.whl", hash = "sha256:a4d3091415f010369ae4ed1fc6b79def9416358877534caf6a0fdd2146c87a3e"},
+    {file = "PyYAML-6.0.2-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:cc1c1159b3d456576af7a3e4d1ba7e6924cb39de8f67111c735f6fc832082774"},
+    {file = "PyYAML-6.0.2-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:1e2120ef853f59c7419231f3bf4e7021f1b936f6ebd222406c3b60212205d2ee"},
+    {file = "PyYAML-6.0.2-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:5d225db5a45f21e78dd9358e58a98702a0302f2659a3c6cd320564b75b86f47c"},
+    {file = "PyYAML-6.0.2-cp311-cp311-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:5ac9328ec4831237bec75defaf839f7d4564be1e6b25ac710bd1a96321cc8317"},
+    {file = "PyYAML-6.0.2-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:3ad2a3decf9aaba3d29c8f537ac4b243e36bef957511b4766cb0057d32b0be85"},
+    {file = "PyYAML-6.0.2-cp311-cp311-musllinux_1_1_aarch64.whl", hash = "sha256:ff3824dc5261f50c9b0dfb3be22b4567a6f938ccce4587b38952d85fd9e9afe4"},
+    {file = "PyYAML-6.0.2-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:797b4f722ffa07cc8d62053e4cff1486fa6dc094105d13fea7b1de7d8bf71c9e"},
+    {file = "PyYAML-6.0.2-cp311-cp311-win32.whl", hash = "sha256:11d8f3dd2b9c1207dcaf2ee0bbbfd5991f571186ec9cc78427ba5bd32afae4b5"},
+    {file = "PyYAML-6.0.2-cp311-cp311-win_amd64.whl", hash = "sha256:e10ce637b18caea04431ce14fabcf5c64a1c61ec9c56b071a4b7ca131ca52d44"},
+    {file = "PyYAML-6.0.2-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:c70c95198c015b85feafc136515252a261a84561b7b1d51e3384e0655ddf25ab"},
+    {file = "PyYAML-6.0.2-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:ce826d6ef20b1bc864f0a68340c8b3287705cae2f8b4b1d932177dcc76721725"},
+    {file = "PyYAML-6.0.2-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1f71ea527786de97d1a0cc0eacd1defc0985dcf6b3f17bb77dcfc8c34bec4dc5"},
+    {file = "PyYAML-6.0.2-cp312-cp312-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:9b22676e8097e9e22e36d6b7bda33190d0d400f345f23d4065d48f4ca7ae0425"},
+    {file = "PyYAML-6.0.2-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:80bab7bfc629882493af4aa31a4cfa43a4c57c83813253626916b8c7ada83476"},
+    {file = "PyYAML-6.0.2-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:0833f8694549e586547b576dcfaba4a6b55b9e96098b36cdc7ebefe667dfed48"},
+    {file = "PyYAML-6.0.2-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:8b9c7197f7cb2738065c481a0461e50ad02f18c78cd75775628afb4d7137fb3b"},
+    {file = "PyYAML-6.0.2-cp312-cp312-win32.whl", hash = "sha256:ef6107725bd54b262d6dedcc2af448a266975032bc85ef0172c5f059da6325b4"},
+    {file = "PyYAML-6.0.2-cp312-cp312-win_amd64.whl", hash = "sha256:7e7401d0de89a9a855c839bc697c079a4af81cf878373abd7dc625847d25cbd8"},
+    {file = "PyYAML-6.0.2-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:efdca5630322a10774e8e98e1af481aad470dd62c3170801852d752aa7a783ba"},
+    {file = "PyYAML-6.0.2-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:50187695423ffe49e2deacb8cd10510bc361faac997de9efef88badc3bb9e2d1"},
+    {file = "PyYAML-6.0.2-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:0ffe8360bab4910ef1b9e87fb812d8bc0a308b0d0eef8c8f44e0254ab3b07133"},
+    {file = "PyYAML-6.0.2-cp313-cp313-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:17e311b6c678207928d649faa7cb0d7b4c26a0ba73d41e99c4fff6b6c3276484"},
+    {file = "PyYAML-6.0.2-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:70b189594dbe54f75ab3a1acec5f1e3faa7e8cf2f1e08d9b561cb41b845f69d5"},
+    {file = "PyYAML-6.0.2-cp313-cp313-musllinux_1_1_aarch64.whl", hash = "sha256:41e4e3953a79407c794916fa277a82531dd93aad34e29c2a514c2c0c5fe971cc"},
+    {file = "PyYAML-6.0.2-cp313-cp313-musllinux_1_1_x86_64.whl", hash = "sha256:68ccc6023a3400877818152ad9a1033e3db8625d899c72eacb5a668902e4d652"},
+    {file = "PyYAML-6.0.2-cp313-cp313-win32.whl", hash = "sha256:bc2fa7c6b47d6bc618dd7fb02ef6fdedb1090ec036abab80d4681424b84c1183"},
+    {file = "PyYAML-6.0.2-cp313-cp313-win_amd64.whl", hash = "sha256:8388ee1976c416731879ac16da0aff3f63b286ffdd57cdeb95f3f2e085687563"},
+    {file = "PyYAML-6.0.2-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:24471b829b3bf607e04e88d79542a9d48bb037c2267d7927a874e6c205ca7e9a"},
+    {file = "PyYAML-6.0.2-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:d7fded462629cfa4b685c5416b949ebad6cec74af5e2d42905d41e257e0869f5"},
+    {file = "PyYAML-6.0.2-cp38-cp38-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:d84a1718ee396f54f3a086ea0a66d8e552b2ab2017ef8b420e92edbc841c352d"},
+    {file = "PyYAML-6.0.2-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9056c1ecd25795207ad294bcf39f2db3d845767be0ea6e6a34d856f006006083"},
+    {file = "PyYAML-6.0.2-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:82d09873e40955485746739bcb8b4586983670466c23382c19cffecbf1fd8706"},
+    {file = "PyYAML-6.0.2-cp38-cp38-win32.whl", hash = "sha256:43fa96a3ca0d6b1812e01ced1044a003533c47f6ee8aca31724f78e93ccc089a"},
+    {file = "PyYAML-6.0.2-cp38-cp38-win_amd64.whl", hash = "sha256:01179a4a8559ab5de078078f37e5c1a30d76bb88519906844fd7bdea1b7729ff"},
+    {file = "PyYAML-6.0.2-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:688ba32a1cffef67fd2e9398a2efebaea461578b0923624778664cc1c914db5d"},
+    {file = "PyYAML-6.0.2-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:a8786accb172bd8afb8be14490a16625cbc387036876ab6ba70912730faf8e1f"},
+    {file = "PyYAML-6.0.2-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:d8e03406cac8513435335dbab54c0d385e4a49e4945d2909a581c83647ca0290"},
+    {file = "PyYAML-6.0.2-cp39-cp39-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:f753120cb8181e736c57ef7636e83f31b9c0d1722c516f7e86cf15b7aa57ff12"},
+    {file = "PyYAML-6.0.2-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:3b1fdb9dc17f5a7677423d508ab4f243a726dea51fa5e70992e59a7411c89d19"},
+    {file = "PyYAML-6.0.2-cp39-cp39-musllinux_1_1_aarch64.whl", hash = "sha256:0b69e4ce7a131fe56b7e4d770c67429700908fc0752af059838b1cfb41960e4e"},
+    {file = "PyYAML-6.0.2-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:a9f8c2e67970f13b16084e04f134610fd1d374bf477b17ec1599185cf611d725"},
+    {file = "PyYAML-6.0.2-cp39-cp39-win32.whl", hash = "sha256:6395c297d42274772abc367baaa79683958044e5d3835486c16da75d2a694631"},
+    {file = "PyYAML-6.0.2-cp39-cp39-win_amd64.whl", hash = "sha256:39693e1f8320ae4f43943590b49779ffb98acb81f788220ea932a6b6c51004d8"},
+    {file = "pyyaml-6.0.2.tar.gz", hash = "sha256:d584d9ec91ad65861cc08d42e834324ef890a082e591037abe114850ff7bbc3e"},
 ]
 
 [[package]]
@@ -2890,6 +2912,58 @@ files = [
 [package.dependencies]
 mpmath = ">=0.19"
 
+[[package]]
+name = "testcontainers"
+version = "4.8.1"
+description = "Python library for throwaway instances of anything that can run in a Docker container"
+optional = false
+python-versions = "<4.0,>=3.9"
+files = [
+    {file = "testcontainers-4.8.1-py3-none-any.whl", hash = "sha256:d8ae43e8fe34060fcd5c3f494e0b7652b7774beabe94568a2283d0881e94d489"},
+    {file = "testcontainers-4.8.1.tar.gz", hash = "sha256:5ded4820b7227ad526857eb3caaafcabce1bbac05d22ad194849b136ffae3cb0"},
+]
+
+[package.dependencies]
+docker = "*"
+typing-extensions = "*"
+urllib3 = "*"
+wrapt = "*"
+
+[package.extras]
+arangodb = ["python-arango (>=7.8,<8.0)"]
+aws = ["boto3", "httpx"]
+azurite = ["azure-storage-blob (>=12.19,<13.0)"]
+chroma = ["chromadb-client"]
+clickhouse = ["clickhouse-driver"]
+cosmosdb = ["azure-cosmos"]
+db2 = ["ibm_db_sa", "sqlalchemy"]
+generic = ["httpx", "redis"]
+google = ["google-cloud-datastore (>=2)", "google-cloud-pubsub (>=2)"]
+influxdb = ["influxdb", "influxdb-client"]
+k3s = ["kubernetes", "pyyaml"]
+keycloak = ["python-keycloak"]
+localstack = ["boto3"]
+mailpit = ["cryptography"]
+minio = ["minio"]
+mongodb = ["pymongo"]
+mssql = ["pymssql", "sqlalchemy"]
+mysql = ["pymysql[rsa]", "sqlalchemy"]
+nats = ["nats-py"]
+neo4j = ["neo4j"]
+opensearch = ["opensearch-py"]
+oracle = ["oracledb", "sqlalchemy"]
+oracle-free = ["oracledb", "sqlalchemy"]
+qdrant = ["qdrant-client"]
+rabbitmq = ["pika"]
+redis = ["redis"]
+registry = ["bcrypt"]
+scylla = ["cassandra-driver (==3.29.1)"]
+selenium = ["selenium"]
+sftp = ["cryptography"]
+test-module-import = ["httpx"]
+trino = ["trino"]
+weaviate = ["weaviate-client (>=4.5.4,<5.0.0)"]
+
 [[package]]
 name = "toml"
 version = "0.10.2"
@@ -2912,6 +2986,20 @@ files = [
     {file = "tomli-2.0.1.tar.gz", hash = "sha256:de526c12914f0c550d15924c62d72abc48d6fe7364aa87328337a31007fe8a4f"},
 ]
 
+[[package]]
+name = "types-jwcrypto"
+version = "1.5.0.20240925"
+description = "Typing stubs for jwcrypto"
+optional = false
+python-versions = ">=3.8"
+files = [
+    {file = "types-jwcrypto-1.5.0.20240925.tar.gz", hash = "sha256:50e17b790378c96239344476c7bd13b52d0c7eeb6d16c2d53723e48cc6bbf4fe"},
+    {file = "types_jwcrypto-1.5.0.20240925-py3-none-any.whl", hash = "sha256:2d12a2d528240d326075e896aafec7056b9136bf3207fa6ccf3fcb8fbf9e11a1"},
+]
+
+[package.dependencies]
+cryptography = "*"
+
 [[package]]
 name = "types-psutil"
 version = "5.9.5.12"
@@ -2945,6 +3033,17 @@ files = [
     {file = "types_pytest_lazy_fixture-0.6.3.3-py3-none-any.whl", hash = "sha256:a56a55649147ff960ff79d4b2c781a4f769351abc1876873f3116d0bd0c96353"},
 ]
 
+[[package]]
+name = "types-pyyaml"
+version = "6.0.12.20240917"
+description = "Typing stubs for PyYAML"
+optional = false
+python-versions = ">=3.8"
+files = [
+    {file = "types-PyYAML-6.0.12.20240917.tar.gz", hash = "sha256:d1405a86f9576682234ef83bcb4e6fff7c9305c8b1fbad5e0bcd4f7dbdc9c587"},
+    {file = "types_PyYAML-6.0.12.20240917-py3-none-any.whl", hash = "sha256:392b267f1c0fe6022952462bf5d6523f31e37f6cea49b14cee7ad634b6301570"},
+]
+
 [[package]]
 name = "types-requests"
 version = "2.31.0.0"
@@ -3019,22 +3118,6 @@ brotli = ["brotli (==1.0.9)", "brotli (>=1.0.9)", "brotlicffi (>=0.8.0)", "brotl
 secure = ["certifi", "cryptography (>=1.3.4)", "idna (>=2.0.0)", "ipaddress", "pyOpenSSL (>=0.14)", "urllib3-secure-extra"]
 socks = ["PySocks (>=1.5.6,!=1.5.7,<2.0)"]
 
-[[package]]
-name = "websocket-client"
-version = "1.3.3"
-description = "WebSocket client for Python with low level API options"
-optional = false
-python-versions = ">=3.7"
-files = [
-    {file = "websocket-client-1.3.3.tar.gz", hash = "sha256:d58c5f284d6a9bf8379dab423259fe8f85b70d5fa5d2916d5791a84594b122b1"},
-    {file = "websocket_client-1.3.3-py3-none-any.whl", hash = "sha256:5d55652dc1d0b3c734f044337d929aaf83f4f9138816ec680c1aefefb4dc4877"},
-]
-
-[package.extras]
-docs = ["Sphinx (>=3.4)", "sphinx-rtd-theme (>=0.5)"]
-optional = ["python-socks", "wsaccel"]
-test = ["websockets"]
-
 [[package]]
 name = "websockets"
 version = "12.0"
@@ -3406,4 +3489,4 @@ cffi = ["cffi (>=1.11)"]
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.9"
-content-hash = "0f4804119f417edf8e1fbd6d715d2e8d70ad731334fa9570304a2203f83339cf"
+content-hash = "13bfc7479aacfe051abb92252b8ddc2e0c429f4607b2d9d8c4b353d2f75c1927"

proxy/src/serverless/error.rs

@@ -0,0 +1,5 @@
+use http::StatusCode;
+
+pub trait HttpCodeError {
+    fn get_http_status_code(&self) -> StatusCode;
+}

proxy/src/serverless/http_conn_pool.rs

@@ -294,6 +294,11 @@ pub(crate) fn poll_http2_client(
                 conn_id,
                 aux: aux.clone(),
             });
+            Metrics::get()
+                .proxy
+                .http_pool_opened_connections
+                .get_metric()
+                .inc();
 
             Arc::downgrade(&pool)
         }
@@ -306,7 +311,7 @@ pub(crate) fn poll_http2_client(
             let res = connection.await;
             match res {
                 Ok(()) => info!("connection closed"),
-                Err(e) => error!(%session_id, "connection error: {}", e),
+                Err(e) => error!(%session_id, "connection error: {e:?}"),
             }
 
             // remove from connection pool

proxy/src/serverless/mod.rs

@@ -6,6 +6,7 @@ mod backend;
 pub mod cancel_set;
 mod conn_pool;
 mod conn_pool_lib;
+mod error;
 mod http_conn_pool;
 mod http_util;
 mod json;

proxy/src/serverless/sql_over_http.rs

@@ -28,6 +28,7 @@ use uuid::Uuid;
 use super::backend::{LocalProxyConnError, PoolingBackend};
 use super::conn_pool::{AuthData, ConnInfoWithAuth};
 use super::conn_pool_lib::{self, ConnInfo};
+use super::error::HttpCodeError;
 use super::http_util::json_response;
 use super::json::{json_to_pg_text, pg_text_row_to_json, JsonConversionError};
 use super::local_conn_pool;
@@ -238,7 +239,6 @@ fn get_conn_info(
     Ok(ConnInfoWithAuth { conn_info, auth })
 }
 
-// TODO: return different http error codes
 pub(crate) async fn handle(
     config: &'static ProxyConfig,
     ctx: RequestMonitoring,
@@ -319,9 +319,8 @@ pub(crate) async fn handle(
                 "forwarding error to user"
             );
 
-            // TODO: this shouldn't always be bad request.
             json_response(
-                StatusCode::BAD_REQUEST,
+                e.get_http_status_code(),
                 json!({
                     "message": message,
                     "code": code,
@@ -405,6 +404,25 @@ impl UserFacingError for SqlOverHttpError {
     }
 }
 
+impl HttpCodeError for SqlOverHttpError {
+    fn get_http_status_code(&self) -> StatusCode {
+        match self {
+            SqlOverHttpError::ReadPayload(_) => StatusCode::BAD_REQUEST,
+            SqlOverHttpError::ConnectCompute(h) => match h.get_error_kind() {
+                ErrorKind::User => StatusCode::BAD_REQUEST,
+                _ => StatusCode::INTERNAL_SERVER_ERROR,
+            },
+            SqlOverHttpError::ConnInfo(_) => StatusCode::BAD_REQUEST,
+            SqlOverHttpError::RequestTooLarge(_) => StatusCode::PAYLOAD_TOO_LARGE,
+            SqlOverHttpError::ResponseTooLarge(_) => StatusCode::INSUFFICIENT_STORAGE,
+            SqlOverHttpError::InvalidIsolationLevel => StatusCode::BAD_REQUEST,
+            SqlOverHttpError::Postgres(_) => StatusCode::BAD_REQUEST,
+            SqlOverHttpError::JsonConversion(_) => StatusCode::INTERNAL_SERVER_ERROR,
+            SqlOverHttpError::Cancelled(_) => StatusCode::INTERNAL_SERVER_ERROR,
+        }
+    }
+}
+
 #[derive(Debug, thiserror::Error)]
 pub(crate) enum ReadPayloadError {
     #[error("could not read the HTTP request body: {0}")]

pyproject.toml

@@ -42,6 +42,13 @@ pytest-repeat = "^0.9.3"
 websockets = "^12.0"
 clickhouse-connect = "^0.7.16"
 kafka-python = "^2.0.2"
+jwcrypto = "^1.5.6"
+h2 = "^4.1.0"
+types-jwcrypto = "^1.5.0.20240925"
+pyyaml = "^6.0.2"
+types-pyyaml = "^6.0.12.20240917"
+testcontainers = "^4.8.1"
+jsonnet = "^0.20.0"
 
 [tool.poetry.group.dev.dependencies]
 mypy = "==1.3.0"
@@ -70,12 +77,14 @@ strict = true
 
 [[tool.mypy.overrides]]
 module = [
+    "_jsonnet.*",
     "asyncpg.*",
     "pg8000.*",
     "allure.*",
     "allure_commons.*",
     "allure_pytest.*",
     "kafka.*",
+    "testcontainers.*",
 ]
 ignore_missing_imports = true
 

safekeeper/src/json_ctrl.rs

@@ -7,7 +7,6 @@
 //!
 
 use anyhow::Context;
-use bytes::Bytes;
 use postgres_backend::QueryError;
 use serde::{Deserialize, Serialize};
 use tokio::io::{AsyncRead, AsyncWrite};
@@ -176,7 +175,7 @@ pub async fn append_logical_message(
             truncate_lsn: msg.truncate_lsn,
             proposer_uuid: [0u8; 16],
         },
-        wal_data: Bytes::from(wal_data),
+        wal_data,
     });
 
     let response = tli.process_msg(&append_request).await?;

safekeeper/tests/walproposer_sim/simulation.rs

@@ -151,8 +151,7 @@ impl WalProposer {
         for _ in 0..cnt {
             self.disk
                 .lock()
-                .insert_logical_message("prefix", b"message")
-                .expect("failed to generate logical message");
+                .insert_logical_message(c"prefix", b"message");
         }
 
         let end_lsn = self.disk.lock().flush_rec_ptr();

safekeeper/tests/walproposer_sim/walproposer_disk.rs

@@ -1,24 +1,7 @@
-use std::{ffi::CString, sync::Arc};
+use std::{ffi::CStr, sync::Arc};
 
-use byteorder::{LittleEndian, WriteBytesExt};
-use crc32c::crc32c_append;
 use parking_lot::{Mutex, MutexGuard};
-use postgres_ffi::{
-    pg_constants::{
-        RM_LOGICALMSG_ID, XLOG_LOGICAL_MESSAGE, XLP_LONG_HEADER, XLR_BLOCK_ID_DATA_LONG,
-        XLR_BLOCK_ID_DATA_SHORT,
-    },
-    v16::{
-        wal_craft_test_export::{XLogLongPageHeaderData, XLogPageHeaderData, XLOG_PAGE_MAGIC},
-        xlog_utils::{
-            XLogSegNoOffsetToRecPtr, XlLogicalMessage, XLOG_RECORD_CRC_OFFS,
-            XLOG_SIZE_OF_XLOG_LONG_PHD, XLOG_SIZE_OF_XLOG_RECORD, XLOG_SIZE_OF_XLOG_SHORT_PHD,
-            XLP_FIRST_IS_CONTRECORD,
-        },
-        XLogRecord,
-    },
-    WAL_SEGMENT_SIZE, XLOG_BLCKSZ,
-};
+use postgres_ffi::v16::wal_generator::WalGenerator;
 use utils::lsn::Lsn;
 
 use super::block_storage::BlockStorage;
@@ -35,6 +18,7 @@ impl DiskWalProposer {
                 internal_available_lsn: Lsn(0),
                 prev_lsn: Lsn(0),
                 disk: BlockStorage::new(),
+                wal_generator: WalGenerator::new(),
             }),
         })
     }
@@ -51,6 +35,8 @@ pub struct State {
     prev_lsn: Lsn,
     // actual WAL storage
     disk: BlockStorage,
+    // WAL record generator
+    wal_generator: WalGenerator,
 }
 
 impl State {
@@ -66,6 +52,9 @@ impl State {
     /// Update the internal available LSN to the given value.
     pub fn reset_to(&mut self, lsn: Lsn) {
         self.internal_available_lsn = lsn;
+        self.prev_lsn = Lsn(0); // Safekeeper doesn't care if this is omitted
+        self.wal_generator.lsn = self.internal_available_lsn;
+        self.wal_generator.prev_lsn = self.prev_lsn;
     }
 
     /// Get current LSN.
@@ -73,242 +62,11 @@ impl State {
         self.internal_available_lsn
     }
 
-    /// Generate a new WAL record at the current LSN.
-    pub fn insert_logical_message(&mut self, prefix: &str, msg: &[u8]) -> anyhow::Result<()> {
-        let prefix_cstr = CString::new(prefix)?;
-        let prefix_bytes = prefix_cstr.as_bytes_with_nul();
-
-        let lm = XlLogicalMessage {
-            db_id: 0,
-            transactional: 0,
-            prefix_size: prefix_bytes.len() as ::std::os::raw::c_ulong,
-            message_size: msg.len() as ::std::os::raw::c_ulong,
-        };
-
-        let record_bytes = lm.encode();
-        let rdatas: Vec<&[u8]> = vec![&record_bytes, prefix_bytes, msg];
-        insert_wal_record(self, rdatas, RM_LOGICALMSG_ID, XLOG_LOGICAL_MESSAGE)
-    }
-}
-
-fn insert_wal_record(
-    state: &mut State,
-    rdatas: Vec<&[u8]>,
-    rmid: u8,
-    info: u8,
-) -> anyhow::Result<()> {
-    // bytes right after the header, in the same rdata block
-    let mut scratch = Vec::new();
-    let mainrdata_len: usize = rdatas.iter().map(|rdata| rdata.len()).sum();
-
-    if mainrdata_len > 0 {
-        if mainrdata_len > 255 {
-            scratch.push(XLR_BLOCK_ID_DATA_LONG);
-            // TODO: verify endiness
-            let _ = scratch.write_u32::<LittleEndian>(mainrdata_len as u32);
-        } else {
-            scratch.push(XLR_BLOCK_ID_DATA_SHORT);
-            scratch.push(mainrdata_len as u8);
-        }
-    }
-
-    let total_len: u32 = (XLOG_SIZE_OF_XLOG_RECORD + scratch.len() + mainrdata_len) as u32;
-    let size = maxalign(total_len);
-    assert!(size as usize > XLOG_SIZE_OF_XLOG_RECORD);
-
-    let start_bytepos = recptr_to_bytepos(state.internal_available_lsn);
-    let end_bytepos = start_bytepos + size as u64;
-
-    let start_recptr = bytepos_to_recptr(start_bytepos);
-    let end_recptr = bytepos_to_recptr(end_bytepos);
-
-    assert!(recptr_to_bytepos(start_recptr) == start_bytepos);
-    assert!(recptr_to_bytepos(end_recptr) == end_bytepos);
-
-    let mut crc = crc32c_append(0, &scratch);
-    for rdata in &rdatas {
-        crc = crc32c_append(crc, rdata);
-    }
-
-    let mut header = XLogRecord {
-        xl_tot_len: total_len,
-        xl_xid: 0,
-        xl_prev: state.prev_lsn.0,
-        xl_info: info,
-        xl_rmid: rmid,
-        __bindgen_padding_0: [0u8; 2usize],
-        xl_crc: crc,
-    };
-
-    // now we have the header and can finish the crc
-    let header_bytes = header.encode()?;
-    let crc = crc32c_append(crc, &header_bytes[0..XLOG_RECORD_CRC_OFFS]);
-    header.xl_crc = crc;
-
-    let mut header_bytes = header.encode()?.to_vec();
-    assert!(header_bytes.len() == XLOG_SIZE_OF_XLOG_RECORD);
-
-    header_bytes.extend_from_slice(&scratch);
-
-    // finish rdatas
-    let mut rdatas = rdatas;
-    rdatas.insert(0, &header_bytes);
-
-    write_walrecord_to_disk(state, total_len as u64, rdatas, start_recptr, end_recptr)?;
-
-    state.internal_available_lsn = end_recptr;
-    state.prev_lsn = start_recptr;
-    Ok(())
-}
-
-fn write_walrecord_to_disk(
-    state: &mut State,
-    total_len: u64,
-    rdatas: Vec<&[u8]>,
-    start: Lsn,
-    end: Lsn,
-) -> anyhow::Result<()> {
-    let mut curr_ptr = start;
-    let mut freespace = insert_freespace(curr_ptr);
-    let mut written: usize = 0;
-
-    assert!(freespace >= size_of::<u32>());
-
-    for mut rdata in rdatas {
-        while rdata.len() >= freespace {
-            assert!(
-                curr_ptr.segment_offset(WAL_SEGMENT_SIZE) >= XLOG_SIZE_OF_XLOG_SHORT_PHD
-                    || freespace == 0
-            );
-
-            state.write(curr_ptr.0, &rdata[..freespace]);
-            rdata = &rdata[freespace..];
-            written += freespace;
-            curr_ptr = Lsn(curr_ptr.0 + freespace as u64);
-
-            let mut new_page = XLogPageHeaderData {
-                xlp_magic: XLOG_PAGE_MAGIC as u16,
-                xlp_info: XLP_BKP_REMOVABLE,
-                xlp_tli: 1,
-                xlp_pageaddr: curr_ptr.0,
-                xlp_rem_len: (total_len - written as u64) as u32,
-                ..Default::default() // Put 0 in padding fields.
-            };
-            if new_page.xlp_rem_len > 0 {
-                new_page.xlp_info |= XLP_FIRST_IS_CONTRECORD;
-            }
-
-            if curr_ptr.segment_offset(WAL_SEGMENT_SIZE) == 0 {
-                new_page.xlp_info |= XLP_LONG_HEADER;
-                let long_page = XLogLongPageHeaderData {
-                    std: new_page,
-                    xlp_sysid: 0,
-                    xlp_seg_size: WAL_SEGMENT_SIZE as u32,
-                    xlp_xlog_blcksz: XLOG_BLCKSZ as u32,
-                };
-                let header_bytes = long_page.encode()?;
-                assert!(header_bytes.len() == XLOG_SIZE_OF_XLOG_LONG_PHD);
-                state.write(curr_ptr.0, &header_bytes);
-                curr_ptr = Lsn(curr_ptr.0 + header_bytes.len() as u64);
-            } else {
-                let header_bytes = new_page.encode()?;
-                assert!(header_bytes.len() == XLOG_SIZE_OF_XLOG_SHORT_PHD);
-                state.write(curr_ptr.0, &header_bytes);
-                curr_ptr = Lsn(curr_ptr.0 + header_bytes.len() as u64);
-            }
-            freespace = insert_freespace(curr_ptr);
-        }
-
-        assert!(
-            curr_ptr.segment_offset(WAL_SEGMENT_SIZE) >= XLOG_SIZE_OF_XLOG_SHORT_PHD
-                || rdata.is_empty()
-        );
-        state.write(curr_ptr.0, rdata);
-        curr_ptr = Lsn(curr_ptr.0 + rdata.len() as u64);
-        written += rdata.len();
-        freespace -= rdata.len();
-    }
-
-    assert!(written == total_len as usize);
-    curr_ptr.0 = maxalign(curr_ptr.0);
-    assert!(curr_ptr == end);
-    Ok(())
-}
-
-fn maxalign<T>(size: T) -> T
-where
-    T: std::ops::BitAnd<Output = T>
-        + std::ops::Add<Output = T>
-        + std::ops::Not<Output = T>
-        + From<u8>,
-{
-    (size + T::from(7)) & !T::from(7)
-}
-
-fn insert_freespace(ptr: Lsn) -> usize {
-    if ptr.block_offset() == 0 {
-        0
-    } else {
-        (XLOG_BLCKSZ as u64 - ptr.block_offset()) as usize
-    }
-}
-
-const XLP_BKP_REMOVABLE: u16 = 0x0004;
-const USABLE_BYTES_IN_PAGE: u64 = (XLOG_BLCKSZ - XLOG_SIZE_OF_XLOG_SHORT_PHD) as u64;
-const USABLE_BYTES_IN_SEGMENT: u64 = ((WAL_SEGMENT_SIZE / XLOG_BLCKSZ) as u64
-    * USABLE_BYTES_IN_PAGE)
-    - (XLOG_SIZE_OF_XLOG_RECORD - XLOG_SIZE_OF_XLOG_SHORT_PHD) as u64;
-
-fn bytepos_to_recptr(bytepos: u64) -> Lsn {
-    let fullsegs = bytepos / USABLE_BYTES_IN_SEGMENT;
-    let mut bytesleft = bytepos % USABLE_BYTES_IN_SEGMENT;
-
-    let seg_offset = if bytesleft < (XLOG_BLCKSZ - XLOG_SIZE_OF_XLOG_SHORT_PHD) as u64 {
-        // fits on first page of segment
-        bytesleft + XLOG_SIZE_OF_XLOG_SHORT_PHD as u64
-    } else {
-        // account for the first page on segment with long header
-        bytesleft -= (XLOG_BLCKSZ - XLOG_SIZE_OF_XLOG_SHORT_PHD) as u64;
-        let fullpages = bytesleft / USABLE_BYTES_IN_PAGE;
-        bytesleft %= USABLE_BYTES_IN_PAGE;
-
-        XLOG_BLCKSZ as u64
-            + fullpages * XLOG_BLCKSZ as u64
-            + bytesleft
-            + XLOG_SIZE_OF_XLOG_SHORT_PHD as u64
-    };
-
-    Lsn(XLogSegNoOffsetToRecPtr(
-        fullsegs,
-        seg_offset as u32,
-        WAL_SEGMENT_SIZE,
-    ))
-}
-
-fn recptr_to_bytepos(ptr: Lsn) -> u64 {
-    let fullsegs = ptr.segment_number(WAL_SEGMENT_SIZE);
-    let offset = ptr.segment_offset(WAL_SEGMENT_SIZE) as u64;
-
-    let fullpages = offset / XLOG_BLCKSZ as u64;
-    let offset = offset % XLOG_BLCKSZ as u64;
-
-    if fullpages == 0 {
-        fullsegs * USABLE_BYTES_IN_SEGMENT
-            + if offset > 0 {
-                assert!(offset >= XLOG_SIZE_OF_XLOG_SHORT_PHD as u64);
-                offset - XLOG_SIZE_OF_XLOG_SHORT_PHD as u64
-            } else {
-                0
-            }
-    } else {
-        fullsegs * USABLE_BYTES_IN_SEGMENT
-            + (XLOG_BLCKSZ - XLOG_SIZE_OF_XLOG_SHORT_PHD) as u64
-            + (fullpages - 1) * USABLE_BYTES_IN_PAGE
-            + if offset > 0 {
-                assert!(offset >= XLOG_SIZE_OF_XLOG_SHORT_PHD as u64);
-                offset - XLOG_SIZE_OF_XLOG_SHORT_PHD as u64
-            } else {
-                0
-            }
+    /// Inserts a logical record in the WAL at the current LSN.
+    pub fn insert_logical_message(&mut self, prefix: &CStr, msg: &[u8]) {
+        let record = self.wal_generator.generate_logical_message(prefix, msg);
+        self.disk.write(self.internal_available_lsn.into(), &record);
+        self.prev_lsn = self.internal_available_lsn;
+        self.internal_available_lsn += record.len() as u64;
     }
 }

test_runner/conftest.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 pytest_plugins = (
     "fixtures.pg_version",
     "fixtures.parametrize",
+    "fixtures.h2server",
     "fixtures.httpserver",
     "fixtures.compute_reconfigure",
     "fixtures.storage_controller_proxy",

test_runner/fixtures/h2server.py

@@ -0,0 +1,198 @@
+"""
+https://python-hyper.org/projects/hyper-h2/en/stable/asyncio-example.html
+
+auth-broker -> local-proxy needs a h2 connection, so we need a h2 server :)
+"""
+
+import asyncio
+import collections
+import io
+import json
+from collections.abc import AsyncIterable
+
+import pytest_asyncio
+from h2.config import H2Configuration
+from h2.connection import H2Connection
+from h2.errors import ErrorCodes
+from h2.events import (
+    ConnectionTerminated,
+    DataReceived,
+    RemoteSettingsChanged,
+    RequestReceived,
+    StreamEnded,
+    StreamReset,
+    WindowUpdated,
+)
+from h2.exceptions import ProtocolError, StreamClosedError
+from h2.settings import SettingCodes
+
+RequestData = collections.namedtuple("RequestData", ["headers", "data"])
+
+
+class H2Server:
+    def __init__(self, host, port) -> None:
+        self.host = host
+        self.port = port
+
+
+class H2Protocol(asyncio.Protocol):
+    def __init__(self):
+        config = H2Configuration(client_side=False, header_encoding="utf-8")
+        self.conn = H2Connection(config=config)
+        self.transport = None
+        self.stream_data = {}
+        self.flow_control_futures = {}
+
+    def connection_made(self, transport: asyncio.Transport):  # type: ignore[override]
+        self.transport = transport
+        self.conn.initiate_connection()
+        self.transport.write(self.conn.data_to_send())
+
+    def connection_lost(self, _exc):
+        for future in self.flow_control_futures.values():
+            future.cancel()
+        self.flow_control_futures = {}
+
+    def data_received(self, data: bytes):
+        assert self.transport is not None
+        try:
+            events = self.conn.receive_data(data)
+        except ProtocolError:
+            self.transport.write(self.conn.data_to_send())
+            self.transport.close()
+        else:
+            self.transport.write(self.conn.data_to_send())
+            for event in events:
+                if isinstance(event, RequestReceived):
+                    self.request_received(event.headers, event.stream_id)
+                elif isinstance(event, DataReceived):
+                    self.receive_data(event.data, event.stream_id)
+                elif isinstance(event, StreamEnded):
+                    self.stream_complete(event.stream_id)
+                elif isinstance(event, ConnectionTerminated):
+                    self.transport.close()
+                elif isinstance(event, StreamReset):
+                    self.stream_reset(event.stream_id)
+                elif isinstance(event, WindowUpdated):
+                    self.window_updated(event.stream_id, event.delta)
+                elif isinstance(event, RemoteSettingsChanged):
+                    if SettingCodes.INITIAL_WINDOW_SIZE in event.changed_settings:
+                        self.window_updated(None, 0)
+
+                self.transport.write(self.conn.data_to_send())
+
+    def request_received(self, headers: list[tuple[str, str]], stream_id: int):
+        headers_map = collections.OrderedDict(headers)
+
+        # Store off the request data.
+        request_data = RequestData(headers_map, io.BytesIO())
+        self.stream_data[stream_id] = request_data
+
+    def stream_complete(self, stream_id: int):
+        """
+        When a stream is complete, we can send our response.
+        """
+        try:
+            request_data = self.stream_data[stream_id]
+        except KeyError:
+            # Just return, we probably 405'd this already
+            return
+
+        headers = request_data.headers
+        body = request_data.data.getvalue().decode("utf-8")
+
+        data = json.dumps({"headers": headers, "body": body}, indent=4).encode("utf8")
+
+        response_headers = (
+            (":status", "200"),
+            ("content-type", "application/json"),
+            ("content-length", str(len(data))),
+        )
+        self.conn.send_headers(stream_id, response_headers)
+        asyncio.ensure_future(self.send_data(data, stream_id))
+
+    def receive_data(self, data: bytes, stream_id: int):
+        """
+        We've received some data on a stream. If that stream is one we're
+        expecting data on, save it off. Otherwise, reset the stream.
+        """
+        try:
+            stream_data = self.stream_data[stream_id]
+        except KeyError:
+            self.conn.reset_stream(stream_id, error_code=ErrorCodes.PROTOCOL_ERROR)
+        else:
+            stream_data.data.write(data)
+
+    def stream_reset(self, stream_id):
+        """
+        A stream reset was sent. Stop sending data.
+        """
+        if stream_id in self.flow_control_futures:
+            future = self.flow_control_futures.pop(stream_id)
+            future.cancel()
+
+    async def send_data(self, data, stream_id):
+        """
+        Send data according to the flow control rules.
+        """
+        while data:
+            while self.conn.local_flow_control_window(stream_id) < 1:
+                try:
+                    await self.wait_for_flow_control(stream_id)
+                except asyncio.CancelledError:
+                    return
+
+            chunk_size = min(
+                self.conn.local_flow_control_window(stream_id),
+                len(data),
+                self.conn.max_outbound_frame_size,
+            )
+
+            try:
+                self.conn.send_data(
+                    stream_id, data[:chunk_size], end_stream=(chunk_size == len(data))
+                )
+            except (StreamClosedError, ProtocolError):
+                # The stream got closed and we didn't get told. We're done
+                # here.
+                break
+
+            assert self.transport is not None
+            self.transport.write(self.conn.data_to_send())
+            data = data[chunk_size:]
+
+    async def wait_for_flow_control(self, stream_id):
+        """
+        Waits for a Future that fires when the flow control window is opened.
+        """
+        f: asyncio.Future[None] = asyncio.Future()
+        self.flow_control_futures[stream_id] = f
+        await f
+
+    def window_updated(self, stream_id, delta):
+        """
+        A window update frame was received. Unblock some number of flow control
+        Futures.
+        """
+        if stream_id and stream_id in self.flow_control_futures:
+            f = self.flow_control_futures.pop(stream_id)
+            f.set_result(delta)
+        elif not stream_id:
+            for f in self.flow_control_futures.values():
+                f.set_result(delta)
+
+            self.flow_control_futures = {}
+
+
+@pytest_asyncio.fixture(scope="function")
+async def http2_echoserver() -> AsyncIterable[H2Server]:
+    loop = asyncio.get_event_loop()
+    serve = await loop.create_server(H2Protocol, "127.0.0.1", 0)
+    (host, port) = serve.sockets[0].getsockname()
+
+    asyncio.create_task(serve.wait_closed())
+
+    server = H2Server(host, port)
+    yield server
+
+    serve.close()

test_runner/fixtures/neon_fixtures.py

@@ -35,11 +35,13 @@ import toml
 from _pytest.config import Config
 from _pytest.config.argparsing import Parser
 from _pytest.fixtures import FixtureRequest
+from jwcrypto import jwk
 
 # Type-related stuff
 from psycopg2.extensions import connection as PgConnection
 from psycopg2.extensions import cursor as PgCursor
 from psycopg2.extensions import make_dsn, parse_dsn
+from pytest_httpserver import HTTPServer
 from urllib3.util.retry import Retry
 
 from fixtures import overlayfs
@@ -53,6 +55,7 @@ from fixtures.common_types import (
     TimelineId,
 )
 from fixtures.endpoint.http import EndpointHttpClient
+from fixtures.h2server import H2Server
 from fixtures.log_helper import log
 from fixtures.metrics import Metrics, MetricsGetter, parse_metrics
 from fixtures.neon_cli import NeonLocalCli, Pagectl
@@ -3080,6 +3083,31 @@ class PSQL:
         )
 
 
+def generate_proxy_tls_certs(common_name: str, key_path: Path, crt_path: Path):
+    if not key_path.exists():
+        r = subprocess.run(
+            [
+                "openssl",
+                "req",
+                "-new",
+                "-x509",
+                "-days",
+                "365",
+                "-nodes",
+                "-text",
+                "-out",
+                str(crt_path),
+                "-keyout",
+                str(key_path),
+                "-subj",
+                f"/CN={common_name}",
+                "-addext",
+                f"subjectAltName = DNS:{common_name}",
+            ]
+        )
+        assert r.returncode == 0
+
+
 class NeonProxy(PgProtocol):
     link_auth_uri: str = "http://dummy-uri"
 
@@ -3178,29 +3206,7 @@ class NeonProxy(PgProtocol):
         # generate key of it doesn't exist
         crt_path = self.test_output_dir / "proxy.crt"
         key_path = self.test_output_dir / "proxy.key"
-
-        if not key_path.exists():
-            r = subprocess.run(
-                [
-                    "openssl",
-                    "req",
-                    "-new",
-                    "-x509",
-                    "-days",
-                    "365",
-                    "-nodes",
-                    "-text",
-                    "-out",
-                    str(crt_path),
-                    "-keyout",
-                    str(key_path),
-                    "-subj",
-                    "/CN=*.localtest.me",
-                    "-addext",
-                    "subjectAltName = DNS:*.localtest.me",
-                ]
-            )
-            assert r.returncode == 0
+        generate_proxy_tls_certs("*.localtest.me", key_path, crt_path)
 
         args = [
             str(self.neon_binpath / "proxy"),
@@ -3380,6 +3386,125 @@ class NeonProxy(PgProtocol):
         assert out == "ok"
 
 
+class NeonAuthBroker:
+    class ControlPlane:
+        def __init__(self, endpoint: str):
+            self.endpoint = endpoint
+
+        def extra_args(self) -> list[str]:
+            args = [
+                *["--auth-backend", "console"],
+                *["--auth-endpoint", self.endpoint],
+            ]
+            return args
+
+    def __init__(
+        self,
+        neon_binpath: Path,
+        test_output_dir: Path,
+        http_port: int,
+        mgmt_port: int,
+        external_http_port: int,
+        auth_backend: NeonAuthBroker.ControlPlane,
+    ):
+        self.domain = "apiauth.localtest.me"  # resolves to 127.0.0.1
+        self.host = "127.0.0.1"
+        self.http_port = http_port
+        self.external_http_port = external_http_port
+        self.neon_binpath = neon_binpath
+        self.test_output_dir = test_output_dir
+        self.mgmt_port = mgmt_port
+        self.auth_backend = auth_backend
+        self.http_timeout_seconds = 15
+        self._popen: Optional[subprocess.Popen[bytes]] = None
+
+    def start(self) -> NeonAuthBroker:
+        assert self._popen is None
+
+        # generate key of it doesn't exist
+        crt_path = self.test_output_dir / "proxy.crt"
+        key_path = self.test_output_dir / "proxy.key"
+        generate_proxy_tls_certs("apiauth.localtest.me", key_path, crt_path)
+
+        args = [
+            str(self.neon_binpath / "proxy"),
+            *["--http", f"{self.host}:{self.http_port}"],
+            *["--mgmt", f"{self.host}:{self.mgmt_port}"],
+            *["--wss", f"{self.host}:{self.external_http_port}"],
+            *["-c", str(crt_path)],
+            *["-k", str(key_path)],
+            *["--sql-over-http-pool-opt-in", "false"],
+            *["--is-auth-broker", "true"],
+            *self.auth_backend.extra_args(),
+        ]
+
+        logfile = open(self.test_output_dir / "proxy.log", "w")
+        self._popen = subprocess.Popen(args, stdout=logfile, stderr=logfile)
+        self._wait_until_ready()
+        return self
+
+    # Sends SIGTERM to the proxy if it has been started
+    def terminate(self):
+        if self._popen:
+            self._popen.terminate()
+
+    # Waits for proxy to exit if it has been opened with a default timeout of
+    # two seconds. Raises subprocess.TimeoutExpired if the proxy does not exit in time.
+    def wait_for_exit(self, timeout=2):
+        if self._popen:
+            self._popen.wait(timeout=timeout)
+
+    @backoff.on_exception(backoff.expo, requests.exceptions.RequestException, max_time=10)
+    def _wait_until_ready(self):
+        assert (
+            self._popen and self._popen.poll() is None
+        ), "Proxy exited unexpectedly. Check test log."
+        requests.get(f"http://{self.host}:{self.http_port}/v1/status")
+
+    async def query(self, query, args, **kwargs):
+        user = kwargs["user"]
+        token = kwargs["token"]
+        expected_code = kwargs.get("expected_code")
+
+        log.info(f"Executing http query: {query}")
+
+        connstr = f"postgresql://{user}@{self.domain}/postgres"
+        async with httpx.AsyncClient(verify=str(self.test_output_dir / "proxy.crt")) as client:
+            response = await client.post(
+                f"https://{self.domain}:{self.external_http_port}/sql",
+                json={"query": query, "params": args},
+                headers={
+                    "Neon-Connection-String": connstr,
+                    "Authorization": f"Bearer {token}",
+                },
+            )
+
+            if expected_code is not None:
+                assert response.status_code == expected_code, f"response: {response.json()}"
+            return response.json()
+
+    def get_metrics(self) -> str:
+        request_result = requests.get(f"http://{self.host}:{self.http_port}/metrics")
+        return request_result.text
+
+    def __enter__(self) -> NeonAuthBroker:
+        return self
+
+    def __exit__(
+        self,
+        _exc_type: Optional[type[BaseException]],
+        _exc_value: Optional[BaseException],
+        _traceback: Optional[TracebackType],
+    ):
+        if self._popen is not None:
+            self._popen.terminate()
+            try:
+                self._popen.wait(timeout=5)
+            except subprocess.TimeoutExpired:
+                log.warning("failed to gracefully terminate proxy; killing")
+                self._popen.kill()
+
+
 @pytest.fixture(scope="function")
 def link_proxy(
     port_distributor: PortDistributor, neon_binpath: Path, test_output_dir: Path
@@ -3444,6 +3569,74 @@ def static_proxy(
         yield proxy
 
 
+@pytest.fixture(scope="function")
+def neon_authorize_jwk() -> jwk.JWK:
+    kid = str(uuid.uuid4())
+    key = jwk.JWK.generate(kty="RSA", size=2048, alg="RS256", use="sig", kid=kid)
+    assert isinstance(key, jwk.JWK)
+    return key
+
+
+@pytest.fixture(scope="function")
+def static_auth_broker(
+    port_distributor: PortDistributor,
+    neon_binpath: Path,
+    test_output_dir: Path,
+    httpserver: HTTPServer,
+    neon_authorize_jwk: jwk.JWK,
+    http2_echoserver: H2Server,
+) -> Iterable[NeonAuthBroker]:
+    """Neon Auth Broker that routes to a mocked local_proxy and a mocked cplane HTTP API."""
+
+    local_proxy_addr = f"{http2_echoserver.host}:{http2_echoserver.port}"
+
+    # return local_proxy addr on ProxyWakeCompute.
+    httpserver.expect_request("/cplane/proxy_wake_compute").respond_with_json(
+        {
+            "address": local_proxy_addr,
+            "aux": {
+                "endpoint_id": "ep-foo-bar-1234",
+                "branch_id": "br-foo-bar",
+                "project_id": "foo-bar",
+            },
+        }
+    )
+
+    # return jwks mock addr on GetEndpointJwks
+    httpserver.expect_request(re.compile("^/cplane/endpoints/.+/jwks$")).respond_with_json(
+        {
+            "jwks": [
+                {
+                    "id": "foo",
+                    "jwks_url": httpserver.url_for("/authorize/jwks.json"),
+                    "provider_name": "test",
+                    "jwt_audience": None,
+                    "role_names": ["anonymous", "authenticated"],
+                }
+            ]
+        }
+    )
+
+    # return static fixture jwks.
+    jwk = neon_authorize_jwk.export_public(as_dict=True)
+    httpserver.expect_request("/authorize/jwks.json").respond_with_json({"keys": [jwk]})
+
+    mgmt_port = port_distributor.get_port()
+    http_port = port_distributor.get_port()
+    external_http_port = port_distributor.get_port()
+
+    with NeonAuthBroker(
+        neon_binpath=neon_binpath,
+        test_output_dir=test_output_dir,
+        http_port=http_port,
+        mgmt_port=mgmt_port,
+        external_http_port=external_http_port,
+        auth_backend=NeonAuthBroker.ControlPlane(httpserver.url_for("/cplane")),
+    ) as proxy:
+        proxy.start()
+        yield proxy
+
+
 class Endpoint(PgProtocol, LogUtils):
     """An object representing a Postgres compute endpoint managed by the control plane."""
 

test_runner/fixtures/paths.py

@@ -21,6 +21,8 @@ if TYPE_CHECKING:
     from typing import Optional
 
 
+BASE_DIR = Path(__file__).parents[2]
+COMPUTE_CONFIG_DIR = BASE_DIR / "compute" / "etc"
 DEFAULT_OUTPUT_DIR: str = "test_output"
 
 
@@ -64,18 +66,17 @@ def get_test_repo_dir(request: FixtureRequest, top_output_dir: Path) -> Path:
 @pytest.fixture(scope="session")
 def base_dir() -> Iterator[Path]:
     # find the base directory (currently this is the git root)
-    base_dir = Path(__file__).parents[2]
-    log.info(f"base_dir is {base_dir}")
+    log.info(f"base_dir is {BASE_DIR}")
 
-    yield base_dir
+    yield BASE_DIR
 
 
 @pytest.fixture(scope="session")
-def compute_config_dir(base_dir: Path) -> Iterator[Path]:
+def compute_config_dir() -> Iterator[Path]:
     """
     Retrieve the path to the compute configuration directory.
     """
-    yield base_dir / "compute" / "etc"
+    yield COMPUTE_CONFIG_DIR
 
 
 @pytest.fixture(scope="function")

test_runner/regress/test_auth_broker.py

@@ -0,0 +1,37 @@
+import json
+
+import pytest
+from fixtures.neon_fixtures import NeonAuthBroker
+from jwcrypto import jwk, jwt
+
+
+@pytest.mark.asyncio
+async def test_auth_broker_happy(
+    static_auth_broker: NeonAuthBroker,
+    neon_authorize_jwk: jwk.JWK,
+):
+    """
+    Signs a JWT and uses it to authorize a query to local_proxy.
+    """
+
+    token = jwt.JWT(
+        header={"kid": neon_authorize_jwk.key_id, "alg": "RS256"}, claims={"sub": "user1"}
+    )
+    token.make_signed_token(neon_authorize_jwk)
+    res = await static_auth_broker.query("foo", ["arg1"], user="anonymous", token=token.serialize())
+
+    # local proxy mock just echos back the request
+    # check that we forward the correct data
+
+    assert (
+        res["headers"]["authorization"] == f"Bearer {token.serialize()}"
+    ), "JWT should be forwarded"
+
+    assert (
+        "anonymous" in res["headers"]["neon-connection-string"]
+    ), "conn string should be forwarded"
+
+    assert json.loads(res["body"]) == {
+        "query": "foo",
+        "params": ["arg1"],
+    }, "Query body should be forwarded"

test_runner/regress/test_broken_timeline.py

@@ -103,7 +103,6 @@ def test_timeline_init_break_before_checkpoint(neon_env_builder: NeonEnvBuilder)
     env.pageserver.allowed_errors.extend(
         [
             ".*Failed to process timeline dir contents.*Timeline has no ancestor and no layer files.*",
-            ".*Timeline got dropped without initializing, cleaning its files.*",
         ]
     )
 
@@ -145,7 +144,6 @@ def test_timeline_init_break_before_checkpoint_recreate(
     env.pageserver.allowed_errors.extend(
         [
             ".*Failed to process timeline dir contents.*Timeline has no ancestor and no layer files.*",
-            ".*Timeline got dropped without initializing, cleaning its files.*",
             ".*Failed to load index_part from remote storage, failed creation?.*",
         ]
     )

test_runner/regress/test_compute_metrics.py

@@ -1,9 +1,453 @@
 from __future__ import annotations
 
-from fixtures.neon_fixtures import NeonEnv
+import enum
+import os
+import shutil
+from pathlib import Path
+from typing import TYPE_CHECKING, cast
+
+# Docs are available at https://jsonnet.org/ref/bindings.html#python_api
+import _jsonnet
+import pytest
+import requests
+import yaml
+from fixtures.log_helper import log
+from fixtures.paths import BASE_DIR, COMPUTE_CONFIG_DIR
+
+if TYPE_CHECKING:
+    from types import TracebackType
+    from typing import Optional, TypedDict, Union
+
+    from fixtures.neon_fixtures import NeonEnv
+    from fixtures.pg_version import PgVersion
+    from fixtures.port_distributor import PortDistributor
+
+    class Metric(TypedDict):
+        metric_name: str
+        type: str
+        help: str
+        key_labels: Optional[list[str]]
+        values: Optional[list[str]]
+        query: Optional[str]
+        query_ref: Optional[str]
+
+    class Collector(TypedDict):
+        collector_name: str
+        metrics: list[Metric]
+        queries: Optional[list[Query]]
+
+    class Query(TypedDict):
+        query_name: str
+        query: str
 
 
-def test_compute_metrics(neon_simple_env: NeonEnv):
+JSONNET_IMPORT_CACHE: dict[str, bytes] = {}
+JSONNET_PATH: list[Path] = [BASE_DIR / "compute" / "jsonnet", COMPUTE_CONFIG_DIR]
+
+
+def __import_callback(dir: str, rel: str) -> tuple[str, bytes]:
+    """
+    dir: The directory of the Jsonnet file which tried to import a file
+    rel: The actual import path from Jsonnet
+    """
+    if not rel:
+        raise RuntimeError("Empty filename")
+
+    full_path: Optional[str] = None
+    if os.path.isabs(rel):
+        full_path = rel
+    else:
+        for p in (dir, *JSONNET_PATH):
+            assert isinstance(p, (str, Path)), "for mypy"
+            full_path = os.path.join(p, rel)
+
+            assert isinstance(full_path, str), "for mypy"
+            if not os.path.exists(full_path):
+                full_path = None
+                continue
+
+            break
+
+        if not full_path:
+            raise RuntimeError(f"Could not resolve import ({rel}) in {dir}")
+
+    if os.path.isdir(full_path):
+        raise RuntimeError(f"Attempted to import directory: {full_path}")
+
+    if full_path not in JSONNET_IMPORT_CACHE:
+        with open(full_path, encoding="utf-8") as f:
+            JSONNET_IMPORT_CACHE[full_path] = f.read().encode()
+
+    return full_path, JSONNET_IMPORT_CACHE[full_path]
+
+
+def jsonnet_evaluate_file(
+    jsonnet_file: Union[str, Path],
+    ext_vars: Optional[Union[str, dict[str, str]]] = None,
+    tla_vars: Optional[Union[str, dict[str, str]]] = None,
+) -> str:
+    return cast(
+        "str",
+        _jsonnet.evaluate_file(
+            str(jsonnet_file),
+            ext_vars=ext_vars,
+            tla_vars=tla_vars,
+            import_callback=__import_callback,
+        ),
+    )
+
+
+def evaluate_collector(jsonnet_file: Path, pg_version: PgVersion) -> str:
+    return jsonnet_evaluate_file(jsonnet_file, ext_vars={"pg_version": str(pg_version)})
+
+
+def evaluate_config(
+    jsonnet_file: Path, collector_name: str, collector_file: Union[str, Path], connstr: str
+) -> str:
+    return jsonnet_evaluate_file(
+        jsonnet_file,
+        tla_vars={
+            "collector_name": collector_name,
+            "collector_file": str(collector_file),
+            "connection_string": connstr,
+        },
+    )
+
+
+@enum.unique
+class SqlExporterProcess(str, enum.Enum):
+    COMPUTE = "compute"
+    AUTOSCALING = "autoscaling"
+
+
+@pytest.mark.parametrize(
+    "collector_name",
+    ["neon_collector", "neon_collector_autoscaling"],
+    ids=[SqlExporterProcess.COMPUTE, SqlExporterProcess.AUTOSCALING],
+)
+def test_sql_exporter_metrics_smoke(
+    pg_version: PgVersion,
+    neon_simple_env: NeonEnv,
+    compute_config_dir: Path,
+    collector_name: str,
+):
+    """
+    This is a smoke test to ensure the metrics SQL queries for sql_exporter
+    work without errors.
+    """
+    env = neon_simple_env
+
+    endpoint = env.endpoints.create("main")
+    endpoint.respec(skip_pg_catalog_updates=False)
+    endpoint.start()
+
+    # Extract all the SQL queries from the sql_exporter config files, and run
+    # them.
+    collector = cast(
+        "Collector",
+        yaml.safe_load(
+            jsonnet_evaluate_file(
+                str(compute_config_dir / f"{collector_name}.jsonnet"),
+                ext_vars={"pg_version": pg_version},
+            )
+        ),
+    )
+
+    for metric in collector["metrics"]:
+        query = metric.get("query")
+        if query is not None:
+            log.info("Checking query for metric %s in %s", metric["metric_name"], collector_name)
+            endpoint.safe_psql(query)
+
+    queries = collector.get("queries")
+    if queries is not None:
+        # This variable is named q because mypy is too silly to understand it is
+        # different from the query above.
+        #
+        # query: Optional[str]
+        # q: Metric
+        for q in queries:
+            log.info("Checking query %s in %s", q["query_name"], collector_name)
+            endpoint.safe_psql(q["query"])
+
+
+class SqlExporterRunner:
+    def __init__(self, test_output_dir: Path, sql_exporter_port: int) -> None:
+        self._log_file_name = test_output_dir / "sql_exporter.stderr"
+        self._sql_exporter_port = sql_exporter_port
+
+        log.info(f"Starting sql_exporter at http://localhost:{self._sql_exporter_port}")
+
+    def start(self) -> None:
+        raise NotImplementedError()
+
+    def stop(self) -> None:
+        raise NotImplementedError()
+
+    def __enter__(self) -> SqlExporterRunner:
+        self.start()
+
+        return self
+
+    def __exit__(
+        self,
+        exc_type: Optional[type[BaseException]],
+        exc: Optional[BaseException],
+        tb: Optional[TracebackType],
+    ):
+        self.stop()
+
+
+SQL_EXPORTER = shutil.which("sql_exporter")
+
+if SQL_EXPORTER is None:
+    from testcontainers.core.container import DockerContainer
+    from testcontainers.core.waiting_utils import wait_for_logs
+    from typing_extensions import override
+
+    class SqlExporterContainer(DockerContainer):  # type: ignore
+        def __init__(
+            self, logs_dir: Path, config_file: Path, collector_file: Path, port: int
+        ) -> None:
+            # NOTE: Keep the version the same as in
+            # compute/Dockerfile.compute-node and Dockerfile.build-tools.
+            #
+            # The "host" network mode allows sql_exporter to talk to the
+            # endpoint which is running on the host.
+            super().__init__("docker.io/burningalchemist/sql_exporter:0.13.1", network_mode="host")
+
+            self.__logs_dir = logs_dir
+            self.__port = port
+
+            config_file_name = config_file.name
+            collector_file_name = collector_file.name
+
+            self.with_command(f"-config.file=/etc/{config_file_name} -web.listen-address=:{port}")
+
+            container_config_file = f"/etc/{config_file_name}"
+            container_collector_file = f"/etc/{collector_file_name}"
+            log.info(
+                "Mapping %s to %s in sql_exporter container", config_file, container_config_file
+            )
+            log.info(
+                "Mapping %s to %s in sql_exporter container",
+                collector_file,
+                container_collector_file,
+            )
+
+            # NOTE: z allows Podman to work with SELinux. Please don't change it.
+            # Ideally this would be a ro (read-only) mount, but I couldn't seem to
+            # get it to work.
+            self.with_volume_mapping(str(config_file), container_config_file, "z")
+            self.with_volume_mapping(str(collector_file), container_collector_file, "z")
+
+        @override
+        def start(self) -> SqlExporterContainer:
+            super().start()
+
+            log.info("Waiting for sql_exporter to be ready")
+            wait_for_logs(
+                self,
+                rf'level=info msg="Listening on" address=\[::\]:{self.__port}',
+                timeout=5,
+            )
+
+            return self
+
+    class SqlExporterContainerRunner(SqlExporterRunner):
+        def __init__(
+            self,
+            test_output_dir: Path,
+            config_file: Path,
+            collector_file: Path,
+            sql_exporter_port: int,
+        ) -> None:
+            super().__init__(test_output_dir, sql_exporter_port)
+
+            self.__container = SqlExporterContainer(
+                test_output_dir, config_file, collector_file, sql_exporter_port
+            )
+
+        @override
+        def start(self) -> None:
+            self.__container.start()
+
+        @override
+        def stop(self) -> None:
+            try:
+                # sql_exporter doesn't print anything to stdout
+                with open(self._log_file_name, "w", encoding="utf-8") as f:
+                    f.write(self.__container.get_logs()[1].decode())
+            except Exception:
+                log.exception("Failed to write sql_exporter logs")
+
+            # Stop the container *after* getting the logs
+            self.__container.stop()
+
+else:
+    import subprocess
+    import time
+    from signal import Signals
+
+    from typing_extensions import override
+
+    if TYPE_CHECKING:
+        from collections.abc import Mapping
+
+    class SqlExporterNativeRunner(SqlExporterRunner):
+        def __init__(
+            self,
+            test_output_dir: Path,
+            config_file: Path,
+            collector_file: Path,
+            sql_exporter_port: int,
+        ) -> None:
+            super().__init__(test_output_dir, sql_exporter_port)
+
+            self.__config_file = config_file
+            self.__collector_file = collector_file
+            self.__proc: subprocess.Popen[str]
+
+        @override
+        def start(self) -> None:
+            assert SQL_EXPORTER is not None
+
+            log_file = open(self._log_file_name, "w", encoding="utf-8")
+            self.__proc = subprocess.Popen(
+                [
+                    os.path.realpath(SQL_EXPORTER),
+                    f"-config.file={self.__config_file}",
+                    f"-web.listen-address=:{self._sql_exporter_port}",
+                ],
+                # If PGSERVICEFILE is set, sql_exporter won't launch.
+                env=cast("Mapping[str, str]", {}),
+                stderr=log_file,
+                bufsize=0,
+                text=True,
+            )
+
+            log.info("Waiting for sql_exporter to be ready")
+
+            with open(self._log_file_name, encoding="utf-8") as f:
+                started = time.time()
+                while True:
+                    if time.time() - started > 5:
+                        self.__proc.kill()
+                        raise RuntimeError("sql_exporter did not start up properly")
+
+                    line = f.readline()
+                    if not line:
+                        time.sleep(0.5)
+                        continue
+
+                    if (
+                        f'level=info msg="Listening on" address=[::]:{self._sql_exporter_port}'
+                        in line
+                    ):
+                        break
+
+        @override
+        def stop(self) -> None:
+            self.__proc.send_signal(Signals.SIGINT)
+            self.__proc.wait()
+
+
+@pytest.mark.parametrize(
+    "exporter",
+    [SqlExporterProcess.COMPUTE, SqlExporterProcess.AUTOSCALING],
+)
+def test_sql_exporter_metrics_e2e(
+    pg_version: PgVersion,
+    neon_simple_env: NeonEnv,
+    test_output_dir: Path,
+    compute_config_dir: Path,
+    exporter: SqlExporterProcess,
+    port_distributor: PortDistributor,
+):
+    """
+    This is a full E2E test of the sql_exporter setup to make sure it works
+    without error.
+
+    If you use Podman instead of Docker, you may run into issues. If you run
+    rootful Podman, you may need to add a ~/.testcontainers.properties file
+    with the following content:
+
+        ryuk.container.privileged=true
+
+    If you are not running rootful Podman, set the following environment
+    variable:
+
+        TESTCONTAINERS_RYUK_DISABLED=true
+
+    Note that you will need the Podman socket to be running. On a systemd-based
+    system, that command will look something like:
+
+        # Use `enable --now` to start the socket on login and immediately.
+        systemctl --user start podman.socket
+
+    Whether you use the user service manager or the system service manager is
+    up to you, but may have implications on the above ryuk related steps. Note
+    that you may also need the docker(1) Podman frontend. I am unsure if the
+    docker Python package supports Podman natively.
+    """
+    env = neon_simple_env
+
+    endpoint = env.endpoints.create("main")
+    endpoint.respec(skip_pg_catalog_updates=False)
+    endpoint.start()
+
+    if exporter == SqlExporterProcess.COMPUTE:
+        stem_suffix = ""
+    elif exporter == SqlExporterProcess.AUTOSCALING:
+        stem_suffix = "_autoscaling"
+
+    # Write the collector file
+    collector_file = test_output_dir / f"neon_collector{stem_suffix}.yml"
+    with open(collector_file, "w", encoding="utf-8") as o:
+        collector = evaluate_collector(
+            compute_config_dir / f"neon_collector{stem_suffix}.jsonnet", pg_version
+        )
+        o.write(collector)
+
+    conn_options = endpoint.conn_options()
+    pg_host = conn_options["host"]
+    pg_port = conn_options["port"]
+    pg_user = conn_options["user"]
+    pg_dbname = conn_options["dbname"]
+    pg_application_name = f"sql_exporter{stem_suffix}"
+    connstr = f"postgresql://{pg_user}@{pg_host}:{pg_port}/{pg_dbname}?sslmode=disable&application_name={pg_application_name}"
+
+    def escape_go_filepath_match_characters(s: str) -> str:
+        """
+        Unfortunately sql_exporter doesn't use plain file paths, so we need to
+        escape special characters. pytest encodes the parameters of a test using
+        [ and ], so we need to escape them with backslashes.
+        See https://pkg.go.dev/path/filepath#Match.
+        """
+        return s.replace("[", r"\[").replace("]", r"\]")
+
+    # Write the config file
+    config_file = test_output_dir / f"sql_exporter{stem_suffix}.yml"
+    with open(config_file, "w", encoding="utf-8") as o:
+        config = evaluate_config(
+            compute_config_dir / "sql_exporter.jsonnet",
+            collector_name=collector_file.stem,
+            collector_file=escape_go_filepath_match_characters(str(collector_file))
+            if SQL_EXPORTER
+            else collector_file.name,
+            connstr=connstr,
+        )
+        o.write(config)
+
+    sql_exporter_port = port_distributor.get_port()
+    with (SqlExporterNativeRunner if SQL_EXPORTER else SqlExporterContainerRunner)(
+        test_output_dir, config_file, collector_file, sql_exporter_port
+    ) as _runner:
+        resp = requests.get(f"http://localhost:{sql_exporter_port}/metrics")
+        resp.raise_for_status()
+
+
+def test_perf_counters(neon_simple_env: NeonEnv):
     """
     Test compute metrics, exposed in the neon_backend_perf_counters and
     neon_perf_counters views

test_runner/regress/test_import.py

@@ -91,7 +91,6 @@ def test_import_from_vanilla(test_output_dir, pg_bin, vanilla_pg, neon_env_build
         [
             ".*Failed to import basebackup.*",
             ".*unexpected non-zero bytes after the tar archive.*",
-            ".*Timeline got dropped without initializing, cleaning its files.*",
             ".*InternalServerError.*timeline not found.*",
             ".*InternalServerError.*Tenant .* not found.*",
             ".*InternalServerError.*Timeline .* not found.*",

test_runner/regress/test_proxy.py

@@ -561,7 +561,7 @@ def test_sql_over_http_pool_dos(static_proxy: NeonProxy):
 
     # query generates a million rows - should hit the 10MB reponse limit quickly
     response = query(
-        400,
+        507,
         "select * from generate_series(1, 5000) a cross join generate_series(1, 5000) b cross join (select 'foo'::foo) c;",
     )
     assert "response is too large (max is 10485760 bytes)" in response["message"]

test_runner/regress/test_tenant_delete.py

@@ -146,8 +146,6 @@ def test_long_timeline_create_cancelled_by_tenant_delete(neon_env_builder: NeonE
 
     env.pageserver.allowed_errors.extend(
         [
-            # happens with the cancellation bailing flushing loop earlier, leaving disk_consistent_lsn at zero
-            ".*Timeline got dropped without initializing, cleaning its files",
             # the response hit_pausable_failpoint_and_later_fail
             f".*Error processing HTTP request: InternalServerError\\(new timeline {env.initial_tenant}/{env.initial_timeline} has invalid disk_consistent_lsn",
         ]

test_runner/regress/test_timeline_archive.py

@@ -213,6 +213,13 @@ def test_timeline_offloading(neon_env_builder: NeonEnvBuilder, manual_offload: b
     wait_until(30, 1, leaf_offloaded)
     wait_until(30, 1, parent_offloaded)
 
+    # Offloaded child timelines should still prevent deletion
+    with pytest.raises(
+        PageserverApiException,
+        match=f".* timeline which has child timelines: \\[{leaf_timeline_id}\\]",
+    ):
+        ps_http.timeline_delete(tenant_id, parent_timeline_id)
+
     ps_http.timeline_archival_config(
         tenant_id,
         grandparent_timeline_id,

test_runner/regress/test_timeline_detach_ancestor.py

@@ -9,7 +9,7 @@ from queue import Empty, Queue
 from threading import Barrier
 
 import pytest
-from fixtures.common_types import Lsn, TimelineId
+from fixtures.common_types import Lsn, TimelineArchivalState, TimelineId
 from fixtures.log_helper import log
 from fixtures.neon_fixtures import (
     LogCursor,
@@ -634,7 +634,13 @@ def test_timeline_ancestor_detach_errors(neon_env_builder: NeonEnvBuilder, shard
     shards = 2 if sharded else 1
 
     neon_env_builder.num_pageservers = shards
-    env = neon_env_builder.init_start(initial_tenant_shard_count=shards if sharded else None)
+    env = neon_env_builder.init_start(
+        initial_tenant_shard_count=shards if sharded else None,
+        initial_tenant_conf={
+            # turn off gc, we want to do manual offloading here.
+            "gc_period": "0s",
+        },
+    )
 
     pageservers = dict((int(p.id), p) for p in env.pageservers)
 
@@ -656,7 +662,9 @@ def test_timeline_ancestor_detach_errors(neon_env_builder: NeonEnvBuilder, shard
         client.detach_ancestor(env.initial_tenant, env.initial_timeline)
     assert info.value.status_code == 409
 
-    _ = env.create_branch("first_branch")
+    early_branch = env.create_branch("early_branch")
+
+    first_branch = env.create_branch("first_branch")
 
     second_branch = env.create_branch("second_branch", ancestor_branch_name="first_branch")
 
@@ -665,6 +673,29 @@ def test_timeline_ancestor_detach_errors(neon_env_builder: NeonEnvBuilder, shard
         client.detach_ancestor(env.initial_tenant, second_branch)
     assert info.value.status_code == 400
 
+    client.timeline_archival_config(
+        env.initial_tenant, second_branch, TimelineArchivalState.ARCHIVED
+    )
+
+    client.timeline_archival_config(
+        env.initial_tenant, early_branch, TimelineArchivalState.ARCHIVED
+    )
+
+    with pytest.raises(PageserverApiException, match=f".*archived: {early_branch}") as info:
+        client.detach_ancestor(env.initial_tenant, first_branch)
+    assert info.value.status_code == 400
+
+    if not sharded:
+        client.timeline_offload(env.initial_tenant, early_branch)
+
+    client.timeline_archival_config(
+        env.initial_tenant, first_branch, TimelineArchivalState.ARCHIVED
+    )
+
+    with pytest.raises(PageserverApiException, match=f".*archived: {first_branch}") as info:
+        client.detach_ancestor(env.initial_tenant, first_branch)
+    assert info.value.status_code == 400
+
 
 def test_sharded_timeline_detach_ancestor(neon_env_builder: NeonEnvBuilder):
     """

test_runner/stubs/h2/README.md

@@ -0,0 +1 @@
+generated via `poetry run stubgen -p h2 -o test_runner/stubs`

test_runner/stubs/h2/config.pyi

@@ -0,0 +1,42 @@
+from _typeshed import Incomplete
+
+class _BooleanConfigOption:
+    name: Incomplete
+    attr_name: Incomplete
+    def __init__(self, name) -> None: ...
+    def __get__(self, instance, owner): ...
+    def __set__(self, instance, value) -> None: ...
+
+class DummyLogger:
+    def __init__(self, *vargs) -> None: ...
+    def debug(self, *vargs, **kwargs) -> None: ...
+    def trace(self, *vargs, **kwargs) -> None: ...
+
+class OutputLogger:
+    file: Incomplete
+    trace_level: Incomplete
+    def __init__(self, file: Incomplete | None = ..., trace_level: bool = ...) -> None: ...
+    def debug(self, fmtstr, *args) -> None: ...
+    def trace(self, fmtstr, *args) -> None: ...
+
+class H2Configuration:
+    client_side: Incomplete
+    validate_outbound_headers: Incomplete
+    normalize_outbound_headers: Incomplete
+    validate_inbound_headers: Incomplete
+    normalize_inbound_headers: Incomplete
+    logger: Incomplete
+    def __init__(
+        self,
+        client_side: bool = ...,
+        header_encoding: Incomplete | None = ...,
+        validate_outbound_headers: bool = ...,
+        normalize_outbound_headers: bool = ...,
+        validate_inbound_headers: bool = ...,
+        normalize_inbound_headers: bool = ...,
+        logger: Incomplete | None = ...,
+    ) -> None: ...
+    @property
+    def header_encoding(self): ...
+    @header_encoding.setter
+    def header_encoding(self, value) -> None: ...

test_runner/stubs/h2/connection.pyi

@@ -0,0 +1,142 @@
+from enum import Enum, IntEnum
+
+from _typeshed import Incomplete
+
+from .config import H2Configuration as H2Configuration
+from .errors import ErrorCodes as ErrorCodes
+from .events import AlternativeServiceAvailable as AlternativeServiceAvailable
+from .events import ConnectionTerminated as ConnectionTerminated
+from .events import PingAckReceived as PingAckReceived
+from .events import PingReceived as PingReceived
+from .events import PriorityUpdated as PriorityUpdated
+from .events import RemoteSettingsChanged as RemoteSettingsChanged
+from .events import SettingsAcknowledged as SettingsAcknowledged
+from .events import UnknownFrameReceived as UnknownFrameReceived
+from .events import WindowUpdated as WindowUpdated
+from .exceptions import DenialOfServiceError as DenialOfServiceError
+from .exceptions import FlowControlError as FlowControlError
+from .exceptions import FrameTooLargeError as FrameTooLargeError
+from .exceptions import NoAvailableStreamIDError as NoAvailableStreamIDError
+from .exceptions import NoSuchStreamError as NoSuchStreamError
+from .exceptions import ProtocolError as ProtocolError
+from .exceptions import RFC1122Error as RFC1122Error
+from .exceptions import StreamClosedError as StreamClosedError
+from .exceptions import StreamIDTooLowError as StreamIDTooLowError
+from .exceptions import TooManyStreamsError as TooManyStreamsError
+from .frame_buffer import FrameBuffer as FrameBuffer
+from .settings import SettingCodes as SettingCodes
+from .settings import Settings as Settings
+from .stream import H2Stream as H2Stream
+from .stream import StreamClosedBy as StreamClosedBy
+from .utilities import guard_increment_window as guard_increment_window
+from .windows import WindowManager as WindowManager
+
+class ConnectionState(Enum):
+    IDLE: int
+    CLIENT_OPEN: int
+    SERVER_OPEN: int
+    CLOSED: int
+
+class ConnectionInputs(Enum):
+    SEND_HEADERS: int
+    SEND_PUSH_PROMISE: int
+    SEND_DATA: int
+    SEND_GOAWAY: int
+    SEND_WINDOW_UPDATE: int
+    SEND_PING: int
+    SEND_SETTINGS: int
+    SEND_RST_STREAM: int
+    SEND_PRIORITY: int
+    RECV_HEADERS: int
+    RECV_PUSH_PROMISE: int
+    RECV_DATA: int
+    RECV_GOAWAY: int
+    RECV_WINDOW_UPDATE: int
+    RECV_PING: int
+    RECV_SETTINGS: int
+    RECV_RST_STREAM: int
+    RECV_PRIORITY: int
+    SEND_ALTERNATIVE_SERVICE: int
+    RECV_ALTERNATIVE_SERVICE: int
+
+class AllowedStreamIDs(IntEnum):
+    EVEN: int
+    ODD: int
+
+class H2ConnectionStateMachine:
+    state: Incomplete
+    def __init__(self) -> None: ...
+    def process_input(self, input_): ...
+
+class H2Connection:
+    DEFAULT_MAX_OUTBOUND_FRAME_SIZE: int
+    DEFAULT_MAX_INBOUND_FRAME_SIZE: Incomplete
+    HIGHEST_ALLOWED_STREAM_ID: Incomplete
+    MAX_WINDOW_INCREMENT: Incomplete
+    DEFAULT_MAX_HEADER_LIST_SIZE: Incomplete
+    MAX_CLOSED_STREAMS: Incomplete
+    state_machine: Incomplete
+    streams: Incomplete
+    highest_inbound_stream_id: int
+    highest_outbound_stream_id: int
+    encoder: Incomplete
+    decoder: Incomplete
+    config: Incomplete
+    local_settings: Incomplete
+    remote_settings: Incomplete
+    outbound_flow_control_window: Incomplete
+    max_outbound_frame_size: Incomplete
+    max_inbound_frame_size: Incomplete
+    incoming_buffer: Incomplete
+    def __init__(self, config: Incomplete | None = ...) -> None: ...
+    @property
+    def open_outbound_streams(self): ...
+    @property
+    def open_inbound_streams(self): ...
+    @property
+    def inbound_flow_control_window(self): ...
+    def initiate_connection(self) -> None: ...
+    def initiate_upgrade_connection(self, settings_header: Incomplete | None = ...): ...
+    def get_next_available_stream_id(self): ...
+    def send_headers(
+        self,
+        stream_id,
+        headers,
+        end_stream: bool = ...,
+        priority_weight: Incomplete | None = ...,
+        priority_depends_on: Incomplete | None = ...,
+        priority_exclusive: Incomplete | None = ...,
+    ) -> None: ...
+    def send_data(
+        self, stream_id, data, end_stream: bool = ..., pad_length: Incomplete | None = ...
+    ) -> None: ...
+    def end_stream(self, stream_id) -> None: ...
+    def increment_flow_control_window(
+        self, increment, stream_id: Incomplete | None = ...
+    ) -> None: ...
+    def push_stream(self, stream_id, promised_stream_id, request_headers) -> None: ...
+    def ping(self, opaque_data) -> None: ...
+    def reset_stream(self, stream_id, error_code: int = ...) -> None: ...
+    def close_connection(
+        self,
+        error_code: int = ...,
+        additional_data: Incomplete | None = ...,
+        last_stream_id: Incomplete | None = ...,
+    ) -> None: ...
+    def update_settings(self, new_settings) -> None: ...
+    def advertise_alternative_service(
+        self, field_value, origin: Incomplete | None = ..., stream_id: Incomplete | None = ...
+    ) -> None: ...
+    def prioritize(
+        self,
+        stream_id,
+        weight: Incomplete | None = ...,
+        depends_on: Incomplete | None = ...,
+        exclusive: Incomplete | None = ...,
+    ) -> None: ...
+    def local_flow_control_window(self, stream_id): ...
+    def remote_flow_control_window(self, stream_id): ...
+    def acknowledge_received_data(self, acknowledged_size, stream_id) -> None: ...
+    def data_to_send(self, amount: Incomplete | None = ...): ...
+    def clear_outbound_data_buffer(self) -> None: ...
+    def receive_data(self, data): ...

test_runner/stubs/h2/errors.pyi

@@ -0,0 +1,17 @@
+import enum
+
+class ErrorCodes(enum.IntEnum):
+    NO_ERROR: int
+    PROTOCOL_ERROR: int
+    INTERNAL_ERROR: int
+    FLOW_CONTROL_ERROR: int
+    SETTINGS_TIMEOUT: int
+    STREAM_CLOSED: int
+    FRAME_SIZE_ERROR: int
+    REFUSED_STREAM: int
+    CANCEL: int
+    COMPRESSION_ERROR: int
+    CONNECT_ERROR: int
+    ENHANCE_YOUR_CALM: int
+    INADEQUATE_SECURITY: int
+    HTTP_1_1_REQUIRED: int

test_runner/stubs/h2/events.pyi

@@ -0,0 +1,106 @@
+from _typeshed import Incomplete
+
+from .settings import ChangedSetting as ChangedSetting
+
+class Event: ...
+
+class RequestReceived(Event):
+    stream_id: Incomplete
+    headers: Incomplete
+    stream_ended: Incomplete
+    priority_updated: Incomplete
+    def __init__(self) -> None: ...
+
+class ResponseReceived(Event):
+    stream_id: Incomplete
+    headers: Incomplete
+    stream_ended: Incomplete
+    priority_updated: Incomplete
+    def __init__(self) -> None: ...
+
+class TrailersReceived(Event):
+    stream_id: Incomplete
+    headers: Incomplete
+    stream_ended: Incomplete
+    priority_updated: Incomplete
+    def __init__(self) -> None: ...
+
+class _HeadersSent(Event): ...
+class _ResponseSent(_HeadersSent): ...
+class _RequestSent(_HeadersSent): ...
+class _TrailersSent(_HeadersSent): ...
+class _PushedRequestSent(_HeadersSent): ...
+
+class InformationalResponseReceived(Event):
+    stream_id: Incomplete
+    headers: Incomplete
+    priority_updated: Incomplete
+    def __init__(self) -> None: ...
+
+class DataReceived(Event):
+    stream_id: Incomplete
+    data: Incomplete
+    flow_controlled_length: Incomplete
+    stream_ended: Incomplete
+    def __init__(self) -> None: ...
+
+class WindowUpdated(Event):
+    stream_id: Incomplete
+    delta: Incomplete
+    def __init__(self) -> None: ...
+
+class RemoteSettingsChanged(Event):
+    changed_settings: Incomplete
+    def __init__(self) -> None: ...
+    @classmethod
+    def from_settings(cls, old_settings, new_settings): ...
+
+class PingReceived(Event):
+    ping_data: Incomplete
+    def __init__(self) -> None: ...
+
+class PingAckReceived(Event):
+    ping_data: Incomplete
+    def __init__(self) -> None: ...
+
+class StreamEnded(Event):
+    stream_id: Incomplete
+    def __init__(self) -> None: ...
+
+class StreamReset(Event):
+    stream_id: Incomplete
+    error_code: Incomplete
+    remote_reset: bool
+    def __init__(self) -> None: ...
+
+class PushedStreamReceived(Event):
+    pushed_stream_id: Incomplete
+    parent_stream_id: Incomplete
+    headers: Incomplete
+    def __init__(self) -> None: ...
+
+class SettingsAcknowledged(Event):
+    changed_settings: Incomplete
+    def __init__(self) -> None: ...
+
+class PriorityUpdated(Event):
+    stream_id: Incomplete
+    weight: Incomplete
+    depends_on: Incomplete
+    exclusive: Incomplete
+    def __init__(self) -> None: ...
+
+class ConnectionTerminated(Event):
+    error_code: Incomplete
+    last_stream_id: Incomplete
+    additional_data: Incomplete
+    def __init__(self) -> None: ...
+
+class AlternativeServiceAvailable(Event):
+    origin: Incomplete
+    field_value: Incomplete
+    def __init__(self) -> None: ...
+
+class UnknownFrameReceived(Event):
+    frame: Incomplete
+    def __init__(self) -> None: ...

test_runner/stubs/h2/exceptions.pyi

@@ -0,0 +1,48 @@
+from _typeshed import Incomplete
+
+class H2Error(Exception): ...
+
+class ProtocolError(H2Error):
+    error_code: Incomplete
+
+class FrameTooLargeError(ProtocolError):
+    error_code: Incomplete
+
+class FrameDataMissingError(ProtocolError):
+    error_code: Incomplete
+
+class TooManyStreamsError(ProtocolError): ...
+
+class FlowControlError(ProtocolError):
+    error_code: Incomplete
+
+class StreamIDTooLowError(ProtocolError):
+    stream_id: Incomplete
+    max_stream_id: Incomplete
+    def __init__(self, stream_id, max_stream_id) -> None: ...
+
+class NoAvailableStreamIDError(ProtocolError): ...
+
+class NoSuchStreamError(ProtocolError):
+    stream_id: Incomplete
+    def __init__(self, stream_id) -> None: ...
+
+class StreamClosedError(NoSuchStreamError):
+    stream_id: Incomplete
+    error_code: Incomplete
+    def __init__(self, stream_id) -> None: ...
+
+class InvalidSettingsValueError(ProtocolError, ValueError):
+    error_code: Incomplete
+    def __init__(self, msg, error_code) -> None: ...
+
+class InvalidBodyLengthError(ProtocolError):
+    expected_length: Incomplete
+    actual_length: Incomplete
+    def __init__(self, expected, actual) -> None: ...
+
+class UnsupportedFrameError(ProtocolError): ...
+class RFC1122Error(H2Error): ...
+
+class DenialOfServiceError(ProtocolError):
+    error_code: Incomplete

test_runner/stubs/h2/frame_buffer.pyi

@@ -0,0 +1,19 @@
+from .exceptions import (
+    FrameDataMissingError as FrameDataMissingError,
+)
+from .exceptions import (
+    FrameTooLargeError as FrameTooLargeError,
+)
+from .exceptions import (
+    ProtocolError as ProtocolError,
+)
+
+CONTINUATION_BACKLOG: int
+
+class FrameBuffer:
+    data: bytes
+    max_frame_size: int
+    def __init__(self, server: bool = ...) -> None: ...
+    def add_data(self, data) -> None: ...
+    def __iter__(self): ...
+    def __next__(self): ...

test_runner/stubs/h2/settings.pyi

@@ -0,0 +1,61 @@
+import enum
+from collections.abc import MutableMapping
+from typing import Any
+
+from _typeshed import Incomplete
+from h2.errors import ErrorCodes as ErrorCodes
+from h2.exceptions import InvalidSettingsValueError as InvalidSettingsValueError
+
+class SettingCodes(enum.IntEnum):
+    HEADER_TABLE_SIZE: Incomplete
+    ENABLE_PUSH: Incomplete
+    MAX_CONCURRENT_STREAMS: Incomplete
+    INITIAL_WINDOW_SIZE: Incomplete
+    MAX_FRAME_SIZE: Incomplete
+    MAX_HEADER_LIST_SIZE: Incomplete
+    ENABLE_CONNECT_PROTOCOL: Incomplete
+
+class ChangedSetting:
+    setting: Incomplete
+    original_value: Incomplete
+    new_value: Incomplete
+    def __init__(self, setting, original_value, new_value) -> None: ...
+
+class Settings(MutableMapping[str, Any]):
+    def __init__(self, client: bool = ..., initial_values: Incomplete | None = ...) -> None: ...
+    def acknowledge(self): ...
+    @property
+    def header_table_size(self): ...
+    @header_table_size.setter
+    def header_table_size(self, value) -> None: ...
+    @property
+    def enable_push(self): ...
+    @enable_push.setter
+    def enable_push(self, value) -> None: ...
+    @property
+    def initial_window_size(self): ...
+    @initial_window_size.setter
+    def initial_window_size(self, value) -> None: ...
+    @property
+    def max_frame_size(self): ...
+    @max_frame_size.setter
+    def max_frame_size(self, value) -> None: ...
+    @property
+    def max_concurrent_streams(self): ...
+    @max_concurrent_streams.setter
+    def max_concurrent_streams(self, value) -> None: ...
+    @property
+    def max_header_list_size(self): ...
+    @max_header_list_size.setter
+    def max_header_list_size(self, value) -> None: ...
+    @property
+    def enable_connect_protocol(self): ...
+    @enable_connect_protocol.setter
+    def enable_connect_protocol(self, value) -> None: ...
+    def __getitem__(self, key): ...
+    def __setitem__(self, key, value) -> None: ...
+    def __delitem__(self, key) -> None: ...
+    def __iter__(self): ...
+    def __len__(self) -> int: ...
+    def __eq__(self, other): ...
+    def __ne__(self, other): ...

test_runner/stubs/h2/stream.pyi

@@ -0,0 +1,184 @@
+from enum import Enum, IntEnum
+
+from _typeshed import Incomplete
+
+from .errors import ErrorCodes as ErrorCodes
+from .events import (
+    AlternativeServiceAvailable as AlternativeServiceAvailable,
+)
+from .events import (
+    DataReceived as DataReceived,
+)
+from .events import (
+    InformationalResponseReceived as InformationalResponseReceived,
+)
+from .events import (
+    PushedStreamReceived as PushedStreamReceived,
+)
+from .events import (
+    RequestReceived as RequestReceived,
+)
+from .events import (
+    ResponseReceived as ResponseReceived,
+)
+from .events import (
+    StreamEnded as StreamEnded,
+)
+from .events import (
+    StreamReset as StreamReset,
+)
+from .events import (
+    TrailersReceived as TrailersReceived,
+)
+from .events import (
+    WindowUpdated as WindowUpdated,
+)
+from .exceptions import (
+    FlowControlError as FlowControlError,
+)
+from .exceptions import (
+    InvalidBodyLengthError as InvalidBodyLengthError,
+)
+from .exceptions import (
+    ProtocolError as ProtocolError,
+)
+from .exceptions import (
+    StreamClosedError as StreamClosedError,
+)
+from .utilities import (
+    HeaderValidationFlags as HeaderValidationFlags,
+)
+from .utilities import (
+    authority_from_headers as authority_from_headers,
+)
+from .utilities import (
+    extract_method_header as extract_method_header,
+)
+from .utilities import (
+    guard_increment_window as guard_increment_window,
+)
+from .utilities import (
+    is_informational_response as is_informational_response,
+)
+from .utilities import (
+    normalize_inbound_headers as normalize_inbound_headers,
+)
+from .utilities import (
+    normalize_outbound_headers as normalize_outbound_headers,
+)
+from .utilities import (
+    validate_headers as validate_headers,
+)
+from .utilities import (
+    validate_outbound_headers as validate_outbound_headers,
+)
+from .windows import WindowManager as WindowManager
+
+class StreamState(IntEnum):
+    IDLE: int
+    RESERVED_REMOTE: int
+    RESERVED_LOCAL: int
+    OPEN: int
+    HALF_CLOSED_REMOTE: int
+    HALF_CLOSED_LOCAL: int
+    CLOSED: int
+
+class StreamInputs(Enum):
+    SEND_HEADERS: int
+    SEND_PUSH_PROMISE: int
+    SEND_RST_STREAM: int
+    SEND_DATA: int
+    SEND_WINDOW_UPDATE: int
+    SEND_END_STREAM: int
+    RECV_HEADERS: int
+    RECV_PUSH_PROMISE: int
+    RECV_RST_STREAM: int
+    RECV_DATA: int
+    RECV_WINDOW_UPDATE: int
+    RECV_END_STREAM: int
+    RECV_CONTINUATION: int
+    SEND_INFORMATIONAL_HEADERS: int
+    RECV_INFORMATIONAL_HEADERS: int
+    SEND_ALTERNATIVE_SERVICE: int
+    RECV_ALTERNATIVE_SERVICE: int
+    UPGRADE_CLIENT: int
+    UPGRADE_SERVER: int
+
+class StreamClosedBy(Enum):
+    SEND_END_STREAM: int
+    RECV_END_STREAM: int
+    SEND_RST_STREAM: int
+    RECV_RST_STREAM: int
+
+STREAM_OPEN: Incomplete
+
+class H2StreamStateMachine:
+    state: Incomplete
+    stream_id: Incomplete
+    client: Incomplete
+    headers_sent: Incomplete
+    trailers_sent: Incomplete
+    headers_received: Incomplete
+    trailers_received: Incomplete
+    stream_closed_by: Incomplete
+    def __init__(self, stream_id) -> None: ...
+    def process_input(self, input_): ...
+    def request_sent(self, previous_state): ...
+    def response_sent(self, previous_state): ...
+    def request_received(self, previous_state): ...
+    def response_received(self, previous_state): ...
+    def data_received(self, previous_state): ...
+    def window_updated(self, previous_state): ...
+    def stream_half_closed(self, previous_state): ...
+    def stream_ended(self, previous_state): ...
+    def stream_reset(self, previous_state): ...
+    def send_new_pushed_stream(self, previous_state): ...
+    def recv_new_pushed_stream(self, previous_state): ...
+    def send_push_promise(self, previous_state): ...
+    def recv_push_promise(self, previous_state): ...
+    def send_end_stream(self, previous_state) -> None: ...
+    def send_reset_stream(self, previous_state) -> None: ...
+    def reset_stream_on_error(self, previous_state) -> None: ...
+    def recv_on_closed_stream(self, previous_state) -> None: ...
+    def send_on_closed_stream(self, previous_state) -> None: ...
+    def recv_push_on_closed_stream(self, previous_state) -> None: ...
+    def send_push_on_closed_stream(self, previous_state) -> None: ...
+    def send_informational_response(self, previous_state): ...
+    def recv_informational_response(self, previous_state): ...
+    def recv_alt_svc(self, previous_state): ...
+    def send_alt_svc(self, previous_state) -> None: ...
+
+class H2Stream:
+    state_machine: Incomplete
+    stream_id: Incomplete
+    max_outbound_frame_size: Incomplete
+    request_method: Incomplete
+    outbound_flow_control_window: Incomplete
+    config: Incomplete
+    def __init__(self, stream_id, config, inbound_window_size, outbound_window_size) -> None: ...
+    @property
+    def inbound_flow_control_window(self): ...
+    @property
+    def open(self): ...
+    @property
+    def closed(self): ...
+    @property
+    def closed_by(self): ...
+    def upgrade(self, client_side) -> None: ...
+    def send_headers(self, headers, encoder, end_stream: bool = ...): ...
+    def push_stream_in_band(self, related_stream_id, headers, encoder): ...
+    def locally_pushed(self): ...
+    def send_data(self, data, end_stream: bool = ..., pad_length: Incomplete | None = ...): ...
+    def end_stream(self): ...
+    def advertise_alternative_service(self, field_value): ...
+    def increase_flow_control_window(self, increment): ...
+    def receive_push_promise_in_band(self, promised_stream_id, headers, header_encoding): ...
+    def remotely_pushed(self, pushed_headers): ...
+    def receive_headers(self, headers, end_stream, header_encoding): ...
+    def receive_data(self, data, end_stream, flow_control_len): ...
+    def receive_window_update(self, increment): ...
+    def receive_continuation(self) -> None: ...
+    def receive_alt_svc(self, frame): ...
+    def reset_stream(self, error_code: int = ...): ...
+    def stream_reset(self, frame): ...
+    def acknowledge_received_data(self, acknowledged_size): ...

test_runner/stubs/h2/utilities.pyi

@@ -0,0 +1,25 @@
+from typing import NamedTuple
+
+from _typeshed import Incomplete
+
+from .exceptions import FlowControlError as FlowControlError
+from .exceptions import ProtocolError as ProtocolError
+
+UPPER_RE: Incomplete
+CONNECTION_HEADERS: Incomplete
+
+def extract_method_header(headers): ...
+def is_informational_response(headers): ...
+def guard_increment_window(current, increment): ...
+def authority_from_headers(headers): ...
+
+class HeaderValidationFlags(NamedTuple):
+    is_client: Incomplete
+    is_trailer: Incomplete
+    is_response_header: Incomplete
+    is_push_promise: Incomplete
+
+def validate_headers(headers, hdr_validation_flags): ...
+def normalize_outbound_headers(headers, hdr_validation_flags): ...
+def normalize_inbound_headers(headers, hdr_validation_flags): ...
+def validate_outbound_headers(headers, hdr_validation_flags): ...

test_runner/stubs/h2/windows.pyi

@@ -0,0 +1,13 @@
+from _typeshed import Incomplete
+
+from .exceptions import FlowControlError as FlowControlError
+
+LARGEST_FLOW_CONTROL_WINDOW: Incomplete
+
+class WindowManager:
+    max_window_size: Incomplete
+    current_window_size: Incomplete
+    def __init__(self, max_window_size) -> None: ...
+    def window_consumed(self, size) -> None: ...
+    def window_opened(self, size) -> None: ...
+    def process_bytes(self, size): ...