alternative

Apart from sticking closer to the comment above the function,
this reduces the diff in the next patch.

.github/workflows/build_and_test.yml

@@ -341,7 +341,9 @@ jobs:
 
       - name: Run rust tests
         run: |
-          ${cov_prefix} cargo nextest run $CARGO_FLAGS $CARGO_FEATURES
+          for io_engine in std-fs tokio-epoll-uring ; do
+            NEON_PAGESERVER_UNIT_TEST_VIRTUAL_FILE_IOENGINE=$io_engine ${cov_prefix} cargo nextest run $CARGO_FLAGS $CARGO_FEATURES
+          done
 
           # Run separate tests for real S3
           export ENABLE_REAL_S3_REMOTE_STORAGE=nonempty
@@ -426,6 +428,7 @@ jobs:
       matrix:
         build_type: [ debug, release ]
         pg_version: [ v14, v15, v16 ]
+        pageserver_virtual_file_io_engine: [ std-fs, tokio-epoll-uring ]
     steps:
       - name: Checkout
         uses: actions/checkout@v3
@@ -448,6 +451,7 @@ jobs:
           TEST_RESULT_CONNSTR: ${{ secrets.REGRESS_TEST_RESULT_CONNSTR_NEW }}
           CHECK_ONDISK_DATA_COMPATIBILITY: nonempty
           BUILD_TAG: ${{ needs.tag.outputs.build-tag }}
+          PAGESERVER_VIRTUAL_FILE_IO_ENGINE: ${{ matrix.pageserver_virtual_file_io_engine }}
 
       - name: Merge and upload coverage data
         if: matrix.build_type == 'debug' && matrix.pg_version == 'v14'
@@ -466,6 +470,7 @@ jobs:
       matrix:
         pytest_split_group: [ 1, 2, 3, 4 ]
         build_type: [ release ]
+        pageserver_virtual_file_io_engine: [ std-fs, tokio-epoll-uring ]
     steps:
       - name: Checkout
         uses: actions/checkout@v3
@@ -482,6 +487,7 @@ jobs:
           VIP_VAP_ACCESS_TOKEN: "${{ secrets.VIP_VAP_ACCESS_TOKEN }}"
           PERF_TEST_RESULT_CONNSTR: "${{ secrets.PERF_TEST_RESULT_CONNSTR }}"
           TEST_RESULT_CONNSTR: "${{ secrets.REGRESS_TEST_RESULT_CONNSTR_NEW }}"
+          PAGESERVER_VIRTUAL_FILE_IO_ENGINE: "${{ matrix.pageserver_virtual_file_io_engine }}"
       # XXX: no coverage data handling here, since benchmarks are run on release builds,
       # while coverage is currently collected for the debug ones
 

Cargo.lock

@@ -2520,6 +2520,16 @@ dependencies = [
  "windows-sys 0.48.0",
 ]
 
+[[package]]
+name = "io-uring"
+version = "0.6.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "460648e47a07a43110fbfa2e0b14afb2be920093c31e5dccc50e49568e099762"
+dependencies = [
+ "bitflags 1.3.2",
+ "libc",
+]
+
 [[package]]
 name = "ipnet"
 version = "2.9.0"
@@ -3211,19 +3221,18 @@ name = "pagebench"
 version = "0.1.0"
 dependencies = [
  "anyhow",
- "camino",
  "clap",
  "futures",
  "hdrhistogram",
  "humantime",
  "humantime-serde",
+ "pageserver",
  "pageserver_api",
  "pageserver_client",
  "rand 0.8.5",
  "serde",
  "serde_json",
  "tokio",
- "tokio-util",
  "tracing",
  "utils",
  "workspace_hack",
@@ -3316,6 +3325,7 @@ dependencies = [
  "tenant_size_model",
  "thiserror",
  "tokio",
+ "tokio-epoll-uring",
  "tokio-io-timeout",
  "tokio-postgres",
  "tokio-stream",
@@ -5335,18 +5345,18 @@ dependencies = [
 
 [[package]]
 name = "thiserror"
-version = "1.0.40"
+version = "1.0.47"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "978c9a314bd8dc99be594bc3c175faaa9794be04a5a5e153caba6915336cebac"
+checksum = "97a802ec30afc17eee47b2855fc72e0c4cd62be9b4efe6591edde0ec5bd68d8f"
 dependencies = [
  "thiserror-impl",
 ]
 
 [[package]]
 name = "thiserror-impl"
-version = "1.0.40"
+version = "1.0.47"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "f9456a42c5b0d803c8cd86e73dd7cc9edd429499f37a3550d286d5e86720569f"
+checksum = "6bb623b56e39ab7dcd4b1b98bb6c8f8d907ed255b18de254088016b27a8ee19b"
 dependencies = [
  "proc-macro2",
  "quote",
@@ -5470,6 +5480,21 @@ dependencies = [
  "windows-sys 0.48.0",
 ]
 
+[[package]]
+name = "tokio-epoll-uring"
+version = "0.1.0"
+source = "git+https://github.com/neondatabase/tokio-epoll-uring.git?branch=problame/macos-usage-alternative#f3983bcce513d80f938ca6f1e9eb038e7803afeb"
+dependencies = [
+ "futures",
+ "once_cell",
+ "scopeguard",
+ "thiserror",
+ "tokio",
+ "tokio-util",
+ "tracing",
+ "uring-common",
+]
+
 [[package]]
 name = "tokio-io-timeout"
 version = "1.2.0"
@@ -6021,6 +6046,15 @@ dependencies = [
  "webpki-roots 0.23.1",
 ]
 
+[[package]]
+name = "uring-common"
+version = "0.1.0"
+source = "git+https://github.com/neondatabase/tokio-epoll-uring.git?branch=problame/macos-usage-alternative#f3983bcce513d80f938ca6f1e9eb038e7803afeb"
+dependencies = [
+ "io-uring",
+ "libc",
+]
+
 [[package]]
 name = "url"
 version = "2.3.1"

Cargo.toml

@@ -173,6 +173,12 @@ x509-parser = "0.15"
 env_logger = "0.10"
 log = "0.4"
 
+## tokio-epoll-uring
+#tokio-epoll-uring = { path = "../tokio-epoll-uring/tokio-epoll-uring" }
+#uring-common = { path = "../tokio-epoll-uring/uring-common" }
+tokio-epoll-uring = { git = "https://github.com/neondatabase/tokio-epoll-uring.git" , branch = "problame/macos-usage-alternative" }
+#uring-common = { git = "https://github.com/neondatabase/tokio-epoll-uring.git" , branch = "problame/macos-usage" }
+
 ## Libraries from neondatabase/ git forks, ideally with changes to be upstreamed
 postgres = { git = "https://github.com/neondatabase/rust-postgres.git", branch="neon" }
 postgres-native-tls = { git = "https://github.com/neondatabase/rust-postgres.git", branch="neon" }

compute_tools/src/bin/compute_ctl.rs

@@ -231,9 +231,9 @@ fn main() -> Result<()> {
     let compute = Arc::new(compute_node);
 
     // If this is a pooled VM, prewarm before starting HTTP server and becoming
-    // available for binding. Prewarming helps Postgres start quicker later,
+    // available for binding. Prewarming helps postgres start quicker later,
     // because QEMU will already have it's memory allocated from the host, and
-    // the necessary binaries will already be cached.
+    // the necessary binaries will alreaady be cached.
     if !spec_set {
         compute.prewarm_postgres()?;
     }
@@ -276,11 +276,6 @@ fn main() -> Result<()> {
 
     state.status = ComputeStatus::Init;
     compute.state_changed.notify_all();
-
-    info!(
-        "running compute with features: {:?}",
-        state.pspec.as_ref().unwrap().spec.features
-    );
     drop(state);
 
     // Launch remaining service threads
@@ -293,7 +288,7 @@ fn main() -> Result<()> {
     let pg = match compute.start_compute(extension_server_port) {
         Ok(pg) => Some(pg),
         Err(err) => {
-            error!("could not start the compute node: {:#}", err);
+            error!("could not start the compute node: {:?}", err);
             let mut state = compute.state.lock().unwrap();
             state.error = Some(format!("{:?}", err));
             state.status = ComputeStatus::Failed;

compute_tools/src/compute.rs

@@ -20,7 +20,7 @@ use futures::StreamExt;
 use postgres::{Client, NoTls};
 use tokio;
 use tokio_postgres;
-use tracing::{debug, error, info, instrument, warn};
+use tracing::{error, info, instrument, warn};
 use utils::id::{TenantId, TimelineId};
 use utils::lsn::Lsn;
 
@@ -280,7 +280,7 @@ fn create_neon_superuser(spec: &ComputeSpec, client: &mut Client) -> Result<()>
             $$;"#,
         roles_decl, database_decl,
     );
-    info!("Neon superuser created: {}", inlinify(&query));
+    info!("Neon superuser created:\n{}", inlinify(&query));
     client
         .simple_query(&query)
         .map_err(|e| anyhow::anyhow!(e).context(query))?;
@@ -964,16 +964,6 @@ impl ComputeNode {
         Ok(pg_process)
     }
 
-    /// Update the `last_active` in the shared state, but ensure that it's a more recent one.
-    pub fn update_last_active(&self, last_active: Option<DateTime<Utc>>) {
-        let mut state = self.state.lock().unwrap();
-        // NB: `Some(<DateTime>)` is always greater than `None`.
-        if last_active > state.last_active {
-            state.last_active = last_active;
-            debug!("set the last compute activity time to: {:?}", last_active);
-        }
-    }
-
     // Look for core dumps and collect backtraces.
     //
     // EKS worker nodes have following core dump settings:

compute_tools/src/monitor.rs

@@ -3,118 +3,88 @@ use std::{thread, time::Duration};
 
 use chrono::{DateTime, Utc};
 use postgres::{Client, NoTls};
-use tracing::{debug, error, info, warn};
+use tracing::{debug, info, warn};
 
 use crate::compute::ComputeNode;
-use compute_api::responses::ComputeStatus;
-use compute_api::spec::ComputeFeature;
 
 const MONITOR_CHECK_INTERVAL: Duration = Duration::from_millis(500);
 
 // Spin in a loop and figure out the last activity time in the Postgres.
 // Then update it in the shared state. This function never errors out.
-// NB: the only expected panic is at `Mutex` unwrap(), all other errors
-// should be handled gracefully.
+// XXX: the only expected panic is at `RwLock` unwrap().
 fn watch_compute_activity(compute: &ComputeNode) {
     // Suppose that `connstr` doesn't change
     let connstr = compute.connstr.as_str();
-
-    // During startup and configuration we connect to every Postgres database,
-    // but we don't want to count this as some user activity. So wait until
-    // the compute fully started before monitoring activity.
-    wait_for_postgres_start(compute);
-
     // Define `client` outside of the loop to reuse existing connection if it's active.
     let mut client = Client::connect(connstr, NoTls);
 
-    let mut sleep = false;
-    let mut prev_active_time: Option<f64> = None;
-    let mut prev_sessions: Option<i64> = None;
-
-    if compute.has_feature(ComputeFeature::ActivityMonitorExperimental) {
-        info!("starting experimental activity monitor for {}", connstr);
-    } else {
-        info!("starting activity monitor for {}", connstr);
-    }
+    info!("watching Postgres activity at {}", connstr);
 
     loop {
-        // We use `continue` a lot, so it's more convenient to sleep at the top of the loop.
-        // But skip the first sleep, so we can connect to Postgres immediately.
-        if sleep {
-            // Should be outside of the mutex lock to allow others to read while we sleep.
-            thread::sleep(MONITOR_CHECK_INTERVAL);
-        } else {
-            sleep = true;
-        }
+        // Should be outside of the write lock to allow others to read while we sleep.
+        thread::sleep(MONITOR_CHECK_INTERVAL);
 
         match &mut client {
             Ok(cli) => {
                 if cli.is_closed() {
-                    info!("connection to Postgres is closed, trying to reconnect");
+                    info!("connection to postgres closed, trying to reconnect");
 
                     // Connection is closed, reconnect and try again.
                     client = Client::connect(connstr, NoTls);
                     continue;
                 }
 
-                // This is a new logic, only enable if the feature flag is set.
-                // TODO: remove this once we are sure that it works OR drop it altogether.
-                if compute.has_feature(ComputeFeature::ActivityMonitorExperimental) {
-                    // First, check if the total active time or sessions across all databases has changed.
-                    // If it did, it means that user executed some queries. In theory, it can even go down if
-                    // some databases were dropped, but it's still a user activity.
-                    match get_database_stats(cli) {
-                        Ok((active_time, sessions)) => {
-                            let mut detected_activity = false;
+                // Get all running client backends except ourself, use RFC3339 DateTime format.
+                let backends = cli
+                    .query(
+                        "SELECT state, to_char(state_change, 'YYYY-MM-DD\"T\"HH24:MI:SS.US\"Z\"') AS state_change
+                         FROM pg_stat_activity
+                         WHERE backend_type = 'client backend'
+                            AND pid != pg_backend_pid()
+                            AND usename != 'cloud_admin';", // XXX: find a better way to filter other monitors?
+                        &[],
+                    );
+                let mut last_active = compute.state.lock().unwrap().last_active;
 
-                            prev_active_time = match prev_active_time {
-                                Some(prev_active_time) => {
-                                    if active_time != prev_active_time {
-                                        detected_activity = true;
-                                    }
-                                    Some(active_time)
-                                }
-                                None => Some(active_time),
-                            };
-                            prev_sessions = match prev_sessions {
-                                Some(prev_sessions) => {
-                                    if sessions != prev_sessions {
-                                        detected_activity = true;
-                                    }
-                                    Some(sessions)
-                                }
-                                None => Some(sessions),
-                            };
+                if let Ok(backs) = backends {
+                    let mut idle_backs: Vec<DateTime<Utc>> = vec![];
 
-                            if detected_activity {
-                                // Update the last active time and continue, we don't need to
-                                // check backends state change.
-                                compute.update_last_active(Some(Utc::now()));
-                                continue;
+                    for b in backs.into_iter() {
+                        let state: String = match b.try_get("state") {
+                            Ok(state) => state,
+                            Err(_) => continue,
+                        };
+
+                        if state == "idle" {
+                            let change: String = match b.try_get("state_change") {
+                                Ok(state_change) => state_change,
+                                Err(_) => continue,
+                            };
+                            let change = DateTime::parse_from_rfc3339(&change);
+                            match change {
+                                Ok(t) => idle_backs.push(t.with_timezone(&Utc)),
+                                Err(e) => {
+                                    info!("cannot parse backend state_change DateTime: {}", e);
+                                    continue;
+                                }
                             }
+                        } else {
+                            // Found non-idle backend, so the last activity is NOW.
+                            // Save it and exit the for loop. Also clear the idle backend
+                            // `state_change` timestamps array as it doesn't matter now.
+                            last_active = Some(Utc::now());
+                            idle_backs.clear();
+                            break;
                         }
-                        Err(e) => {
-                            error!("could not get database statistics: {}", e);
-                            continue;
-                        }
+                    }
+
+                    // Get idle backend `state_change` with the max timestamp.
+                    if let Some(last) = idle_backs.iter().max() {
+                        last_active = Some(*last);
                     }
                 }
 
-                // Second, if database statistics is the same, check all backends state change,
-                // maybe there is some with more recent activity. `get_backends_state_change()`
-                // can return None or stale timestamp, so it's `compute.update_last_active()`
-                // responsibility to check if the new timestamp is more recent than the current one.
-                // This helps us to discover new sessions, that did nothing yet.
-                match get_backends_state_change(cli) {
-                    Ok(last_active) => {
-                        compute.update_last_active(last_active);
-                    }
-                    Err(e) => {
-                        error!("could not get backends state change: {}", e);
-                    }
-                }
-
-                // Finally, if there are existing (logical) walsenders, do not suspend.
+                // If there are existing (logical) walsenders, do not suspend.
                 //
                 // walproposer doesn't currently show up in pg_stat_replication,
                 // but protect if it will be
@@ -123,12 +93,11 @@ fn watch_compute_activity(compute: &ComputeNode) {
                     Ok(r) => match r.try_get::<&str, i64>("count") {
                         Ok(num_ws) => {
                             if num_ws > 0 {
-                                compute.update_last_active(Some(Utc::now()));
-                                continue;
+                                last_active = Some(Utc::now());
                             }
                         }
                         Err(e) => {
-                            warn!("failed to parse walsenders count: {:?}", e);
+                            warn!("failed to parse ws count: {:?}", e);
                             continue;
                         }
                     },
@@ -137,31 +106,17 @@ fn watch_compute_activity(compute: &ComputeNode) {
                         continue;
                     }
                 }
-                //
-                // Do not suspend compute if autovacuum is running
-                //
-                let autovacuum_count_query = "select count(*) from pg_stat_activity where backend_type = 'autovacuum worker'";
-                match cli.query_one(autovacuum_count_query, &[]) {
-                    Ok(r) => match r.try_get::<&str, i64>("count") {
-                        Ok(num_workers) => {
-                            if num_workers > 0 {
-                                compute.update_last_active(Some(Utc::now()));
-                                continue;
-                            }
-                        }
-                        Err(e) => {
-                            warn!("failed to parse autovacuum workers count: {:?}", e);
-                            continue;
-                        }
-                    },
-                    Err(e) => {
-                        warn!("failed to get list of autovacuum workers: {:?}", e);
-                        continue;
-                    }
+
+                // Update the last activity in the shared state if we got a more recent one.
+                let mut state = compute.state.lock().unwrap();
+                // NB: `Some(<DateTime>)` is always greater than `None`.
+                if last_active > state.last_active {
+                    state.last_active = last_active;
+                    debug!("set the last compute activity time to: {:?}", last_active);
                 }
             }
             Err(e) => {
-                debug!("could not connect to Postgres: {}, retrying", e);
+                debug!("cannot connect to postgres: {}, retrying", e);
 
                 // Establish a new connection and try again.
                 client = Client::connect(connstr, NoTls);
@@ -170,124 +125,12 @@ fn watch_compute_activity(compute: &ComputeNode) {
     }
 }
 
-// Hang on condition variable waiting until the compute status is `Running`.
-fn wait_for_postgres_start(compute: &ComputeNode) {
-    let mut state = compute.state.lock().unwrap();
-    while state.status != ComputeStatus::Running {
-        info!("compute is not running, waiting before monitoring activity");
-        state = compute.state_changed.wait(state).unwrap();
-
-        if state.status == ComputeStatus::Running {
-            break;
-        }
-    }
-}
-
-// Figure out the total active time and sessions across all non-system databases.
-// Returned tuple is `(active_time, sessions)`.
-// It can return `0.0` active time or `0` sessions, which means no user databases exist OR
-// it was a start with skipped `pg_catalog` updates and user didn't do any queries
-// (or open any sessions) yet.
-fn get_database_stats(cli: &mut Client) -> anyhow::Result<(f64, i64)> {
-    // Filter out `postgres` database as `compute_ctl` and other monitoring tools
-    // like `postgres_exporter` use it to query Postgres statistics.
-    // Use explicit 8 bytes type casts to match Rust types.
-    let stats = cli.query_one(
-        "SELECT coalesce(sum(active_time), 0.0)::float8 AS total_active_time,
-            coalesce(sum(sessions), 0)::bigint AS total_sessions
-        FROM pg_stat_database
-        WHERE datname NOT IN (
-                'postgres',
-                'template0',
-                'template1'
-            );",
-        &[],
-    );
-    let stats = match stats {
-        Ok(stats) => stats,
-        Err(e) => {
-            return Err(anyhow::anyhow!("could not query active_time: {}", e));
-        }
-    };
-
-    let active_time: f64 = match stats.try_get("total_active_time") {
-        Ok(active_time) => active_time,
-        Err(e) => return Err(anyhow::anyhow!("could not get total_active_time: {}", e)),
-    };
-
-    let sessions: i64 = match stats.try_get("total_sessions") {
-        Ok(sessions) => sessions,
-        Err(e) => return Err(anyhow::anyhow!("could not get total_sessions: {}", e)),
-    };
-
-    Ok((active_time, sessions))
-}
-
-// Figure out the most recent state change time across all client backends.
-// If there is currently active backend, timestamp will be `Utc::now()`.
-// It can return `None`, which means no client backends exist or we were
-// unable to parse the timestamp.
-fn get_backends_state_change(cli: &mut Client) -> anyhow::Result<Option<DateTime<Utc>>> {
-    let mut last_active: Option<DateTime<Utc>> = None;
-    // Get all running client backends except ourself, use RFC3339 DateTime format.
-    let backends = cli.query(
-        "SELECT state, to_char(state_change, 'YYYY-MM-DD\"T\"HH24:MI:SS.US\"Z\"') AS state_change
-                FROM pg_stat_activity
-                    WHERE backend_type = 'client backend'
-                    AND pid != pg_backend_pid()
-                    AND usename != 'cloud_admin';", // XXX: find a better way to filter other monitors?
-        &[],
-    );
-
-    match backends {
-        Ok(backs) => {
-            let mut idle_backs: Vec<DateTime<Utc>> = vec![];
-
-            for b in backs.into_iter() {
-                let state: String = match b.try_get("state") {
-                    Ok(state) => state,
-                    Err(_) => continue,
-                };
-
-                if state == "idle" {
-                    let change: String = match b.try_get("state_change") {
-                        Ok(state_change) => state_change,
-                        Err(_) => continue,
-                    };
-                    let change = DateTime::parse_from_rfc3339(&change);
-                    match change {
-                        Ok(t) => idle_backs.push(t.with_timezone(&Utc)),
-                        Err(e) => {
-                            info!("cannot parse backend state_change DateTime: {}", e);
-                            continue;
-                        }
-                    }
-                } else {
-                    // Found non-idle backend, so the last activity is NOW.
-                    // Return immediately, no need to check other backends.
-                    return Ok(Some(Utc::now()));
-                }
-            }
-
-            // Get idle backend `state_change` with the max timestamp.
-            if let Some(last) = idle_backs.iter().max() {
-                last_active = Some(*last);
-            }
-        }
-        Err(e) => {
-            return Err(anyhow::anyhow!("could not query backends: {}", e));
-        }
-    }
-
-    Ok(last_active)
-}
-
 /// Launch a separate compute monitor thread and return its `JoinHandle`.
-pub fn launch_monitor(compute: &Arc<ComputeNode>) -> thread::JoinHandle<()> {
-    let compute = Arc::clone(compute);
+pub fn launch_monitor(state: &Arc<ComputeNode>) -> thread::JoinHandle<()> {
+    let state = Arc::clone(state);
 
     thread::Builder::new()
         .name("compute-monitor".into())
-        .spawn(move || watch_compute_activity(&compute))
+        .spawn(move || watch_compute_activity(&state))
         .expect("cannot launch compute monitor thread")
 }

compute_tools/src/pg_helpers.rs

@@ -456,7 +456,7 @@ pub fn handle_postgres_logs(stderr: std::process::ChildStderr) -> JoinHandle<()>
 /// - no new lines were written for the last second
 async fn handle_postgres_logs_async(stderr: tokio::process::ChildStderr) -> Result<()> {
     let mut lines = tokio::io::BufReader::new(stderr).lines();
-    let timeout_duration = Duration::from_millis(100);
+    let timeout_duration = Duration::from_secs(1);
     let ts_regex =
         regex::Regex::new(r"^\d+-\d{2}-\d{2} \d{2}:\d{2}:\d{2}").expect("regex is valid");
 

compute_tools/src/spec.rs

@@ -190,20 +190,18 @@ pub fn handle_roles(spec: &ComputeSpec, client: &mut Client) -> Result<()> {
 
     // Print a list of existing Postgres roles (only in debug mode)
     if span_enabled!(Level::INFO) {
-        let mut vec = Vec::new();
+        info!("postgres roles:");
         for r in &existing_roles {
-            vec.push(format!(
-                "{}:{}",
+            info!(
+                "    - {}:{}",
                 r.name,
                 if r.encrypted_password.is_some() {
                     "[FILTERED]"
                 } else {
                     "(null)"
                 }
-            ));
+            );
         }
-
-        info!("postgres roles (total {}): {:?}", vec.len(), vec);
     }
 
     // Process delta operations first
@@ -241,10 +239,7 @@ pub fn handle_roles(spec: &ComputeSpec, client: &mut Client) -> Result<()> {
     // Refresh Postgres roles info to handle possible roles renaming
     let existing_roles: Vec<Role> = get_existing_roles(&mut xact)?;
 
-    info!(
-        "handling cluster spec roles (total {})",
-        spec.cluster.roles.len()
-    );
+    info!("cluster spec roles:");
     for role in &spec.cluster.roles {
         let name = &role.name;
         // XXX: with a limited number of roles it is fine, but consider making it a HashMap
@@ -307,7 +302,7 @@ pub fn handle_roles(spec: &ComputeSpec, client: &mut Client) -> Result<()> {
                     "CREATE ROLE {} INHERIT CREATEROLE CREATEDB BYPASSRLS REPLICATION IN ROLE neon_superuser",
                     name.pg_quote()
                 );
-                info!("running role create query: '{}'", &query);
+                info!("role create query: '{}'", &query);
                 query.push_str(&role.to_pg_options());
                 xact.execute(query.as_str(), &[])?;
             }
@@ -324,7 +319,7 @@ pub fn handle_roles(spec: &ComputeSpec, client: &mut Client) -> Result<()> {
                 RoleAction::Create => " -> create",
                 RoleAction::Update => " -> update",
             };
-            info!(" - {}:{}{}", name, pwd, action_str);
+            info!("   - {}:{}{}", name, pwd, action_str);
         }
     }
 
@@ -433,11 +428,10 @@ pub fn handle_databases(spec: &ComputeSpec, client: &mut Client) -> Result<()> {
 
     // Print a list of existing Postgres databases (only in debug mode)
     if span_enabled!(Level::INFO) {
-        let mut vec = Vec::new();
+        info!("postgres databases:");
         for (dbname, db) in &existing_dbs {
-            vec.push(format!("{}:{}", dbname, db.owner));
+            info!("    {}:{}", dbname, db.owner);
         }
-        info!("postgres databases (total {}): {:?}", vec.len(), vec);
     }
 
     // Process delta operations first
@@ -509,10 +503,7 @@ pub fn handle_databases(spec: &ComputeSpec, client: &mut Client) -> Result<()> {
     // Refresh Postgres databases info to handle possible renames
     let existing_dbs = get_existing_dbs(client)?;
 
-    info!(
-        "handling cluster spec databases (total {})",
-        spec.cluster.databases.len()
-    );
+    info!("cluster spec databases:");
     for db in &spec.cluster.databases {
         let name = &db.name;
         let pg_db = existing_dbs.get(name);
@@ -571,7 +562,7 @@ pub fn handle_databases(spec: &ComputeSpec, client: &mut Client) -> Result<()> {
                 DatabaseAction::Create => " -> create",
                 DatabaseAction::Update => " -> update",
             };
-            info!(" - {}:{}{}", db.name, db.owner, action_str);
+            info!("   - {}:{}{}", db.name, db.owner, action_str);
         }
     }
 

control_plane/src/bin/neon_local.rs

@@ -6,11 +6,11 @@
 //! rely on `neon_local` to set up the environment for each test.
 //!
 use anyhow::{anyhow, bail, Context, Result};
-use clap::{value_parser, Arg, ArgAction, ArgMatches, Command, ValueEnum};
+use clap::{value_parser, Arg, ArgAction, ArgMatches, Command};
 use compute_api::spec::ComputeMode;
 use control_plane::attachment_service::AttachmentService;
 use control_plane::endpoint::ComputeControlPlane;
-use control_plane::local_env::{InitForceMode, LocalEnv};
+use control_plane::local_env::LocalEnv;
 use control_plane::pageserver::{PageServerNode, PAGESERVER_REMOTE_STORAGE_DIR};
 use control_plane::safekeeper::SafekeeperNode;
 use control_plane::tenant_migration::migrate_tenant;
@@ -338,7 +338,7 @@ fn handle_init(init_match: &ArgMatches) -> anyhow::Result<LocalEnv> {
 
     let mut env =
         LocalEnv::parse_config(&toml_file).context("Failed to create neon configuration")?;
-    let force = init_match.get_one("force").expect("we set a default value");
+    let force = init_match.get_flag("force");
     env.init(pg_version, force)
         .context("Failed to initialize neon repository")?;
 
@@ -1266,15 +1266,9 @@ fn cli() -> Command {
         .required(false);
 
     let force_arg = Arg::new("force")
-        .value_parser(value_parser!(InitForceMode))
+        .value_parser(value_parser!(bool))
         .long("force")
-        .default_value(
-            InitForceMode::MustNotExist
-                .to_possible_value()
-                .unwrap()
-                .get_name()
-                .to_owned(),
-        )
+        .action(ArgAction::SetTrue)
         .help("Force initialization even if the repository is not empty")
         .required(false);
 

control_plane/src/local_env.rs

@@ -5,7 +5,6 @@
 
 use anyhow::{bail, ensure, Context};
 
-use clap::ValueEnum;
 use postgres_backend::AuthType;
 use reqwest::Url;
 use serde::{Deserialize, Serialize};
@@ -163,31 +162,6 @@ impl Default for SafekeeperConf {
     }
 }
 
-#[derive(Clone, Copy)]
-pub enum InitForceMode {
-    MustNotExist,
-    EmptyDirOk,
-    RemoveAllContents,
-}
-
-impl ValueEnum for InitForceMode {
-    fn value_variants<'a>() -> &'a [Self] {
-        &[
-            Self::MustNotExist,
-            Self::EmptyDirOk,
-            Self::RemoveAllContents,
-        ]
-    }
-
-    fn to_possible_value(&self) -> Option<clap::builder::PossibleValue> {
-        Some(clap::builder::PossibleValue::new(match self {
-            InitForceMode::MustNotExist => "must-not-exist",
-            InitForceMode::EmptyDirOk => "empty-dir-ok",
-            InitForceMode::RemoveAllContents => "remove-all-contents",
-        }))
-    }
-}
-
 impl SafekeeperConf {
     /// Compute is served by port on which only tenant scoped tokens allowed, if
     /// it is configured.
@@ -410,7 +384,7 @@ impl LocalEnv {
     //
     // Initialize a new Neon repository
     //
-    pub fn init(&mut self, pg_version: u32, force: &InitForceMode) -> anyhow::Result<()> {
+    pub fn init(&mut self, pg_version: u32, force: bool) -> anyhow::Result<()> {
         // check if config already exists
         let base_path = &self.base_data_dir;
         ensure!(
@@ -419,34 +393,25 @@ impl LocalEnv {
         );
 
         if base_path.exists() {
-            match force {
-                InitForceMode::MustNotExist => {
-                    bail!(
-                        "directory '{}' already exists. Perhaps already initialized?",
-                        base_path.display()
-                    );
-                }
-                InitForceMode::EmptyDirOk => {
-                    if let Some(res) = std::fs::read_dir(base_path)?.next() {
-                        res.context("check if directory is empty")?;
-                        anyhow::bail!("directory not empty: {base_path:?}");
-                    }
-                }
-                InitForceMode::RemoveAllContents => {
-                    println!("removing all contents of '{}'", base_path.display());
-                    // instead of directly calling `remove_dir_all`, we keep the original dir but removing
-                    // all contents inside. This helps if the developer symbol links another directory (i.e.,
-                    // S3 local SSD) to the `.neon` base directory.
-                    for entry in std::fs::read_dir(base_path)? {
-                        let entry = entry?;
-                        let path = entry.path();
-                        if path.is_dir() {
-                            fs::remove_dir_all(&path)?;
-                        } else {
-                            fs::remove_file(&path)?;
-                        }
+            if force {
+                println!("removing all contents of '{}'", base_path.display());
+                // instead of directly calling `remove_dir_all`, we keep the original dir but removing
+                // all contents inside. This helps if the developer symbol links another directory (i.e.,
+                // S3 local SSD) to the `.neon` base directory.
+                for entry in std::fs::read_dir(base_path)? {
+                    let entry = entry?;
+                    let path = entry.path();
+                    if path.is_dir() {
+                        fs::remove_dir_all(&path)?;
+                    } else {
+                        fs::remove_file(&path)?;
                     }
                 }
+            } else {
+                bail!(
+                    "directory '{}' already exists. Perhaps already initialized? (Hint: use --force to remove all contents)",
+                    base_path.display()
+                );
             }
         }
 

libs/compute_api/src/spec.rs

@@ -82,13 +82,10 @@ pub struct ComputeSpec {
 #[serde(rename_all = "snake_case")]
 pub enum ComputeFeature {
     // XXX: Add more feature flags here.
-    /// Enable the experimental activity monitor logic, which uses `pg_stat_database` to
-    /// track short-lived connections as user activity.
-    ActivityMonitorExperimental,
 
-    /// This is a special feature flag that is used to represent unknown feature flags.
-    /// Basically all unknown to enum flags are represented as this one. See unit test
-    /// `parse_unknown_features()` for more details.
+    // This is a special feature flag that is used to represent unknown feature flags.
+    // Basically all unknown to enum flags are represented as this one. See unit test
+    // `parse_unknown_features()` for more details.
     #[serde(other)]
     UnknownFeature,
 }
@@ -285,23 +282,4 @@ mod tests {
         assert!(spec.features.contains(&ComputeFeature::UnknownFeature));
         assert_eq!(spec.features, vec![ComputeFeature::UnknownFeature; 2]);
     }
-
-    #[test]
-    fn parse_known_features() {
-        // Test that we can properly parse known feature flags.
-        let file = File::open("tests/cluster_spec.json").unwrap();
-        let mut json: serde_json::Value = serde_json::from_reader(file).unwrap();
-        let ob = json.as_object_mut().unwrap();
-
-        // Add known feature flags.
-        let features = vec!["activity_monitor_experimental"];
-        ob.insert("features".into(), features.into());
-
-        let spec: ComputeSpec = serde_json::from_value(json).unwrap();
-
-        assert_eq!(
-            spec.features,
-            vec![ComputeFeature::ActivityMonitorExperimental]
-        );
-    }
 }

libs/pageserver_api/src/key.rs

@@ -3,8 +3,6 @@ use byteorder::{ByteOrder, BE};
 use serde::{Deserialize, Serialize};
 use std::fmt;
 
-use crate::reltag::{BlockNumber, RelTag};
-
 /// Key used in the Repository kv-store.
 ///
 /// The Repository treats this as an opaque struct, but see the code in pgdatadir_mapping.rs
@@ -148,22 +146,6 @@ pub fn is_rel_block_key(key: &Key) -> bool {
     key.field1 == 0x00 && key.field4 != 0 && key.field6 != 0xffffffff
 }
 
-/// Guaranteed to return `Ok()` if [[is_rel_block_key]] returns `true` for `key`.
-pub fn key_to_rel_block(key: Key) -> anyhow::Result<(RelTag, BlockNumber)> {
-    Ok(match key.field1 {
-        0x00 => (
-            RelTag {
-                spcnode: key.field2,
-                dbnode: key.field3,
-                relnode: key.field4,
-                forknum: key.field5,
-            },
-            key.field6,
-        ),
-        _ => anyhow::bail!("unexpected value kind 0x{:02x}", key.field1),
-    })
-}
-
 impl std::str::FromStr for Key {
     type Err = anyhow::Error;
 

libs/pageserver_api/src/models.rs

@@ -368,8 +368,6 @@ pub struct TenantInfo {
     /// If a layer is present in both local FS and S3, it counts only once.
     pub current_physical_size: Option<u64>, // physical size is only included in `tenant_status` endpoint
     pub attachment_status: TenantAttachmentStatus,
-    #[serde(skip_serializing_if = "Option::is_none")]
-    pub generation: Option<u32>,
 }
 
 #[derive(Serialize, Deserialize, Clone)]
@@ -912,7 +910,6 @@ mod tests {
             state: TenantState::Active,
             current_physical_size: Some(42),
             attachment_status: TenantAttachmentStatus::Attached,
-            generation: None,
         };
         let expected_active = json!({
             "id": original_active.id.to_string(),
@@ -933,7 +930,6 @@ mod tests {
             },
             current_physical_size: Some(42),
             attachment_status: TenantAttachmentStatus::Attached,
-            generation: None,
         };
         let expected_broken = json!({
             "id": original_broken.id.to_string(),

libs/pageserver_api/src/reltag.rs

@@ -32,9 +32,6 @@ pub struct RelTag {
     pub relnode: Oid,
 }
 
-/// Block number within a relation or SLRU. This matches PostgreSQL's BlockNumber type.
-pub type BlockNumber = u32;
-
 impl PartialOrd for RelTag {
     fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
         Some(self.cmp(other))

pageserver/Cargo.toml

@@ -61,6 +61,7 @@ sync_wrapper.workspace = true
 tokio-tar.workspace = true
 thiserror.workspace = true
 tokio = { workspace = true, features = ["process", "sync", "fs", "rt", "io-util", "time"] }
+tokio-epoll-uring.workspace = true
 tokio-io-timeout.workspace = true
 tokio-postgres.workspace = true
 tokio-stream.workspace = true

pageserver/client/src/mgmt_api.rs

@@ -221,16 +221,4 @@ impl Client {
             .await
             .map_err(Error::ReceiveBody)
     }
-
-    pub async fn tenant_reset(&self, tenant_shard_id: TenantShardId) -> Result<()> {
-        let uri = format!(
-            "{}/v1/tenant/{}/reset",
-            self.mgmt_api_endpoint, tenant_shard_id
-        );
-        self.request(Method::POST, &uri, ())
-            .await?
-            .json()
-            .await
-            .map_err(Error::ReceiveBody)
-    }
 }

pageserver/ctl/src/layer_map_analyzer.rs

@@ -18,7 +18,7 @@ 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::virtual_file::VirtualFile;
+use pageserver::virtual_file::{self, VirtualFile};
 
 use utils::{bin_ser::BeSer, lsn::Lsn};
 
@@ -142,7 +142,7 @@ pub(crate) async fn main(cmd: &AnalyzeLayerMapCmd) -> Result<()> {
     let ctx = RequestContext::new(TaskKind::DebugTool, DownloadBehavior::Error);
 
     // Initialize virtual_file (file desriptor cache) and page cache which are needed to access layer persistent B-Tree.
-    pageserver::virtual_file::init(10);
+    pageserver::virtual_file::init(10, virtual_file::IoEngineKind::StdFs);
     pageserver::page_cache::init(100);
 
     let mut total_delta_layers = 0usize;

pageserver/ctl/src/layers.rs

@@ -59,7 +59,7 @@ 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::init(10, virtual_file::IoEngineKind::StdFs);
     page_cache::init(100);
     let file = FileBlockReader::new(VirtualFile::open(path).await?);
     let summary_blk = file.read_blk(0, ctx).await?;
@@ -187,7 +187,7 @@ pub(crate) async fn main(cmd: &LayerCmd) -> Result<()> {
             new_tenant_id,
             new_timeline_id,
         } => {
-            pageserver::virtual_file::init(10);
+            pageserver::virtual_file::init(10, virtual_file::IoEngineKind::StdFs);
             pageserver::page_cache::init(100);
 
             let ctx = RequestContext::new(TaskKind::DebugTool, DownloadBehavior::Error);

pageserver/ctl/src/main.rs

@@ -123,7 +123,7 @@ fn read_pg_control_file(control_file_path: &Utf8Path) -> anyhow::Result<()> {
 
 async fn print_layerfile(path: &Utf8Path) -> anyhow::Result<()> {
     // Basic initialization of things that don't change after startup
-    virtual_file::init(10);
+    virtual_file::init(10, virtual_file::IoEngineKind::StdFs);
     page_cache::init(100);
     let ctx = RequestContext::new(TaskKind::DebugTool, DownloadBehavior::Error);
     dump_layerfile_from_path(path, true, &ctx).await

pageserver/pagebench/Cargo.toml

@@ -8,7 +8,6 @@ license.workspace = true
 
 [dependencies]
 anyhow.workspace = true
-camino.workspace = true
 clap.workspace = true
 futures.workspace = true
 hdrhistogram.workspace = true
@@ -19,8 +18,8 @@ serde.workspace = true
 serde_json.workspace = true
 tracing.workspace = true
 tokio.workspace = true
-tokio-util.workspace = true
 
+pageserver = { path = ".." }
 pageserver_client.workspace = true
 pageserver_api.workspace = true
 utils = { path = "../../libs/utils/" }

pageserver/pagebench/src/cmd/getpage_latest_lsn.rs

@@ -1,11 +1,11 @@
 use anyhow::Context;
-use camino::Utf8PathBuf;
 use futures::future::join_all;
-use pageserver_api::key::{is_rel_block_key, key_to_rel_block, Key};
+use pageserver::pgdatadir_mapping::key_to_rel_block;
+use pageserver::repository;
+use pageserver_api::key::is_rel_block_key;
 use pageserver_api::keyspace::KeySpaceAccum;
 use pageserver_api::models::PagestreamGetPageRequest;
 
-use tokio_util::sync::CancellationToken;
 use utils::id::TenantTimelineId;
 use utils::lsn::Lsn;
 
@@ -14,7 +14,7 @@ use tokio::sync::Barrier;
 use tokio::task::JoinSet;
 use tracing::{info, instrument};
 
-use std::collections::{HashMap, HashSet};
+use std::collections::HashMap;
 use std::future::Future;
 use std::num::NonZeroUsize;
 use std::pin::Pin;
@@ -45,12 +45,6 @@ pub(crate) struct Args {
     req_latest_probability: f64,
     #[clap(long)]
     limit_to_first_n_targets: Option<usize>,
-    /// For large pageserver installations, enumerating the keyspace takes a lot of time.
-    /// If specified, the specified path is used to maintain a cache of the keyspace enumeration result.
-    /// The cache is tagged and auto-invalided by the tenant/timeline ids only.
-    /// It doesn't get invalidated if the keyspace changes under the hood, e.g., due to new ingested data or compaction.
-    #[clap(long)]
-    keyspace_cache: Option<Utf8PathBuf>,
     targets: Option<Vec<TenantTimelineId>>,
 }
 
@@ -65,7 +59,7 @@ impl LiveStats {
     }
 }
 
-#[derive(Clone, serde::Serialize, serde::Deserialize)]
+#[derive(Clone)]
 struct KeyRange {
     timeline: TenantTimelineId,
     timeline_lsn: Lsn,
@@ -113,107 +107,63 @@ async fn main_impl(
     )
     .await?;
 
-    #[derive(serde::Deserialize)]
-    struct KeyspaceCacheDe {
-        tag: Vec<TenantTimelineId>,
-        data: Vec<KeyRange>,
-    }
-    #[derive(serde::Serialize)]
-    struct KeyspaceCacheSer<'a> {
-        tag: &'a [TenantTimelineId],
-        data: &'a [KeyRange],
-    }
-    let cache = args
-        .keyspace_cache
-        .as_ref()
-        .map(|keyspace_cache_file| {
-            let contents = match std::fs::read(keyspace_cache_file) {
-                Err(e) if e.kind() == std::io::ErrorKind::NotFound => {
-                    return anyhow::Ok(None);
-                }
-                x => x.context("read keyspace cache file")?,
-            };
-            let cache: KeyspaceCacheDe =
-                serde_json::from_slice(&contents).context("deserialize cache file")?;
-            let tag_ok = HashSet::<TenantTimelineId>::from_iter(cache.tag.into_iter())
-                == HashSet::from_iter(timelines.iter().cloned());
-            info!("keyspace cache file matches tag: {tag_ok}");
-            anyhow::Ok(if tag_ok { Some(cache.data) } else { None })
-        })
-        .transpose()?
-        .flatten();
-    let all_ranges: Vec<KeyRange> = if let Some(cached) = cache {
-        info!("using keyspace cache file");
-        cached
-    } else {
-        let mut js = JoinSet::new();
-        for timeline in &timelines {
-            js.spawn({
-                let mgmt_api_client = Arc::clone(&mgmt_api_client);
-                let timeline = *timeline;
-                async move {
-                    let partitioning = mgmt_api_client
-                        .keyspace(timeline.tenant_id, timeline.timeline_id)
-                        .await?;
-                    let lsn = partitioning.at_lsn;
-                    let start = Instant::now();
-                    let mut filtered = KeySpaceAccum::new();
-                    // let's hope this is inlined and vectorized...
-                    // TODO: turn this loop into a is_rel_block_range() function.
-                    for r in partitioning.keys.ranges.iter() {
-                        let mut i = r.start;
-                        while i != r.end {
-                            if is_rel_block_key(&i) {
-                                filtered.add_key(i);
-                            }
-                            i = i.next();
+    let mut js = JoinSet::new();
+    for timeline in &timelines {
+        js.spawn({
+            let mgmt_api_client = Arc::clone(&mgmt_api_client);
+            let timeline = *timeline;
+            async move {
+                let partitioning = mgmt_api_client
+                    .keyspace(timeline.tenant_id, timeline.timeline_id)
+                    .await?;
+                let lsn = partitioning.at_lsn;
+                let start = Instant::now();
+                let mut filtered = KeySpaceAccum::new();
+                // let's hope this is inlined and vectorized...
+                // TODO: turn this loop into a is_rel_block_range() function.
+                for r in partitioning.keys.ranges.iter() {
+                    let mut i = r.start;
+                    while i != r.end {
+                        if is_rel_block_key(&i) {
+                            filtered.add_key(i);
                         }
+                        i = i.next();
                     }
-                    let filtered = filtered.to_keyspace();
-                    let filter_duration = start.elapsed();
-
-                    anyhow::Ok((
-                        filter_duration,
-                        filtered.ranges.into_iter().map(move |r| KeyRange {
-                            timeline,
-                            timeline_lsn: lsn,
-                            start: r.start.to_i128(),
-                            end: r.end.to_i128(),
-                        }),
-                    ))
                 }
-            });
-        }
-        let mut total_filter_duration = Duration::from_secs(0);
-        let mut all_ranges: Vec<KeyRange> = Vec::new();
-        while let Some(res) = js.join_next().await {
-            let (filter_duration, range) = res.unwrap().unwrap();
-            all_ranges.extend(range);
-            total_filter_duration += filter_duration;
-        }
-        info!("filter duration: {}", total_filter_duration.as_secs_f64());
-        if let Some(cachefile) = args.keyspace_cache.as_ref() {
-            let cache = KeyspaceCacheSer {
-                tag: &timelines,
-                data: &all_ranges,
-            };
-            let bytes = serde_json::to_vec(&cache).context("serialize keyspace for cache file")?;
-            std::fs::write(cachefile, bytes).context("write keyspace cache file to disk")?;
-            info!("successfully wrote keyspace cache file");
-        }
-        all_ranges
-    };
+                let filtered = filtered.to_keyspace();
+                let filter_duration = start.elapsed();
+
+                anyhow::Ok((
+                    filter_duration,
+                    filtered.ranges.into_iter().map(move |r| KeyRange {
+                        timeline,
+                        timeline_lsn: lsn,
+                        start: r.start.to_i128(),
+                        end: r.end.to_i128(),
+                    }),
+                ))
+            }
+        });
+    }
+    let mut total_filter_duration = Duration::from_secs(0);
+    let mut all_ranges: Vec<KeyRange> = Vec::new();
+    while let Some(res) = js.join_next().await {
+        let (filter_duration, range) = res.unwrap().unwrap();
+        all_ranges.extend(range);
+        total_filter_duration += filter_duration;
+    }
+    info!("filter duration: {}", total_filter_duration.as_secs_f64());
 
     let live_stats = Arc::new(LiveStats::default());
 
     let num_client_tasks = timelines.len();
     let num_live_stats_dump = 1;
     let num_work_sender_tasks = 1;
-    let num_main_impl = 1;
 
     let start_work_barrier = Arc::new(tokio::sync::Barrier::new(
-        num_client_tasks + num_live_stats_dump + num_work_sender_tasks + num_main_impl,
+        num_client_tasks + num_live_stats_dump + num_work_sender_tasks,
     ));
+    let all_work_done_barrier = Arc::new(tokio::sync::Barrier::new(num_client_tasks));
 
     tokio::spawn({
         let stats = Arc::clone(&live_stats);
@@ -233,143 +183,126 @@ async fn main_impl(
         }
     });
 
-    let cancel = CancellationToken::new();
-
-    let mut work_senders: HashMap<TenantTimelineId, _> = HashMap::new();
+    let mut work_senders = HashMap::new();
     let mut tasks = Vec::new();
     for tl in &timelines {
         let (sender, receiver) = tokio::sync::mpsc::channel(10); // TODO: not sure what the implications of this are
-        work_senders.insert(*tl, sender);
+        work_senders.insert(tl, sender);
         tasks.push(tokio::spawn(client(
             args,
             *tl,
             Arc::clone(&start_work_barrier),
             receiver,
+            Arc::clone(&all_work_done_barrier),
             Arc::clone(&live_stats),
-            cancel.clone(),
         )));
     }
 
-    let work_sender: Pin<Box<dyn Send + Future<Output = ()>>> = {
-        let start_work_barrier = start_work_barrier.clone();
-        let cancel = cancel.clone();
-        match args.per_target_rate_limit {
-            None => Box::pin(async move {
+    let work_sender: Pin<Box<dyn Send + Future<Output = ()>>> = match args.per_target_rate_limit {
+        None => Box::pin(async move {
+            let weights = rand::distributions::weighted::WeightedIndex::new(
+                all_ranges.iter().map(|v| v.len()),
+            )
+            .unwrap();
+
+            start_work_barrier.wait().await;
+
+            loop {
+                let (timeline, req) = {
+                    let mut rng = rand::thread_rng();
+                    let r = &all_ranges[weights.sample(&mut rng)];
+                    let key: i128 = rng.gen_range(r.start..r.end);
+                    let key = repository::Key::from_i128(key);
+                    let (rel_tag, block_no) =
+                        key_to_rel_block(key).expect("we filter non-rel-block keys out above");
+                    (
+                        r.timeline,
+                        PagestreamGetPageRequest {
+                            latest: rng.gen_bool(args.req_latest_probability),
+                            lsn: r.timeline_lsn,
+                            rel: rel_tag,
+                            blkno: block_no,
+                        },
+                    )
+                };
+                let sender = work_senders.get(&timeline).unwrap();
+                // TODO: what if this blocks?
+                sender.send(req).await.ok().unwrap();
+            }
+        }),
+        Some(rps_limit) => Box::pin(async move {
+            let period = Duration::from_secs_f64(1.0 / (rps_limit as f64));
+
+            let make_timeline_task: &dyn Fn(
+                TenantTimelineId,
+            )
+                -> Pin<Box<dyn Send + Future<Output = ()>>> = &|timeline| {
+                let sender = work_senders.get(&timeline).unwrap();
+                let ranges: Vec<KeyRange> = all_ranges
+                    .iter()
+                    .filter(|r| r.timeline == timeline)
+                    .cloned()
+                    .collect();
                 let weights = rand::distributions::weighted::WeightedIndex::new(
-                    all_ranges.iter().map(|v| v.len()),
+                    ranges.iter().map(|v| v.len()),
                 )
                 .unwrap();
 
-                start_work_barrier.wait().await;
-
-                while !cancel.is_cancelled() {
-                    let (timeline, req) = {
-                        let mut rng = rand::thread_rng();
-                        let r = &all_ranges[weights.sample(&mut rng)];
-                        let key: i128 = rng.gen_range(r.start..r.end);
-                        let key = Key::from_i128(key);
-                        let (rel_tag, block_no) =
-                            key_to_rel_block(key).expect("we filter non-rel-block keys out above");
-                        (
-                            r.timeline,
+                Box::pin(async move {
+                    let mut ticker = tokio::time::interval(period);
+                    ticker.set_missed_tick_behavior(
+                        /* TODO review this choice */
+                        tokio::time::MissedTickBehavior::Burst,
+                    );
+                    loop {
+                        ticker.tick().await;
+                        let req = {
+                            let mut rng = rand::thread_rng();
+                            let r = &ranges[weights.sample(&mut rng)];
+                            let key: i128 = rng.gen_range(r.start..r.end);
+                            let key = repository::Key::from_i128(key);
+                            assert!(is_rel_block_key(&key));
+                            let (rel_tag, block_no) = key_to_rel_block(key)
+                                .expect("we filter non-rel-block keys out above");
                             PagestreamGetPageRequest {
                                 latest: rng.gen_bool(args.req_latest_probability),
                                 lsn: r.timeline_lsn,
                                 rel: rel_tag,
                                 blkno: block_no,
-                            },
-                        )
-                    };
-                    let sender = work_senders.get(&timeline).unwrap();
-                    // TODO: what if this blocks?
-                    if sender.send(req).await.is_err() {
-                        assert!(cancel.is_cancelled(), "client has gone away unexpectedly");
-                    }
-                }
-            }),
-            Some(rps_limit) => Box::pin(async move {
-                let period = Duration::from_secs_f64(1.0 / (rps_limit as f64));
-                let make_timeline_task: &dyn Fn(
-                    TenantTimelineId,
-                )
-                    -> Pin<Box<dyn Send + Future<Output = ()>>> = &|timeline| {
-                    let sender = work_senders.get(&timeline).unwrap();
-                    let ranges: Vec<KeyRange> = all_ranges
-                        .iter()
-                        .filter(|r| r.timeline == timeline)
-                        .cloned()
-                        .collect();
-                    let weights = rand::distributions::weighted::WeightedIndex::new(
-                        ranges.iter().map(|v| v.len()),
-                    )
-                    .unwrap();
-
-                    let cancel = cancel.clone();
-                    Box::pin(async move {
-                        let mut ticker = tokio::time::interval(period);
-                        ticker.set_missed_tick_behavior(
-                            /* TODO review this choice */
-                            tokio::time::MissedTickBehavior::Burst,
-                        );
-                        while !cancel.is_cancelled() {
-                            ticker.tick().await;
-                            let req = {
-                                let mut rng = rand::thread_rng();
-                                let r = &ranges[weights.sample(&mut rng)];
-                                let key: i128 = rng.gen_range(r.start..r.end);
-                                let key = Key::from_i128(key);
-                                assert!(is_rel_block_key(&key));
-                                let (rel_tag, block_no) = key_to_rel_block(key)
-                                    .expect("we filter non-rel-block keys out above");
-                                PagestreamGetPageRequest {
-                                    latest: rng.gen_bool(args.req_latest_probability),
-                                    lsn: r.timeline_lsn,
-                                    rel: rel_tag,
-                                    blkno: block_no,
-                                }
-                            };
-                            if sender.send(req).await.is_err() {
-                                assert!(cancel.is_cancelled(), "client has gone away unexpectedly");
                             }
-                        }
-                    })
-                };
+                        };
+                        sender.send(req).await.ok().unwrap();
+                    }
+                })
+            };
 
-                let tasks: Vec<_> = work_senders
-                    .keys()
-                    .map(|tl| make_timeline_task(*tl))
-                    .collect();
+            let tasks: Vec<_> = work_senders
+                .keys()
+                .map(|tl| make_timeline_task(**tl))
+                .collect();
 
-                start_work_barrier.wait().await;
+            start_work_barrier.wait().await;
 
-                join_all(tasks).await;
-            }),
-        }
+            join_all(tasks).await;
+        }),
     };
 
-    let work_sender_task = tokio::spawn(work_sender);
-
     if let Some(runtime) = args.runtime {
-        info!("waiting for everything to become ready");
-        start_work_barrier.wait().await;
-        info!("work started");
-        tokio::time::sleep(runtime.into()).await;
-        info!("runtime over, signalling cancellation");
-        cancel.cancel();
-        work_sender_task.await.unwrap();
-        info!("work sender exited");
+        match tokio::time::timeout(runtime.into(), work_sender).await {
+            Ok(()) => unreachable!("work sender never terminates"),
+            Err(_timeout) => {
+                // this implicitly drops the work_senders, making all the clients exit
+            }
+        }
     } else {
-        work_sender_task.await.unwrap();
+        work_sender.await;
         unreachable!("work sender never terminates");
     }
 
-    info!("joining clients");
     for t in tasks {
         t.await.unwrap();
     }
 
-    info!("all clients stopped");
-
     let output = Output {
         total: {
             let mut agg_stats = request_stats::Stats::new();
@@ -393,9 +326,11 @@ async fn client(
     timeline: TenantTimelineId,
     start_work_barrier: Arc<Barrier>,
     mut work: tokio::sync::mpsc::Receiver<PagestreamGetPageRequest>,
+    all_work_done_barrier: Arc<Barrier>,
     live_stats: Arc<LiveStats>,
-    cancel: CancellationToken,
 ) {
+    start_work_barrier.wait().await;
+
     let client = pageserver_client::page_service::Client::new(args.page_service_connstring.clone())
         .await
         .unwrap();
@@ -404,18 +339,12 @@ async fn client(
         .await
         .unwrap();
 
-    start_work_barrier.wait().await;
-
-    while let Some(req) =
-        tokio::select! { work = work.recv() => { work } , _ = cancel.cancelled() => { return; } }
-    {
+    while let Some(req) = work.recv().await {
         let start = Instant::now();
-
-        let res = tokio::select! {
-            res = client.getpage(req) => { res },
-            _ = cancel.cancelled() => { return; }
-        };
-        res.with_context(|| format!("getpage for {timeline}"))
+        client
+            .getpage(req)
+            .await
+            .with_context(|| format!("getpage for {timeline}"))
             .unwrap();
         let elapsed = start.elapsed();
         live_stats.inc();
@@ -423,4 +352,6 @@ async fn client(
             stats.borrow().lock().unwrap().observe(elapsed).unwrap();
         });
     }
+
+    all_work_done_barrier.wait().await;
 }

pageserver/src/bin/pageserver.rs

@@ -130,7 +130,7 @@ fn main() -> anyhow::Result<()> {
     let scenario = failpoint_support::init();
 
     // Basic initialization of things that don't change after startup
-    virtual_file::init(conf.max_file_descriptors);
+    virtual_file::init(conf.max_file_descriptors, conf.virtual_file_io_engine);
     page_cache::init(conf.page_cache_size);
 
     start_pageserver(launch_ts, conf).context("Failed to start pageserver")?;

pageserver/src/config.rs

@@ -36,6 +36,7 @@ use crate::tenant::config::TenantConfOpt;
 use crate::tenant::{
     TENANTS_SEGMENT_NAME, TENANT_DELETED_MARKER_FILE_NAME, TIMELINES_SEGMENT_NAME,
 };
+use crate::virtual_file;
 use crate::{
     IGNORED_TENANT_FILE_NAME, METADATA_FILE_NAME, TENANT_CONFIG_NAME, TENANT_HEATMAP_BASENAME,
     TENANT_LOCATION_CONFIG_NAME, TIMELINE_DELETE_MARK_SUFFIX, TIMELINE_UNINIT_MARK_SUFFIX,
@@ -43,6 +44,8 @@ use crate::{
 
 use self::defaults::DEFAULT_CONCURRENT_TENANT_WARMUP;
 
+use self::defaults::DEFAULT_VIRTUAL_FILE_IO_ENGINE;
+
 pub mod defaults {
     use crate::tenant::config::defaults::*;
     use const_format::formatcp;
@@ -79,6 +82,8 @@ pub mod defaults {
 
     pub const DEFAULT_INGEST_BATCH_SIZE: u64 = 100;
 
+    pub const DEFAULT_VIRTUAL_FILE_IO_ENGINE: &str = "std-fs";
+
     ///
     /// Default built-in configuration file.
     ///
@@ -114,6 +119,8 @@ pub mod defaults {
 
 #ingest_batch_size = {DEFAULT_INGEST_BATCH_SIZE}
 
+#virtual_file_io_engine = '{DEFAULT_VIRTUAL_FILE_IO_ENGINE}'
+
 [tenant_config]
 #checkpoint_distance = {DEFAULT_CHECKPOINT_DISTANCE} # in bytes
 #checkpoint_timeout = {DEFAULT_CHECKPOINT_TIMEOUT}
@@ -247,6 +254,8 @@ pub struct PageServerConf {
 
     /// Maximum number of WAL records to be ingested and committed at the same time
     pub ingest_batch_size: u64,
+
+    pub virtual_file_io_engine: virtual_file::IoEngineKind,
 }
 
 /// We do not want to store this in a PageServerConf because the latter may be logged
@@ -331,6 +340,8 @@ struct PageServerConfigBuilder {
     secondary_download_concurrency: BuilderValue<usize>,
 
     ingest_batch_size: BuilderValue<u64>,
+
+    virtual_file_io_engine: BuilderValue<virtual_file::IoEngineKind>,
 }
 
 impl Default for PageServerConfigBuilder {
@@ -406,6 +417,8 @@ impl Default for PageServerConfigBuilder {
             secondary_download_concurrency: Set(DEFAULT_SECONDARY_DOWNLOAD_CONCURRENCY),
 
             ingest_batch_size: Set(DEFAULT_INGEST_BATCH_SIZE),
+
+            virtual_file_io_engine: Set(DEFAULT_VIRTUAL_FILE_IO_ENGINE.parse().unwrap()),
         }
     }
 }
@@ -562,6 +575,10 @@ impl PageServerConfigBuilder {
         self.ingest_batch_size = BuilderValue::Set(ingest_batch_size)
     }
 
+    pub fn virtual_file_io_engine(&mut self, value: virtual_file::IoEngineKind) {
+        self.virtual_file_io_engine = BuilderValue::Set(value);
+    }
+
     pub fn build(self) -> anyhow::Result<PageServerConf> {
         let concurrent_tenant_warmup = self
             .concurrent_tenant_warmup
@@ -669,6 +686,9 @@ impl PageServerConfigBuilder {
             ingest_batch_size: self
                 .ingest_batch_size
                 .ok_or(anyhow!("missing ingest_batch_size"))?,
+            virtual_file_io_engine: self
+                .virtual_file_io_engine
+                .ok_or(anyhow!("missing virtual_file_io_engine"))?,
         })
     }
 }
@@ -920,6 +940,9 @@ impl PageServerConf {
                     builder.secondary_download_concurrency(parse_toml_u64(key, item)? as usize)
                 },
                 "ingest_batch_size" => builder.ingest_batch_size(parse_toml_u64(key, item)?),
+                "virtual_file_io_engine" => {
+                    builder.virtual_file_io_engine(parse_toml_from_str("virtual_file_io_engine", item)?)
+                }
                 _ => bail!("unrecognized pageserver option '{key}'"),
             }
         }
@@ -993,6 +1016,7 @@ impl PageServerConf {
             heatmap_upload_concurrency: defaults::DEFAULT_HEATMAP_UPLOAD_CONCURRENCY,
             secondary_download_concurrency: defaults::DEFAULT_SECONDARY_DOWNLOAD_CONCURRENCY,
             ingest_batch_size: defaults::DEFAULT_INGEST_BATCH_SIZE,
+            virtual_file_io_engine: DEFAULT_VIRTUAL_FILE_IO_ENGINE.parse().unwrap(),
         }
     }
 }
@@ -1224,6 +1248,7 @@ background_task_maximum_delay = '334 s'
                 heatmap_upload_concurrency: defaults::DEFAULT_HEATMAP_UPLOAD_CONCURRENCY,
                 secondary_download_concurrency: defaults::DEFAULT_SECONDARY_DOWNLOAD_CONCURRENCY,
                 ingest_batch_size: defaults::DEFAULT_INGEST_BATCH_SIZE,
+                virtual_file_io_engine: DEFAULT_VIRTUAL_FILE_IO_ENGINE.parse().unwrap(),
             },
             "Correct defaults should be used when no config values are provided"
         );
@@ -1287,6 +1312,7 @@ background_task_maximum_delay = '334 s'
                 heatmap_upload_concurrency: defaults::DEFAULT_HEATMAP_UPLOAD_CONCURRENCY,
                 secondary_download_concurrency: defaults::DEFAULT_SECONDARY_DOWNLOAD_CONCURRENCY,
                 ingest_batch_size: 100,
+                virtual_file_io_engine: DEFAULT_VIRTUAL_FILE_IO_ENGINE.parse().unwrap(),
             },
             "Should be able to parse all basic config values correctly"
         );

pageserver/src/consumption_metrics.rs

@@ -267,7 +267,7 @@ async fn calculate_synthetic_size_worker(
             }
         };
 
-        for (tenant_shard_id, tenant_state, _gen) in tenants {
+        for (tenant_shard_id, tenant_state) in tenants {
             if tenant_state != TenantState::Active {
                 continue;
             }

pageserver/src/consumption_metrics/metrics.rs

@@ -196,7 +196,7 @@ pub(super) async fn collect_all_metrics(
         }
     };
 
-    let tenants = futures::stream::iter(tenants).filter_map(|(id, state, _)| async move {
+    let tenants = futures::stream::iter(tenants).filter_map(|(id, state)| async move {
         if state != TenantState::Active || !id.is_zero() {
             None
         } else {

pageserver/src/disk_usage_eviction_task.rs

@@ -633,7 +633,7 @@ async fn collect_eviction_candidates(
 
     let mut candidates = Vec::new();
 
-    for (tenant_id, _state, _gen) in &tenants {
+    for (tenant_id, _state) in &tenants {
         if cancel.is_cancelled() {
             return Ok(EvictionCandidates::Cancelled);
         }

pageserver/src/http/routes.rs

@@ -874,12 +874,11 @@ async fn tenant_list_handler(
             ApiError::ResourceUnavailable("Tenant map is initializing or shutting down".into())
         })?
         .iter()
-        .map(|(id, state, gen)| TenantInfo {
+        .map(|(id, state)| TenantInfo {
             id: *id,
             state: state.clone(),
             current_physical_size: None,
             attachment_status: state.attachment_status(),
-            generation: (*gen).into(),
         })
         .collect::<Vec<TenantInfo>>();
 
@@ -909,7 +908,6 @@ async fn tenant_status(
                 state: state.clone(),
                 current_physical_size: Some(current_physical_size),
                 attachment_status: state.attachment_status(),
-                generation: tenant.generation().into(),
             },
             timelines: tenant.list_timeline_ids(),
         })

pageserver/src/lib.rs

@@ -1,3 +1,4 @@
+#![recursion_limit = "300"]
 #![deny(clippy::undocumented_unsafe_blocks)]
 
 mod auth;

pageserver/src/metrics.rs

@@ -941,6 +941,7 @@ pub(crate) static STORAGE_IO_SIZE: Lazy<IntGaugeVec> = Lazy::new(|| {
     .expect("failed to define a metric")
 });
 
+#[cfg(not(test))]
 pub(crate) mod virtual_file_descriptor_cache {
     use super::*;
 

pageserver/src/page_service.rs

@@ -1500,8 +1500,7 @@ impl From<GetActiveTenantError> for QueryError {
             GetActiveTenantError::WaitForActiveTimeout { .. } => QueryError::Disconnected(
                 ConnectionError::Io(io::Error::new(io::ErrorKind::TimedOut, e.to_string())),
             ),
-            GetActiveTenantError::Cancelled
-            | GetActiveTenantError::WillNotBecomeActive(TenantState::Stopping { .. }) => {
+            GetActiveTenantError::WillNotBecomeActive(TenantState::Stopping { .. }) => {
                 QueryError::Shutdown
             }
             e => QueryError::Other(anyhow::anyhow!(e)),

pageserver/src/pgdatadir_mapping.rs

@@ -14,7 +14,7 @@ use crate::walrecord::NeonWalRecord;
 use anyhow::{ensure, Context};
 use bytes::{Buf, Bytes};
 use pageserver_api::key::is_rel_block_key;
-use pageserver_api::reltag::{BlockNumber, RelTag, SlruKind};
+use pageserver_api::reltag::{RelTag, SlruKind};
 use postgres_ffi::relfile_utils::{FSM_FORKNUM, VISIBILITYMAP_FORKNUM};
 use postgres_ffi::BLCKSZ;
 use postgres_ffi::{Oid, TimestampTz, TransactionId};
@@ -27,6 +27,9 @@ use tracing::{debug, trace, warn};
 use utils::bin_ser::DeserializeError;
 use utils::{bin_ser::BeSer, lsn::Lsn};
 
+/// Block number within a relation or SLRU. This matches PostgreSQL's BlockNumber type.
+pub type BlockNumber = u32;
+
 #[derive(Debug)]
 pub enum LsnForTimestamp {
     /// Found commits both before and after the given timestamp
@@ -1860,6 +1863,21 @@ pub fn is_inherited_key(key: Key) -> bool {
     key != AUX_FILES_KEY
 }
 
+/// Guaranteed to return `Ok()` if [[is_rel_block_key]] returns `true` for `key`.
+pub fn key_to_rel_block(key: Key) -> anyhow::Result<(RelTag, BlockNumber)> {
+    Ok(match key.field1 {
+        0x00 => (
+            RelTag {
+                spcnode: key.field2,
+                dbnode: key.field3,
+                relnode: key.field4,
+                forknum: key.field5,
+            },
+            key.field6,
+        ),
+        _ => anyhow::bail!("unexpected value kind 0x{:02x}", key.field1),
+    })
+}
 pub fn is_rel_fsm_block_key(key: Key) -> bool {
     key.field1 == 0x00 && key.field4 != 0 && key.field5 == FSM_FORKNUM && key.field6 != 0xffffffff
 }

pageserver/src/tenant.rs

@@ -13,7 +13,6 @@
 
 use anyhow::{bail, Context};
 use camino::Utf8Path;
-use camino::Utf8PathBuf;
 use enumset::EnumSet;
 use futures::stream::FuturesUnordered;
 use futures::FutureExt;
@@ -1931,10 +1930,6 @@ impl Tenant {
         self.current_state() == TenantState::Active
     }
 
-    pub fn generation(&self) -> Generation {
-        self.generation
-    }
-
     /// Changes tenant status to active, unless shutdown was already requested.
     ///
     /// `background_jobs_can_start` is an optional barrier set to a value during pageserver startup
@@ -3182,55 +3177,6 @@ impl Tenant {
         .await
     }
 
-    async fn upload_initdb(
-        &self,
-        timelines_path: &Utf8PathBuf,
-        pgdata_path: &Utf8PathBuf,
-        timeline_id: &TimelineId,
-    ) -> anyhow::Result<()> {
-        let Some(storage) = &self.remote_storage else {
-            // No remote storage?  No upload.
-            return Ok(());
-        };
-
-        let temp_path = timelines_path.join(format!(
-            "{INITDB_PATH}.upload-{timeline_id}.{TEMP_FILE_SUFFIX}"
-        ));
-
-        scopeguard::defer! {
-            if let Err(e) = fs::remove_file(&temp_path) {
-                error!("Failed to remove temporary initdb archive '{temp_path}': {e}");
-            }
-        }
-
-        let (pgdata_zstd, tar_zst_size) =
-            import_datadir::create_tar_zst(pgdata_path, &temp_path).await?;
-
-        pausable_failpoint!("before-initdb-upload");
-
-        backoff::retry(
-            || async {
-                self::remote_timeline_client::upload_initdb_dir(
-                    storage,
-                    &self.tenant_shard_id.tenant_id,
-                    timeline_id,
-                    pgdata_zstd.try_clone().await?,
-                    tar_zst_size,
-                    &self.cancel,
-                )
-                .await
-            },
-            |_| false,
-            3,
-            u32::MAX,
-            "persist_initdb_tar_zst",
-            backoff::Cancel::new(self.cancel.clone(), || anyhow::anyhow!("Cancelled")),
-        )
-        .await?;
-
-        Ok(())
-    }
-
     /// - run initdb to init temporary instance and get bootstrap data
     /// - after initialization completes, tar up the temp dir and upload it to S3.
     ///
@@ -3280,26 +3226,66 @@ impl Tenant {
                 )
                 .await
                 .context("download initdb tar")?;
-
-            scopeguard::defer! {
-                if let Err(e) = fs::remove_file(&initdb_tar_zst_path) {
-                    error!("Failed to remove temporary initdb archive '{initdb_tar_zst_path}': {e}");
-                }
-            }
-
             let buf_read =
                 BufReader::with_capacity(remote_timeline_client::BUFFER_SIZE, initdb_tar_zst);
             import_datadir::extract_tar_zst(&pgdata_path, buf_read)
                 .await
                 .context("extract initdb tar")?;
+
+            tokio::fs::remove_file(&initdb_tar_zst_path)
+                .await
+                .or_else(|e| {
+                    if e.kind() == std::io::ErrorKind::NotFound {
+                        // If something else already removed the file, ignore the error
+                        Ok(())
+                    } else {
+                        Err(e)
+                    }
+                })
+                .with_context(|| format!("tempfile removal {initdb_tar_zst_path}"))?;
         } else {
             // Init temporarily repo to get bootstrap data, this creates a directory in the `initdb_path` path
             run_initdb(self.conf, &pgdata_path, pg_version, &self.cancel).await?;
 
             // Upload the created data dir to S3
-            if self.tenant_shard_id().is_zero() {
-                self.upload_initdb(&timelines_path, &pgdata_path, &timeline_id)
-                    .await?;
+            if let Some(storage) = &self.remote_storage {
+                let temp_path = timelines_path.join(format!(
+                    "{INITDB_PATH}.upload-{timeline_id}.{TEMP_FILE_SUFFIX}"
+                ));
+
+                let (pgdata_zstd, tar_zst_size) =
+                    import_datadir::create_tar_zst(&pgdata_path, &temp_path).await?;
+                backoff::retry(
+                    || async {
+                        self::remote_timeline_client::upload_initdb_dir(
+                            storage,
+                            &self.tenant_shard_id.tenant_id,
+                            &timeline_id,
+                            pgdata_zstd.try_clone().await?,
+                            tar_zst_size,
+                            &self.cancel,
+                        )
+                        .await
+                    },
+                    |_| false,
+                    3,
+                    u32::MAX,
+                    "persist_initdb_tar_zst",
+                    backoff::Cancel::new(self.cancel.clone(), || anyhow::anyhow!("Cancelled")),
+                )
+                .await?;
+
+                tokio::fs::remove_file(&temp_path)
+                    .await
+                    .or_else(|e| {
+                        if e.kind() == std::io::ErrorKind::NotFound {
+                            // If something else already removed the file, ignore the error
+                            Ok(())
+                        } else {
+                            Err(e)
+                        }
+                    })
+                    .with_context(|| format!("tempfile removal {temp_path}"))?;
             }
         }
         let pgdata_lsn = import_datadir::get_lsn_from_controlfile(&pgdata_path)?.align();

pageserver/src/tenant/block_io.rs

@@ -5,10 +5,10 @@
 use super::ephemeral_file::EphemeralFile;
 use super::storage_layer::delta_layer::{Adapter, DeltaLayerInner};
 use crate::context::RequestContext;
-use crate::page_cache::{self, PageReadGuard, ReadBufResult, PAGE_SZ};
+use crate::page_cache::{self, PageReadGuard, PageWriteGuard, ReadBufResult, PAGE_SZ};
 use crate::virtual_file::VirtualFile;
 use bytes::Bytes;
-use std::ops::{Deref, DerefMut};
+use std::ops::Deref;
 
 /// This is implemented by anything that can read 8 kB (PAGE_SZ)
 /// blocks, using the page cache
@@ -39,6 +39,8 @@ pub enum BlockLease<'a> {
     EphemeralFileMutableTail(&'a [u8; PAGE_SZ]),
     #[cfg(test)]
     Arc(std::sync::Arc<[u8; PAGE_SZ]>),
+    #[cfg(test)]
+    Vec(Vec<u8>),
 }
 
 impl From<PageReadGuard<'static>> for BlockLease<'static> {
@@ -63,6 +65,13 @@ impl<'a> Deref for BlockLease<'a> {
             BlockLease::EphemeralFileMutableTail(v) => v,
             #[cfg(test)]
             BlockLease::Arc(v) => v.deref(),
+            #[cfg(test)]
+            BlockLease::Vec(v) => {
+                let v: &Vec<u8> = v;
+                assert_eq!(v.len(), PAGE_SZ, "caller must ensure that v has PAGE_SZ");
+                // Safety: see above assertion.
+                unsafe { &*(v.as_ptr() as *const [u8; PAGE_SZ]) }
+            }
         }
     }
 }
@@ -169,10 +178,14 @@ impl FileBlockReader {
     }
 
     /// Read a page from the underlying file into given buffer.
-    async fn fill_buffer(&self, buf: &mut [u8], blkno: u32) -> Result<(), std::io::Error> {
+    async fn fill_buffer(
+        &self,
+        buf: PageWriteGuard<'static>,
+        blkno: u32,
+    ) -> Result<PageWriteGuard<'static>, std::io::Error> {
         assert!(buf.len() == PAGE_SZ);
         self.file
-            .read_exact_at(buf, blkno as u64 * PAGE_SZ as u64)
+            .read_exact_at_page(buf, blkno as u64 * PAGE_SZ as u64)
             .await
     }
     /// Read a block.
@@ -196,9 +209,9 @@ impl FileBlockReader {
                 )
             })? {
             ReadBufResult::Found(guard) => Ok(guard.into()),
-            ReadBufResult::NotFound(mut write_guard) => {
+            ReadBufResult::NotFound(write_guard) => {
                 // Read the page from disk into the buffer
-                self.fill_buffer(write_guard.deref_mut(), blknum).await?;
+                let write_guard = self.fill_buffer(write_guard, blknum).await?;
                 Ok(write_guard.mark_valid().into())
             }
         }

pageserver/src/tenant/delete.rs

@@ -542,7 +542,6 @@ impl DeleteTenantFlow {
         )
         .await?;
 
-        pausable_failpoint!("tenant-delete-before-cleanup-remaining-fs-traces-pausable");
         fail::fail_point!("tenant-delete-before-cleanup-remaining-fs-traces", |_| {
             Err(anyhow::anyhow!(
                 "failpoint: tenant-delete-before-cleanup-remaining-fs-traces"

pageserver/src/tenant/ephemeral_file.rs

@@ -5,11 +5,11 @@ use crate::config::PageServerConf;
 use crate::context::RequestContext;
 use crate::page_cache::{self, PAGE_SZ};
 use crate::tenant::block_io::{BlockCursor, BlockLease, BlockReader};
-use crate::virtual_file::VirtualFile;
+use crate::virtual_file::{self, VirtualFile};
 use camino::Utf8PathBuf;
 use pageserver_api::shard::TenantShardId;
 use std::cmp::min;
-use std::fs::OpenOptions;
+
 use std::io::{self, ErrorKind};
 use std::ops::DerefMut;
 use std::sync::atomic::AtomicU64;
@@ -47,7 +47,10 @@ impl EphemeralFile {
 
         let file = VirtualFile::open_with_options(
             &filename,
-            OpenOptions::new().read(true).write(true).create(true),
+            virtual_file::OpenOptions::new()
+                .read(true)
+                .write(true)
+                .create(true),
         )
         .await?;
 
@@ -89,11 +92,10 @@ impl EphemeralFile {
                 page_cache::ReadBufResult::Found(guard) => {
                     return Ok(BlockLease::PageReadGuard(guard))
                 }
-                page_cache::ReadBufResult::NotFound(mut write_guard) => {
-                    let buf: &mut [u8] = write_guard.deref_mut();
-                    debug_assert_eq!(buf.len(), PAGE_SZ);
-                    self.file
-                        .read_exact_at(&mut buf[..], blknum as u64 * PAGE_SZ as u64)
+                page_cache::ReadBufResult::NotFound(write_guard) => {
+                    let write_guard = self
+                        .file
+                        .read_exact_at_page(write_guard, blknum as u64 * PAGE_SZ as u64)
                         .await?;
                     let read_guard = write_guard.mark_valid();
                     return Ok(BlockLease::PageReadGuard(read_guard));

pageserver/src/tenant/mgr.rs

@@ -1629,8 +1629,8 @@ pub(crate) enum TenantMapListError {
 ///
 /// Get list of tenants, for the mgmt API
 ///
-pub(crate) async fn list_tenants(
-) -> Result<Vec<(TenantShardId, TenantState, Generation)>, TenantMapListError> {
+pub(crate) async fn list_tenants() -> Result<Vec<(TenantShardId, TenantState)>, TenantMapListError>
+{
     let tenants = TENANTS.read().unwrap();
     let m = match &*tenants {
         TenantsMap::Initializing => return Err(TenantMapListError::Initializing),
@@ -1638,9 +1638,7 @@ pub(crate) async fn list_tenants(
     };
     Ok(m.iter()
         .filter_map(|(id, tenant)| match tenant {
-            TenantSlot::Attached(tenant) => {
-                Some((*id, tenant.current_state(), tenant.generation()))
-            }
+            TenantSlot::Attached(tenant) => Some((*id, tenant.current_state())),
             TenantSlot::Secondary(_) => None,
             TenantSlot::InProgress(_) => None,
         })

pageserver/src/tenant/storage_layer/delta_layer.rs

@@ -36,7 +36,7 @@ use crate::tenant::block_io::{BlockBuf, BlockCursor, BlockLease, BlockReader, Fi
 use crate::tenant::disk_btree::{DiskBtreeBuilder, DiskBtreeReader, VisitDirection};
 use crate::tenant::storage_layer::{Layer, ValueReconstructResult, ValueReconstructState};
 use crate::tenant::Timeline;
-use crate::virtual_file::VirtualFile;
+use crate::virtual_file::{self, VirtualFile};
 use crate::{walrecord, TEMP_FILE_SUFFIX};
 use crate::{DELTA_FILE_MAGIC, STORAGE_FORMAT_VERSION};
 use anyhow::{bail, ensure, Context, Result};
@@ -649,7 +649,7 @@ impl DeltaLayer {
     {
         let file = VirtualFile::open_with_options(
             path,
-            &*std::fs::OpenOptions::new().read(true).write(true),
+            virtual_file::OpenOptions::new().read(true).write(true),
         )
         .await
         .with_context(|| format!("Failed to open file '{}'", path))?;

pageserver/src/tenant/storage_layer/image_layer.rs

@@ -34,7 +34,7 @@ use crate::tenant::storage_layer::{
     LayerAccessStats, ValueReconstructResult, ValueReconstructState,
 };
 use crate::tenant::Timeline;
-use crate::virtual_file::VirtualFile;
+use crate::virtual_file::{self, VirtualFile};
 use crate::{IMAGE_FILE_MAGIC, STORAGE_FORMAT_VERSION, TEMP_FILE_SUFFIX};
 use anyhow::{bail, ensure, Context, Result};
 use bytes::Bytes;
@@ -327,7 +327,7 @@ impl ImageLayer {
     {
         let file = VirtualFile::open_with_options(
             path,
-            &*std::fs::OpenOptions::new().read(true).write(true),
+            virtual_file::OpenOptions::new().read(true).write(true),
         )
         .await
         .with_context(|| format!("Failed to open file '{}'", path))?;
@@ -492,11 +492,15 @@ impl ImageLayerWriterInner {
             },
         );
         info!("new image layer {path}");
-        let mut file = VirtualFile::open_with_options(
-            &path,
-            std::fs::OpenOptions::new().write(true).create_new(true),
-        )
-        .await?;
+        let mut file = {
+            VirtualFile::open_with_options(
+                &path,
+                virtual_file::OpenOptions::new()
+                    .write(true)
+                    .create_new(true),
+            )
+            .await?
+        };
         // make room for the header block
         file.seek(SeekFrom::Start(PAGE_SZ as u64)).await?;
         let blob_writer = BlobWriter::new(file, PAGE_SZ as u64);

pageserver/src/tenant/timeline.rs

@@ -65,10 +65,11 @@ use crate::keyspace::{KeyPartitioning, KeySpace, KeySpaceRandomAccum};
 use crate::metrics::{
     TimelineMetrics, MATERIALIZED_PAGE_CACHE_HIT, MATERIALIZED_PAGE_CACHE_HIT_DIRECT,
 };
+use crate::pgdatadir_mapping::LsnForTimestamp;
 use crate::pgdatadir_mapping::{is_inherited_key, is_rel_fsm_block_key, is_rel_vm_block_key};
-use crate::pgdatadir_mapping::{CalculateLogicalSizeError, LsnForTimestamp};
+use crate::pgdatadir_mapping::{BlockNumber, CalculateLogicalSizeError};
 use crate::tenant::config::{EvictionPolicy, TenantConfOpt};
-use pageserver_api::reltag::{BlockNumber, RelTag};
+use pageserver_api::reltag::RelTag;
 use pageserver_api::shard::ShardIndex;
 
 use postgres_connection::PgConnectionConfig;

pageserver/src/virtual_file.rs

@@ -11,17 +11,27 @@
 //! src/backend/storage/file/fd.c
 //!
 use crate::metrics::{StorageIoOperation, STORAGE_IO_SIZE, STORAGE_IO_TIME_METRIC};
+
+use crate::page_cache::PageWriteGuard;
 use crate::tenant::TENANTS_SEGMENT_NAME;
 use camino::{Utf8Path, Utf8PathBuf};
 use once_cell::sync::OnceCell;
-use std::fs::{self, File, OpenOptions};
+use std::fs::{self, File};
 use std::io::{Error, ErrorKind, Seek, SeekFrom};
+use tokio_epoll_uring::IoBufMut;
+
+use std::os::fd::{AsRawFd, FromRawFd, IntoRawFd, OwnedFd, RawFd};
 use std::os::unix::fs::FileExt;
 use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
 use tokio::sync::{RwLock, RwLockReadGuard, RwLockWriteGuard};
 use tokio::time::Instant;
 use utils::fs_ext;
 
+mod io_engine;
+mod open_options;
+pub use io_engine::IoEngineKind;
+pub(crate) use open_options::*;
+
 ///
 /// A virtual file descriptor. You can use this just like std::fs::File, but internally
 /// the underlying file is closed if the system is low on file descriptors,
@@ -104,7 +114,38 @@ struct SlotInner {
     tag: u64,
 
     /// the underlying file
-    file: Option<File>,
+    file: Option<OwnedFd>,
+}
+
+/// Impl of [`tokio_epoll_uring::IoBuf`] and [`tokio_epoll_uring::IoBufMut`] for [`PageWriteGuard`].
+struct PageWriteGuardBuf {
+    page: PageWriteGuard<'static>,
+    init_up_to: usize,
+}
+// Safety: the [`PageWriteGuard`] gives us exclusive ownership of the page cache slot,
+// and the location remains stable even if [`Self`] or the [`PageWriteGuard`] is moved.
+unsafe impl tokio_epoll_uring::IoBuf for PageWriteGuardBuf {
+    fn stable_ptr(&self) -> *const u8 {
+        self.page.as_ptr()
+    }
+    fn bytes_init(&self) -> usize {
+        self.init_up_to
+    }
+    fn bytes_total(&self) -> usize {
+        self.page.len()
+    }
+}
+// Safety: see above, plus: the ownership of [`PageWriteGuard`] means exclusive access,
+// hence it's safe to hand out the `stable_mut_ptr()`.
+unsafe impl tokio_epoll_uring::IoBufMut for PageWriteGuardBuf {
+    fn stable_mut_ptr(&mut self) -> *mut u8 {
+        self.page.as_mut_ptr()
+    }
+
+    unsafe fn set_init(&mut self, pos: usize) {
+        assert!(pos <= self.page.len());
+        self.init_up_to = pos;
+    }
 }
 
 impl OpenFiles {
@@ -272,6 +313,10 @@ macro_rules! with_file {
         let $ident = $this.lock_file().await?;
         observe_duration!($op, $($body)*)
     }};
+    ($this:expr, $op:expr, | mut $ident:ident | $($body:tt)*) => {{
+        let mut $ident = $this.lock_file().await?;
+        observe_duration!($op, $($body)*)
+    }};
 }
 
 impl VirtualFile {
@@ -315,7 +360,9 @@ impl VirtualFile {
         // NB: there is also StorageIoOperation::OpenAfterReplace which is for the case
         // where our caller doesn't get to use the returned VirtualFile before its
         // slot gets re-used by someone else.
-        let file = observe_duration!(StorageIoOperation::Open, open_options.open(path))?;
+        let file = observe_duration!(StorageIoOperation::Open, {
+            open_options.open(path.as_std_path()).await?
+        });
 
         // Strip all options other than read and write.
         //
@@ -388,15 +435,13 @@ impl VirtualFile {
 
     /// Call File::sync_all() on the underlying File.
     pub async fn sync_all(&self) -> Result<(), Error> {
-        with_file!(self, StorageIoOperation::Fsync, |file| file
-            .as_ref()
-            .sync_all())
+        with_file!(self, StorageIoOperation::Fsync, |file_guard| file_guard
+            .with_std_file(|std_file| std_file.sync_all()))
     }
 
     pub async fn metadata(&self) -> Result<fs::Metadata, Error> {
-        with_file!(self, StorageIoOperation::Metadata, |file| file
-            .as_ref()
-            .metadata())
+        with_file!(self, StorageIoOperation::Metadata, |file_guard| file_guard
+            .with_std_file(|std_file| std_file.metadata()))
     }
 
     /// Helper function internal to `VirtualFile` that looks up the underlying File,
@@ -405,7 +450,7 @@ impl VirtualFile {
     ///
     /// We are doing it via a macro as Rust doesn't support async closures that
     /// take on parameters with lifetimes.
-    async fn lock_file(&self) -> Result<FileGuard<'_>, Error> {
+    async fn lock_file(&self) -> Result<FileGuard, Error> {
         let open_files = get_open_files();
 
         let mut handle_guard = {
@@ -451,10 +496,9 @@ impl VirtualFile {
         // NB: we use StorageIoOperation::OpenAferReplace for this to distinguish this
         // case from StorageIoOperation::Open. This helps with identifying thrashing
         // of the virtual file descriptor cache.
-        let file = observe_duration!(
-            StorageIoOperation::OpenAfterReplace,
-            self.open_options.open(&self.path)
-        )?;
+        let file = observe_duration!(StorageIoOperation::OpenAfterReplace, {
+            self.open_options.open(self.path.as_std_path()).await?
+        });
 
         // Store the File in the slot and update the handle in the VirtualFile
         // to point to it.
@@ -479,9 +523,8 @@ impl VirtualFile {
                 self.pos = offset;
             }
             SeekFrom::End(offset) => {
-                self.pos = with_file!(self, StorageIoOperation::Seek, |file| file
-                    .as_ref()
-                    .seek(SeekFrom::End(offset)))?
+                self.pos = with_file!(self, StorageIoOperation::Seek, |mut file_guard| file_guard
+                    .with_std_file_mut(|std_file| std_file.seek(SeekFrom::End(offset))))?
             }
             SeekFrom::Current(offset) => {
                 let pos = self.pos as i128 + offset as i128;
@@ -501,24 +544,48 @@ impl VirtualFile {
     }
 
     // Copied from https://doc.rust-lang.org/1.72.0/src/std/os/unix/fs.rs.html#117-135
-    pub async fn read_exact_at(&self, mut buf: &mut [u8], mut offset: u64) -> Result<(), Error> {
-        while !buf.is_empty() {
-            match self.read_at(buf, offset).await {
-                Ok(0) => {
-                    return Err(Error::new(
-                        std::io::ErrorKind::UnexpectedEof,
-                        "failed to fill whole buffer",
-                    ))
-                }
+    pub async fn read_exact_at<B>(&self, buf: B, mut offset: u64) -> Result<B, Error>
+    where
+        B: IoBufMut + Send,
+    {
+        use tokio_epoll_uring::BoundedBuf;
+        let mut buf: tokio_epoll_uring::Slice<B> = buf.slice_full();
+        while buf.bytes_total() != 0 {
+            let res;
+            (buf, res) = self.read_at(buf, offset).await;
+            match res {
+                Ok(0) => break,
                 Ok(n) => {
-                    buf = &mut buf[n..];
+                    buf = buf.slice(n..);
                     offset += n as u64;
                 }
                 Err(ref e) if e.kind() == std::io::ErrorKind::Interrupted => {}
                 Err(e) => return Err(e),
             }
         }
-        Ok(())
+        if !buf.is_empty() {
+            Err(std::io::Error::new(
+                std::io::ErrorKind::UnexpectedEof,
+                "failed to fill whole buffer",
+            ))
+        } else {
+            Ok(buf.into_inner())
+        }
+    }
+
+    /// Like [`Self::read_exact_at`] but for [`PageWriteGuard`].
+    pub async fn read_exact_at_page(
+        &self,
+        page: PageWriteGuard<'static>,
+        offset: u64,
+    ) -> Result<PageWriteGuard<'static>, Error> {
+        let buf = PageWriteGuardBuf {
+            page,
+            init_up_to: 0,
+        };
+        let res = self.read_exact_at(buf, offset).await;
+        res.map(|PageWriteGuardBuf { page, .. }| page)
+            .map_err(|e| Error::new(ErrorKind::Other, e))
     }
 
     // Copied from https://doc.rust-lang.org/1.72.0/src/std/os/unix/fs.rs.html#219-235
@@ -568,22 +635,30 @@ impl VirtualFile {
         Ok(n)
     }
 
-    pub async fn read_at(&self, buf: &mut [u8], offset: u64) -> Result<usize, Error> {
-        let result = with_file!(self, StorageIoOperation::Read, |file| file
-            .as_ref()
-            .read_at(buf, offset));
-        if let Ok(size) = result {
-            STORAGE_IO_SIZE
-                .with_label_values(&["read", &self.tenant_id, &self.timeline_id])
-                .add(size as i64);
-        }
-        result
+    pub(crate) async fn read_at<B>(&self, buf: B, offset: u64) -> (B, Result<usize, Error>)
+    where
+        B: tokio_epoll_uring::BoundedBufMut + Send,
+    {
+        let file_guard = match self.lock_file().await {
+            Ok(file_guard) => file_guard,
+            Err(e) => return (buf, Err(e)),
+        };
+
+        observe_duration!(StorageIoOperation::Read, {
+            let ((_file_guard, buf), res) = io_engine::get().read_at(file_guard, offset, buf).await;
+            if let Ok(size) = res {
+                STORAGE_IO_SIZE
+                    .with_label_values(&["read", &self.tenant_id, &self.timeline_id])
+                    .add(size as i64);
+            }
+            (buf, res)
+        })
     }
 
     async fn write_at(&self, buf: &[u8], offset: u64) -> Result<usize, Error> {
-        let result = with_file!(self, StorageIoOperation::Write, |file| file
-            .as_ref()
-            .write_at(buf, offset));
+        let result = with_file!(self, StorageIoOperation::Write, |file_guard| {
+            file_guard.with_std_file(|std_file| std_file.write_at(buf, offset))
+        });
         if let Ok(size) = result {
             STORAGE_IO_SIZE
                 .with_label_values(&["write", &self.tenant_id, &self.timeline_id])
@@ -593,18 +668,54 @@ impl VirtualFile {
     }
 }
 
-struct FileGuard<'a> {
-    slot_guard: RwLockReadGuard<'a, SlotInner>,
+struct FileGuard {
+    slot_guard: RwLockReadGuard<'static, SlotInner>,
 }
 
-impl<'a> AsRef<File> for FileGuard<'a> {
-    fn as_ref(&self) -> &File {
+impl AsRef<OwnedFd> for FileGuard {
+    fn as_ref(&self) -> &OwnedFd {
         // This unwrap is safe because we only create `FileGuard`s
         // if we know that the file is Some.
         self.slot_guard.file.as_ref().unwrap()
     }
 }
 
+impl FileGuard {
+    /// Soft deprecation: we'll move VirtualFile to async APIs and remove this function eventually.
+    fn with_std_file<F, R>(&self, with: F) -> R
+    where
+        F: FnOnce(&File) -> R,
+    {
+        // SAFETY:
+        // - lifetime of the fd: `file` doesn't outlive the OwnedFd stored in `self`.
+        // - `&` usage below: `self` is `&`, hence Rust typesystem guarantees there are is no `&mut`
+        let file = unsafe { File::from_raw_fd(self.as_ref().as_raw_fd()) };
+        let res = with(&file);
+        let _ = file.into_raw_fd();
+        res
+    }
+    /// Soft deprecation: we'll move VirtualFile to async APIs and remove this function eventually.
+    fn with_std_file_mut<F, R>(&mut self, with: F) -> R
+    where
+        F: FnOnce(&mut File) -> R,
+    {
+        // SAFETY:
+        // - lifetime of the fd: `file` doesn't outlive the OwnedFd stored in `self`.
+        // - &mut usage below: `self` is `&mut`, hence this call is the only task/thread that has control over the underlying fd
+        let mut file = unsafe { File::from_raw_fd(self.as_ref().as_raw_fd()) };
+        let res = with(&mut file);
+        let _ = file.into_raw_fd();
+        res
+    }
+}
+
+impl tokio_epoll_uring::IoFd for FileGuard {
+    unsafe fn as_fd(&self) -> RawFd {
+        let owned_fd: &OwnedFd = self.as_ref();
+        owned_fd.as_raw_fd()
+    }
+}
+
 #[cfg(test)]
 impl VirtualFile {
     pub(crate) async fn read_blk(
@@ -612,16 +723,19 @@ impl VirtualFile {
         blknum: u32,
     ) -> Result<crate::tenant::block_io::BlockLease<'_>, std::io::Error> {
         use crate::page_cache::PAGE_SZ;
-        let mut buf = [0; PAGE_SZ];
-        self.read_exact_at(&mut buf, blknum as u64 * (PAGE_SZ as u64))
+        let buf = vec![0; PAGE_SZ];
+        let buf = self
+            .read_exact_at(buf, blknum as u64 * (PAGE_SZ as u64))
             .await?;
-        Ok(std::sync::Arc::new(buf).into())
+        Ok(crate::tenant::block_io::BlockLease::Vec(buf))
     }
 
     async fn read_to_end(&mut self, buf: &mut Vec<u8>) -> Result<(), Error> {
+        let mut tmp = vec![0; 128];
         loop {
-            let mut tmp = [0; 128];
-            match self.read_at(&mut tmp, self.pos).await {
+            let res;
+            (tmp, res) = self.read_at(tmp, self.pos).await;
+            match res {
                 Ok(0) => return Ok(()),
                 Ok(n) => {
                     self.pos += n as u64;
@@ -697,10 +811,12 @@ impl OpenFiles {
 /// Initialize the virtual file module. This must be called once at page
 /// server startup.
 ///
-pub fn init(num_slots: usize) {
+#[cfg(not(test))]
+pub fn init(num_slots: usize, engine: IoEngineKind) {
     if OPEN_FILES.set(OpenFiles::new(num_slots)).is_err() {
         panic!("virtual_file::init called twice");
     }
+    io_engine::init(engine);
     crate::metrics::virtual_file_descriptor_cache::SIZE_MAX.set(num_slots as u64);
 }
 
@@ -745,10 +861,10 @@ mod tests {
     }
 
     impl MaybeVirtualFile {
-        async fn read_exact_at(&self, buf: &mut [u8], offset: u64) -> Result<(), Error> {
+        async fn read_exact_at(&self, mut buf: Vec<u8>, offset: u64) -> Result<Vec<u8>, Error> {
             match self {
                 MaybeVirtualFile::VirtualFile(file) => file.read_exact_at(buf, offset).await,
-                MaybeVirtualFile::File(file) => file.read_exact_at(buf, offset),
+                MaybeVirtualFile::File(file) => file.read_exact_at(&mut buf, offset).map(|()| buf),
             }
         }
         async fn write_all_at(&self, buf: &[u8], offset: u64) -> Result<(), Error> {
@@ -790,14 +906,14 @@ mod tests {
 
         // Helper function to slurp a portion of a file into a string
         async fn read_string_at(&mut self, pos: u64, len: usize) -> Result<String, Error> {
-            let mut buf = vec![0; len];
-            self.read_exact_at(&mut buf, pos).await?;
+            let buf = vec![0; len];
+            let buf = self.read_exact_at(buf, pos).await?;
             Ok(String::from_utf8(buf).unwrap())
         }
     }
 
     #[tokio::test]
-    async fn test_virtual_files() -> Result<(), Error> {
+    async fn test_virtual_files() -> anyhow::Result<()> {
         // The real work is done in the test_files() helper function. This
         // allows us to run the same set of tests against a native File, and
         // VirtualFile. We trust the native Files and wouldn't need to test them,
@@ -813,14 +929,17 @@ mod tests {
     }
 
     #[tokio::test]
-    async fn test_physical_files() -> Result<(), Error> {
+    async fn test_physical_files() -> anyhow::Result<()> {
         test_files("physical_files", |path, open_options| async move {
-            Ok(MaybeVirtualFile::File(open_options.open(path)?))
+            Ok(MaybeVirtualFile::File({
+                let owned_fd = open_options.open(path.as_std_path()).await?;
+                File::from(owned_fd)
+            }))
         })
         .await
     }
 
-    async fn test_files<OF, FT>(testname: &str, openfunc: OF) -> Result<(), Error>
+    async fn test_files<OF, FT>(testname: &str, openfunc: OF) -> anyhow::Result<()>
     where
         OF: Fn(Utf8PathBuf, OpenOptions) -> FT,
         FT: Future<Output = Result<MaybeVirtualFile, std::io::Error>>,
@@ -964,11 +1083,11 @@ mod tests {
         for _threadno in 0..THREADS {
             let files = files.clone();
             let hdl = rt.spawn(async move {
-                let mut buf = [0u8; SIZE];
+                let mut buf = vec![0u8; SIZE];
                 let mut rng = rand::rngs::OsRng;
                 for _ in 1..1000 {
                     let f = &files[rng.gen_range(0..files.len())];
-                    f.read_exact_at(&mut buf, 0).await.unwrap();
+                    buf = f.read_exact_at(buf, 0).await.unwrap();
                     assert!(buf == SAMPLE);
                 }
             });

pageserver/src/virtual_file/io_engine.rs

@@ -0,0 +1,111 @@
+//! [`super::VirtualFile`] supports different IO engines.
+//!
+//! The [`IoEngineKind`] enum identifies them.
+//!
+//! The choice of IO engine is global.
+//! Initialize using [`init`].
+//!
+//! Then use [`get`] and  [`super::OpenOptions`].
+
+#[derive(
+    Copy,
+    Clone,
+    PartialEq,
+    Eq,
+    Hash,
+    strum_macros::EnumString,
+    strum_macros::Display,
+    serde_with::DeserializeFromStr,
+    serde_with::SerializeDisplay,
+    Debug,
+)]
+#[strum(serialize_all = "kebab-case")]
+pub enum IoEngineKind {
+    StdFs,
+    #[cfg(target_os = "linux")]
+    TokioEpollUring,
+}
+
+static IO_ENGINE: once_cell::sync::OnceCell<IoEngineKind> = once_cell::sync::OnceCell::new();
+
+#[cfg(not(test))]
+pub(super) fn init(engine: IoEngineKind) {
+    if IO_ENGINE.set(engine).is_err() {
+        panic!("called twice");
+    }
+}
+
+pub(super) fn get() -> &'static IoEngineKind {
+    #[cfg(test)]
+    {
+        let env_var_name = "NEON_PAGESERVER_UNIT_TEST_VIRTUAL_FILE_IOENGINE";
+        IO_ENGINE.get_or_init(|| match std::env::var(env_var_name) {
+            Ok(v) => match v.parse::<IoEngineKind>() {
+                Ok(engine_kind) => engine_kind,
+                Err(e) => {
+                    panic!("invalid VirtualFile io engine for env var {env_var_name}: {e:#}: {v:?}")
+                }
+            },
+            Err(std::env::VarError::NotPresent) => {
+                crate::config::defaults::DEFAULT_VIRTUAL_FILE_IO_ENGINE
+                    .parse()
+                    .unwrap()
+            }
+            Err(std::env::VarError::NotUnicode(_)) => {
+                panic!("env var {env_var_name} is not unicode");
+            }
+        })
+    }
+    #[cfg(not(test))]
+    IO_ENGINE.get().unwrap()
+}
+
+use std::os::unix::prelude::FileExt;
+
+use super::FileGuard;
+
+impl IoEngineKind {
+    pub(super) async fn read_at<B>(
+        &self,
+        file_guard: FileGuard,
+        offset: u64,
+        mut buf: B,
+    ) -> ((FileGuard, B), std::io::Result<usize>)
+    where
+        B: tokio_epoll_uring::BoundedBufMut + Send,
+    {
+        match self {
+            IoEngineKind::StdFs => {
+                // SAFETY: `dst` only lives at most as long as this match arm, during which buf remains valid memory.
+                let dst = unsafe {
+                    std::slice::from_raw_parts_mut(buf.stable_mut_ptr(), buf.bytes_total())
+                };
+                let res = file_guard.with_std_file(|std_file| std_file.read_at(dst, offset));
+                if let Ok(nbytes) = &res {
+                    assert!(*nbytes <= buf.bytes_total());
+                    // SAFETY: see above assertion
+                    unsafe {
+                        buf.set_init(*nbytes);
+                    }
+                }
+                #[allow(dropping_references)]
+                drop(dst);
+                ((file_guard, buf), res)
+            }
+            #[cfg(target_os = "linux")]
+            IoEngineKind::TokioEpollUring => {
+                let system = tokio_epoll_uring::thread_local_system().await;
+                let (resources, res) = system.read(file_guard, offset, buf).await;
+                (
+                    resources,
+                    res.map_err(|e| match e {
+                        tokio_epoll_uring::Error::Op(e) => e,
+                        tokio_epoll_uring::Error::System(system) => {
+                            std::io::Error::new(std::io::ErrorKind::Other, system)
+                        }
+                    }),
+                )
+            }
+        }
+    }
+}

pageserver/src/virtual_file/open_options.rs

@@ -0,0 +1,138 @@
+//! Enum-dispatch to the `OpenOptions` type of the respective [`super::IoEngineKind`];
+
+use super::IoEngineKind;
+use std::{os::fd::OwnedFd, path::Path};
+
+#[derive(Debug, Clone)]
+pub enum OpenOptions {
+    StdFs(std::fs::OpenOptions),
+    #[cfg(target_os = "linux")]
+    TokioEpollUring(tokio_epoll_uring::ops::open_at::OpenOptions),
+}
+
+impl Default for OpenOptions {
+    fn default() -> Self {
+        match super::io_engine::get() {
+            IoEngineKind::StdFs => Self::StdFs(std::fs::OpenOptions::new()),
+            #[cfg(target_os = "linux")]
+            IoEngineKind::TokioEpollUring => {
+                Self::TokioEpollUring(tokio_epoll_uring::ops::open_at::OpenOptions::new())
+            }
+        }
+    }
+}
+
+impl OpenOptions {
+    pub fn new() -> OpenOptions {
+        Self::default()
+    }
+
+    pub fn read(&mut self, read: bool) -> &mut OpenOptions {
+        match self {
+            OpenOptions::StdFs(x) => {
+                let _ = x.read(read);
+            }
+            #[cfg(target_os = "linux")]
+            OpenOptions::TokioEpollUring(x) => {
+                let _ = x.read(read);
+            }
+        }
+        self
+    }
+
+    pub fn write(&mut self, write: bool) -> &mut OpenOptions {
+        match self {
+            OpenOptions::StdFs(x) => {
+                let _ = x.write(write);
+            }
+            #[cfg(target_os = "linux")]
+            OpenOptions::TokioEpollUring(x) => {
+                let _ = x.write(write);
+            }
+        }
+        self
+    }
+
+    pub fn create(&mut self, create: bool) -> &mut OpenOptions {
+        match self {
+            OpenOptions::StdFs(x) => {
+                let _ = x.create(create);
+            }
+            #[cfg(target_os = "linux")]
+            OpenOptions::TokioEpollUring(x) => {
+                let _ = x.create(create);
+            }
+        }
+        self
+    }
+
+    pub fn create_new(&mut self, create_new: bool) -> &mut OpenOptions {
+        match self {
+            OpenOptions::StdFs(x) => {
+                let _ = x.create_new(create_new);
+            }
+            #[cfg(target_os = "linux")]
+            OpenOptions::TokioEpollUring(x) => {
+                let _ = x.create_new(create_new);
+            }
+        }
+        self
+    }
+
+    pub fn truncate(&mut self, truncate: bool) -> &mut OpenOptions {
+        match self {
+            OpenOptions::StdFs(x) => {
+                let _ = x.truncate(truncate);
+            }
+            #[cfg(target_os = "linux")]
+            OpenOptions::TokioEpollUring(x) => {
+                let _ = x.truncate(truncate);
+            }
+        }
+        self
+    }
+
+    pub(in crate::virtual_file) async fn open(&self, path: &Path) -> std::io::Result<OwnedFd> {
+        match self {
+            OpenOptions::StdFs(x) => x.open(path).map(|file| file.into()),
+            #[cfg(target_os = "linux")]
+            OpenOptions::TokioEpollUring(x) => {
+                let system = tokio_epoll_uring::thread_local_system().await;
+                system.open(path, x).await.map_err(|e| match e {
+                    tokio_epoll_uring::Error::Op(e) => e,
+                    tokio_epoll_uring::Error::System(system) => {
+                        std::io::Error::new(std::io::ErrorKind::Other, system)
+                    }
+                })
+            }
+        }
+    }
+}
+
+impl std::os::unix::prelude::OpenOptionsExt for OpenOptions {
+    fn mode(&mut self, mode: u32) -> &mut OpenOptions {
+        match self {
+            OpenOptions::StdFs(x) => {
+                let _ = x.mode(mode);
+            }
+            #[cfg(target_os = "linux")]
+            OpenOptions::TokioEpollUring(x) => {
+                let _ = x.mode(mode);
+            }
+        }
+        self
+    }
+
+    fn custom_flags(&mut self, flags: i32) -> &mut OpenOptions {
+        match self {
+            OpenOptions::StdFs(x) => {
+                let _ = x.custom_flags(flags);
+            }
+            #[cfg(target_os = "linux")]
+            OpenOptions::TokioEpollUring(x) => {
+                let _ = x.custom_flags(flags);
+            }
+        }
+        self
+    }
+}

pageserver/src/walingest.rs

@@ -38,7 +38,7 @@ use crate::tenant::PageReconstructError;
 use crate::tenant::Timeline;
 use crate::walrecord::*;
 use crate::ZERO_PAGE;
-use pageserver_api::reltag::{BlockNumber, RelTag, SlruKind};
+use pageserver_api::reltag::{RelTag, SlruKind};
 use postgres_ffi::pg_constants;
 use postgres_ffi::relfile_utils::{FSM_FORKNUM, INIT_FORKNUM, MAIN_FORKNUM, VISIBILITYMAP_FORKNUM};
 use postgres_ffi::v14::nonrelfile_utils::mx_offset_to_member_segment;

pageserver/src/walredo.rs

@@ -47,11 +47,9 @@ use crate::metrics::{
     WAL_REDO_PROCESS_LAUNCH_DURATION_HISTOGRAM, WAL_REDO_RECORDS_HISTOGRAM,
     WAL_REDO_RECORD_COUNTER, WAL_REDO_TIME,
 };
-use crate::pgdatadir_mapping::key_to_slru_block;
+use crate::pgdatadir_mapping::{key_to_rel_block, key_to_slru_block};
 use crate::repository::Key;
 use crate::walrecord::NeonWalRecord;
-
-use pageserver_api::key::key_to_rel_block;
 use pageserver_api::reltag::{RelTag, SlruKind};
 use postgres_ffi::pg_constants;
 use postgres_ffi::relfile_utils::VISIBILITYMAP_FORKNUM;

poetry.lock

@@ -1118,13 +1118,13 @@ files = [
 
 [[package]]
 name = "jinja2"
-version = "3.1.3"
+version = "3.1.2"
 description = "A very fast and expressive template engine."
 optional = false
 python-versions = ">=3.7"
 files = [
-    {file = "Jinja2-3.1.3-py3-none-any.whl", hash = "sha256:7d6d50dd97d52cbc355597bd845fabfbac3f551e1f99619e39a35ce8c370b5fa"},
-    {file = "Jinja2-3.1.3.tar.gz", hash = "sha256:ac8bd6544d4bb2c9792bf3a159e80bba8fda7f07e81bc3aed565432d5925ba90"},
+    {file = "Jinja2-3.1.2-py3-none-any.whl", hash = "sha256:6088930bfe239f0e6710546ab9c19c9ef35e29792895fed6e6e31a023a182a61"},
+    {file = "Jinja2-3.1.2.tar.gz", hash = "sha256:31351a702a408a9e7595a8fc6150fc3f43bb6bf7e319770cbc0db9df9437e852"},
 ]
 
 [package.dependencies]
@@ -2421,16 +2421,6 @@ files = [
     {file = "wrapt-1.14.1-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:8ad85f7f4e20964db4daadcab70b47ab05c7c1cf2a7c1e51087bfaa83831854c"},
     {file = "wrapt-1.14.1-cp310-cp310-win32.whl", hash = "sha256:a9a52172be0b5aae932bef82a79ec0a0ce87288c7d132946d645eba03f0ad8a8"},
     {file = "wrapt-1.14.1-cp310-cp310-win_amd64.whl", hash = "sha256:6d323e1554b3d22cfc03cd3243b5bb815a51f5249fdcbb86fda4bf62bab9e164"},
-    {file = "wrapt-1.14.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:ecee4132c6cd2ce5308e21672015ddfed1ff975ad0ac8d27168ea82e71413f55"},
-    {file = "wrapt-1.14.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:2020f391008ef874c6d9e208b24f28e31bcb85ccff4f335f15a3251d222b92d9"},
-    {file = "wrapt-1.14.1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:2feecf86e1f7a86517cab34ae6c2f081fd2d0dac860cb0c0ded96d799d20b335"},
-    {file = "wrapt-1.14.1-cp311-cp311-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:240b1686f38ae665d1b15475966fe0472f78e71b1b4903c143a842659c8e4cb9"},
-    {file = "wrapt-1.14.1-cp311-cp311-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:a9008dad07d71f68487c91e96579c8567c98ca4c3881b9b113bc7b33e9fd78b8"},
-    {file = "wrapt-1.14.1-cp311-cp311-musllinux_1_1_aarch64.whl", hash = "sha256:6447e9f3ba72f8e2b985a1da758767698efa72723d5b59accefd716e9e8272bf"},
-    {file = "wrapt-1.14.1-cp311-cp311-musllinux_1_1_i686.whl", hash = "sha256:acae32e13a4153809db37405f5eba5bac5fbe2e2ba61ab227926a22901051c0a"},
-    {file = "wrapt-1.14.1-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:49ef582b7a1152ae2766557f0550a9fcbf7bbd76f43fbdc94dd3bf07cc7168be"},
-    {file = "wrapt-1.14.1-cp311-cp311-win32.whl", hash = "sha256:358fe87cc899c6bb0ddc185bf3dbfa4ba646f05b1b0b9b5a27c2cb92c2cea204"},
-    {file = "wrapt-1.14.1-cp311-cp311-win_amd64.whl", hash = "sha256:26046cd03936ae745a502abf44dac702a5e6880b2b01c29aea8ddf3353b68224"},
     {file = "wrapt-1.14.1-cp35-cp35m-manylinux1_i686.whl", hash = "sha256:43ca3bbbe97af00f49efb06e352eae40434ca9d915906f77def219b88e85d907"},
     {file = "wrapt-1.14.1-cp35-cp35m-manylinux1_x86_64.whl", hash = "sha256:6b1a564e6cb69922c7fe3a678b9f9a3c54e72b469875aa8018f18b4d1dd1adf3"},
     {file = "wrapt-1.14.1-cp35-cp35m-manylinux2010_i686.whl", hash = "sha256:00b6d4ea20a906c0ca56d84f93065b398ab74b927a7a3dbd470f6fc503f95dc3"},
@@ -2668,4 +2658,4 @@ cffi = ["cffi (>=1.11)"]
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.9"
-content-hash = "9cf2734cafd5b6963165d398f1b24621193d5284d0bc7cc26a720a014f523860"
+content-hash = "35c237fe6a9278b2dc65b06ed96bde5afb9e393d52c01b00c59acf1df3a8d482"

pyproject.toml

@@ -14,7 +14,7 @@ requests = "^2.31.0"
 pytest-xdist = "^3.3.1"
 asyncpg = "^0.29.0"
 aiopg = "^1.4.0"
-Jinja2 = "^3.1.3"
+Jinja2 = "^3.0.2"
 types-requests = "^2.31.0.0"
 types-psycopg2 = "^2.9.21.10"
 boto3 = "^1.34.11"

safekeeper/src/control_file.rs

@@ -13,16 +13,13 @@ use std::time::Instant;
 
 use crate::control_file_upgrade::upgrade_control_file;
 use crate::metrics::PERSIST_CONTROL_FILE_SECONDS;
-use crate::state::TimelinePersistentState;
+use crate::safekeeper::{SafeKeeperState, SK_FORMAT_VERSION, SK_MAGIC};
 use utils::{bin_ser::LeSer, id::TenantTimelineId};
 
 use crate::SafeKeeperConf;
 
 use std::convert::TryInto;
 
-pub const SK_MAGIC: u32 = 0xcafeceefu32;
-pub const SK_FORMAT_VERSION: u32 = 7;
-
 // contains persistent metadata for safekeeper
 const CONTROL_FILE_NAME: &str = "safekeeper.control";
 // needed to atomically update the state using `rename`
@@ -32,9 +29,9 @@ pub const CHECKSUM_SIZE: usize = std::mem::size_of::<u32>();
 /// Storage should keep actual state inside of it. It should implement Deref
 /// trait to access state fields and have persist method for updating that state.
 #[async_trait::async_trait]
-pub trait Storage: Deref<Target = TimelinePersistentState> {
+pub trait Storage: Deref<Target = SafeKeeperState> {
     /// Persist safekeeper state on disk and update internal state.
-    async fn persist(&mut self, s: &TimelinePersistentState) -> Result<()>;
+    async fn persist(&mut self, s: &SafeKeeperState) -> Result<()>;
 
     /// Timestamp of last persist.
     fn last_persist_at(&self) -> Instant;
@@ -47,7 +44,7 @@ pub struct FileStorage {
     conf: SafeKeeperConf,
 
     /// Last state persisted to disk.
-    state: TimelinePersistentState,
+    state: SafeKeeperState,
     /// Not preserved across restarts.
     last_persist_at: Instant,
 }
@@ -71,7 +68,7 @@ impl FileStorage {
     pub fn create_new(
         timeline_dir: Utf8PathBuf,
         conf: &SafeKeeperConf,
-        state: TimelinePersistentState,
+        state: SafeKeeperState,
     ) -> Result<FileStorage> {
         let store = FileStorage {
             timeline_dir,
@@ -84,7 +81,7 @@ impl FileStorage {
     }
 
     /// Check the magic/version in the on-disk data and deserialize it, if possible.
-    fn deser_sk_state(buf: &mut &[u8]) -> Result<TimelinePersistentState> {
+    fn deser_sk_state(buf: &mut &[u8]) -> Result<SafeKeeperState> {
         // Read the version independent part
         let magic = ReadBytesExt::read_u32::<LittleEndian>(buf)?;
         if magic != SK_MAGIC {
@@ -96,7 +93,7 @@ impl FileStorage {
         }
         let version = ReadBytesExt::read_u32::<LittleEndian>(buf)?;
         if version == SK_FORMAT_VERSION {
-            let res = TimelinePersistentState::des(buf)?;
+            let res = SafeKeeperState::des(buf)?;
             return Ok(res);
         }
         // try to upgrade
@@ -107,15 +104,13 @@ impl FileStorage {
     pub fn load_control_file_conf(
         conf: &SafeKeeperConf,
         ttid: &TenantTimelineId,
-    ) -> Result<TimelinePersistentState> {
+    ) -> Result<SafeKeeperState> {
         let path = conf.timeline_dir(ttid).join(CONTROL_FILE_NAME);
         Self::load_control_file(path)
     }
 
     /// Read in the control file.
-    pub fn load_control_file<P: AsRef<Path>>(
-        control_file_path: P,
-    ) -> Result<TimelinePersistentState> {
+    pub fn load_control_file<P: AsRef<Path>>(control_file_path: P) -> Result<SafeKeeperState> {
         let mut control_file = std::fs::OpenOptions::new()
             .read(true)
             .write(true)
@@ -158,7 +153,7 @@ impl FileStorage {
 }
 
 impl Deref for FileStorage {
-    type Target = TimelinePersistentState;
+    type Target = SafeKeeperState;
 
     fn deref(&self) -> &Self::Target {
         &self.state
@@ -170,7 +165,7 @@ impl Storage for FileStorage {
     /// Persists state durably to the underlying storage.
     ///
     /// For a description, see <https://lwn.net/Articles/457667/>.
-    async fn persist(&mut self, s: &TimelinePersistentState) -> Result<()> {
+    async fn persist(&mut self, s: &SafeKeeperState) -> Result<()> {
         let _timer = PERSIST_CONTROL_FILE_SECONDS.start_timer();
 
         // write data to safekeeper.control.partial
@@ -247,7 +242,7 @@ impl Storage for FileStorage {
 mod test {
     use super::FileStorage;
     use super::*;
-    use crate::SafeKeeperConf;
+    use crate::{safekeeper::SafeKeeperState, SafeKeeperConf};
     use anyhow::Result;
     use utils::{id::TenantTimelineId, lsn::Lsn};
 
@@ -262,7 +257,7 @@ mod test {
     async fn load_from_control_file(
         conf: &SafeKeeperConf,
         ttid: &TenantTimelineId,
-    ) -> Result<(FileStorage, TimelinePersistentState)> {
+    ) -> Result<(FileStorage, SafeKeeperState)> {
         fs::create_dir_all(conf.timeline_dir(ttid))
             .await
             .expect("failed to create timeline dir");
@@ -275,11 +270,11 @@ mod test {
     async fn create(
         conf: &SafeKeeperConf,
         ttid: &TenantTimelineId,
-    ) -> Result<(FileStorage, TimelinePersistentState)> {
+    ) -> Result<(FileStorage, SafeKeeperState)> {
         fs::create_dir_all(conf.timeline_dir(ttid))
             .await
             .expect("failed to create timeline dir");
-        let state = TimelinePersistentState::empty();
+        let state = SafeKeeperState::empty();
         let timeline_dir = conf.timeline_dir(ttid);
         let storage = FileStorage::create_new(timeline_dir, conf, state.clone())?;
         Ok((storage, state))

safekeeper/src/control_file_upgrade.rs

@@ -1,7 +1,6 @@
 //! Code to deal with safekeeper control file upgrades
-use crate::{
-    safekeeper::{AcceptorState, PgUuid, ServerInfo, Term, TermHistory, TermLsn},
-    state::{PersistedPeers, TimelinePersistentState},
+use crate::safekeeper::{
+    AcceptorState, PersistedPeers, PgUuid, SafeKeeperState, ServerInfo, Term, TermHistory, TermLsn,
 };
 use anyhow::{bail, Result};
 use pq_proto::SystemId;
@@ -138,7 +137,7 @@ pub struct SafeKeeperStateV4 {
     pub peers: PersistedPeers,
 }
 
-pub fn upgrade_control_file(buf: &[u8], version: u32) -> Result<TimelinePersistentState> {
+pub fn upgrade_control_file(buf: &[u8], version: u32) -> Result<SafeKeeperState> {
     // migrate to storing full term history
     if version == 1 {
         info!("reading safekeeper control file version {}", version);
@@ -150,7 +149,7 @@ pub fn upgrade_control_file(buf: &[u8], version: u32) -> Result<TimelinePersiste
                 lsn: Lsn(0),
             }]),
         };
-        return Ok(TimelinePersistentState {
+        return Ok(SafeKeeperState {
             tenant_id: oldstate.server.tenant_id,
             timeline_id: oldstate.server.timeline_id,
             acceptor_state: ac,
@@ -177,7 +176,7 @@ pub fn upgrade_control_file(buf: &[u8], version: u32) -> Result<TimelinePersiste
             system_id: oldstate.server.system_id,
             wal_seg_size: oldstate.server.wal_seg_size,
         };
-        return Ok(TimelinePersistentState {
+        return Ok(SafeKeeperState {
             tenant_id: oldstate.server.tenant_id,
             timeline_id: oldstate.server.timeline_id,
             acceptor_state: oldstate.acceptor_state,
@@ -200,7 +199,7 @@ pub fn upgrade_control_file(buf: &[u8], version: u32) -> Result<TimelinePersiste
             system_id: oldstate.server.system_id,
             wal_seg_size: oldstate.server.wal_seg_size,
         };
-        return Ok(TimelinePersistentState {
+        return Ok(SafeKeeperState {
             tenant_id: oldstate.server.tenant_id,
             timeline_id: oldstate.server.timeline_id,
             acceptor_state: oldstate.acceptor_state,
@@ -223,7 +222,7 @@ pub fn upgrade_control_file(buf: &[u8], version: u32) -> Result<TimelinePersiste
             system_id: oldstate.server.system_id,
             wal_seg_size: oldstate.server.wal_seg_size,
         };
-        return Ok(TimelinePersistentState {
+        return Ok(SafeKeeperState {
             tenant_id: oldstate.tenant_id,
             timeline_id: oldstate.timeline_id,
             acceptor_state: oldstate.acceptor_state,
@@ -239,7 +238,7 @@ pub fn upgrade_control_file(buf: &[u8], version: u32) -> Result<TimelinePersiste
         });
     } else if version == 5 {
         info!("reading safekeeper control file version {}", version);
-        let mut oldstate = TimelinePersistentState::des(&buf[..buf.len()])?;
+        let mut oldstate = SafeKeeperState::des(&buf[..buf.len()])?;
         if oldstate.timeline_start_lsn != Lsn(0) {
             return Ok(oldstate);
         }
@@ -252,7 +251,7 @@ pub fn upgrade_control_file(buf: &[u8], version: u32) -> Result<TimelinePersiste
         return Ok(oldstate);
     } else if version == 6 {
         info!("reading safekeeper control file version {}", version);
-        let mut oldstate = TimelinePersistentState::des(&buf[..buf.len()])?;
+        let mut oldstate = SafeKeeperState::des(&buf[..buf.len()])?;
         if oldstate.server.pg_version != 0 {
             return Ok(oldstate);
         }

safekeeper/src/copy_timeline.rs

@@ -14,7 +14,7 @@ use utils::{id::TenantTimelineId, lsn::Lsn};
 use crate::{
     control_file::{FileStorage, Storage},
     pull_timeline::{create_temp_timeline_dir, load_temp_timeline, validate_temp_timeline},
-    state::TimelinePersistentState,
+    safekeeper::SafeKeeperState,
     timeline::{Timeline, TimelineError},
     wal_backup::copy_s3_segments,
     wal_storage::{wal_file_paths, WalReader},
@@ -137,7 +137,7 @@ pub async fn handle_request(request: Request) -> Result<()> {
     )
     .await?;
 
-    let mut new_state = TimelinePersistentState::new(
+    let mut new_state = SafeKeeperState::new(
         &request.destination_ttid,
         state.server.clone(),
         vec![],
@@ -160,7 +160,7 @@ pub async fn handle_request(request: Request) -> Result<()> {
 
 async fn copy_disk_segments(
     conf: &SafeKeeperConf,
-    persisted_state: &TimelinePersistentState,
+    persisted_state: &SafeKeeperState,
     wal_seg_size: usize,
     source_ttid: &TenantTimelineId,
     start_lsn: Lsn,

safekeeper/src/debug_dump.rs

@@ -22,13 +22,14 @@ use utils::id::TenantTimelineId;
 use utils::id::{TenantId, TimelineId};
 use utils::lsn::Lsn;
 
+use crate::safekeeper::SafeKeeperState;
+use crate::safekeeper::SafekeeperMemState;
 use crate::safekeeper::TermHistory;
+use crate::SafeKeeperConf;
+
 use crate::send_wal::WalSenderState;
-use crate::state::TimelineMemState;
-use crate::state::TimelinePersistentState;
 use crate::wal_storage::WalReader;
 use crate::GlobalTimelines;
-use crate::SafeKeeperConf;
 
 /// Various filters that influence the resulting JSON output.
 #[derive(Debug, Serialize, Deserialize, Clone)]
@@ -142,7 +143,7 @@ pub struct Config {
 pub struct Timeline {
     pub tenant_id: TenantId,
     pub timeline_id: TimelineId,
-    pub control_file: Option<TimelinePersistentState>,
+    pub control_file: Option<SafeKeeperState>,
     pub memory: Option<Memory>,
     pub disk_content: Option<DiskContent>,
 }
@@ -157,7 +158,7 @@ pub struct Memory {
     pub num_computes: u32,
     pub last_removed_segno: XLogSegNo,
     pub epoch_start_lsn: Lsn,
-    pub mem_state: TimelineMemState,
+    pub mem_state: SafekeeperMemState,
 
     // PhysicalStorage state.
     pub write_lsn: Lsn,

safekeeper/src/http/routes.rs

@@ -160,7 +160,7 @@ async fn timeline_status_handler(request: Request<Body>) -> Result<Response<Body
         commit_lsn: inmem.commit_lsn,
         backup_lsn: inmem.backup_lsn,
         peer_horizon_lsn: inmem.peer_horizon_lsn,
-        remote_consistent_lsn: inmem.remote_consistent_lsn,
+        remote_consistent_lsn: tli.get_walsenders().get_remote_consistent_lsn(),
         peers: tli.get_peers(conf).await,
         walsenders: tli.get_walsenders().get_all(),
         walreceivers: tli.get_walreceivers().get_all(),

safekeeper/src/json_ctrl.rs

@@ -21,8 +21,7 @@ use crate::safekeeper::{AcceptorProposerMessage, AppendResponse, ServerInfo};
 use crate::safekeeper::{
     AppendRequest, AppendRequestHeader, ProposerAcceptorMessage, ProposerElected,
 };
-use crate::safekeeper::{Term, TermHistory, TermLsn};
-use crate::state::TimelinePersistentState;
+use crate::safekeeper::{SafeKeeperState, Term, TermHistory, TermLsn};
 use crate::timeline::Timeline;
 use crate::GlobalTimelines;
 use postgres_backend::PostgresBackend;
@@ -57,7 +56,7 @@ pub struct AppendLogicalMessage {
 #[derive(Debug, Serialize)]
 struct AppendResult {
     // safekeeper state after append
-    state: TimelinePersistentState,
+    state: SafeKeeperState,
     // info about new record in the WAL
     inserted_wal: InsertedWAL,
 }

safekeeper/src/lib.rs

@@ -28,7 +28,6 @@ pub mod recovery;
 pub mod remove_wal;
 pub mod safekeeper;
 pub mod send_wal;
-pub mod state;
 pub mod timeline;
 pub mod wal_backup;
 pub mod wal_service;

safekeeper/src/metrics.rs

@@ -21,7 +21,7 @@ use utils::pageserver_feedback::PageserverFeedback;
 use utils::{id::TenantTimelineId, lsn::Lsn};
 
 use crate::{
-    state::{TimelineMemState, TimelinePersistentState},
+    safekeeper::{SafeKeeperState, SafekeeperMemState},
     GlobalTimelines,
 };
 
@@ -308,10 +308,11 @@ pub struct FullTimelineInfo {
     pub last_removed_segno: XLogSegNo,
 
     pub epoch_start_lsn: Lsn,
-    pub mem_state: TimelineMemState,
-    pub persisted_state: TimelinePersistentState,
+    pub mem_state: SafekeeperMemState,
+    pub persisted_state: SafeKeeperState,
 
     pub flush_lsn: Lsn,
+    pub remote_consistent_lsn: Lsn,
 
     pub wal_storage: WalStorageMetrics,
 }
@@ -607,7 +608,7 @@ impl Collector for TimelineCollector {
                 .set(tli.mem_state.peer_horizon_lsn.into());
             self.remote_consistent_lsn
                 .with_label_values(labels)
-                .set(tli.mem_state.remote_consistent_lsn.into());
+                .set(tli.remote_consistent_lsn.into());
             self.timeline_active
                 .with_label_values(labels)
                 .set(tli.timeline_is_active as u64);

safekeeper/src/safekeeper.rs

@@ -18,16 +18,17 @@ use tracing::*;
 use crate::control_file;
 use crate::send_wal::HotStandbyFeedback;
 
-use crate::state::TimelineState;
 use crate::wal_storage;
 use pq_proto::SystemId;
 use utils::pageserver_feedback::PageserverFeedback;
 use utils::{
     bin_ser::LeSer,
-    id::{NodeId, TenantId, TimelineId},
+    id::{NodeId, TenantId, TenantTimelineId, TimelineId},
     lsn::Lsn,
 };
 
+pub const SK_MAGIC: u32 = 0xcafeceefu32;
+pub const SK_FORMAT_VERSION: u32 = 7;
 const SK_PROTOCOL_VERSION: u32 = 2;
 pub const UNKNOWN_SERVER_VERSION: u32 = 0;
 
@@ -221,7 +222,7 @@ pub struct PersistedPeerInfo {
 }
 
 impl PersistedPeerInfo {
-    pub fn new() -> Self {
+    fn new() -> Self {
         Self {
             backup_lsn: Lsn::INVALID,
             term: INVALID_TERM,
@@ -231,10 +232,111 @@ impl PersistedPeerInfo {
     }
 }
 
-// make clippy happy
-impl Default for PersistedPeerInfo {
-    fn default() -> Self {
-        Self::new()
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
+pub struct PersistedPeers(pub Vec<(NodeId, PersistedPeerInfo)>);
+
+/// Persistent information stored on safekeeper node
+/// On disk data is prefixed by magic and format version and followed by checksum.
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
+pub struct SafeKeeperState {
+    #[serde(with = "hex")]
+    pub tenant_id: TenantId,
+    #[serde(with = "hex")]
+    pub timeline_id: TimelineId,
+    /// persistent acceptor state
+    pub acceptor_state: AcceptorState,
+    /// information about server
+    pub server: ServerInfo,
+    /// Unique id of the last *elected* proposer we dealt with. Not needed
+    /// for correctness, exists for monitoring purposes.
+    #[serde(with = "hex")]
+    pub proposer_uuid: PgUuid,
+    /// Since which LSN this timeline generally starts. Safekeeper might have
+    /// joined later.
+    pub timeline_start_lsn: Lsn,
+    /// Since which LSN safekeeper has (had) WAL for this timeline.
+    /// All WAL segments next to one containing local_start_lsn are
+    /// filled with data from the beginning.
+    pub local_start_lsn: Lsn,
+    /// Part of WAL acknowledged by quorum *and available locally*. Always points
+    /// to record boundary.
+    pub commit_lsn: Lsn,
+    /// LSN that points to the end of the last backed up segment. Useful to
+    /// persist to avoid finding out offloading progress on boot.
+    pub backup_lsn: Lsn,
+    /// Minimal LSN which may be needed for recovery of some safekeeper (end_lsn
+    /// of last record streamed to everyone). Persisting it helps skipping
+    /// recovery in walproposer, generally we compute it from peers. In
+    /// walproposer proto called 'truncate_lsn'. Updates are currently drived
+    /// only by walproposer.
+    pub peer_horizon_lsn: Lsn,
+    /// LSN of the oldest known checkpoint made by pageserver and successfully
+    /// pushed to s3. We don't remove WAL beyond it. Persisted only for
+    /// informational purposes, we receive it from pageserver (or broker).
+    pub remote_consistent_lsn: Lsn,
+    // Peers and their state as we remember it. Knowing peers themselves is
+    // fundamental; but state is saved here only for informational purposes and
+    // obviously can be stale. (Currently not saved at all, but let's provision
+    // place to have less file version upgrades).
+    pub peers: PersistedPeers,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+// In memory safekeeper state. Fields mirror ones in `SafeKeeperState`; values
+// are not flushed yet.
+pub struct SafekeeperMemState {
+    pub commit_lsn: Lsn,
+    pub backup_lsn: Lsn,
+    pub peer_horizon_lsn: Lsn,
+    #[serde(with = "hex")]
+    pub proposer_uuid: PgUuid,
+}
+
+impl SafeKeeperState {
+    pub fn new(
+        ttid: &TenantTimelineId,
+        server_info: ServerInfo,
+        peers: Vec<NodeId>,
+        commit_lsn: Lsn,
+        local_start_lsn: Lsn,
+    ) -> SafeKeeperState {
+        SafeKeeperState {
+            tenant_id: ttid.tenant_id,
+            timeline_id: ttid.timeline_id,
+            acceptor_state: AcceptorState {
+                term: 0,
+                term_history: TermHistory::empty(),
+            },
+            server: server_info,
+            proposer_uuid: [0; 16],
+            timeline_start_lsn: Lsn(0),
+            local_start_lsn,
+            commit_lsn,
+            backup_lsn: local_start_lsn,
+            peer_horizon_lsn: local_start_lsn,
+            remote_consistent_lsn: Lsn(0),
+            peers: PersistedPeers(
+                peers
+                    .iter()
+                    .map(|p| (*p, PersistedPeerInfo::new()))
+                    .collect(),
+            ),
+        }
+    }
+
+    #[cfg(test)]
+    pub fn empty() -> Self {
+        SafeKeeperState::new(
+            &TenantTimelineId::empty(),
+            ServerInfo {
+                pg_version: UNKNOWN_SERVER_VERSION, /* Postgres server version */
+                system_id: 0,                       /* Postgres system identifier */
+                wal_seg_size: 0,
+            },
+            vec![],
+            Lsn::INVALID,
+            Lsn::INVALID,
+        )
     }
 }
 
@@ -481,7 +583,9 @@ pub struct SafeKeeper<CTRL: control_file::Storage, WAL: wal_storage::Storage> {
     /// determines epoch switch point.
     pub epoch_start_lsn: Lsn,
 
-    pub state: TimelineState<CTRL>, // persistent state storage
+    pub inmem: SafekeeperMemState, // in memory part
+    pub state: CTRL,               // persistent state storage
+
     pub wal_store: WAL,
 
     node_id: NodeId, // safekeeper's node id
@@ -508,7 +612,13 @@ where
 
         Ok(SafeKeeper {
             epoch_start_lsn: Lsn(0),
-            state: TimelineState::new(state),
+            inmem: SafekeeperMemState {
+                commit_lsn: state.commit_lsn,
+                backup_lsn: state.backup_lsn,
+                peer_horizon_lsn: state.peer_horizon_lsn,
+                proposer_uuid: state.proposer_uuid,
+            },
+            state,
             wal_store,
             node_id,
         })
@@ -616,12 +726,12 @@ where
                 );
             }
 
-            let mut state = self.state.start_change();
+            let mut state = self.state.clone();
             state.server.system_id = msg.system_id;
             if msg.pg_version != UNKNOWN_SERVER_VERSION {
                 state.server.pg_version = msg.pg_version;
             }
-            self.state.finish_change(&state).await?;
+            self.state.persist(&state).await?;
         }
 
         info!(
@@ -656,15 +766,15 @@ where
             term: self.state.acceptor_state.term,
             vote_given: false as u64,
             flush_lsn: self.flush_lsn(),
-            truncate_lsn: self.state.inmem.peer_horizon_lsn,
+            truncate_lsn: self.inmem.peer_horizon_lsn,
             term_history: self.get_term_history(),
             timeline_start_lsn: self.state.timeline_start_lsn,
         };
         if self.state.acceptor_state.term < msg.term {
-            let mut state = self.state.start_change();
+            let mut state = self.state.clone();
             state.acceptor_state.term = msg.term;
             // persist vote before sending it out
-            self.state.finish_change(&state).await?;
+            self.state.persist(&state).await?;
 
             resp.term = self.state.acceptor_state.term;
             resp.vote_given = true as u64;
@@ -693,9 +803,9 @@ where
     ) -> Result<Option<AcceptorProposerMessage>> {
         info!("received ProposerElected {:?}", msg);
         if self.state.acceptor_state.term < msg.term {
-            let mut state = self.state.start_change();
+            let mut state = self.state.clone();
             state.acceptor_state.term = msg.term;
-            self.state.finish_change(&state).await?;
+            self.state.persist(&state).await?;
         }
 
         // If our term is higher, ignore the message (next feedback will inform the compute)
@@ -715,10 +825,10 @@ where
         }
         // Otherwise we must never attempt to truncate committed data.
         assert!(
-            msg.start_streaming_at >= self.state.inmem.commit_lsn,
+            msg.start_streaming_at >= self.inmem.commit_lsn,
             "attempt to truncate committed data: start_streaming_at={}, commit_lsn={}",
             msg.start_streaming_at,
-            self.state.inmem.commit_lsn
+            self.inmem.commit_lsn
         );
 
         // TODO: cross check divergence point, check if msg.start_streaming_at corresponds to
@@ -729,7 +839,7 @@ where
 
         // and now adopt term history from proposer
         {
-            let mut state = self.state.start_change();
+            let mut state = self.state.clone();
 
             // Here we learn initial LSN for the first time, set fields
             // interested in that.
@@ -753,13 +863,13 @@ where
             // NB: on new clusters, this happens at the same time as
             // timeline_start_lsn initialization, it is taken outside to provide
             // upgrade.
-            state.commit_lsn = max(state.commit_lsn, state.timeline_start_lsn);
+            self.inmem.commit_lsn = max(self.inmem.commit_lsn, state.timeline_start_lsn);
 
             // Initializing backup_lsn is useful to avoid making backup think it should upload 0 segment.
-            state.backup_lsn = max(state.backup_lsn, state.timeline_start_lsn);
+            self.inmem.backup_lsn = max(self.inmem.backup_lsn, state.timeline_start_lsn);
 
             state.acceptor_state.term_history = msg.term_history.clone();
-            self.state.finish_change(&state).await?;
+            self.persist_control_file(state).await?;
         }
 
         info!("start receiving WAL since {:?}", msg.start_streaming_at);
@@ -782,41 +892,63 @@ where
     async fn update_commit_lsn(&mut self, mut candidate: Lsn) -> Result<()> {
         // Both peers and walproposer communicate this value, we might already
         // have a fresher (higher) version.
-        candidate = max(candidate, self.state.inmem.commit_lsn);
+        candidate = max(candidate, self.inmem.commit_lsn);
         let commit_lsn = min(candidate, self.flush_lsn());
         assert!(
-            commit_lsn >= self.state.inmem.commit_lsn,
+            commit_lsn >= self.inmem.commit_lsn,
             "commit_lsn monotonicity violated: old={} new={}",
-            self.state.inmem.commit_lsn,
+            self.inmem.commit_lsn,
             commit_lsn
         );
 
-        self.state.inmem.commit_lsn = commit_lsn;
+        self.inmem.commit_lsn = commit_lsn;
 
         // If new commit_lsn reached epoch switch, force sync of control
         // file: walproposer in sync mode is very interested when this
         // happens. Note: this is for sync-safekeepers mode only, as
         // otherwise commit_lsn might jump over epoch_start_lsn.
         if commit_lsn >= self.epoch_start_lsn && self.state.commit_lsn < self.epoch_start_lsn {
-            self.state.flush().await?;
+            self.persist_control_file(self.state.clone()).await?;
         }
 
         Ok(())
     }
 
+    /// Persist in-memory state of control file to disk.
+    //
+    // TODO: passing inmem_remote_consistent_lsn everywhere is ugly, better
+    // separate state completely and give Arc to all those who need it.
+    pub async fn persist_inmem(&mut self, inmem_remote_consistent_lsn: Lsn) -> Result<()> {
+        let mut state = self.state.clone();
+        state.remote_consistent_lsn = inmem_remote_consistent_lsn;
+        self.persist_control_file(state).await
+    }
+
+    /// Persist in-memory state to the disk, taking other data from state.
+    async fn persist_control_file(&mut self, mut state: SafeKeeperState) -> Result<()> {
+        state.commit_lsn = self.inmem.commit_lsn;
+        state.backup_lsn = self.inmem.backup_lsn;
+        state.peer_horizon_lsn = self.inmem.peer_horizon_lsn;
+        state.proposer_uuid = self.inmem.proposer_uuid;
+        self.state.persist(&state).await
+    }
+
     /// Persist control file if there is something to save and enough time
     /// passed after the last save.
-    pub async fn maybe_persist_inmem_control_file(&mut self) -> Result<()> {
+    pub async fn maybe_persist_inmem_control_file(
+        &mut self,
+        inmem_remote_consistent_lsn: Lsn,
+    ) -> Result<()> {
         const CF_SAVE_INTERVAL: Duration = Duration::from_secs(300);
-        if self.state.pers.last_persist_at().elapsed() < CF_SAVE_INTERVAL {
+        if self.state.last_persist_at().elapsed() < CF_SAVE_INTERVAL {
             return Ok(());
         }
-        let need_persist = self.state.inmem.commit_lsn > self.state.commit_lsn
-            || self.state.inmem.backup_lsn > self.state.backup_lsn
-            || self.state.inmem.peer_horizon_lsn > self.state.peer_horizon_lsn
-            || self.state.inmem.remote_consistent_lsn > self.state.remote_consistent_lsn;
+        let need_persist = self.inmem.commit_lsn > self.state.commit_lsn
+            || self.inmem.backup_lsn > self.state.backup_lsn
+            || self.inmem.peer_horizon_lsn > self.state.peer_horizon_lsn
+            || inmem_remote_consistent_lsn > self.state.remote_consistent_lsn;
         if need_persist {
-            self.state.flush().await?;
+            self.persist_inmem(inmem_remote_consistent_lsn).await?;
             trace!("saved control file: {CF_SAVE_INTERVAL:?} passed");
         }
         Ok(())
@@ -842,7 +974,7 @@ where
         // Now we know that we are in the same term as the proposer,
         // processing the message.
 
-        self.state.inmem.proposer_uuid = msg.h.proposer_uuid;
+        self.inmem.proposer_uuid = msg.h.proposer_uuid;
 
         // do the job
         if !msg.wal_data.is_empty() {
@@ -866,16 +998,15 @@ where
         // - if we make safekeepers always send persistent value,
         //   any compute restart would pull it down.
         // Thus, take max before adopting.
-        self.state.inmem.peer_horizon_lsn =
-            max(self.state.inmem.peer_horizon_lsn, msg.h.truncate_lsn);
+        self.inmem.peer_horizon_lsn = max(self.inmem.peer_horizon_lsn, msg.h.truncate_lsn);
 
         // Update truncate and commit LSN in control file.
         // To avoid negative impact on performance of extra fsync, do it only
-        // when commit_lsn delta exceeds WAL segment size.
-        if self.state.commit_lsn + (self.state.server.wal_seg_size as u64)
-            < self.state.inmem.commit_lsn
+        // when truncate_lsn delta exceeds WAL segment size.
+        if self.state.peer_horizon_lsn + (self.state.server.wal_seg_size as u64)
+            < self.inmem.peer_horizon_lsn
         {
-            self.state.flush().await?;
+            self.persist_control_file(self.state.clone()).await?;
         }
 
         trace!(
@@ -917,27 +1048,27 @@ where
             }
         }
 
-        self.state.inmem.backup_lsn = max(Lsn(sk_info.backup_lsn), self.state.inmem.backup_lsn);
-        sync_control_file |= self.state.backup_lsn + (self.state.server.wal_seg_size as u64)
-            < self.state.inmem.backup_lsn;
+        let new_backup_lsn = max(Lsn(sk_info.backup_lsn), self.inmem.backup_lsn);
+        sync_control_file |=
+            self.state.backup_lsn + (self.state.server.wal_seg_size as u64) < new_backup_lsn;
+        self.inmem.backup_lsn = new_backup_lsn;
 
-        self.state.inmem.remote_consistent_lsn = max(
-            Lsn(sk_info.remote_consistent_lsn),
-            self.state.inmem.remote_consistent_lsn,
-        );
+        // value in sk_info should be maximized over our local in memory value.
+        let new_remote_consistent_lsn = Lsn(sk_info.remote_consistent_lsn);
+        assert!(self.state.remote_consistent_lsn <= new_remote_consistent_lsn);
         sync_control_file |= self.state.remote_consistent_lsn
             + (self.state.server.wal_seg_size as u64)
-            < self.state.inmem.remote_consistent_lsn;
+            < new_remote_consistent_lsn;
 
-        self.state.inmem.peer_horizon_lsn = max(
-            Lsn(sk_info.peer_horizon_lsn),
-            self.state.inmem.peer_horizon_lsn,
-        );
+        let new_peer_horizon_lsn = max(Lsn(sk_info.peer_horizon_lsn), self.inmem.peer_horizon_lsn);
         sync_control_file |= self.state.peer_horizon_lsn + (self.state.server.wal_seg_size as u64)
-            < self.state.inmem.peer_horizon_lsn;
+            < new_peer_horizon_lsn;
+        self.inmem.peer_horizon_lsn = new_peer_horizon_lsn;
 
         if sync_control_file {
-            self.state.flush().await?;
+            let mut state = self.state.clone();
+            state.remote_consistent_lsn = new_remote_consistent_lsn;
+            self.persist_control_file(state).await?;
         }
         Ok(())
     }
@@ -965,20 +1096,17 @@ mod tests {
     use postgres_ffi::WAL_SEGMENT_SIZE;
 
     use super::*;
-    use crate::{
-        state::{PersistedPeers, TimelinePersistentState},
-        wal_storage::Storage,
-    };
+    use crate::wal_storage::Storage;
     use std::{ops::Deref, str::FromStr, time::Instant};
 
     // fake storage for tests
     struct InMemoryState {
-        persisted_state: TimelinePersistentState,
+        persisted_state: SafeKeeperState,
     }
 
     #[async_trait::async_trait]
     impl control_file::Storage for InMemoryState {
-        async fn persist(&mut self, s: &TimelinePersistentState) -> Result<()> {
+        async fn persist(&mut self, s: &SafeKeeperState) -> Result<()> {
             self.persisted_state = s.clone();
             Ok(())
         }
@@ -989,15 +1117,15 @@ mod tests {
     }
 
     impl Deref for InMemoryState {
-        type Target = TimelinePersistentState;
+        type Target = SafeKeeperState;
 
         fn deref(&self) -> &Self::Target {
             &self.persisted_state
         }
     }
 
-    fn test_sk_state() -> TimelinePersistentState {
-        let mut state = TimelinePersistentState::empty();
+    fn test_sk_state() -> SafeKeeperState {
+        let mut state = SafeKeeperState::empty();
         state.server.wal_seg_size = WAL_SEGMENT_SIZE as u32;
         state.tenant_id = TenantId::from([1u8; 16]);
         state.timeline_id = TimelineId::from([1u8; 16]);
@@ -1054,7 +1182,7 @@ mod tests {
         }
 
         // reboot...
-        let state = sk.state.deref().clone();
+        let state = sk.state.persisted_state.clone();
         let storage = InMemoryState {
             persisted_state: state,
         };
@@ -1193,7 +1321,7 @@ mod tests {
         use utils::Hex;
         let tenant_id = TenantId::from_str("cf0480929707ee75372337efaa5ecf96").unwrap();
         let timeline_id = TimelineId::from_str("112ded66422aa5e953e5440fa5427ac4").unwrap();
-        let state = TimelinePersistentState {
+        let state = SafeKeeperState {
             tenant_id,
             timeline_id,
             acceptor_state: AcceptorState {
@@ -1277,7 +1405,7 @@ mod tests {
 
         assert_eq!(Hex(&ser), Hex(&expected));
 
-        let deser = TimelinePersistentState::des(&ser).unwrap();
+        let deser = SafeKeeperState::des(&ser).unwrap();
 
         assert_eq!(deser, state);
     }

safekeeper/src/send_wal.rs

@@ -19,6 +19,7 @@ use serde::{Deserialize, Serialize};
 use tokio::io::{AsyncRead, AsyncWrite};
 use utils::failpoint_support;
 use utils::id::TenantTimelineId;
+use utils::lsn::AtomicLsn;
 use utils::pageserver_feedback::PageserverFeedback;
 
 use std::cmp::{max, min};
@@ -89,12 +90,16 @@ pub struct StandbyFeedback {
 
 /// WalSenders registry. Timeline holds it (wrapped in Arc).
 pub struct WalSenders {
+    /// Lsn maximized over all walsenders *and* peer data, so might be higher
+    /// than what we receive from replicas.
+    remote_consistent_lsn: AtomicLsn,
     mutex: Mutex<WalSendersShared>,
 }
 
 impl WalSenders {
-    pub fn new() -> Arc<WalSenders> {
+    pub fn new(remote_consistent_lsn: Lsn) -> Arc<WalSenders> {
         Arc::new(WalSenders {
+            remote_consistent_lsn: AtomicLsn::from(remote_consistent_lsn),
             mutex: Mutex::new(WalSendersShared::new()),
         })
     }
@@ -152,6 +157,7 @@ impl WalSenders {
         let mut shared = self.mutex.lock();
         shared.get_slot_mut(id).feedback = ReplicationFeedback::Pageserver(*feedback);
         shared.update_ps_feedback();
+        self.update_remote_consistent_lsn(shared.agg_ps_feedback.remote_consistent_lsn);
     }
 
     /// Record standby reply.
@@ -196,6 +202,18 @@ impl WalSenders {
         }
     }
 
+    /// Get remote_consistent_lsn maximized across all walsenders and peers.
+    pub fn get_remote_consistent_lsn(self: &Arc<WalSenders>) -> Lsn {
+        self.remote_consistent_lsn.load()
+    }
+
+    /// Update maximized remote_consistent_lsn, return new (potentially) value.
+    pub fn update_remote_consistent_lsn(self: &Arc<WalSenders>, candidate: Lsn) -> Lsn {
+        self.remote_consistent_lsn
+            .fetch_max(candidate)
+            .max(candidate)
+    }
+
     /// Unregister walsender.
     fn unregister(self: &Arc<WalSenders>, id: WalSenderId) {
         let mut shared = self.mutex.lock();
@@ -426,11 +444,7 @@ impl SafekeeperPostgresHandler {
             wal_reader,
             send_buf: [0; MAX_SEND_SIZE],
         };
-        let mut reply_reader = ReplyReader {
-            reader,
-            ws_guard,
-            tli,
-        };
+        let mut reply_reader = ReplyReader { reader, ws_guard };
 
         let res = tokio::select! {
             // todo: add read|write .context to these errors
@@ -624,18 +638,17 @@ impl<IO: AsyncRead + AsyncWrite + Unpin> WalSender<'_, IO> {
 struct ReplyReader<IO> {
     reader: PostgresBackendReader<IO>,
     ws_guard: Arc<WalSenderGuard>,
-    tli: Arc<Timeline>,
 }
 
 impl<IO: AsyncRead + AsyncWrite + Unpin> ReplyReader<IO> {
     async fn run(&mut self) -> Result<(), CopyStreamHandlerEnd> {
         loop {
             let msg = self.reader.read_copy_message().await?;
-            self.handle_feedback(&msg).await?
+            self.handle_feedback(&msg)?
         }
     }
 
-    async fn handle_feedback(&mut self, msg: &Bytes) -> anyhow::Result<()> {
+    fn handle_feedback(&mut self, msg: &Bytes) -> anyhow::Result<()> {
         match msg.first().cloned() {
             Some(HOT_STANDBY_FEEDBACK_TAG_BYTE) => {
                 // Note: deserializing is on m[1..] because we skip the tag byte.
@@ -662,9 +675,6 @@ impl<IO: AsyncRead + AsyncWrite + Unpin> ReplyReader<IO> {
                 self.ws_guard
                     .walsenders
                     .record_ps_feedback(self.ws_guard.id, &ps_feedback);
-                self.tli
-                    .update_remote_consistent_lsn(ps_feedback.remote_consistent_lsn)
-                    .await;
                 // in principle new remote_consistent_lsn could allow to
                 // deactivate the timeline, but we check that regularly through
                 // broker updated, not need to do it here

safekeeper/src/state.rs

@@ -1,197 +0,0 @@
-//! Defines per timeline data stored persistently (SafeKeeperPersistentState)
-//! and its wrapper with in memory layer (SafekeeperState).
-
-use std::ops::Deref;
-
-use anyhow::Result;
-use serde::{Deserialize, Serialize};
-use utils::{
-    id::{NodeId, TenantId, TenantTimelineId, TimelineId},
-    lsn::Lsn,
-};
-
-use crate::{
-    control_file,
-    safekeeper::{AcceptorState, PersistedPeerInfo, PgUuid, ServerInfo, TermHistory},
-};
-
-/// Persistent information stored on safekeeper node about timeline.
-/// On disk data is prefixed by magic and format version and followed by checksum.
-#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
-pub struct TimelinePersistentState {
-    #[serde(with = "hex")]
-    pub tenant_id: TenantId,
-    #[serde(with = "hex")]
-    pub timeline_id: TimelineId,
-    /// persistent acceptor state
-    pub acceptor_state: AcceptorState,
-    /// information about server
-    pub server: ServerInfo,
-    /// Unique id of the last *elected* proposer we dealt with. Not needed
-    /// for correctness, exists for monitoring purposes.
-    #[serde(with = "hex")]
-    pub proposer_uuid: PgUuid,
-    /// Since which LSN this timeline generally starts. Safekeeper might have
-    /// joined later.
-    pub timeline_start_lsn: Lsn,
-    /// Since which LSN safekeeper has (had) WAL for this timeline.
-    /// All WAL segments next to one containing local_start_lsn are
-    /// filled with data from the beginning.
-    pub local_start_lsn: Lsn,
-    /// Part of WAL acknowledged by quorum *and available locally*. Always points
-    /// to record boundary.
-    pub commit_lsn: Lsn,
-    /// LSN that points to the end of the last backed up segment. Useful to
-    /// persist to avoid finding out offloading progress on boot.
-    pub backup_lsn: Lsn,
-    /// Minimal LSN which may be needed for recovery of some safekeeper (end_lsn
-    /// of last record streamed to everyone). Persisting it helps skipping
-    /// recovery in walproposer, generally we compute it from peers. In
-    /// walproposer proto called 'truncate_lsn'. Updates are currently drived
-    /// only by walproposer.
-    pub peer_horizon_lsn: Lsn,
-    /// LSN of the oldest known checkpoint made by pageserver and successfully
-    /// pushed to s3. We don't remove WAL beyond it. Persisted only for
-    /// informational purposes, we receive it from pageserver (or broker).
-    pub remote_consistent_lsn: Lsn,
-    // Peers and their state as we remember it. Knowing peers themselves is
-    // fundamental; but state is saved here only for informational purposes and
-    // obviously can be stale. (Currently not saved at all, but let's provision
-    // place to have less file version upgrades).
-    pub peers: PersistedPeers,
-}
-
-#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
-pub struct PersistedPeers(pub Vec<(NodeId, PersistedPeerInfo)>);
-
-impl TimelinePersistentState {
-    pub fn new(
-        ttid: &TenantTimelineId,
-        server_info: ServerInfo,
-        peers: Vec<NodeId>,
-        commit_lsn: Lsn,
-        local_start_lsn: Lsn,
-    ) -> TimelinePersistentState {
-        TimelinePersistentState {
-            tenant_id: ttid.tenant_id,
-            timeline_id: ttid.timeline_id,
-            acceptor_state: AcceptorState {
-                term: 0,
-                term_history: TermHistory::empty(),
-            },
-            server: server_info,
-            proposer_uuid: [0; 16],
-            timeline_start_lsn: Lsn(0),
-            local_start_lsn,
-            commit_lsn,
-            backup_lsn: local_start_lsn,
-            peer_horizon_lsn: local_start_lsn,
-            remote_consistent_lsn: Lsn(0),
-            peers: PersistedPeers(
-                peers
-                    .iter()
-                    .map(|p| (*p, PersistedPeerInfo::new()))
-                    .collect(),
-            ),
-        }
-    }
-
-    #[cfg(test)]
-    pub fn empty() -> Self {
-        use crate::safekeeper::UNKNOWN_SERVER_VERSION;
-
-        TimelinePersistentState::new(
-            &TenantTimelineId::empty(),
-            ServerInfo {
-                pg_version: UNKNOWN_SERVER_VERSION, /* Postgres server version */
-                system_id: 0,                       /* Postgres system identifier */
-                wal_seg_size: 0,
-            },
-            vec![],
-            Lsn::INVALID,
-            Lsn::INVALID,
-        )
-    }
-}
-
-#[derive(Debug, Clone, Serialize, Deserialize)]
-// In memory safekeeper state. Fields mirror ones in `SafeKeeperPersistentState`; values
-// are not flushed yet.
-pub struct TimelineMemState {
-    pub commit_lsn: Lsn,
-    pub backup_lsn: Lsn,
-    pub peer_horizon_lsn: Lsn,
-    pub remote_consistent_lsn: Lsn,
-    #[serde(with = "hex")]
-    pub proposer_uuid: PgUuid,
-}
-
-/// Safekeeper persistent state plus in memory layer, to avoid frequent fsyncs
-/// when we update fields like commit_lsn which don't need immediate
-/// persistence. Provides transactional like API to atomically update the state.
-///
-/// Implements Deref into *persistent* part.
-pub struct TimelineState<CTRL: control_file::Storage> {
-    pub inmem: TimelineMemState,
-    pub pers: CTRL, // persistent
-}
-
-impl<CTRL> TimelineState<CTRL>
-where
-    CTRL: control_file::Storage,
-{
-    pub fn new(state: CTRL) -> Self {
-        TimelineState {
-            inmem: TimelineMemState {
-                commit_lsn: state.commit_lsn,
-                backup_lsn: state.backup_lsn,
-                peer_horizon_lsn: state.peer_horizon_lsn,
-                remote_consistent_lsn: state.remote_consistent_lsn,
-                proposer_uuid: state.proposer_uuid,
-            },
-            pers: state,
-        }
-    }
-
-    /// Start atomic change. Returns SafeKeeperPersistentState with in memory
-    /// values applied; the protocol is to 1) change returned struct as desired
-    /// 2) atomically persist it with finish_change.
-    pub fn start_change(&self) -> TimelinePersistentState {
-        let mut s = self.pers.clone();
-        s.commit_lsn = self.inmem.commit_lsn;
-        s.backup_lsn = self.inmem.backup_lsn;
-        s.peer_horizon_lsn = self.inmem.peer_horizon_lsn;
-        s.remote_consistent_lsn = self.inmem.remote_consistent_lsn;
-        s.proposer_uuid = self.inmem.proposer_uuid;
-        s
-    }
-
-    /// Persist given state. c.f. start_change.
-    pub async fn finish_change(&mut self, s: &TimelinePersistentState) -> Result<()> {
-        self.pers.persist(s).await?;
-        // keep in memory values up to date
-        self.inmem.commit_lsn = s.commit_lsn;
-        self.inmem.backup_lsn = s.backup_lsn;
-        self.inmem.peer_horizon_lsn = s.peer_horizon_lsn;
-        self.inmem.remote_consistent_lsn = s.remote_consistent_lsn;
-        self.inmem.proposer_uuid = s.proposer_uuid;
-        Ok(())
-    }
-
-    /// Flush in memory values.
-    pub async fn flush(&mut self) -> Result<()> {
-        let s = self.start_change();
-        self.finish_change(&s).await
-    }
-}
-
-impl<CTRL> Deref for TimelineState<CTRL>
-where
-    CTRL: control_file::Storage,
-{
-    type Target = TimelinePersistentState;
-
-    fn deref(&self) -> &Self::Target {
-        &self.pers
-    }
-}

safekeeper/src/timeline.rs

@@ -28,11 +28,10 @@ use storage_broker::proto::TenantTimelineId as ProtoTenantTimelineId;
 use crate::receive_wal::WalReceivers;
 use crate::recovery::{recovery_main, Donor, RecoveryNeededInfo};
 use crate::safekeeper::{
-    AcceptorProposerMessage, ProposerAcceptorMessage, SafeKeeper, ServerInfo, Term, TermLsn,
-    INVALID_TERM,
+    AcceptorProposerMessage, ProposerAcceptorMessage, SafeKeeper, SafeKeeperState,
+    SafekeeperMemState, ServerInfo, Term, TermLsn, INVALID_TERM,
 };
 use crate::send_wal::WalSenders;
-use crate::state::{TimelineMemState, TimelinePersistentState};
 use crate::{control_file, safekeeper::UNKNOWN_SERVER_VERSION};
 
 use crate::metrics::FullTimelineInfo;
@@ -122,7 +121,7 @@ impl SharedState {
     fn create_new(
         conf: &SafeKeeperConf,
         ttid: &TenantTimelineId,
-        state: TimelinePersistentState,
+        state: SafeKeeperState,
     ) -> Result<Self> {
         if state.server.wal_seg_size == 0 {
             bail!(TimelineError::UninitializedWalSegSize(*ttid));
@@ -176,28 +175,30 @@ impl SharedState {
         })
     }
 
-    fn is_active(&self, num_computes: usize) -> bool {
+    fn is_active(&self, num_computes: usize, remote_consistent_lsn: Lsn) -> bool {
         self.is_wal_backup_required(num_computes)
             // FIXME: add tracking of relevant pageservers and check them here individually,
             // otherwise migration won't work (we suspend too early).
-            || self.sk.state.inmem.remote_consistent_lsn < self.sk.state.inmem.commit_lsn
+            || remote_consistent_lsn < self.sk.inmem.commit_lsn
     }
 
     /// Mark timeline active/inactive and return whether s3 offloading requires
     /// start/stop action. If timeline is deactivated, control file is persisted
     /// as maintenance task does that only for active timelines.
-    async fn update_status(&mut self, num_computes: usize, ttid: TenantTimelineId) -> bool {
-        let is_active = self.is_active(num_computes);
+    async fn update_status(
+        &mut self,
+        num_computes: usize,
+        remote_consistent_lsn: Lsn,
+        ttid: TenantTimelineId,
+    ) -> bool {
+        let is_active = self.is_active(num_computes, remote_consistent_lsn);
         if self.active != is_active {
             info!(
                 "timeline {} active={} now, remote_consistent_lsn={}, commit_lsn={}",
-                ttid,
-                is_active,
-                self.sk.state.inmem.remote_consistent_lsn,
-                self.sk.state.inmem.commit_lsn
+                ttid, is_active, remote_consistent_lsn, self.sk.inmem.commit_lsn
             );
             if !is_active {
-                if let Err(e) = self.sk.state.flush().await {
+                if let Err(e) = self.sk.persist_inmem(remote_consistent_lsn).await {
                     warn!("control file save in update_status failed: {:?}", e);
                 }
             }
@@ -211,8 +212,8 @@ impl SharedState {
         let seg_size = self.get_wal_seg_size();
         num_computes > 0 ||
         // Currently only the whole segment is offloaded, so compare segment numbers.
-            (self.sk.state.inmem.commit_lsn.segment_number(seg_size) >
-             self.sk.state.inmem.backup_lsn.segment_number(seg_size))
+            (self.sk.inmem.commit_lsn.segment_number(seg_size) >
+             self.sk.inmem.backup_lsn.segment_number(seg_size))
     }
 
     /// Is current state of s3 offloading is not what it ought to be?
@@ -226,7 +227,7 @@ impl SharedState {
             };
             trace!(
                 "timeline {} s3 offloading action {} pending: num_computes={}, commit_lsn={}, backup_lsn={}",
-                self.sk.state.timeline_id, action_pending, num_computes, self.sk.state.inmem.commit_lsn, self.sk.state.inmem.backup_lsn
+                self.sk.state.timeline_id, action_pending, num_computes, self.sk.inmem.commit_lsn, self.sk.inmem.backup_lsn
             );
         }
         res
@@ -247,6 +248,7 @@ impl SharedState {
         &self,
         ttid: &TenantTimelineId,
         conf: &SafeKeeperConf,
+        remote_consistent_lsn: Lsn,
     ) -> SafekeeperTimelineInfo {
         SafekeeperTimelineInfo {
             safekeeper_id: conf.my_id.0,
@@ -258,15 +260,15 @@ impl SharedState {
             last_log_term: self.sk.get_epoch(),
             flush_lsn: self.sk.flush_lsn().0,
             // note: this value is not flushed to control file yet and can be lost
-            commit_lsn: self.sk.state.inmem.commit_lsn.0,
-            remote_consistent_lsn: self.sk.state.inmem.remote_consistent_lsn.0,
-            peer_horizon_lsn: self.sk.state.inmem.peer_horizon_lsn.0,
+            commit_lsn: self.sk.inmem.commit_lsn.0,
+            remote_consistent_lsn: remote_consistent_lsn.0,
+            peer_horizon_lsn: self.sk.inmem.peer_horizon_lsn.0,
             safekeeper_connstr: conf
                 .advertise_pg_addr
                 .to_owned()
                 .unwrap_or(conf.listen_pg_addr.clone()),
             http_connstr: conf.listen_http_addr.to_owned(),
-            backup_lsn: self.sk.state.inmem.backup_lsn.0,
+            backup_lsn: self.sk.inmem.backup_lsn.0,
             local_start_lsn: self.sk.state.local_start_lsn.0,
             availability_zone: conf.availability_zone.clone(),
         }
@@ -364,6 +366,7 @@ impl Timeline {
         let _enter = info_span!("load_timeline", timeline = %ttid.timeline_id).entered();
 
         let shared_state = SharedState::restore(conf, &ttid)?;
+        let rcl = shared_state.sk.state.remote_consistent_lsn;
         let (commit_lsn_watch_tx, commit_lsn_watch_rx) =
             watch::channel(shared_state.sk.state.commit_lsn);
         let (term_flush_lsn_watch_tx, term_flush_lsn_watch_rx) = watch::channel(TermLsn::from((
@@ -380,7 +383,7 @@ impl Timeline {
             term_flush_lsn_watch_tx,
             term_flush_lsn_watch_rx,
             mutex: Mutex::new(shared_state),
-            walsenders: WalSenders::new(),
+            walsenders: WalSenders::new(rcl),
             walreceivers: WalReceivers::new(),
             cancellation_rx,
             cancellation_tx,
@@ -401,8 +404,7 @@ impl Timeline {
         let (term_flush_lsn_watch_tx, term_flush_lsn_watch_rx) =
             watch::channel(TermLsn::from((INVALID_TERM, Lsn::INVALID)));
         let (cancellation_tx, cancellation_rx) = watch::channel(false);
-        let state =
-            TimelinePersistentState::new(&ttid, server_info, vec![], commit_lsn, local_start_lsn);
+        let state = SafeKeeperState::new(&ttid, server_info, vec![], commit_lsn, local_start_lsn);
 
         Ok(Timeline {
             ttid,
@@ -412,7 +414,7 @@ impl Timeline {
             term_flush_lsn_watch_tx,
             term_flush_lsn_watch_rx,
             mutex: Mutex::new(SharedState::create_new(conf, &ttid, state)?),
-            walsenders: WalSenders::new(),
+            walsenders: WalSenders::new(Lsn(0)),
             walreceivers: WalReceivers::new(),
             cancellation_rx,
             cancellation_tx,
@@ -446,7 +448,7 @@ impl Timeline {
         fs::create_dir_all(&self.timeline_dir).await?;
 
         // Write timeline to disk and start background tasks.
-        if let Err(e) = shared_state.sk.state.flush().await {
+        if let Err(e) = shared_state.sk.persist_inmem(Lsn::INVALID).await {
             // Bootstrap failed, cancel timeline and remove timeline directory.
             self.cancel(shared_state);
 
@@ -521,7 +523,11 @@ impl Timeline {
 
     async fn update_status(&self, shared_state: &mut SharedState) -> bool {
         shared_state
-            .update_status(self.walreceivers.get_num(), self.ttid)
+            .update_status(
+                self.walreceivers.get_num(),
+                self.get_walsenders().get_remote_consistent_lsn(),
+                self.ttid,
+            )
             .await
     }
 
@@ -552,8 +558,8 @@ impl Timeline {
         }
         let shared_state = self.write_shared_state().await;
         if self.walreceivers.get_num() == 0 {
-            return shared_state.sk.state.inmem.commit_lsn == Lsn(0) || // no data at all yet
-            reported_remote_consistent_lsn >= shared_state.sk.state.inmem.commit_lsn;
+            return shared_state.sk.inmem.commit_lsn == Lsn(0) || // no data at all yet
+            reported_remote_consistent_lsn >= shared_state.sk.inmem.commit_lsn;
         }
         false
     }
@@ -617,7 +623,7 @@ impl Timeline {
                 resp.pageserver_feedback = ps_feedback;
             }
 
-            commit_lsn = shared_state.sk.state.inmem.commit_lsn;
+            commit_lsn = shared_state.sk.inmem.commit_lsn;
             term_flush_lsn =
                 TermLsn::from((shared_state.sk.get_term(), shared_state.sk.flush_lsn()));
         }
@@ -641,14 +647,14 @@ impl Timeline {
     }
 
     /// Returns state of the timeline.
-    pub async fn get_state(&self) -> (TimelineMemState, TimelinePersistentState) {
+    pub async fn get_state(&self) -> (SafekeeperMemState, SafeKeeperState) {
         let state = self.write_shared_state().await;
-        (state.sk.state.inmem.clone(), state.sk.state.clone())
+        (state.sk.inmem.clone(), state.sk.state.clone())
     }
 
     /// Returns latest backup_lsn.
     pub async fn get_wal_backup_lsn(&self) -> Lsn {
-        self.write_shared_state().await.sk.state.inmem.backup_lsn
+        self.write_shared_state().await.sk.inmem.backup_lsn
     }
 
     /// Sets backup_lsn to the given value.
@@ -658,7 +664,7 @@ impl Timeline {
         }
 
         let mut state = self.write_shared_state().await;
-        state.sk.state.inmem.backup_lsn = max(state.sk.state.inmem.backup_lsn, backup_lsn);
+        state.sk.inmem.backup_lsn = max(state.sk.inmem.backup_lsn, backup_lsn);
         // we should check whether to shut down offloader, but this will be done
         // soon by peer communication anyway.
         Ok(())
@@ -667,11 +673,21 @@ impl Timeline {
     /// Get safekeeper info for broadcasting to broker and other peers.
     pub async fn get_safekeeper_info(&self, conf: &SafeKeeperConf) -> SafekeeperTimelineInfo {
         let shared_state = self.write_shared_state().await;
-        shared_state.get_safekeeper_info(&self.ttid, conf)
+        shared_state.get_safekeeper_info(
+            &self.ttid,
+            conf,
+            self.walsenders.get_remote_consistent_lsn(),
+        )
     }
 
     /// Update timeline state with peer safekeeper data.
-    pub async fn record_safekeeper_info(&self, sk_info: SafekeeperTimelineInfo) -> Result<()> {
+    pub async fn record_safekeeper_info(&self, mut sk_info: SafekeeperTimelineInfo) -> Result<()> {
+        // Update local remote_consistent_lsn in memory (in .walsenders) and in
+        // sk_info to pass it down to control file.
+        sk_info.remote_consistent_lsn = self
+            .walsenders
+            .update_remote_consistent_lsn(Lsn(sk_info.remote_consistent_lsn))
+            .0;
         let is_wal_backup_action_pending: bool;
         let commit_lsn: Lsn;
         {
@@ -680,7 +696,7 @@ impl Timeline {
             let peer_info = PeerInfo::from_sk_info(&sk_info, Instant::now());
             shared_state.peers_info.upsert(&peer_info);
             is_wal_backup_action_pending = self.update_status(&mut shared_state).await;
-            commit_lsn = shared_state.sk.state.inmem.commit_lsn;
+            commit_lsn = shared_state.sk.inmem.commit_lsn;
         }
         self.commit_lsn_watch_tx.send(commit_lsn)?;
         // Wake up wal backup launcher, if it is time to stop the offloading.
@@ -690,13 +706,6 @@ impl Timeline {
         Ok(())
     }
 
-    /// Update in memory remote consistent lsn.
-    pub async fn update_remote_consistent_lsn(&self, candidate: Lsn) {
-        let mut shared_state = self.write_shared_state().await;
-        shared_state.sk.state.inmem.remote_consistent_lsn =
-            max(shared_state.sk.state.inmem.remote_consistent_lsn, candidate);
-    }
-
     pub async fn get_peers(&self, conf: &SafeKeeperConf) -> Vec<PeerInfo> {
         let shared_state = self.write_shared_state().await;
         shared_state.get_peers(conf.heartbeat_timeout)
@@ -827,10 +836,11 @@ impl Timeline {
     /// to date so that storage nodes restart doesn't cause many pageserver ->
     /// safekeeper reconnections.
     pub async fn maybe_persist_control_file(&self) -> Result<()> {
+        let remote_consistent_lsn = self.walsenders.get_remote_consistent_lsn();
         self.write_shared_state()
             .await
             .sk
-            .maybe_persist_inmem_control_file()
+            .maybe_persist_inmem_control_file(remote_consistent_lsn)
             .await
     }
 
@@ -852,9 +862,10 @@ impl Timeline {
                 num_computes: self.walreceivers.get_num() as u32,
                 last_removed_segno: state.last_removed_segno,
                 epoch_start_lsn: state.sk.epoch_start_lsn,
-                mem_state: state.sk.state.inmem.clone(),
+                mem_state: state.sk.inmem.clone(),
                 persisted_state: state.sk.state.clone(),
                 flush_lsn: state.sk.wal_store.flush_lsn(),
+                remote_consistent_lsn: self.get_walsenders().get_remote_consistent_lsn(),
                 wal_storage: state.sk.wal_store.get_metrics(),
             })
         } else {
@@ -878,7 +889,7 @@ impl Timeline {
             num_computes: self.walreceivers.get_num() as u32,
             last_removed_segno: state.last_removed_segno,
             epoch_start_lsn: state.sk.epoch_start_lsn,
-            mem_state: state.sk.state.inmem.clone(),
+            mem_state: state.sk.inmem.clone(),
             write_lsn,
             write_record_lsn,
             flush_lsn,

safekeeper/src/wal_storage.rs

@@ -23,7 +23,7 @@ use tokio::io::{AsyncReadExt, AsyncSeekExt};
 use tracing::*;
 
 use crate::metrics::{time_io_closure, WalStorageMetrics, REMOVED_WAL_SEGMENTS};
-use crate::state::TimelinePersistentState;
+use crate::safekeeper::SafeKeeperState;
 use crate::wal_backup::read_object;
 use crate::SafeKeeperConf;
 use postgres_ffi::waldecoder::WalStreamDecoder;
@@ -125,7 +125,7 @@ impl PhysicalStorage {
         ttid: &TenantTimelineId,
         timeline_dir: Utf8PathBuf,
         conf: &SafeKeeperConf,
-        state: &TimelinePersistentState,
+        state: &SafeKeeperState,
     ) -> Result<PhysicalStorage> {
         let wal_seg_size = state.server.wal_seg_size as usize;
 
@@ -525,7 +525,7 @@ impl WalReader {
     pub fn new(
         workdir: Utf8PathBuf,
         timeline_dir: Utf8PathBuf,
-        state: &TimelinePersistentState,
+        state: &SafeKeeperState,
         start_pos: Lsn,
         enable_remote_read: bool,
     ) -> Result<Self> {

scripts/ps_ec2_setup_instance_store

@@ -1,52 +0,0 @@
-#!/usr/bin/env bash
-
-# This script sets up an ext4 partition on an EC2 storage-optimized instance's instance store volume.
-# Unix permission/ownership is set to the calling user (the script does sudo internally.)
-#
-# It's intentionally not idempotent; don't take on that complexity in a bash script.
-
-set -euo pipefail
-set -x
-
-# This seems crude, but, apparently instance store NVMe volumes aren't exposed in the in instance metadata block-device-mapping.
-# https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/block-device-mapping-concepts.html#bdm-instance-metadata
-if [ "$(cat /sys/class/block/nvme1n1/device/model)" != "Amazon EC2 NVMe Instance Storage        " ]; then
-    echo "nvme1n1 is not Amazon EC2 NVMe Instance Storage: '$(cat /sys/class/block/nvme1n1/device/model)'"
-    exit 1
-fi
-
-# NB: we DO NOT warm up all the blocks on the drive as recommended by https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/disk-performance.html
-# The reason is that we don't do that in production either.
-
-# do all the on-disk initialization work now instead of a background kernel thread
-# so that we're ready for benchmarking right after this line
-sudo mkfs.ext4 -E lazy_itable_init=0,lazy_journal_init=0  /dev/nvme1n1
-
-MOUNTPOINT=/instance_store
-sudo mkdir "$MOUNTPOINT"
-sudo mount /dev/nvme1n1 "$MOUNTPOINT"
-sudo chown -R "$(id -u)":"$(id -g)" "$MOUNTPOINT"
-
-TEST_OUTPUT="$MOUNTPOINT/test_output"
-mkdir "$TEST_OUTPUT"
-
-NEON_REPO_DIR="$MOUNTPOINT/repo_dir"
-mkdir "$NEON_REPO_DIR"
-
-cat <<EOF
-SETUP COMPLETE
-
-To run your local neon.git build on the instance store volume,
-run the following commands from the top of the neon.git checkout
-
-    # test suite run
-    export TEST_OUTPUT="$TEST_OUTPUT"
-    DEFAULT_PG_VERSION=15 BUILD_TYPE=release ./scripts/pytest test_runner/performance/test_latency.py
-
-    # for interactive use
-    export NEON_REPO_DIR="$NEON_REPO_DIR"
-    cargo build_testing --release
-    ./target/release/neon_local init --force empty-dir-ok
-EOF
-
-

test_runner/fixtures/broker.py

@@ -54,10 +54,7 @@ class NeonBroker:
             else:
                 break  # success
 
-    def stop(self, immediate: bool = False):
+    def stop(self):
         if self.handle is not None:
-            if immediate:
-                self.handle.kill()
-            else:
-                self.handle.terminate()
+            self.handle.terminate()
             self.handle.wait()

test_runner/fixtures/neon_fixtures.py

@@ -440,6 +440,7 @@ class NeonEnvBuilder:
         preserve_database_files: bool = False,
         initial_tenant: Optional[TenantId] = None,
         initial_timeline: Optional[TimelineId] = None,
+        pageserver_virtual_file_io_engine: Optional[str] = None,
     ):
         self.repo_dir = repo_dir
         self.rust_log_override = rust_log_override
@@ -472,7 +473,8 @@ class NeonEnvBuilder:
         self.test_output_dir = test_output_dir
         self.test_overlay_dir = test_overlay_dir
         self.overlay_mounts_created_by_us: List[Tuple[str, Path]] = []
-        self.config_init_force: Optional[str] = None
+
+        self.pageserver_virtual_file_io_engine: Optional[str] = pageserver_virtual_file_io_engine
 
         assert test_name.startswith(
             "test_"
@@ -754,12 +756,20 @@ class NeonEnvBuilder:
         # Stop all the nodes.
         if self.env:
             log.info("Cleaning up all storage and compute nodes")
-            self.env.stop(
-                immediate=True,
+            self.env.endpoints.stop_all()
+            for sk in self.env.safekeepers:
+                sk.stop(immediate=True)
+
+            for pageserver in self.env.pageservers:
                 # if the test threw an exception, don't check for errors
                 # as a failing assertion would cause the cleanup below to fail
-                ps_assert_metric_no_errors=(exc_type is None),
-            )
+                if exc_type is not None:
+                    pageserver.assert_no_metric_errors()
+
+                pageserver.stop(immediate=True)
+
+            self.env.attachment_service.stop(immediate=True)
+
             cleanup_error = None
 
             if self.scrub_on_exit:
@@ -863,6 +873,8 @@ class NeonEnv:
         self.control_plane_api: str = f"http://127.0.0.1:{attachment_service_port}"
         self.attachment_service: NeonAttachmentService = NeonAttachmentService(self)
 
+        self.pageserver_virtual_file_io_engine = config.pageserver_virtual_file_io_engine
+
         # Create a config file corresponding to the options
         cfg: Dict[str, Any] = {
             "default_tenant_id": str(self.initial_tenant),
@@ -894,6 +906,9 @@ class NeonEnv:
                 "pg_auth_type": pg_auth_type,
                 "http_auth_type": http_auth_type,
             }
+            if self.pageserver_virtual_file_io_engine is not None:
+                ps_cfg["virtual_file_io_engine"] = self.pageserver_virtual_file_io_engine
+
             # Create a corresponding NeonPageserver object
             self.pageservers.append(
                 NeonPageserver(
@@ -928,7 +943,7 @@ class NeonEnv:
             cfg["safekeepers"].append(sk_cfg)
 
         log.info(f"Config: {cfg}")
-        self.neon_cli.init(cfg, force=config.config_init_force)
+        self.neon_cli.init(cfg)
 
     def start(self):
         # Start up broker, pageserver and all safekeepers
@@ -942,20 +957,6 @@ class NeonEnv:
         for safekeeper in self.safekeepers:
             safekeeper.start()
 
-    def stop(self, immediate=False, ps_assert_metric_no_errors=False):
-        """
-        After this method returns, there should be no child processes running.
-        """
-        self.endpoints.stop_all()
-        for sk in self.safekeepers:
-            sk.stop(immediate=immediate)
-        for pageserver in self.pageservers:
-            if ps_assert_metric_no_errors:
-                pageserver.assert_no_metric_errors()
-            pageserver.stop(immediate=immediate)
-        self.attachment_service.stop(immediate=immediate)
-        self.broker.stop(immediate=immediate)
-
     @property
     def pageserver(self) -> NeonPageserver:
         """
@@ -1047,6 +1048,7 @@ def _shared_simple_env(
     neon_binpath: Path,
     pg_distrib_dir: Path,
     pg_version: PgVersion,
+    pageserver_virtual_file_io_engine: str,
 ) -> Iterator[NeonEnv]:
     """
     # Internal fixture backing the `neon_simple_env` fixture. If TEST_SHARED_FIXTURES
@@ -1075,6 +1077,7 @@ def _shared_simple_env(
         preserve_database_files=pytestconfig.getoption("--preserve-database-files"),
         test_name=request.node.name,
         test_output_dir=test_output_dir,
+        pageserver_virtual_file_io_engine=pageserver_virtual_file_io_engine,
     ) as builder:
         env = builder.init_start()
 
@@ -1112,6 +1115,7 @@ def neon_env_builder(
     run_id: uuid.UUID,
     request: FixtureRequest,
     test_overlay_dir: Path,
+    pageserver_virtual_file_io_engine: str,
 ) -> Iterator[NeonEnvBuilder]:
     """
     Fixture to create a Neon environment for test.
@@ -1140,6 +1144,7 @@ def neon_env_builder(
         broker=default_broker,
         run_id=run_id,
         preserve_database_files=pytestconfig.getoption("--preserve-database-files"),
+        pageserver_virtual_file_io_engine=pageserver_virtual_file_io_engine,
         test_name=request.node.name,
         test_output_dir=test_output_dir,
         test_overlay_dir=test_overlay_dir,
@@ -1430,7 +1435,6 @@ class NeonCli(AbstractNeonCli):
     def init(
         self,
         config: Dict[str, Any],
-        force: Optional[str] = None,
     ) -> "subprocess.CompletedProcess[str]":
         with tempfile.NamedTemporaryFile(mode="w+") as tmp:
             tmp.write(toml.dumps(config))
@@ -1438,9 +1442,6 @@ class NeonCli(AbstractNeonCli):
 
             cmd = ["init", f"--config={tmp.name}", "--pg-version", self.env.pg_version]
 
-            if force is not None:
-                cmd.extend(["--force", force])
-
             storage = self.env.pageserver_remote_storage
 
             append_pageserver_param_overrides(

test_runner/fixtures/parametrize.py

@@ -8,7 +8,7 @@ from _pytest.python import Metafunc
 from fixtures.pg_version import PgVersion
 
 """
-Dynamically parametrize tests by Postgres version and build type (debug/release/remote)
+Dynamically parametrize tests by Postgres version, build type (debug/release/remote), and possibly by other parameters
 """
 
 
@@ -31,11 +31,12 @@ def build_type(request: FixtureRequest) -> Optional[str]:
     return None
 
 
-def pytest_generate_tests(metafunc: Metafunc):
-    # Do not parametrize performance tests yet, we need to prepare grafana charts first
-    if "test_runner/performance" in metafunc.definition._nodeid:
-        return
+@pytest.fixture(scope="function", autouse=True)
+def pageserver_virtual_file_io_engine(request: FixtureRequest) -> Optional[str]:
+    return None
 
+
+def pytest_generate_tests(metafunc: Metafunc):
     if (v := os.environ.get("DEFAULT_PG_VERSION")) is None:
         pg_versions = [version for version in PgVersion if version != PgVersion.NOT_SET]
     else:
@@ -46,5 +47,14 @@ def pytest_generate_tests(metafunc: Metafunc):
     else:
         build_types = [bt.lower()]
 
+    # A hacky way to parametrize performance tests only for `pageserver_virtual_file_io_engine=tokio-epoll-uring`
+    # And do not change test name for default `pageserver_virtual_file_io_engine=std-fs` to keep perf tests statistics
+    if (io_engine := os.environ.get("PAGESERVER_VIRTUAL_FILE_IO_ENGINE", "")) not in ("", "std-fs"):
+        metafunc.parametrize("pageserver_virtual_file_io_engine", [io_engine])
+
+    # Do not parametrize performance tests yet by Postgres version or build type, we need to prepare grafana charts first
+    if "test_runner/performance" in metafunc.definition._nodeid:
+        return
+
     metafunc.parametrize("build_type", build_types)
     metafunc.parametrize("pg_version", pg_versions, ids=map(lambda v: f"pg{v}", pg_versions))

test_runner/regress/test_tenant_delete.py

@@ -24,9 +24,8 @@ from fixtures.pageserver.utils import (
     wait_until_tenant_state,
 )
 from fixtures.remote_storage import RemoteStorageKind, available_s3_storages, s3_storage
-from fixtures.types import TenantId, TimelineId
+from fixtures.types import TenantId
 from fixtures.utils import run_pg_bench_small, wait_until
-from requests.exceptions import ReadTimeout
 
 
 def test_tenant_delete_smoke(
@@ -554,111 +553,3 @@ def test_tenant_delete_concurrent(
 
     # Zero tenants remain (we deleted the default tenant)
     assert ps_http.get_metric_value("pageserver_tenant_manager_slots") == 0
-
-
-def test_tenant_delete_races_timeline_creation(
-    neon_env_builder: NeonEnvBuilder,
-    pg_bin: PgBin,
-):
-    """
-    Validate that timeline creation executed in parallel with deletion works correctly.
-
-    This is a reproducer for https://github.com/neondatabase/neon/issues/6255
-    """
-    # The remote storage kind doesn't really matter but we use it for iterations calculation below
-    # (and there is no way to reconstruct the used remote storage kind)
-    remote_storage_kind = RemoteStorageKind.MOCK_S3
-    neon_env_builder.enable_pageserver_remote_storage(remote_storage_kind)
-    env = neon_env_builder.init_start(initial_tenant_conf=MANY_SMALL_LAYERS_TENANT_CONFIG)
-    ps_http = env.pageserver.http_client()
-    tenant_id = env.initial_tenant
-
-    # Sometimes it ends with "InternalServerError(Cancelled", sometimes with "InternalServerError(Operation was cancelled"
-    CANCELLED_ERROR = (
-        ".*POST.*Cancelled request finished with an error: InternalServerError\\(.*ancelled"
-    )
-
-    env.pageserver.allowed_errors.extend(
-        [
-            # lucky race with stopping from flushing a layer we fail to schedule any uploads
-            ".*layer flush task.+: could not flush frozen layer: update_metadata_file",
-            # We need the http connection close for successful reproduction
-            ".*POST.*/timeline.* request was dropped before completing",
-            # Timeline creation runs into this error
-            CANCELLED_ERROR,
-        ]
-    )
-
-    BEFORE_INITDB_UPLOAD_FAILPOINT = "before-initdb-upload"
-    DELETE_BEFORE_CLEANUP_FAILPOINT = "tenant-delete-before-cleanup-remaining-fs-traces-pausable"
-
-    # Wait just before the initdb upload
-    ps_http.configure_failpoints((BEFORE_INITDB_UPLOAD_FAILPOINT, "pause"))
-
-    def timeline_create():
-        try:
-            ps_http.timeline_create(env.pg_version, tenant_id, TimelineId.generate(), timeout=1)
-            raise RuntimeError("creation succeeded even though it shouldn't")
-        except ReadTimeout:
-            pass
-
-    Thread(target=timeline_create).start()
-
-    def hit_initdb_upload_failpoint():
-        assert env.pageserver.log_contains(f"at failpoint {BEFORE_INITDB_UPLOAD_FAILPOINT}")
-
-    wait_until(100, 0.1, hit_initdb_upload_failpoint)
-
-    def creation_connection_timed_out():
-        assert env.pageserver.log_contains(
-            "POST.*/timeline.* request was dropped before completing"
-        )
-
-    # Wait so that we hit the timeout and the connection is dropped
-    # (But timeline creation still continues)
-    wait_until(100, 0.1, creation_connection_timed_out)
-
-    ps_http.configure_failpoints((DELETE_BEFORE_CLEANUP_FAILPOINT, "pause"))
-
-    def tenant_delete():
-        ps_http.tenant_delete(tenant_id)
-
-    Thread(target=tenant_delete).start()
-
-    def deletion_arrived():
-        assert env.pageserver.log_contains(
-            f"cfg failpoint: {DELETE_BEFORE_CLEANUP_FAILPOINT} pause"
-        )
-
-    wait_until(100, 0.1, deletion_arrived)
-
-    ps_http.configure_failpoints((DELETE_BEFORE_CLEANUP_FAILPOINT, "off"))
-
-    # Disable the failpoint and wait for deletion to finish
-    ps_http.configure_failpoints((BEFORE_INITDB_UPLOAD_FAILPOINT, "off"))
-
-    iterations = poll_for_remote_storage_iterations(remote_storage_kind)
-    try:
-        tenant_delete_wait_completed(ps_http, tenant_id, iterations)
-    except PageserverApiException:
-        pass
-
-    # Physical deletion should have happened
-    assert_prefix_empty(
-        neon_env_builder,
-        prefix="/".join(
-            (
-                "tenants",
-                str(tenant_id),
-            )
-        ),
-    )
-
-    # Ensure that creation cancelled and deletion didn't end up in broken state or encountered the leftover temp file
-    assert env.pageserver.log_contains(CANCELLED_ERROR)
-    assert not env.pageserver.log_contains(
-        ".*ERROR.*delete_tenant.*Timelines directory is not empty after all timelines deletion"
-    )
-
-    # Zero tenants remain (we deleted the default tenant)
-    assert ps_http.get_metric_value("pageserver_tenant_manager_slots") == 0

vendor/revisions.json

@@ -1,5 +1,5 @@
 {
-    "postgres-v16": "7be4a52d728459b79b59343c57d338c3073059c8",
-    "postgres-v15": "c1c2272f436ed9231f6172f49de219fe71a9280d",
-    "postgres-v14": "82072911287cabb32018cf92c8425fa1c744def4"
+    "postgres-v16": "225071f482774943854c2eec4540757e01171557",
+    "postgres-v15": "a2dc225ddfc8cae1849aa2316f435c58f0333d8c",
+    "postgres-v14": "03358bb0b5e0d33c238710139e768db9e75cfcc8"
 }