diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 0cd5699e2..a968603af 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -281,3 +281,76 @@ jobs: env: RUST_BACKTRACE: full run: cargo test-api smoke docker + + replica_live_s3: + name: Replica Live S3 (${{ matrix.server.name }}) + runs-on: ubuntu-22.04 + strategy: + fail-fast: false + matrix: + server: + # Two S3-compatible implementations: MinIO exercises the full API + # (including ListMultipartUploads for the orphan sweep); RustFS + # pins compatibility with a second, Rust-native implementation. + - name: minio + start: >- + docker run -d --name live-s3 -p 9000:9000 + -e MINIO_ROOT_USER=bencher-test + -e MINIO_ROOT_PASSWORD=bencher-test-secret + minio/minio:RELEASE.2025-09-07T16-13-09Z server /data + - name: rustfs + start: >- + docker run -d --name live-s3 -p 9000:9000 + -e RUSTFS_ACCESS_KEY=bencher-test + -e RUSTFS_SECRET_KEY=bencher-test-secret + rustfs/rustfs:1.0.0-beta.10 + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + - uses: ./.github/actions/setup-rust + with: + cache-key: replica-live-s3 + mold-version: ${{ inputs.mold-version }} + - name: Start S3 server + env: + START_SERVER: ${{ matrix.server.start }} + run: | + ${START_SERVER} + for _ in $(seq 1 30); do + if curl -s -o /dev/null http://127.0.0.1:9000/; then + exit 0 + fi + sleep 2 + done + echo "S3 server did not come up" >&2 + docker logs live-s3 >&2 || true + exit 1 + - name: Create bucket + env: + AWS_ACCESS_KEY_ID: bencher-test + AWS_SECRET_ACCESS_KEY: bencher-test-secret + AWS_DEFAULT_REGION: us-east-1 + run: | + for _ in $(seq 1 15); do + if aws --endpoint-url http://127.0.0.1:9000 s3api create-bucket \ + --bucket bencher-replica-test; then + exit 0 + fi + sleep 2 + done + echo "failed to create the test bucket" >&2 + docker logs live-s3 >&2 || true + exit 1 + - name: Run live S3 tier + env: + RUST_BACKTRACE: "1" + BENCHER_REPLICA_TEST_S3_BUCKET: bencher-replica-test + BENCHER_REPLICA_TEST_S3_ACCESS_KEY_ID: bencher-test + BENCHER_REPLICA_TEST_S3_SECRET_ACCESS_KEY: bencher-test-secret + BENCHER_REPLICA_TEST_S3_ENDPOINT: http://127.0.0.1:9000 + BENCHER_REPLICA_TEST_S3_REGION: us-east-1 + BENCHER_REPLICA_TEST_S3_PREFIX: ci-live-s3 + run: | + cargo nextest run -p bencher_replica --features plus,testing \ + --run-ignored ignored-only --no-fail-fast -E 'binary(live_s3)' diff --git a/Cargo.lock b/Cargo.lock index b4b013b5b..2e858e50d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -323,14 +323,18 @@ dependencies = [ "bencher_endpoint", "bencher_json", "bencher_otel", + "bencher_replica", "bencher_schema", + "camino", "diesel", "dropshot", "http 1.4.1", + "rusqlite", "schemars", "serde", "serde_json", "slog", + "tempfile", "tokio", ] @@ -1116,7 +1120,9 @@ dependencies = [ "bencher_logger", "bencher_otel", "bencher_otel_provider", + "bencher_replica", "bencher_schema", + "camino", "dropshot", "futures-concurrency", "futures-util", @@ -1140,8 +1146,10 @@ dependencies = [ "bencher_license", "bencher_oci_storage", "bencher_rbac", + "bencher_replica", "bencher_schema", "bencher_token", + "camino", "diesel", "dropshot", "hex", @@ -1559,6 +1567,33 @@ dependencies = [ "thiserror 2.0.18", ] +[[package]] +name = "bencher_replica" +version = "0.6.10" +dependencies = [ + "async-compression", + "aws-credential-types", + "aws-sdk-s3", + "bencher_json", + "bencher_otel", + "bytes", + "camino", + "futures", + "hex", + "pretty_assertions", + "rand 0.10.1", + "rusqlite", + "serde", + "serde_json", + "sha2 0.11.0", + "slog", + "tempfile", + "thiserror 2.0.18", + "tokio", + "uuid", + "zstd", +] + [[package]] name = "bencher_rootfs" version = "0.6.10" @@ -2159,6 +2194,8 @@ dependencies = [ "compression-core", "flate2", "memchr", + "zstd", + "zstd-safe", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index ba5100aff..961c0dd6b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -67,6 +67,7 @@ bencher_otel_provider = { path = "plus/bencher_otel_provider" } bencher_oci_storage = { path = "plus/bencher_oci_storage" } bencher_rate_limiter = { path = "plus/bencher_rate_limiter" } bencher_recaptcha = { path = "plus/bencher_recaptcha" } +bencher_replica = { path = "plus/bencher_replica" } # plus - runner bencher_runner = { path = "plus/bencher_runner" } bencher_oci = { path = "plus/bencher_oci" } diff --git a/docker/bench.Dockerfile b/docker/bench.Dockerfile index ce403ca91..15a945118 100644 --- a/docker/bench.Dockerfile +++ b/docker/bench.Dockerfile @@ -61,6 +61,7 @@ RUN cargo init --lib bencher_otel RUN cargo init --lib bencher_otel_provider RUN cargo init --lib bencher_rate_limiter RUN cargo init --lib bencher_recaptcha +RUN cargo init --lib bencher_replica RUN cargo init --lib bencher_rootfs RUN cargo init --lib bencher_runner RUN cargo init --lib bencher_init diff --git a/lib/api_server/Cargo.toml b/lib/api_server/Cargo.toml index 193317f48..fec51db33 100644 --- a/lib/api_server/Cargo.toml +++ b/lib/api_server/Cargo.toml @@ -42,8 +42,12 @@ slog.workspace = true [dev-dependencies] bencher_api_tests.workspace = true -bencher_json = { workspace = true, features = ["server", "schema"] } +bencher_json = { workspace = true, features = ["server", "schema", "test-clock"] } +bencher_replica = { workspace = true, features = ["plus", "testing"] } +camino.workspace = true http.workspace = true +rusqlite.workspace = true +tempfile.workspace = true tokio = { workspace = true, features = ["rt-multi-thread", "macros"] } [lints] diff --git a/lib/api_server/src/config.rs b/lib/api_server/src/config.rs index 989debe8f..ed2fe5902 100644 --- a/lib/api_server/src/config.rs +++ b/lib/api_server/src/config.rs @@ -27,6 +27,7 @@ pub async fn server_config_options( /// View server configuration /// /// View the API server configuration. +/// Secrets in the configuration are masked in the response. /// The user must be an admin on the server to use this route. #[endpoint { method = GET, @@ -53,7 +54,7 @@ async fn get_one_inner(log: &Logger) -> Result { ) })? .unwrap_or_default() - .into()) + .sanitized()) } #[endpoint { diff --git a/lib/api_server/src/lib.rs b/lib/api_server/src/lib.rs index 69a512a2f..93ee32203 100644 --- a/lib/api_server/src/lib.rs +++ b/lib/api_server/src/lib.rs @@ -2,8 +2,16 @@ #[cfg(test)] use bencher_api_tests as _; #[cfg(test)] +use bencher_replica as _; +#[cfg(test)] +use camino as _; +#[cfg(test)] use http as _; #[cfg(test)] +use rusqlite as _; +#[cfg(test)] +use tempfile as _; +#[cfg(test)] use tokio as _; mod backup; diff --git a/lib/api_server/src/stats.rs b/lib/api_server/src/stats.rs index 46dd4750e..bb835c15d 100644 --- a/lib/api_server/src/stats.rs +++ b/lib/api_server/src/stats.rs @@ -4,7 +4,7 @@ use bencher_json::{ }; use bencher_schema::{ auth_conn, - context::{ApiContext, DbConnection}, + context::{ApiContext, DbConnection, configure_standalone_connection}, error::{forbidden_error, issue_error, not_found_error}, model::{ server::QueryServer, @@ -49,8 +49,23 @@ pub async fn server_stats_get( async fn get_one_inner(log: &Logger, context: &ApiContext) -> Result { let query_server = QueryServer::get_server(auth_conn!(context))?; - let conn = DbConnection::establish(context.database.path.to_string_lossy().as_ref()) + let mut conn = DbConnection::establish(context.database.path.to_string_lossy().as_ref()) .map_err(not_found_error)?; + // Route this standalone connection through the shared configuration so it + // does not bypass the busy_timeout / autocheckpoint settings the pools and + // the writer use: under replication a stray write must never checkpoint. + configure_standalone_connection( + &mut conn, + context.database.busy_timeout, + context.database.replicated, + ) + .map_err(|e| { + issue_error( + "Failed to configure server stats connection", + "Failed to configure the server stats database connection PRAGMAs", + e, + ) + })?; query_server .get_stats(log.clone(), conn, context.is_bencher_cloud) .await diff --git a/lib/api_server/tests/replica.rs b/lib/api_server/tests/replica.rs new file mode 100644 index 000000000..242944de3 --- /dev/null +++ b/lib/api_server/tests/replica.rs @@ -0,0 +1,215 @@ +#![cfg(feature = "plus")] +#![expect(unused_crate_dependencies, reason = "integration test file")] +//! Integration tests for in-process replication (`bencher_replica`) against +//! a full API server: real Dropshot server, real writes through the API, +//! step-driven replication engine, restore-and-compare verification. + +#[cfg(test)] +mod cases { + use bencher_api_tests::TestServer; + use bencher_json::system::config::{JsonReplication, ReplicationTarget}; + use bencher_replica::testing::assert_replica_equivalent; + use bencher_replica::{ + CheckpointOutcome, EngineState, ReplicaConfig, RestoreOutcome, SyncEngine, + restore_if_missing, + }; + use camino::{Utf8Path, Utf8PathBuf}; + + fn dir_path(tmp: &tempfile::TempDir) -> Utf8PathBuf { + Utf8Path::from_path(tmp.path()) + .expect("tempdir path is UTF-8") + .to_path_buf() + } + + fn replica_json(root: &Utf8Path) -> JsonReplication { + JsonReplication { + target: ReplicationTarget::File { + path: root.to_path_buf().into_std_path_buf(), + }, + sync_interval_secs: None, + checkpoint_interval_secs: None, + min_checkpoint_pages: None, + snapshot_interval_secs: None, + snapshot_throttle_mib: None, + retention_generations: None, + verification_interval_secs: None, + shutdown_sync_timeout_secs: None, + } + } + + fn logger() -> slog::Logger { + slog::Logger::root(slog::Discard, slog::o!()) + } + + // A fixed, injected clock. The repo rule is that time-based tests inject a + // clock instead of racing the wall clock, and a frozen clock is also + // strictly safer here: the SyncEngine's only clock-driven work is the + // due-ness of periodic checkpoints, verification, and snapshots. With time + // frozen, none of those become due, so these tests exercise exactly the + // steps they drive explicitly (the bootstrap snapshot, ship, and + // checkpoint). The initial-generation snapshot is driven by + // `pending_new_generation` at construction, not the clock, so freezing time + // does not gate `until_streaming` into a hang. + fn test_clock() -> bencher_json::Clock { + bencher_json::Clock::Custom(std::sync::Arc::new(|| bencher_json::DateTime::TEST)) + } + + /// Drive the engine through the fresh-replica bootstrap snapshot. + async fn until_streaming(engine: &mut SyncEngine) { + for _ in 0u8..64 { + if engine.state() == EngineState::Streaming { + return; + } + engine.sync_once().await.expect("bootstrap sync"); + } + panic!( + "engine never reached Streaming; state: {:?}", + engine.state() + ); + } + + // Real API writes replicate, and the replica restores to an equivalent + // database. + #[tokio::test] + async fn server_writes_replicate_and_restore() { + let replica_tmp = tempfile::tempdir().expect("replica tempdir"); + let replica_root = dir_path(&replica_tmp); + let (server, mut engine) = + TestServer::new_with_replica(replica_json(&replica_root), Some(test_clock()), false) + .await; + until_streaming(&mut engine).await; + + // Writes through the real API. + let _admin = server.signup("Replica Admin", "replica@example.com").await; + let _user = server + .signup("Replica User", "replicauser@example.com") + .await; + engine.sync_once().await.expect("ship API writes"); + + // The replica restores to a logically equivalent database. + let restore_tmp = tempfile::tempdir().expect("restore tempdir"); + let target_db = dir_path(&restore_tmp).join("restored.db"); + let config = ReplicaConfig::try_from(replica_json(&replica_root)).expect("config"); + let outcome = restore_if_missing(&logger(), &config, &target_db) + .await + .expect("restore"); + assert!( + matches!(outcome, RestoreOutcome::Restored { .. }), + "expected Restored, got {outcome:?}" + ); + assert_replica_equivalent( + Utf8Path::from_path(server.db_path()).expect("db path is UTF-8"), + &target_db, + ); + } + + // The startup handshake: a missing database file is rebuilt from the + // replica, and the signed-up user is present in the restored database. + #[tokio::test] + async fn server_reboot_missing_db_auto_restores() { + let replica_tmp = tempfile::tempdir().expect("replica tempdir"); + let replica_root = dir_path(&replica_tmp); + let (server, mut engine) = + TestServer::new_with_replica(replica_json(&replica_root), Some(test_clock()), false) + .await; + until_streaming(&mut engine).await; + let _admin = server.signup("Reboot Admin", "reboot@example.com").await; + engine.sync_once().await.expect("ship API writes"); + drop(engine); + drop(server); + + // "Reboot": the volume is gone; the same restore-if-missing handshake + // that main.rs runs rebuilds the database before any connection opens. + let boot_tmp = tempfile::tempdir().expect("boot tempdir"); + let new_db = dir_path(&boot_tmp).join("bencher.db"); + let config = ReplicaConfig::try_from(replica_json(&replica_root)).expect("config"); + let outcome = restore_if_missing(&logger(), &config, &new_db) + .await + .expect("restore"); + assert!( + matches!(outcome, RestoreOutcome::Restored { .. }), + "expected Restored, got {outcome:?}" + ); + let conn = rusqlite::Connection::open(&new_db).expect("open restored db"); + let users: i64 = conn + .query_row( + "SELECT count(*) FROM user WHERE email = 'reboot@example.com'", + [], + |row| row.get(0), + ) + .expect("query restored user"); + assert_eq!(users, 1, "signed-up user must survive the reboot restore"); + } + + // A fresh server over an empty replica is a clean fresh start, and the + // existing database is never touched. + #[tokio::test] + async fn server_boot_empty_replica_fresh_start() { + let replica_tmp = tempfile::tempdir().expect("replica tempdir"); + let replica_root = dir_path(&replica_tmp); + let config = ReplicaConfig::try_from(replica_json(&replica_root)).expect("config"); + + // Empty replica, missing database: NoReplica (fresh server). + let boot_tmp = tempfile::tempdir().expect("boot tempdir"); + let new_db = dir_path(&boot_tmp).join("bencher.db"); + let outcome = restore_if_missing(&logger(), &config, &new_db) + .await + .expect("restore"); + assert!( + matches!(outcome, RestoreOutcome::NoReplica), + "expected NoReplica, got {outcome:?}" + ); + assert!(!new_db.exists(), "fresh start must not create a database"); + + // Existing database: Skipped, file untouched. + let (server, mut engine) = + TestServer::new_with_replica(replica_json(&replica_root), Some(test_clock()), false) + .await; + until_streaming(&mut engine).await; + let outcome = restore_if_missing( + &logger(), + &config, + Utf8Path::from_path(server.db_path()).expect("db path is UTF-8"), + ) + .await + .expect("restore"); + assert!( + matches!(outcome, RestoreOutcome::Skipped), + "expected Skipped, got {outcome:?}" + ); + } + + // Shadow mode: the replica ships and restores alongside a (nominal) + // Litestream, but never checkpoints. + #[tokio::test] + async fn server_shadow_mode_ships_without_checkpoints() { + let replica_tmp = tempfile::tempdir().expect("replica tempdir"); + let replica_root = dir_path(&replica_tmp); + let (server, mut engine) = + TestServer::new_with_replica(replica_json(&replica_root), Some(test_clock()), true) + .await; + until_streaming(&mut engine).await; + let _admin = server.signup("Shadow Admin", "shadow@example.com").await; + engine.sync_once().await.expect("ship API writes"); + + // Shadow never checkpoints: Litestream keeps checkpoint ownership. + let outcome = engine.checkpoint_once().await.expect("checkpoint"); + assert_eq!(outcome, CheckpointOutcome::SkippedShadow); + + // The shadow replica still restores to an equivalent database. + let restore_tmp = tempfile::tempdir().expect("restore tempdir"); + let target_db = dir_path(&restore_tmp).join("restored.db"); + let config = ReplicaConfig::try_from(replica_json(&replica_root)).expect("config"); + let outcome = restore_if_missing(&logger(), &config, &target_db) + .await + .expect("restore"); + assert!( + matches!(outcome, RestoreOutcome::Restored { .. }), + "expected Restored, got {outcome:?}" + ); + assert_replica_equivalent( + Utf8Path::from_path(server.db_path()).expect("db path is UTF-8"), + &target_db, + ); + } +} diff --git a/lib/bencher_api_tests/Cargo.toml b/lib/bencher_api_tests/Cargo.toml index 779bef401..32a90a4f3 100644 --- a/lib/bencher_api_tests/Cargo.toml +++ b/lib/bencher_api_tests/Cargo.toml @@ -12,9 +12,12 @@ plus = [ "bencher_api/plus", "bencher_config/plus", "bencher_oci_storage/plus", + "bencher_replica/plus", "bencher_schema/plus", "dep:bencher_license", "dep:bencher_oci_storage", + "dep:bencher_replica", + "dep:camino", "dep:hex", "dep:http", "dep:serde_json", @@ -30,6 +33,8 @@ bencher_endpoint.workspace = true bencher_json = { workspace = true, features = ["server", "schema", "db", "test-clock"] } bencher_license = { workspace = true, optional = true } bencher_oci_storage = { workspace = true, optional = true, features = ["test-clock"] } +bencher_replica = { workspace = true, optional = true, features = ["test-clock"] } +camino = { workspace = true, optional = true } bencher_rbac.workspace = true bencher_schema.workspace = true bencher_token.workspace = true diff --git a/lib/bencher_api_tests/src/context.rs b/lib/bencher_api_tests/src/context.rs index f376738bd..485032395 100644 --- a/lib/bencher_api_tests/src/context.rs +++ b/lib/bencher_api_tests/src/context.rs @@ -4,7 +4,11 @@ use rustls::crypto::aws_lc_rs; use bencher_config::DEFAULT_MAX_BODY_SIZE; use bencher_endpoint::Registrar as _; +#[cfg(feature = "plus")] +use bencher_json::system::config::JsonReplication; use bencher_rbac::init_rbac; +#[cfg(feature = "plus")] +use bencher_replica::{ReplicaConfig, ReplicaDb, SyncEngine}; use bencher_schema::{ ApiContext, context::{Database, DbConnection, Messenger}, @@ -44,13 +48,13 @@ pub struct TestServer { impl TestServer { /// Create a new test server with default settings. pub async fn new() -> Self { - Self::build(None, None, None, None).await + Self::build(None, None, None, None, false).await } /// Create a new test server with custom upload timeout and max body size. #[cfg(feature = "plus")] pub async fn new_with_limits(upload_timeout: u64, max_body_size: u64) -> Self { - Self::build(Some(upload_timeout), Some(max_body_size), None, None).await + Self::build(Some(upload_timeout), Some(max_body_size), None, None, false).await } /// Create a new test server with custom upload timeout, max body size, and injectable clock. @@ -60,13 +64,52 @@ impl TestServer { max_body_size: u64, clock: bencher_json::Clock, ) -> Self { - Self::build(Some(upload_timeout), Some(max_body_size), Some(clock), None).await + Self::build( + Some(upload_timeout), + Some(max_body_size), + Some(clock), + None, + false, + ) + .await } /// Create a new test server with a custom runner self-update base URL. #[cfg(feature = "plus")] pub async fn new_with_runner_update_base_url(base_url: url::Url) -> Self { - Self::build(None, None, None, Some(base_url)).await + Self::build(None, None, None, Some(base_url), false).await + } + + /// Create a test server whose database replicates to the given target, + /// returning the step-driven replication engine alongside it (tests + /// drive `sync_once` and friends deterministically instead of racing a + /// background tick task). The database runs in WAL mode with + /// `wal_autocheckpoint = 0`, matching production replication (I2). + #[cfg(feature = "plus")] + #[expect(clippy::expect_used, reason = "test server setup with fallible init")] + pub async fn new_with_replica( + replica: JsonReplication, + clock: Option, + shadow: bool, + ) -> (Self, SyncEngine) { + let server = Self::build(None, None, clock, None, true).await; + let config = ReplicaConfig::try_from(replica).expect("invalid replica config"); + let context = server.context(); + let log = slog::Logger::root(slog::Discard, slog::o!()); + let engine = SyncEngine::new( + log, + config, + ReplicaDb { + db_path: camino::Utf8PathBuf::from(server.db_path.clone()), + writer: context.database.connection.clone(), + busy_timeout_ms: context.database.busy_timeout, + }, + context.clock.clone(), + shadow, + ) + .await + .expect("failed to build replication engine"); + (server, engine) } #[cfg(feature = "plus")] @@ -80,6 +123,7 @@ impl TestServer { max_body_size: Option, clock: Option, runner_update_base_url: Option, + replicated: bool, ) -> Self { // Create logger early so it can be used for OCI storage let log_config = ConfigLogging::StderrTerminal { @@ -96,6 +140,20 @@ impl TestServer { // Establish connection and run migrations let mut conn = DbConnection::establish(&db_path).expect("Failed to establish database connection"); + if replicated { + // Match the production replication PRAGMAs (invariant I2: the + // replicator is the sole checkpointer). journal_mode is persistent + // in the database header; the rest route through the shared + // standalone-connection configurator so tests exercise the same + // PRAGMAs as production. + diesel::connection::SimpleConnection::batch_execute( + &mut conn, + "PRAGMA journal_mode = WAL", + ) + .expect("Failed to set WAL journal mode"); + bencher_schema::context::configure_standalone_connection(&mut conn, 5_000, true) + .expect("Failed to set replication PRAGMAs"); + } run_migrations(&mut conn).expect("Failed to run migrations"); // Create connection pools @@ -115,6 +173,7 @@ impl TestServer { let database = Database { path: PathBuf::from(&db_path), busy_timeout: 5_000, + replicated, public_pool, auth_pool, connection: Arc::new(Mutex::new(conn)), @@ -173,6 +232,7 @@ impl TestServer { max_body_size: Option, _clock: Option<()>, _runner_update_base_url: Option, + _replicated: bool, ) -> Self { // Create logger early so it can be used for OCI storage let log_config = ConfigLogging::StderrTerminal { @@ -208,6 +268,7 @@ impl TestServer { let database = Database { path: PathBuf::from(&db_path), busy_timeout: 5_000, + replicated: false, public_pool, auth_pool, connection: Arc::new(Mutex::new(conn)), diff --git a/lib/bencher_config/src/config_tx.rs b/lib/bencher_config/src/config_tx.rs index 452d0a6a2..58e7c1c27 100644 --- a/lib/bencher_config/src/config_tx.rs +++ b/lib/bencher_config/src/config_tx.rs @@ -18,7 +18,9 @@ use bencher_json::{ }, }; use bencher_rbac::init_rbac; -use bencher_schema::context::{ApiContext, Database, DbConnection}; +use bencher_schema::context::{ + ApiContext, Database, DbConnection, configure_standalone_connection, +}; #[cfg(feature = "plus")] use bencher_schema::{ context::RateLimiting, @@ -212,18 +214,32 @@ async fn into_context( // synchronous=NORMAL is safe with WAL mode and reduces fsync overhead. let busy_timeout = json_database.busy_timeout.unwrap_or(DEFAULT_BUSY_TIMEOUT); let cache_size = json_database.cache_size.unwrap_or(DEFAULT_CACHE_SIZE); + #[cfg(feature = "plus")] + let replicated = plus.as_ref().is_some_and(|plus| { + // Litestream (external) or bencher_replica (in-process): either way, + // the replication system is the sole checkpointer. + plus.litestream.is_some() || plus.replica.is_some() + }); + #[cfg(not(feature = "plus"))] + let replicated = false; info!( log, "Setting database PRAGMAs (busy_timeout: {busy_timeout}ms, cache_size: {cache_size} KiB)" ); + if replicated { + info!( + log, + "Replication owns checkpoints: disabling auto-checkpoint" + ); + } + // Writer connection only: WAL mode is persistent in the database header + // and set once here; the shared PRAGMAs (busy_timeout, synchronous, + // extended_result_codes, autocheckpoint) come from the single source of + // truth in `bencher_schema`. database_connection .batch_execute("PRAGMA journal_mode = WAL") .map_err(ConfigTxError::Pragma)?; - database_connection - .batch_execute(&format!("PRAGMA busy_timeout = {busy_timeout}")) - .map_err(ConfigTxError::Pragma)?; - database_connection - .batch_execute("PRAGMA synchronous = NORMAL") + configure_standalone_connection(&mut database_connection, busy_timeout, replicated) .map_err(ConfigTxError::Pragma)?; // Writer connection only: the read pools hold many connections, so a // large per-connection cache would multiply memory, and reads are served @@ -231,27 +247,12 @@ async fn into_context( database_connection .batch_execute(&format!("PRAGMA cache_size = -{cache_size}")) .map_err(ConfigTxError::Pragma)?; - // Surfaces `SQLITE_BUSY_SNAPSHOT` (517) etc. distinctly from plain `SQLITE_BUSY` (5) - // in error messages, instead of both rendering as "database is locked". - database_connection - .batch_execute("PRAGMA extended_result_codes = ON") - .map_err(ConfigTxError::Pragma)?; - - #[cfg(feature = "plus")] - if plus - .as_ref() - .and_then(|plus| plus.litestream.as_ref()) - .is_some() - { - info!(log, "Configuring Litestream"); - run_litestream(&mut database_connection)?; - } info!(log, "Running database migrations"); bencher_schema::run_migrations(&mut database_connection)?; - let public_pool = connection_pool(log, &database_path, busy_timeout)?; - let auth_pool = connection_pool(log, &database_path, busy_timeout)?; + let public_pool = connection_pool(log, &database_path, busy_timeout, replicated)?; + let auth_pool = connection_pool(log, &database_path, busy_timeout, replicated)?; let data_store = if let Some(data_store) = json_database.data_store { Some(data_store.try_into().map_err(ConfigTxError::DataStore)?) @@ -262,6 +263,7 @@ async fn into_context( let database = Database { path: json_database.file, busy_timeout, + replicated, public_pool, auth_pool, connection: Arc::new(tokio::sync::Mutex::new(database_connection)), @@ -429,44 +431,25 @@ fn sqlite_tmpdir(log: &Logger, database_path: &Path) -> Result<(), ConfigTxError Ok(()) } -/// Configure litestream-specific PRAGMAs. -/// -/// WAL mode, `busy_timeout`, and `synchronous=NORMAL` are now set unconditionally -/// on all connections. This function only sets the litestream-specific PRAGMA -/// to disable auto-checkpoints (so litestream can manage them). -#[cfg(feature = "plus")] -fn run_litestream(database: &mut DbConnection) -> Result<(), ConfigTxError> { - // Disable auto-checkpoints - // https://litestream.io/tips/#disable-autocheckpoints-for-high-write-load-servers - // https://sqlite.org/wal.html#automatic_checkpoint - database - .batch_execute("PRAGMA wal_autocheckpoint = 0") - .map_err(ConfigTxError::Pragma)?; - - Ok(()) -} - -/// Sets essential `SQLite` PRAGMAs on every new pool connection. -/// -/// `busy_timeout` is per-connection and prevents immediate `SQLITE_BUSY` failures -/// under lock contention. `synchronous = NORMAL` is safe with WAL mode and -/// reduces fsync overhead. +/// Sets essential `SQLite` PRAGMAs on every new pool connection, via the +/// shared `configure_standalone_connection` in `bencher_schema` (the single +/// source of truth these pool connections, the standalone writer above, the +/// stats sweep, and the test harness all share). `wal_autocheckpoint = 0` is +/// defense in depth when replication owns checkpointing: pool connections are +/// read-only by convention, but an accidental write through one must not +/// checkpoint. #[derive(Debug)] struct SqliteConnectionCustomizer { busy_timeout: u32, + replicated: bool, } impl diesel::r2d2::CustomizeConnection for SqliteConnectionCustomizer { fn on_acquire(&self, conn: &mut DbConnection) -> Result<(), diesel::r2d2::Error> { - conn.batch_execute(&format!("PRAGMA busy_timeout = {}", self.busy_timeout)) - .map_err(diesel::r2d2::Error::QueryError)?; - conn.batch_execute("PRAGMA synchronous = NORMAL") - .map_err(diesel::r2d2::Error::QueryError)?; - conn.batch_execute("PRAGMA extended_result_codes = ON") - .map_err(diesel::r2d2::Error::QueryError)?; - Ok(()) + configure_standalone_connection(conn, self.busy_timeout, self.replicated) + .map_err(diesel::r2d2::Error::QueryError) } } @@ -474,6 +457,7 @@ fn connection_pool( log: &Logger, database_path: &str, busy_timeout: u32, + replicated: bool, ) -> Result>, ConfigTxError> { let cpu_count = std::thread::available_parallelism() .map(NonZeroUsize::get) @@ -488,7 +472,10 @@ fn connection_pool( ); let connection_manager = ConnectionManager::new(database_path); - let customizer = SqliteConnectionCustomizer { busy_timeout }; + let customizer = SqliteConnectionCustomizer { + busy_timeout, + replicated, + }; Pool::builder() .max_size(max_size) @@ -624,6 +611,7 @@ async fn spawn_stats(log: &Logger, context: &ApiContext) -> Result<(), ConfigTxE log.clone(), context.database.path.clone(), context.database.busy_timeout, + context.database.replicated, context.stats, context .is_bencher_cloud diff --git a/lib/bencher_config/src/lib.rs b/lib/bencher_config/src/lib.rs index e69f37881..70147a9fc 100644 --- a/lib/bencher_config/src/lib.rs +++ b/lib/bencher_config/src/lib.rs @@ -6,7 +6,7 @@ use std::{ }; use bencher_json::{ - BENCHER_API_PORT, JsonConfig, sanitize_json, + BENCHER_API_PORT, JsonConfig, Sanitize as _, sanitize_json, system::config::{ JsonConsole, JsonDatabase, JsonLogging, JsonSecurity, JsonServer, LogLevel, ServerLog, }, @@ -208,6 +208,19 @@ impl Config { pub fn into_inner(self) -> JsonConfig { self.0 } + + /// Consume the config and return it with every secret masked. + /// + /// Used by the `GET /v0/server/config` endpoint so an admin can inspect the + /// server configuration without exposing plaintext secrets (security, + /// database, SMTP, OAuth, replica, and Litestream credentials). Unlike + /// [`sanitize_json`], this masks unconditionally in both debug and release + /// builds. + pub fn sanitized(self) -> JsonConfig { + let mut json = self.0; + json.sanitize(); + json + } } impl Default for Config { @@ -257,3 +270,143 @@ impl Deref for Config { &self.0 } } + +#[cfg(test)] +mod tests { + use bencher_json::{ + Sanitize as _, Secret, + system::config::{DataStore, JsonSmtp}, + }; + + use super::Config; + + // Mirror `bencher_valid::secret::SANITIZED_SECRET` (not publicly re-exported) + // by deriving the mask at runtime, so this test tracks the crate's mask + // instead of duplicating the literal. + fn mask() -> String { + let mut secret = secret("placeholder"); + secret.sanitize(); + String::from(secret) + } + + fn secret(value: &str) -> Secret { + value.parse().expect("valid secret") + } + + // `GET /v0/server/config` must never leak plaintext secrets. Sanitizing the + // base (non-plus) config masks the security, database, and SMTP secrets. + #[test] + fn sanitized_masks_base_secrets() { + let mut config = Config::default(); + config.0.security.secret_key = secret("PLAINTEXT_SECURITY"); + config.0.database.data_store = Some(DataStore::AwsS3 { + access_key_id: "access-key-id".to_owned(), + secret_access_key: secret("PLAINTEXT_DB"), + access_point: "access-point".to_owned(), + }); + config.0.smtp = Some(JsonSmtp { + hostname: "smtp.example.com".parse().unwrap(), + port: None, + insecure_host: None, + starttls: None, + username: "smtp-user".parse().unwrap(), + secret: secret("PLAINTEXT_SMTP"), + from_name: "Bencher".parse().unwrap(), + from_email: "bencher@example.com".parse().unwrap(), + }); + + let json = config.sanitized(); + let serialized = serde_json::to_string(&json).unwrap(); + + for plaintext in ["PLAINTEXT_SECURITY", "PLAINTEXT_DB", "PLAINTEXT_SMTP"] { + assert!( + !serialized.contains(plaintext), + "leaked {plaintext}: {serialized}" + ); + } + assert!( + serialized.contains(&mask()), + "no mask present: {serialized}" + ); + } + + // The plus config carries the OAuth, Litestream, and replica secrets flagged + // in the review. Sanitizing must mask every one of them. + #[cfg(feature = "plus")] + #[test] + fn sanitized_masks_plus_secrets() { + use bencher_json::system::config::{ + JsonGitHub, JsonGoogle, JsonLitestream, JsonPlus, JsonReplica, JsonReplication, + ReplicationTarget, + }; + + let mut config = Config::default(); + config.0.plus = Some(JsonPlus { + rate_limiting: None, + github: Some(JsonGitHub { + client_id: "github-client-id".parse().unwrap(), + client_secret: secret("PLAINTEXT_GITHUB"), + }), + google: Some(JsonGoogle { + client_id: "google-client-id".parse().unwrap(), + client_secret: secret("PLAINTEXT_GOOGLE"), + }), + litestream: Some(JsonLitestream { + replica: JsonReplica::S3 { + bucket: "litestream-bucket".to_owned(), + path: None, + endpoint: None, + region: None, + access_key_id: "access-key-id".to_owned(), + secret_access_key: secret("PLAINTEXT_LITESTREAM"), + sync_interval: None, + }, + snapshot: None, + validation: None, + checkpoint: None, + metrics_port: None, + }), + replica: Some(JsonReplication { + target: ReplicationTarget::S3 { + bucket: "replica-bucket".to_owned(), + path: None, + endpoint: None, + region: None, + access_key_id: "access-key-id".to_owned(), + secret_access_key: secret("PLAINTEXT_REPLICA"), + }, + sync_interval_secs: None, + checkpoint_interval_secs: None, + min_checkpoint_pages: None, + snapshot_interval_secs: None, + snapshot_throttle_mib: None, + retention_generations: None, + verification_interval_secs: None, + shutdown_sync_timeout_secs: None, + }), + stats: None, + cloud: None, + registry: None, + runners: None, + }); + + let json = config.sanitized(); + let serialized = serde_json::to_string(&json).unwrap(); + + for plaintext in [ + "PLAINTEXT_GITHUB", + "PLAINTEXT_GOOGLE", + "PLAINTEXT_LITESTREAM", + "PLAINTEXT_REPLICA", + ] { + assert!( + !serialized.contains(plaintext), + "leaked {plaintext}: {serialized}" + ); + } + assert!( + serialized.contains(&mask()), + "no mask present: {serialized}" + ); + } +} diff --git a/lib/bencher_json/src/system/config/mod.rs b/lib/bencher_json/src/system/config/mod.rs index cbc020462..ca5ad57af 100644 --- a/lib/bencher_json/src/system/config/mod.rs +++ b/lib/bencher_json/src/system/config/mod.rs @@ -33,6 +33,7 @@ pub use plus::{ DEFAULT_CHUNK_SIZE, DEFAULT_MAX_BODY_SIZE, DEFAULT_UPLOAD_TIMEOUT_SECS, MAX_CHUNK_SIZE, RegistryDataStore, }, + replica::{JsonReplication, ReplicationTarget}, runners::{DEFAULT_HEARTBEAT_TIMEOUT_SECS, DEFAULT_JOB_TIMEOUT_GRACE_PERIOD_SECS}, stats::JsonStats, }; diff --git a/lib/bencher_json/src/system/config/plus/litestream.rs b/lib/bencher_json/src/system/config/plus/litestream.rs index 3dc1ac043..c42b6a865 100644 --- a/lib/bencher_json/src/system/config/plus/litestream.rs +++ b/lib/bencher_json/src/system/config/plus/litestream.rs @@ -20,6 +20,9 @@ pub struct JsonLitestream { /// Checkpoint configuration #[serde(skip_serializing_if = "Option::is_none")] pub checkpoint: Option, + /// Port for the Prometheus metrics endpoint (serialized as the Litestream `addr` setting) + #[serde(skip_serializing_if = "Option::is_none")] + pub metrics_port: Option, } #[derive(Debug, Clone, Serialize, Deserialize)] @@ -142,7 +145,9 @@ mod db { snapshot, validation, checkpoint, + metrics_port, } = self; + let addr = metrics_port.map(|port| format!(":{port}")); let replica = LitestreamReplica::from(replica); let snapshot = snapshot.map(LitestreamSnapshot::from); let validation = validation.map(LitestreamValidation::from); @@ -170,6 +175,7 @@ mod db { level: Some(log_level.into()), }); let litestream = Litestream { + addr, dbs, snapshot, validation, @@ -182,6 +188,9 @@ mod db { #[derive(Debug, Clone, Serialize)] #[serde(rename_all = "kebab-case")] pub struct Litestream { + // https://litestream.io/reference/config/#metrics + #[serde(skip_serializing_if = "Option::is_none")] + pub addr: Option, pub dbs: Vec, #[serde(skip_serializing_if = "Option::is_none")] pub snapshot: Option, @@ -367,6 +376,7 @@ mod db { snapshot: None, validation: None, checkpoint: None, + metrics_port: None, }; let path = PathBuf::from("/path/to/db"); let log_level = LogLevel::Info; @@ -408,6 +418,7 @@ logging: interval: Some("6h".to_owned()), }), checkpoint: None, + metrics_port: None, }; let path = PathBuf::from("/path/to/db"); let log_level = LogLevel::Info; @@ -453,6 +464,7 @@ logging: min_page_count: Some(2000), truncate_page_n: Some(121_359), }), + metrics_port: None, }; let path = PathBuf::from("/path/to/db"); let log_level = LogLevel::Info; @@ -476,6 +488,44 @@ logging: ); } + #[test] + fn into_yaml_with_metrics_port() { + let json_litestream = JsonLitestream { + replica: JsonReplica::S3 { + bucket: "bucket".to_owned(), + path: Some("/path/to/backup".to_owned()), + endpoint: None, + region: None, + access_key_id: "access_key_id".to_owned(), + secret_access_key: "secret_access_key".parse().unwrap(), + sync_interval: None, + }, + snapshot: None, + validation: None, + checkpoint: None, + metrics_port: Some(9090), + }; + let path = PathBuf::from("/path/to/db"); + let log_level = LogLevel::Info; + let yaml = json_litestream.into_yaml(path, log_level).unwrap(); + pretty_assertions::assert_eq!( + yaml, + "addr: :9090 +dbs: +- path: /path/to/db + replica: + type: s3 + bucket: bucket + path: /path/to/backup + access-key-id: access_key_id + secret-access-key: secret_access_key + truncate-page-n: 0 +logging: + level: info +" + ); + } + #[test] fn into_yaml_file() { let json_litestream = JsonLitestream { @@ -486,6 +536,7 @@ logging: snapshot: None, validation: None, checkpoint: None, + metrics_port: None, }; let path = PathBuf::from("/path/to/db"); let log_level = LogLevel::Info; @@ -520,6 +571,7 @@ logging: snapshot: None, validation: None, checkpoint: None, + metrics_port: None, }; let path = PathBuf::from("/path/to/db"); let log_level = LogLevel::Info; diff --git a/lib/bencher_json/src/system/config/plus/mod.rs b/lib/bencher_json/src/system/config/plus/mod.rs index bb301eb08..b7dc548f1 100644 --- a/lib/bencher_json/src/system/config/plus/mod.rs +++ b/lib/bencher_json/src/system/config/plus/mod.rs @@ -11,6 +11,7 @@ pub mod google; pub mod litestream; pub mod rate_limiting; pub mod registry; +pub mod replica; pub mod runners; pub mod stats; @@ -20,6 +21,7 @@ pub use google::JsonGoogle; pub use litestream::JsonLitestream; pub use rate_limiting::JsonRateLimiting; pub use registry::JsonRegistry; +pub use replica::JsonReplication; pub use runners::JsonRunners; pub use stats::JsonStats; @@ -35,6 +37,8 @@ pub struct JsonPlus { #[serde(alias = "disaster_recovery", skip_serializing_if = "Option::is_none")] pub litestream: Option, #[serde(skip_serializing_if = "Option::is_none")] + pub replica: Option, + #[serde(skip_serializing_if = "Option::is_none")] pub stats: Option, #[serde(skip_serializing_if = "Option::is_none")] pub cloud: Option, @@ -49,6 +53,7 @@ impl Sanitize for JsonPlus { self.github.sanitize(); self.google.sanitize(); self.litestream.sanitize(); + self.replica.sanitize(); self.cloud.sanitize(); self.registry.sanitize(); } diff --git a/lib/bencher_json/src/system/config/plus/replica.rs b/lib/bencher_json/src/system/config/plus/replica.rs new file mode 100644 index 000000000..9f5e232ea --- /dev/null +++ b/lib/bencher_json/src/system/config/plus/replica.rs @@ -0,0 +1,227 @@ +use std::path::PathBuf; + +use bencher_valid::{Sanitize, Secret}; +#[cfg(feature = "schema")] +use schemars::JsonSchema; +use serde::{Deserialize, Serialize}; + +/// In-process disaster recovery replication (`bencher_replica`). +/// +/// Streams `SQLite` WAL frames to a replica target, takes periodic snapshots, +/// and restores the latest state at API server startup. Replaces Litestream. +/// When both `litestream` and `replica` are configured, the replica runs in +/// shadow mode: Litestream keeps checkpoint ownership and remains the restore +/// source. +/// +/// Exactly one server may replicate to a given target. There is no fencing +/// between concurrent writers, so pointing two servers at the same target +/// interleaves their lineages and corrupts the replica. +/// +/// `deny_unknown_fields` is deliberate: this is disaster-recovery config, so a +/// misspelled key (e.g. `sync_interval` for `sync_interval_secs`) must fail +/// loudly at load time instead of silently applying a default. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[cfg_attr(feature = "schema", derive(JsonSchema))] +#[serde(deny_unknown_fields)] +pub struct JsonReplication { + /// Replica target + pub target: ReplicationTarget, + /// How often to ship new WAL frames (seconds; default 1) + #[serde(skip_serializing_if = "Option::is_none")] + pub sync_interval_secs: Option, + /// Minimum interval between checkpoints (seconds; default 60) + #[serde(skip_serializing_if = "Option::is_none")] + pub checkpoint_interval_secs: Option, + /// Minimum WAL pages before a checkpoint triggers (default 1000) + #[serde(skip_serializing_if = "Option::is_none")] + pub min_checkpoint_pages: Option, + /// How often to start a new generation with a fresh snapshot + /// (seconds; default 86400) + #[serde(skip_serializing_if = "Option::is_none")] + pub snapshot_interval_secs: Option, + /// Snapshot copy throttle (MiB per second; default 32) + #[serde(skip_serializing_if = "Option::is_none")] + pub snapshot_throttle_mib: Option, + /// Number of generations to retain (default 3) + #[serde(skip_serializing_if = "Option::is_none")] + pub retention_generations: Option, + /// How often to run restore-and-compare verification + /// (seconds; default 86400, 0 disables). Each run restores a full copy of + /// the database next to the live one, so expect a transient doubling of + /// data-volume usage while it runs. + #[serde(skip_serializing_if = "Option::is_none")] + pub verification_interval_secs: Option, + /// Deadline for the final WAL ship at shutdown (seconds; default 4, + /// which fits inside Fly's 6 second kill timeout) + #[serde(skip_serializing_if = "Option::is_none")] + pub shutdown_sync_timeout_secs: Option, +} + +impl Sanitize for JsonReplication { + fn sanitize(&mut self) { + self.target.sanitize(); + } +} + +/// Where the replica lives: a local directory XOR S3-compatible storage. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[cfg_attr(feature = "schema", derive(JsonSchema))] +#[serde(tag = "scheme", rename_all = "snake_case", deny_unknown_fields)] +pub enum ReplicationTarget { + File { + path: PathBuf, + }, + /// S3-compatible object storage. Configure a lifecycle rule on the bucket + /// to abort incomplete multipart uploads: a crashed snapshot upload leaves + /// orphaned parts that otherwise accrue storage cost indefinitely. + S3 { + bucket: String, + #[serde(skip_serializing_if = "Option::is_none")] + path: Option, + /// Endpoint override for S3-compatible services (R2, `MinIO`) + #[serde(skip_serializing_if = "Option::is_none")] + endpoint: Option, + #[serde(skip_serializing_if = "Option::is_none")] + region: Option, + access_key_id: String, + secret_access_key: Secret, + }, +} + +impl Sanitize for ReplicationTarget { + fn sanitize(&mut self) { + match self { + Self::File { .. } => {}, + Self::S3 { + secret_access_key, .. + } => secret_access_key.sanitize(), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn replication_json_round_trip() { + let json = r#"{ + "target": { + "scheme": "s3", + "bucket": "my-bucket", + "path": "replica/prod", + "endpoint": "https://minio.example.com", + "region": "auto", + "access_key_id": "access_key_id", + "secret_access_key": "secret_access_key" + }, + "snapshot_interval_secs": 86400 + }"#; + let replication: JsonReplication = serde_json::from_str(json).unwrap(); + let ReplicationTarget::S3 { + bucket, + path, + endpoint, + region, + .. + } = &replication.target + else { + panic!("expected S3 target"); + }; + pretty_assertions::assert_eq!(bucket, "my-bucket"); + pretty_assertions::assert_eq!(path.as_deref(), Some("replica/prod")); + pretty_assertions::assert_eq!(endpoint.as_deref(), Some("https://minio.example.com")); + pretty_assertions::assert_eq!(region.as_deref(), Some("auto")); + pretty_assertions::assert_eq!(replication.snapshot_interval_secs, Some(86_400)); + pretty_assertions::assert_eq!(replication.sync_interval_secs, None); + } + + // A misspelled top-level field (e.g. `sync_interval` for + // `sync_interval_secs`) must fail loudly, not silently apply a default. + #[test] + fn replication_rejects_unknown_top_level_field() { + let json = r#"{ + "target": { "scheme": "file", "path": "/var/lib/bencher/replica" }, + "sync_interval": 5 + }"#; + let error = serde_json::from_str::(json) + .expect_err("misspelled field must fail deserialization"); + assert!( + error.to_string().contains("sync_interval"), + "error should name the unknown field, got: {error}" + ); + } + + // A misspelled field inside the target section must also fail loudly. + #[test] + fn replication_rejects_unknown_target_field() { + let json = r#"{ + "target": { + "scheme": "s3", + "buckett": "my-bucket", + "access_key_id": "access_key_id", + "secret_access_key": "secret_access_key" + } + }"#; + let error = serde_json::from_str::(json) + .expect_err("misspelled target field must fail deserialization"); + assert!( + error.to_string().contains("buckett"), + "error should name the unknown field, got: {error}" + ); + } + + // The correct config still parses cleanly. + #[test] + fn replication_accepts_known_fields() { + let json = r#"{ + "target": { "scheme": "file", "path": "/var/lib/bencher/replica" }, + "sync_interval_secs": 5, + "verification_interval_secs": 0 + }"#; + let replication: JsonReplication = + serde_json::from_str(json).expect("valid config must parse"); + pretty_assertions::assert_eq!(replication.sync_interval_secs, Some(5)); + pretty_assertions::assert_eq!(replication.verification_interval_secs, Some(0)); + } + + #[test] + fn replication_file_target_round_trip() { + let json = r#"{ "target": { "scheme": "file", "path": "/var/lib/bencher/replica" } }"#; + let replication: JsonReplication = serde_json::from_str(json).unwrap(); + let ReplicationTarget::File { path } = &replication.target else { + panic!("expected File target"); + }; + pretty_assertions::assert_eq!(path, &PathBuf::from("/var/lib/bencher/replica")); + } + + #[test] + fn sanitize_s3_secret() { + let mut replication = JsonReplication { + target: ReplicationTarget::S3 { + bucket: "bucket".to_owned(), + path: None, + endpoint: None, + region: None, + access_key_id: "access_key_id".to_owned(), + secret_access_key: "my_secret".parse().unwrap(), + }, + sync_interval_secs: None, + checkpoint_interval_secs: None, + min_checkpoint_pages: None, + snapshot_interval_secs: None, + snapshot_throttle_mib: None, + retention_generations: None, + verification_interval_secs: None, + shutdown_sync_timeout_secs: None, + }; + replication.sanitize(); + let ReplicationTarget::S3 { + secret_access_key, .. + } = &replication.target + else { + panic!("expected S3 target"); + }; + pretty_assertions::assert_eq!(secret_access_key.as_ref(), "************"); + } +} diff --git a/lib/bencher_schema/src/context/database.rs b/lib/bencher_schema/src/context/database.rs index 351af9c36..476ed39dc 100644 --- a/lib/bencher_schema/src/context/database.rs +++ b/lib/bencher_schema/src/context/database.rs @@ -6,6 +6,7 @@ use std::{ use bencher_json::{Secret, system::config::DataStore as DataStoreConfig}; use camino::Utf8PathBuf; +use diesel::connection::SimpleConnection as _; use diesel::r2d2::{ConnectionManager, Pool, PooledConnection}; use dropshot::HttpError; @@ -13,9 +14,48 @@ use crate::error::issue_error; pub type DbConnection = diesel::SqliteConnection; +/// Configure a standalone (non-pool) `SQLite` connection with the essential +/// PRAGMAs. +/// +/// `busy_timeout` (per connection) prevents immediate `SQLITE_BUSY` failures +/// under lock contention. `synchronous = NORMAL` is safe with WAL mode and +/// reduces fsync overhead. `extended_result_codes` surfaces `SQLITE_BUSY_SNAPSHOT` +/// (517) etc. distinctly from plain `SQLITE_BUSY` (5) in error messages. +/// +/// When `replicated` is true, a replication system (Litestream or the +/// in-process replica) owns checkpointing, so `wal_autocheckpoint = 0` must be +/// set on EVERY connection that can write: an auto-checkpoint from any writer +/// can restart the WAL behind the replicator's back, overwriting frames it has +/// not yet shipped. +/// +/// This is the single source of truth for shared connection PRAGMAs: the +/// primary writer setup and the connection-pool customizer in +/// `bencher_config::config_tx`, the stats sweep in `crate::model::server`, +/// the server-stats endpoint in `api_server::stats`, and the +/// `bencher_api_tests` replication harness all route through this. It +/// intentionally does NOT set `journal_mode` (persistent in the database +/// header) or `cache_size`; the writer sets those two inline on top. +pub fn configure_standalone_connection( + conn: &mut DbConnection, + busy_timeout: u32, + replicated: bool, +) -> diesel::QueryResult<()> { + conn.batch_execute(&format!("PRAGMA busy_timeout = {busy_timeout}"))?; + conn.batch_execute("PRAGMA synchronous = NORMAL")?; + conn.batch_execute("PRAGMA extended_result_codes = ON")?; + if replicated { + conn.batch_execute("PRAGMA wal_autocheckpoint = 0")?; + } + Ok(()) +} + pub struct Database { pub path: PathBuf, pub busy_timeout: u32, + /// True when a replication system (Litestream or the in-process replica) + /// owns WAL checkpointing: every writing connection must run with + /// `wal_autocheckpoint = 0`. + pub replicated: bool, /// The public database connection pool. /// Unauthenticated requests should only use this pool. pub public_pool: Pool>, @@ -279,10 +319,68 @@ impl S3Arn { #[cfg(test)] mod tests { + use diesel::{Connection as _, RunQueryDsl as _}; + use super::*; const TEST_BUCKET: &str = "arn:aws:s3:some-region-1:123456789:accesspoint/my-bucket"; + #[derive(diesel::QueryableByName)] + struct WalAutocheckpoint { + #[diesel(sql_type = diesel::sql_types::Integer)] + wal_autocheckpoint: i32, + } + + fn wal_autocheckpoint(conn: &mut DbConnection) -> i32 { + diesel::sql_query("PRAGMA wal_autocheckpoint") + .get_result::(conn) + .expect("read wal_autocheckpoint") + .wal_autocheckpoint + } + + // A temp-file database (not `:memory:`) so WAL mode and its + // autocheckpoint setting are meaningful. + fn wal_db(dir: &tempfile::TempDir) -> DbConnection { + let db_path = dir.path().join("test.db"); + let mut conn = DbConnection::establish(db_path.to_str().expect("utf8 db path")) + .expect("establish db connection"); + conn.batch_execute("PRAGMA journal_mode = WAL") + .expect("set WAL journal mode"); + conn + } + + // When replication owns checkpointing, the standalone connection must have + // auto-checkpoints disabled. + #[test] + fn configure_standalone_connection_disables_autocheckpoint_when_replicated() { + let dir = tempfile::tempdir().expect("tempdir"); + let mut conn = wal_db(&dir); + assert_ne!( + wal_autocheckpoint(&mut conn), + 0, + "sanity: a fresh WAL db has autocheckpoint enabled" + ); + configure_standalone_connection(&mut conn, 7_000, true).expect("configure replicated"); + assert_eq!( + wal_autocheckpoint(&mut conn), + 0, + "replication must disable autocheckpoint on the standalone connection" + ); + } + + // Without replication the connection keeps SQLite's default autocheckpoint. + #[test] + fn configure_standalone_connection_keeps_autocheckpoint_when_not_replicated() { + let dir = tempfile::tempdir().expect("tempdir"); + let mut conn = wal_db(&dir); + configure_standalone_connection(&mut conn, 3_000, false).expect("configure non-replicated"); + assert_ne!( + wal_autocheckpoint(&mut conn), + 0, + "non-replicated connections keep the default autocheckpoint" + ); + } + #[test] fn s3_arn_from_str_no_path() { let arn = S3Arn::from_str(TEST_BUCKET).unwrap(); diff --git a/lib/bencher_schema/src/context/mod.rs b/lib/bencher_schema/src/context/mod.rs index 99132aec3..61a344a16 100644 --- a/lib/bencher_schema/src/context/mod.rs +++ b/lib/bencher_schema/src/context/mod.rs @@ -34,7 +34,9 @@ mod stats; #[cfg(feature = "plus")] use bencher_recaptcha::RecaptchaClient; -pub use database::{DataStore, DataStoreError, Database, DbConnection}; +pub use database::{ + DataStore, DataStoreError, Database, DbConnection, configure_standalone_connection, +}; #[cfg(feature = "plus")] pub use heartbeat_tasks::HeartbeatTasks; #[cfg(feature = "plus")] @@ -162,7 +164,7 @@ macro_rules! write_conn { // `BEGIN IMMEDIATE` (not the diesel default `BEGIN` deferred): take the writer lock at // transaction start. Deferred begins as a reader and fails the read→write upgrade with // `SQLITE_BUSY_SNAPSHOT` (which bypasses `busy_timeout`) when the WAL has advanced since -// `BEGIN` — e.g. between a Litestream PASSIVE checkpoint and the first `INSERT`. +// `BEGIN`, e.g. between a Litestream PASSIVE checkpoint and the first `INSERT`. // `BEGIN IMMEDIATE` has no read-snapshot to invalidate, so contention waits via `busy_timeout`. // // Does NOT compose: nested calls fail with "cannot start a transaction within a transaction" diff --git a/lib/bencher_schema/src/model/server/plus/mod.rs b/lib/bencher_schema/src/model/server/plus/mod.rs index d838cd51d..b8930de09 100644 --- a/lib/bencher_schema/src/model/server/plus/mod.rs +++ b/lib/bencher_schema/src/model/server/plus/mod.rs @@ -11,7 +11,7 @@ use bencher_json::{ }; use bencher_license::Licensor; use chrono::{Duration, Utc}; -use diesel::{Connection as _, RunQueryDsl as _, connection::SimpleConnection as _}; +use diesel::{Connection as _, RunQueryDsl as _}; use dropshot::HttpError; use slog::Logger; use url::Url; @@ -19,7 +19,9 @@ use url::Url; use crate::resource_not_found_err; use crate::{ context::StatsSettings, - context::{Body, DbConnection, Message, Messenger, ServerStatsBody}, + context::{ + Body, DbConnection, Message, Messenger, ServerStatsBody, configure_standalone_connection, + }, error::{request_timeout_error, resource_conflict_err}, macros::fn_get::fn_get, model::{organization::plan::LicenseUsage, user::QueryUser}, @@ -79,6 +81,7 @@ impl QueryServer { log: Logger, db_path: PathBuf, busy_timeout: u32, + replicated: bool, stats: StatsSettings, messenger: Option, licensor: Licensor, @@ -122,7 +125,8 @@ impl QueryServer { continue; }; - if let Err(e) = configure_standalone_connection(&mut conn, busy_timeout) { + if let Err(e) = configure_standalone_connection(&mut conn, busy_timeout, replicated) + { slog::error!(log, "Failed to configure database connection PRAGMAs: {e}"); continue; } @@ -198,7 +202,9 @@ impl QueryServer { continue; }; - if let Err(e) = configure_standalone_connection(&mut conn, busy_timeout) { + if let Err(e) = + configure_standalone_connection(&mut conn, busy_timeout, replicated) + { slog::error!( log, "Failed to configure database connection PRAGMAs for sending stats: {e}" @@ -328,13 +334,3 @@ impl Default for InsertServer { } } } - -pub(super) fn configure_standalone_connection( - conn: &mut DbConnection, - busy_timeout: u32, -) -> diesel::QueryResult<()> { - conn.batch_execute(&format!("PRAGMA busy_timeout = {busy_timeout}"))?; - conn.batch_execute("PRAGMA synchronous = NORMAL")?; - conn.batch_execute("PRAGMA extended_result_codes = ON")?; - Ok(()) -} diff --git a/plus/bencher_otel/src/api_histogram.rs b/plus/bencher_otel/src/api_histogram.rs index 2dc5966c4..21fb5051a 100644 --- a/plus/bencher_otel/src/api_histogram.rs +++ b/plus/bencher_otel/src/api_histogram.rs @@ -16,6 +16,8 @@ pub enum ApiHistogram { ReportProcessDuration, /// Time spent in the batched DB write transaction per iteration. ReportWriteDuration, + /// Time the replica checkpoint critical section held the `SQLite` write lock. + ReplicaCriticalSectionDuration, } impl ApiHistogram { @@ -27,6 +29,7 @@ impl ApiHistogram { Self::ReportCreateDuration => "report.create.duration", Self::ReportProcessDuration => "report.process.duration", Self::ReportWriteDuration => "report.write.duration", + Self::ReplicaCriticalSectionDuration => "replica.critical_section.duration", } } @@ -46,6 +49,9 @@ impl ApiHistogram { Self::ReportWriteDuration => { "Time spent in the batched DB write transaction per iteration" }, + Self::ReplicaCriticalSectionDuration => { + "Time the replica checkpoint critical section held the SQLite write lock" + }, } } @@ -56,7 +62,8 @@ impl ApiHistogram { | Self::JobCompleteDuration(_) | Self::ReportCreateDuration | Self::ReportProcessDuration - | Self::ReportWriteDuration => "s", + | Self::ReportWriteDuration + | Self::ReplicaCriticalSectionDuration => "s", } } @@ -67,7 +74,8 @@ impl ApiHistogram { | Self::JobCompleteDuration(priority) => vec![priority_attribute(priority)], Self::ReportCreateDuration | Self::ReportProcessDuration - | Self::ReportWriteDuration => Vec::new(), + | Self::ReportWriteDuration + | Self::ReplicaCriticalSectionDuration => Vec::new(), } } } diff --git a/plus/bencher_otel/src/api_meter.rs b/plus/bencher_otel/src/api_meter.rs index 9e71f1fe4..3bf69c010 100644 --- a/plus/bencher_otel/src/api_meter.rs +++ b/plus/bencher_otel/src/api_meter.rs @@ -161,6 +161,26 @@ pub enum ApiCounter { // Self-hosted specific metrics SelfHostedServerStartup(Uuid), SelfHostedServerStats(Uuid), + + // Replica (in-process disaster recovery) metrics + ReplicaInitFailed, + ReplicaSegmentShip, + ReplicaShipFailed, + ReplicaCheckpoint, + ReplicaCheckpointPartial, + ReplicaCheckpointSkipped, + ReplicaCheckpointSkippedUnshipped, + ReplicaCheckpointFailed, + ReplicaSnapshot, + ReplicaSnapshotFailed, + ReplicaGenerationCreate, + ReplicaDivergence, + ReplicaRestore, + ReplicaRestoreSoftStop, + ReplicaVerifyPass, + ReplicaVerifyDivergence, + ReplicaVerifyFailed, + ReplicaShadowUnverifiedBoundary, } impl ApiCounter { @@ -224,6 +244,24 @@ impl ApiCounter { Self::RunnerHeartbeatTimeout | Self::RunnerJobTimeout => "{timeout}", Self::RunnerDisconnect => "{disconnect}", Self::RunnerSelfUpdateSent(_) => "{update}", + + Self::ReplicaSegmentShip => "{segment}", + Self::ReplicaInitFailed | Self::ReplicaShipFailed | Self::ReplicaSnapshotFailed => { + "{error}" + }, + Self::ReplicaCheckpoint + | Self::ReplicaCheckpointPartial + | Self::ReplicaCheckpointSkipped + | Self::ReplicaCheckpointSkippedUnshipped + | Self::ReplicaCheckpointFailed => "{checkpoint}", + Self::ReplicaSnapshot => "{snapshot}", + Self::ReplicaGenerationCreate => "{generation}", + Self::ReplicaDivergence => "{divergence}", + Self::ReplicaRestore | Self::ReplicaRestoreSoftStop => "{restore}", + Self::ReplicaVerifyPass | Self::ReplicaVerifyDivergence | Self::ReplicaVerifyFailed => { + "{verification}" + }, + Self::ReplicaShadowUnverifiedBoundary => "{boundary}", } } @@ -320,9 +358,33 @@ impl ApiCounter { // Self-hosted specific metrics Self::SelfHostedServerStartup(_) => "self_hosted.server.startup", Self::SelfHostedServerStats(_) => "self_hosted.server.stats", + + // Replica metrics + Self::ReplicaInitFailed => "replica.init.failed", + Self::ReplicaSegmentShip => "replica.segment.ship", + Self::ReplicaShipFailed => "replica.ship.failed", + Self::ReplicaCheckpoint => "replica.checkpoint", + Self::ReplicaCheckpointPartial => "replica.checkpoint.partial", + Self::ReplicaCheckpointSkipped => "replica.checkpoint.skipped", + Self::ReplicaCheckpointSkippedUnshipped => "replica.checkpoint.skipped_unshipped", + Self::ReplicaCheckpointFailed => "replica.checkpoint.failed", + Self::ReplicaSnapshot => "replica.snapshot", + Self::ReplicaSnapshotFailed => "replica.snapshot.failed", + Self::ReplicaGenerationCreate => "replica.generation.create", + Self::ReplicaDivergence => "replica.divergence", + Self::ReplicaRestore => "replica.restore", + Self::ReplicaRestoreSoftStop => "replica.restore.soft_stop", + Self::ReplicaVerifyPass => "replica.verify.pass", + Self::ReplicaVerifyDivergence => "replica.verify.divergence", + Self::ReplicaVerifyFailed => "replica.verify.failed", + Self::ReplicaShadowUnverifiedBoundary => "replica.shadow.unverified_boundary", } } + #[expect( + clippy::too_many_lines, + reason = "exhaustive one-line-per-counter metadata match" + )] fn description(&self) -> &'static str { match self { Self::ServerStartup => "Counts the number of server startups", @@ -446,6 +508,32 @@ impl ApiCounter { // Self-hosted specific metrics Self::SelfHostedServerStartup(_) => "Counts the number of self-hosted server startups", Self::SelfHostedServerStats(_) => "Counts the number of self-hosted server stats sent", + + // Replica metrics + Self::ReplicaInitFailed => "Counts replica engine construction failures", + Self::ReplicaSegmentShip => "Counts the number of replica WAL segments shipped", + Self::ReplicaShipFailed => "Counts the number of replica segment ship failures", + Self::ReplicaCheckpoint => "Counts the number of replica checkpoints completed", + Self::ReplicaCheckpointPartial => "Counts the number of partial replica checkpoints", + Self::ReplicaCheckpointSkipped => "Counts the number of skipped replica checkpoints", + Self::ReplicaCheckpointSkippedUnshipped => { + "Counts replica checkpoints skipped due to an unshipped WAL tail" + }, + Self::ReplicaCheckpointFailed => "Counts the number of replica checkpoint failures", + Self::ReplicaSnapshot => "Counts the number of replica snapshots completed", + Self::ReplicaSnapshotFailed => "Counts the number of replica snapshot failures", + Self::ReplicaGenerationCreate => "Counts the number of replica generations created", + Self::ReplicaDivergence => "Counts the number of replica divergences detected", + Self::ReplicaRestore => "Counts the number of replica restores at startup", + Self::ReplicaRestoreSoftStop => { + "Counts restores that dropped replica data (older generation or epoch prefix)" + }, + Self::ReplicaVerifyPass => "Counts the number of replica verification passes", + Self::ReplicaVerifyDivergence => { + "Counts replica verifications that found content divergence" + }, + Self::ReplicaVerifyFailed => "Counts replica verifications that failed to complete", + Self::ReplicaShadowUnverifiedBoundary => "Counts unverified shadow WAL boundaries", } } @@ -493,7 +581,25 @@ impl ApiCounter { | Self::RunnerMinutesBilledFailed | Self::RunnerHeartbeatTimeout | Self::RunnerJobTimeout - | Self::RunnerDisconnect => Vec::new(), + | Self::RunnerDisconnect + | Self::ReplicaInitFailed + | Self::ReplicaSegmentShip + | Self::ReplicaShipFailed + | Self::ReplicaCheckpoint + | Self::ReplicaCheckpointPartial + | Self::ReplicaCheckpointSkipped + | Self::ReplicaCheckpointSkippedUnshipped + | Self::ReplicaCheckpointFailed + | Self::ReplicaSnapshot + | Self::ReplicaSnapshotFailed + | Self::ReplicaGenerationCreate + | Self::ReplicaDivergence + | Self::ReplicaRestore + | Self::ReplicaRestoreSoftStop + | Self::ReplicaVerifyPass + | Self::ReplicaVerifyDivergence + | Self::ReplicaVerifyFailed + | Self::ReplicaShadowUnverifiedBoundary => Vec::new(), Self::Run(priority) | Self::MetricsCreate(priority) | Self::OciBandwidthMax(priority) => { diff --git a/plus/bencher_replica/Cargo.toml b/plus/bencher_replica/Cargo.toml new file mode 100644 index 000000000..da68bf0a4 --- /dev/null +++ b/plus/bencher_replica/Cargo.toml @@ -0,0 +1,70 @@ +[package] +name = "bencher_replica" +version.workspace = true +authors.workspace = true +edition.workspace = true +license-file.workspace = true +publish = false + +[features] +default = [] +test-clock = ["bencher_json?/test-clock"] +# Exposes the `testing` module (fixtures, fault injection, workload generator) +# to downstream integration tests (e.g. lib/api_server). +testing = ["dep:rand", "test-clock"] +# Required due to feature unification +# https://doc.rust-lang.org/cargo/reference/features.html#feature-unification +plus = [ + "dep:async-compression", + "dep:aws-credential-types", + "dep:aws-sdk-s3", + "dep:bencher_json", + "dep:bytes", + "dep:camino", + "dep:hex", + "dep:rusqlite", + "dep:serde", + "dep:serde_json", + "dep:sha2", + "dep:slog", + "dep:thiserror", + "dep:tokio", + "dep:uuid", + "dep:zstd", +] +otel = ["dep:bencher_otel", "bencher_otel/plus"] + +[dependencies] +# Bencher internal crates +bencher_json = { workspace = true, optional = true, features = ["server", "plus"] } +bencher_otel = { workspace = true, optional = true } + +# AWS S3 (same as database backup and OCI storage) +aws-credential-types = { workspace = true, optional = true } +aws-sdk-s3 = { workspace = true, optional = true } + +# External dependencies +async-compression = { workspace = true, optional = true, features = ["tokio", "zstd"] } +bytes = { workspace = true, optional = true } +camino = { workspace = true, optional = true } +hex = { workspace = true, optional = true } +rand = { workspace = true, optional = true } +rusqlite = { workspace = true, optional = true, features = ["backup"] } +serde = { workspace = true, optional = true } +serde_json = { workspace = true, optional = true } +sha2 = { workspace = true, optional = true } +slog = { workspace = true, optional = true } +thiserror = { workspace = true, optional = true } +tokio = { workspace = true, optional = true, features = ["fs", "io-util", "macros", "rt", "sync", "time"] } +uuid = { workspace = true, optional = true, features = ["v4"] } +zstd = { workspace = true, optional = true } + +[dev-dependencies] +futures.workspace = true +pretty_assertions.workspace = true +rand.workspace = true +tempfile.workspace = true +tokio = { workspace = true, features = ["macros", "rt-multi-thread"] } + +[lints] +workspace = true diff --git a/plus/bencher_replica/golden/wal_le_4096.db b/plus/bencher_replica/golden/wal_le_4096.db new file mode 100644 index 000000000..7999f81c9 Binary files /dev/null and b/plus/bencher_replica/golden/wal_le_4096.db differ diff --git a/plus/bencher_replica/golden/wal_le_4096.dump.txt b/plus/bencher_replica/golden/wal_le_4096.dump.txt new file mode 100644 index 000000000..a3fc05d8d --- /dev/null +++ b/plus/bencher_replica/golden/wal_le_4096.dump.txt @@ -0,0 +1,12 @@ +magic little_endian +page_size 4096 +checkpoint_seq 0 +salt 0x1b5b450b 0x7a0670a4 +frame offset=32 page_no=1 db_size=0 commit=false +frame offset=4152 page_no=2 db_size=2 commit=true +frame offset=8272 page_no=2 db_size=2 commit=true +frame offset=12392 page_no=2 db_size=2 commit=true +frame offset=16512 page_no=2 db_size=2 commit=true +end_offset 20632 +commit_count 4 +db_size_pages 2 diff --git a/plus/bencher_replica/golden/wal_le_4096.wal b/plus/bencher_replica/golden/wal_le_4096.wal new file mode 100644 index 000000000..f693df92e Binary files /dev/null and b/plus/bencher_replica/golden/wal_le_4096.wal differ diff --git a/plus/bencher_replica/src/backoff.rs b/plus/bencher_replica/src/backoff.rs new file mode 100644 index 000000000..d4fa08f6d --- /dev/null +++ b/plus/bencher_replica/src/backoff.rs @@ -0,0 +1,126 @@ +//! Capped exponential backoff for storage errors. +//! +//! Pure state machine: it computes delays but never sleeps, so all timing +//! stays under the caller's injected clock. + +use std::time::Duration; + +/// Capped exponential backoff: `base, base*2, base*4, ..., cap`. +#[derive(Debug, Clone)] +pub struct Backoff { + base: Duration, + cap: Duration, + current: Option, +} + +impl Backoff { + /// Default base delay matching Litestream: 1 second. + pub const DEFAULT_BASE: Duration = Duration::from_secs(1); + /// Default delay cap matching Litestream: 5 minutes (300 seconds). + pub const DEFAULT_CAP: Duration = Duration::from_mins(5); + + #[must_use] + pub fn new(base: Duration, cap: Duration) -> Self { + Self { + base, + cap, + current: None, + } + } + + /// The next delay to wait before retrying: `base` on the first failure, + /// doubling each subsequent call, saturating at `cap` (a cap below the + /// base clamps the very first delay too). + /// + /// A `Duration::ZERO` base disables backoff entirely: doubling zero stays + /// zero, so every retry is immediate and the cap never engages. + pub fn next_delay(&mut self) -> Duration { + let next = match self.current { + None => self.base.min(self.cap), + Some(current) => current.saturating_mul(2).min(self.cap), + }; + self.current = Some(next); + next + } + + /// Reset after a success; the next failure starts back at `base`. + pub fn reset(&mut self) { + self.current = None; + } +} + +impl Default for Backoff { + fn default() -> Self { + Self::new(Self::DEFAULT_BASE, Self::DEFAULT_CAP) + } +} + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use pretty_assertions::assert_eq; + + use super::Backoff; + + fn delays(backoff: &mut Backoff, count: usize) -> Vec { + std::iter::repeat_with(|| backoff.next_delay()) + .take(count) + .collect() + } + + #[test] + fn next_delay_follows_exact_doubling_sequence() { + let mut backoff = Backoff::new(Duration::from_secs(1), Duration::from_mins(5)); + let expected: Vec = [1u64, 2, 4, 8, 16, 32, 64, 128, 256, 300, 300, 300] + .into_iter() + .map(Duration::from_secs) + .collect(); + assert_eq!(delays(&mut backoff, 12), expected); + } + + #[test] + fn next_delay_saturates_at_cap() { + let mut backoff = Backoff::new(Duration::from_secs(1), Duration::from_secs(4)); + let expected: Vec = [1u64, 2, 4, 4, 4] + .into_iter() + .map(Duration::from_secs) + .collect(); + assert_eq!(delays(&mut backoff, 5), expected); + } + + #[test] + fn reset_returns_to_base() { + let mut backoff = Backoff::new(Duration::from_secs(1), Duration::from_mins(5)); + let _advanced = delays(&mut backoff, 4); + backoff.reset(); + assert_eq!(backoff.next_delay(), Duration::from_secs(1)); + assert_eq!(backoff.next_delay(), Duration::from_secs(2)); + } + + #[test] + fn defaults_match_litestream() { + assert_eq!(Backoff::DEFAULT_BASE, Duration::from_secs(1)); + // The cap is 300 seconds, spelled in the unit clippy prefers. + assert_eq!(Backoff::DEFAULT_CAP, Duration::from_mins(5)); + assert_eq!(Backoff::DEFAULT_CAP.as_secs(), 300); + let mut backoff = Backoff::default(); + assert_eq!(backoff.next_delay(), Duration::from_secs(1)); + } + + #[test] + fn zero_base_yields_immediate_retries() { + let mut backoff = Backoff::new(Duration::ZERO, Duration::from_mins(5)); + let expected = vec![Duration::ZERO; 6]; + assert_eq!(delays(&mut backoff, 6), expected); + backoff.reset(); + assert_eq!(backoff.next_delay(), Duration::ZERO); + } + + #[test] + fn cap_below_base_clamps_first_delay() { + let mut backoff = Backoff::new(Duration::from_secs(10), Duration::from_secs(3)); + let expected = vec![Duration::from_secs(3); 3]; + assert_eq!(delays(&mut backoff, 3), expected); + } +} diff --git a/plus/bencher_replica/src/checkpoint.rs b/plus/bencher_replica/src/checkpoint.rs new file mode 100644 index 000000000..e0bbc38e4 --- /dev/null +++ b/plus/bencher_replica/src/checkpoint.rs @@ -0,0 +1,289 @@ +//! The checkpoint critical section: a PASSIVE checkpoint under a frozen +//! writer. +//! +//! The exact order, executed while the caller holds the app writer mutex: +//! +//! 1. `BEGIN IMMEDIATE` on the dedicated lock connection (1s busy budget; +//! a busy stray writer yields [`CheckpointOutcome::SkippedBusy`], retried +//! next tick). +//! 2. Re-scan the WAL tail from the shipped position (plain `std::fs`): ANY +//! new committed frame beyond the position defers the checkpoint +//! ([`CheckpointOutcome::SkippedUnshipped`], invariant I1). Frames past +//! the last commit (an uncommitted or rolled-back spill) do not block: +//! they are not durable data. +//! 3. `PRAGMA wal_checkpoint(PASSIVE)` on the second connection. PASSIVE +//! needs no writer lock, so it runs to completion WHILE the lock +//! connection holds `BEGIN IMMEDIATE`: no instant exists between the +//! tail verification and the checkpoint where a stray writer can append +//! frames that would be backfilled unshipped. This closes the race that +//! litestream's release-then-checkpoint ordering leaves open, and makes +//! RESTART/TRUNCATE checkpoints unnecessary everywhere (I3: the next +//! writer restarts a fully backfilled WAL naturally, at a moment when +//! everything is shipped). +//! 4. Full backfill (`busy == 0 && log == checkpointed`) is +//! [`CheckpointOutcome::Completed`]; anything else (e.g. a long reader +//! pinning a mark) is [`CheckpointOutcome::Partial`]: fine, a WAL +//! restart is impossible without full backfill, so nothing unshipped can +//! be overwritten; retry next interval. +//! 5. `ROLLBACK`. +//! +//! [`checkpoint_locked`] is a synchronous `fn` on purpose: it is +//! compile-time proof that no network I/O (no `.await`) is reachable while +//! the `SQLite` write lock is held (invariant I5). + +use std::fs::File; +use std::io::{ErrorKind, Read as _}; +use std::time::Duration; + +use camino::{Utf8Path, Utf8PathBuf}; + +use crate::position::Position; +use crate::wal::{WAL_HEADER_SIZE, WalError, WalScanner, parse_wal_header}; + +/// Busy budget for `BEGIN IMMEDIATE` on the lock connection: stray writers +/// (stats/credit sweeps) hold the write lock only briefly, so 1s is ample; +/// on expiry the checkpoint is skipped and retried next tick. +const LOCK_BUSY_TIMEOUT: Duration = Duration::from_secs(1); + +/// The two dedicated checkpoint connections, opened lazily on first use and +/// reused across checkpoints. +pub(crate) struct CheckpointConns { + /// Holds `BEGIN IMMEDIATE` across the critical section. + lock: rusqlite::Connection, + /// Runs `PRAGMA wal_checkpoint(PASSIVE)`; `busy_timeout = 0` because + /// PASSIVE never needs to wait for anything. + ckpt: rusqlite::Connection, +} + +impl CheckpointConns { + pub(crate) fn open(db_path: &Utf8Path) -> Result { + let lock = rusqlite::Connection::open(db_path)?; + lock.busy_timeout(LOCK_BUSY_TIMEOUT)?; + disable_autocheckpoint(&lock)?; + let ckpt = rusqlite::Connection::open(db_path)?; + ckpt.busy_timeout(Duration::ZERO)?; + disable_autocheckpoint(&ckpt)?; + Ok(Self { lock, ckpt }) + } +} + +/// The outcome of one checkpoint attempt. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum CheckpointOutcome { + /// Every WAL frame was backfilled into the database file; the next + /// writer will restart the WAL legitimately. + Completed, + /// Some frames were backfilled but a reader mark pinned the rest (e.g. + /// the online backup's long read snapshot). Harmless: retried next + /// interval, and no WAL restart can happen without full backfill. + Partial, + /// A stray writer held the `SQLite` write lock past the busy budget; + /// retried next tick. App writers queue on the tokio mutex instead and + /// never see this. + SkippedBusy, + /// Committed frames beyond the shipped position exist; the checkpoint + /// is deferred until they ship (invariant I1). + SkippedUnshipped, + /// Shadow mode never checkpoints: Litestream keeps checkpoint + /// ownership until cutover. + SkippedShadow, +} + +#[derive(Debug, thiserror::Error)] +pub enum CheckpointError { + #[error("Checkpoint connection: {0}")] + Sqlite(#[from] rusqlite::Error), + #[error("Checkpoint WAL scan: {0}")] + Wal(#[from] WalError), + #[error("Failed to read WAL for checkpoint ({path}): {error}")] + WalRead { + path: Utf8PathBuf, + error: std::io::Error, + }, +} + +/// Run the checkpoint critical section. The caller MUST hold the app writer +/// mutex for the whole call, so app writers queue on the tokio mutex (never +/// burning their `busy_timeout`) and only stray connections can contend. +/// +/// Synchronous by design: see the module docs. Always releases the +/// `BEGIN IMMEDIATE` transaction, including on error paths. +pub(crate) fn checkpoint_locked( + conns: &mut CheckpointConns, + wal_path: &Utf8Path, + position: &Position, + max_transaction_bytes: u64, +) -> Result { + match conns.lock.execute_batch("BEGIN IMMEDIATE") { + Ok(()) => {}, + Err(error) if is_busy(&error) => return Ok(CheckpointOutcome::SkippedBusy), + Err(error) => return Err(error.into()), + } + let result = locked_section(conns, wal_path, position, max_transaction_bytes); + // Always release the write lock, even when the section failed. + let rollback = conns.lock.execute_batch("ROLLBACK"); + match (result, rollback) { + (Err(error), _) => Err(error), + (Ok(_), Err(error)) => Err(error.into()), + (Ok(outcome), Ok(())) => Ok(outcome), + } +} + +/// Steps 2 to 4 of the critical section, with the write lock held. +fn locked_section( + conns: &CheckpointConns, + wal_path: &Utf8Path, + position: &Position, + max_transaction_bytes: u64, +) -> Result { + if wal_has_unshipped_commit(wal_path, position, max_transaction_bytes)? { + return Ok(CheckpointOutcome::SkippedUnshipped); + } + let (busy, log, checkpointed): (i64, i64, i64) = + conns + .ckpt + .query_row("PRAGMA wal_checkpoint(PASSIVE)", [], |row| { + Ok((row.get(0)?, row.get(1)?, row.get(2)?)) + })?; + if busy == 0 && log == checkpointed { + Ok(CheckpointOutcome::Completed) + } else { + Ok(CheckpointOutcome::Partial) + } +} + +/// Whether any committed frame exists beyond the shipped position (the WAL +/// tail is frozen while the caller holds `BEGIN IMMEDIATE`). +/// +/// Conservative on every ambiguity: a missing or restarted WAL while the +/// position claims shipped bytes reports "unshipped", deferring the +/// checkpoint so the ship path can resolve the situation first (rebind, +/// epoch transition, or divergence). Without this, an externally truncated +/// WAL could launder a divergence into a legitimate-looking epoch change. +fn wal_has_unshipped_commit( + wal_path: &Utf8Path, + position: &Position, + max_transaction_bytes: u64, +) -> Result { + let read_err = |error| CheckpointError::WalRead { + path: wal_path.to_owned(), + error, + }; + let mut file = match File::open(wal_path) { + Ok(file) => file, + Err(error) if error.kind() == ErrorKind::NotFound => return Ok(position.offset > 0), + Err(error) => return Err(read_err(error)), + }; + let mut raw = [0u8; 32]; + match file.read_exact(&mut raw) { + Ok(()) => {}, + Err(error) if error.kind() == ErrorKind::UnexpectedEof => { + return Ok(position.offset > 0); + }, + Err(error) => return Err(read_err(error)), + } + let header = parse_wal_header(&raw)?; + if header.salt != position.salt { + // A different salt cycle: anything in it is by definition not + // shipped at this position. + return Ok(true); + } + let (offset, checksum) = if position.offset == 0 { + (WAL_HEADER_SIZE, header.checksum) + } else { + (position.offset, position.checksum) + }; + let mut scanner = WalScanner::resume(file, header, offset, checksum)?; + // Discard-mode scan: only the EXISTENCE of a commit past `offset` matters, + // so page bytes are checksum-validated and dropped instead of buffered. + // `max_bytes = 1` stops at the first commit boundary. `max_transaction_bytes` + // bounds the run since the last commit so a multi-GiB uncommitted tail is + // not fully re-read inside the checkpoint critical section: an oversized + // run aborts the scan, which is treated CONSERVATIVELY as "unshipped" + // (defer the checkpoint) so an ambiguous tail is never backfilled. + match scanner.scan_committed_extent(1, max_transaction_bytes) { + Ok(extent) => Ok(extent > offset), + Err(WalError::TransactionTooLarge { .. }) => Ok(true), + Err(error) => Err(error.into()), + } +} + +/// Outcome of pinning a read snapshot at the shipped position. +pub(crate) enum PinOutcome { + /// A read snapshot pinned exactly at the shipped position; dropping the + /// connection releases it. + Pinned(rusqlite::Connection), + /// A stray writer held the `SQLite` write lock past the busy budget; + /// retry later. + Busy, + /// Committed frames beyond the shipped position exist; ship first. + Unshipped, +} + +/// Pin a read snapshot at `position` for verification. +/// +/// The lock connection holds `BEGIN IMMEDIATE` across the pin, so no commit +/// can land between the shipped-position check and the snapshot: the pinned +/// state IS the shipped state. This holds even in shadow mode, where +/// Litestream-driven WAL churn would otherwise slip a commit in and produce a +/// spurious verification failure (a stray writer holding the lock just yields +/// [`PinOutcome::Busy`], which the caller retries). +/// +/// Synchronous by design, like [`checkpoint_locked`]: no `.await` while the +/// write lock is held. +pub(crate) fn pin_locked( + conns: &mut CheckpointConns, + db_path: &Utf8Path, + wal_path: &Utf8Path, + position: &Position, + max_transaction_bytes: u64, +) -> Result { + match conns.lock.execute_batch("BEGIN IMMEDIATE") { + Ok(()) => {}, + Err(error) if is_busy(&error) => return Ok(PinOutcome::Busy), + Err(error) => return Err(error.into()), + } + let result = pin_section(db_path, wal_path, position, max_transaction_bytes); + let rollback = conns.lock.execute_batch("ROLLBACK"); + match (result, rollback) { + (Err(error), _) => Err(error), + (Ok(_), Err(error)) => Err(error.into()), + (Ok(outcome), Ok(())) => Ok(outcome), + } +} + +/// The pin itself, with the write lock held when in exclusive mode. +fn pin_section( + db_path: &Utf8Path, + wal_path: &Utf8Path, + position: &Position, + max_transaction_bytes: u64, +) -> Result { + if wal_has_unshipped_commit(wal_path, position, max_transaction_bytes)? { + return Ok(PinOutcome::Unshipped); + } + let pinned = rusqlite::Connection::open(db_path)?; + pinned.busy_timeout(Duration::ZERO)?; + // A deferred BEGIN plus any read materializes the snapshot at the + // current committed state, which (with the write lock held) is exactly + // the shipped position. + pinned.execute_batch("BEGIN")?; + let _tables: i64 = + pinned.query_row("SELECT count(*) FROM sqlite_master", [], |row| row.get(0))?; + Ok(PinOutcome::Pinned(pinned)) +} + +/// `SQLITE_BUSY` (or the shared-cache `SQLITE_LOCKED`) from a statement. +fn is_busy(error: &rusqlite::Error) -> bool { + matches!( + error.sqlite_error_code(), + Some(rusqlite::ErrorCode::DatabaseBusy | rusqlite::ErrorCode::DatabaseLocked) + ) +} + +/// `wal_autocheckpoint = 0`: the checkpoint connections must never +/// checkpoint on their own (invariant I2). +fn disable_autocheckpoint(conn: &rusqlite::Connection) -> Result<(), rusqlite::Error> { + let _pages: i64 = conn.query_row("PRAGMA wal_autocheckpoint = 0", [], |row| row.get(0))?; + Ok(()) +} diff --git a/plus/bencher_replica/src/config.rs b/plus/bencher_replica/src/config.rs new file mode 100644 index 000000000..8f52dc55a --- /dev/null +++ b/plus/bencher_replica/src/config.rs @@ -0,0 +1,412 @@ +//! Resolved replication configuration: JSON config plus battle-tested +//! defaults (matching Litestream where applicable). + +use std::str::FromStr as _; +use std::time::Duration; + +use bencher_json::system::config::{JsonReplication, ReplicationTarget}; +use camino::Utf8PathBuf; + +use crate::local::LocalStorage; +use crate::s3::S3Storage; +use crate::storage::ReplicaStorage; + +/// How often to ship new WAL frames. +pub const DEFAULT_SYNC_INTERVAL: Duration = Duration::from_secs(1); +/// Minimum interval between checkpoints. +pub const DEFAULT_CHECKPOINT_INTERVAL: Duration = Duration::from_mins(1); +/// Minimum WAL pages before a checkpoint triggers. +pub const DEFAULT_MIN_CHECKPOINT_PAGES: u32 = 1000; +/// How often to start a new generation with a fresh snapshot. +pub const DEFAULT_SNAPSHOT_INTERVAL: Duration = Duration::from_hours(24); +/// Snapshot copy throttle in MiB per second. +pub const DEFAULT_SNAPSHOT_THROTTLE_MIB: u32 = 32; +/// Number of generations to retain. +pub const DEFAULT_RETENTION_GENERATIONS: u32 = 3; +/// How often to run restore-and-compare verification. On by default: it is +/// the backstop for the residual failure modes no local evidence can catch +/// (an external writer churning the database while the replicator is down). +pub const DEFAULT_VERIFICATION_INTERVAL: Duration = Duration::from_hours(24); +/// Deadline for the final WAL ship at shutdown: fits inside Fly's 6 second +/// kill timeout with margin for the connection drain in `server.close()`. +/// This bounds ONLY the ship loop; on a complete drain in sole mode a final +/// checkpoint runs AFTER it (unbounded, so it can exceed the kill window). The +/// checkpoint is crash-safe if truncated: it only backfills already-shipped +/// frames (invariant I1), so an interrupted checkpoint costs a re-snapshot on +/// the next boot, never data. +pub const DEFAULT_SHUTDOWN_SYNC_TIMEOUT: Duration = Duration::from_secs(4); + +#[derive(Debug, thiserror::Error)] +pub enum ReplicaConfigError { + #[error("Replica file target path is not valid UTF-8: {0}")] + NonUtf8Path(std::path::PathBuf), + #[error( + "checkpoint_interval_secs must be greater than 0 (0 would attempt a checkpoint every tick)" + )] + ZeroCheckpointInterval, + #[error( + "snapshot_interval_secs must be greater than 0 (0 would start a new generation every tick)" + )] + ZeroSnapshotInterval, + #[error("Replica S3 bucket must not be empty")] + EmptyS3Bucket, + #[error("Replica S3 access_key_id must not be empty")] + EmptyS3AccessKeyId, + #[error("Replica S3 secret_access_key must not be empty")] + EmptyS3SecretAccessKey, + #[error("Replica S3 endpoint is not a valid URL: {0}")] + InvalidS3Endpoint(bencher_json::ValidError), +} + +/// Resolved configuration for the replicator. +#[derive(Debug, Clone)] +pub struct ReplicaConfig { + pub target: ReplicationTarget, + pub sync_interval: Duration, + pub checkpoint_interval: Duration, + pub min_checkpoint_pages: u32, + pub snapshot_interval: Duration, + pub snapshot_throttle_mib: u32, + pub retention_generations: u32, + /// Restore-and-compare verification interval; `None` disables it + /// (config value 0 opts out; absent means the default). + pub verification_interval: Option, + pub shutdown_sync_timeout: Duration, + /// Largest single WAL transaction (raw bytes) the ship path will accept. + /// A transaction beyond this is structurally unshippable (it would exceed + /// the restore decompression bound and poison every restore), so shipping + /// it fails loudly. Not operator-configurable: it defaults to + /// [`crate::segment::MAX_DECOMPRESSED_BYTES`] and must never exceed it + /// (tests lower it to exercise the poison path). + pub max_transaction_bytes: u64, +} + +impl TryFrom for ReplicaConfig { + type Error = ReplicaConfigError; + + fn try_from(json: JsonReplication) -> Result { + let JsonReplication { + target, + sync_interval_secs, + checkpoint_interval_secs, + min_checkpoint_pages, + snapshot_interval_secs, + snapshot_throttle_mib, + retention_generations, + verification_interval_secs, + shutdown_sync_timeout_secs, + } = json; + // Validate the target up front, at config load time, so a + // misconfiguration fails loudly here instead of as an infinite runtime + // retry loop. + match &target { + ReplicationTarget::File { path } => { + if Utf8PathBuf::from_path_buf(path.clone()).is_err() { + return Err(ReplicaConfigError::NonUtf8Path(path.clone())); + } + }, + ReplicationTarget::S3 { + bucket, + endpoint, + access_key_id, + secret_access_key, + .. + } => { + if bucket.trim().is_empty() { + return Err(ReplicaConfigError::EmptyS3Bucket); + } + if access_key_id.trim().is_empty() { + return Err(ReplicaConfigError::EmptyS3AccessKeyId); + } + // `Secret` deserialization rejects the fully-empty string but + // not a whitespace-only one; catch that here like the fields + // above. + if secret_access_key.as_ref().trim().is_empty() { + return Err(ReplicaConfigError::EmptyS3SecretAccessKey); + } + if let Some(endpoint) = endpoint { + bencher_json::Url::from_str(endpoint) + .map_err(ReplicaConfigError::InvalidS3Endpoint)?; + } + }, + } + // A zero checkpoint or snapshot interval makes the corresponding + // due-ness check true every tick (a checkpoint attempt, or a whole new + // generation, every second). Retention clamps to 1 just below; + // verification 0 legitimately opts out; a zero sync interval is clamped + // to 1 second at runtime in the replicator tick shell (it drives the + // loop period, not a due-ness check, so it cannot be validated to a + // resolved `Duration` here). + if checkpoint_interval_secs == Some(0) { + return Err(ReplicaConfigError::ZeroCheckpointInterval); + } + if snapshot_interval_secs == Some(0) { + return Err(ReplicaConfigError::ZeroSnapshotInterval); + } + Ok(Self { + target, + sync_interval: sync_interval_secs.map_or(DEFAULT_SYNC_INTERVAL, |secs| { + Duration::from_secs(secs.into()) + }), + checkpoint_interval: checkpoint_interval_secs + .map_or(DEFAULT_CHECKPOINT_INTERVAL, |secs| { + Duration::from_secs(secs.into()) + }), + min_checkpoint_pages: min_checkpoint_pages.unwrap_or(DEFAULT_MIN_CHECKPOINT_PAGES), + snapshot_interval: snapshot_interval_secs.map_or(DEFAULT_SNAPSHOT_INTERVAL, |secs| { + Duration::from_secs(secs.into()) + }), + snapshot_throttle_mib: snapshot_throttle_mib.unwrap_or(DEFAULT_SNAPSHOT_THROTTLE_MIB), + retention_generations: retention_generations + .unwrap_or(DEFAULT_RETENTION_GENERATIONS) + .max(1), + verification_interval: match verification_interval_secs { + // Explicit zero opts out of verification entirely. + Some(0) => None, + Some(secs) => Some(Duration::from_secs(secs.into())), + None => Some(DEFAULT_VERIFICATION_INTERVAL), + }, + shutdown_sync_timeout: shutdown_sync_timeout_secs + .map_or(DEFAULT_SHUTDOWN_SYNC_TIMEOUT, |secs| { + Duration::from_secs(secs.into()) + }), + max_transaction_bytes: crate::segment::MAX_DECOMPRESSED_BYTES, + }) + } +} + +impl ReplicaConfig { + /// Build the storage backend for the configured target. + /// + /// The `File` path was validated as UTF-8 in `try_from`, so this cannot + /// fail afterward; a race is impossible because the target is immutable. + #[must_use] + pub fn build_storage(&self) -> ReplicaStorage { + match &self.target { + ReplicationTarget::File { path } => { + let root = Utf8PathBuf::from_path_buf(path.clone()) + .unwrap_or_else(|path| Utf8PathBuf::from(path.to_string_lossy().into_owned())); + ReplicaStorage::Local(LocalStorage::new(root)) + }, + ReplicationTarget::S3 { + bucket, + path, + endpoint, + region, + access_key_id, + secret_access_key, + } => ReplicaStorage::S3(Box::new(S3Storage::new( + bucket.clone(), + path.clone(), + endpoint.clone(), + region.clone(), + access_key_id.clone(), + secret_access_key.as_ref(), + ))), + } + } +} + +#[cfg(test)] +mod tests { + use std::path::PathBuf; + use std::time::Duration; + + use bencher_json::system::config::{JsonReplication, ReplicationTarget}; + use pretty_assertions::assert_eq; + + use super::{ + DEFAULT_CHECKPOINT_INTERVAL, DEFAULT_MIN_CHECKPOINT_PAGES, DEFAULT_RETENTION_GENERATIONS, + DEFAULT_SHUTDOWN_SYNC_TIMEOUT, DEFAULT_SNAPSHOT_INTERVAL, DEFAULT_SNAPSHOT_THROTTLE_MIB, + DEFAULT_SYNC_INTERVAL, DEFAULT_VERIFICATION_INTERVAL, ReplicaConfig, + }; + + fn file_json(overrides: impl FnOnce(&mut JsonReplication)) -> JsonReplication { + let mut json = JsonReplication { + target: ReplicationTarget::File { + path: PathBuf::from("/tmp/replica"), + }, + sync_interval_secs: None, + checkpoint_interval_secs: None, + min_checkpoint_pages: None, + snapshot_interval_secs: None, + snapshot_throttle_mib: None, + retention_generations: None, + verification_interval_secs: None, + shutdown_sync_timeout_secs: None, + }; + overrides(&mut json); + json + } + + fn s3_json(bucket: &str, access_key_id: &str, endpoint: Option<&str>) -> JsonReplication { + JsonReplication { + target: ReplicationTarget::S3 { + bucket: bucket.to_owned(), + path: None, + endpoint: endpoint.map(str::to_owned), + region: None, + access_key_id: access_key_id.to_owned(), + secret_access_key: "secret".parse().expect("valid secret"), + }, + sync_interval_secs: None, + checkpoint_interval_secs: None, + min_checkpoint_pages: None, + snapshot_interval_secs: None, + snapshot_throttle_mib: None, + retention_generations: None, + verification_interval_secs: None, + shutdown_sync_timeout_secs: None, + } + } + + #[test] + fn defaults_applied() { + let config = ReplicaConfig::try_from(file_json(|_| {})).unwrap(); + assert_eq!(config.sync_interval, DEFAULT_SYNC_INTERVAL); + assert_eq!(config.checkpoint_interval, DEFAULT_CHECKPOINT_INTERVAL); + assert_eq!(config.min_checkpoint_pages, DEFAULT_MIN_CHECKPOINT_PAGES); + assert_eq!(config.snapshot_interval, DEFAULT_SNAPSHOT_INTERVAL); + assert_eq!(config.snapshot_throttle_mib, DEFAULT_SNAPSHOT_THROTTLE_MIB); + assert_eq!(config.retention_generations, DEFAULT_RETENTION_GENERATIONS); + assert_eq!( + config.verification_interval, + Some(DEFAULT_VERIFICATION_INTERVAL) + ); + assert_eq!(config.shutdown_sync_timeout, DEFAULT_SHUTDOWN_SYNC_TIMEOUT); + } + + #[test] + fn overrides_applied() { + let config = ReplicaConfig::try_from(file_json(|json| { + json.sync_interval_secs = Some(2); + json.checkpoint_interval_secs = Some(120); + json.min_checkpoint_pages = Some(500); + json.snapshot_interval_secs = Some(3600); + json.snapshot_throttle_mib = Some(64); + json.retention_generations = Some(5); + json.verification_interval_secs = Some(86_400); + json.shutdown_sync_timeout_secs = Some(2); + })) + .unwrap(); + assert_eq!(config.sync_interval, Duration::from_secs(2)); + assert_eq!(config.checkpoint_interval, Duration::from_mins(2)); + assert_eq!(config.min_checkpoint_pages, 500); + assert_eq!(config.snapshot_interval, Duration::from_hours(1)); + assert_eq!(config.snapshot_throttle_mib, 64); + assert_eq!(config.retention_generations, 5); + assert_eq!(config.verification_interval, Some(Duration::from_hours(24))); + assert_eq!(config.shutdown_sync_timeout, Duration::from_secs(2)); + } + + #[test] + fn retention_zero_clamps_to_one() { + let config = ReplicaConfig::try_from(file_json(|json| { + json.retention_generations = Some(0); + })) + .unwrap(); + assert_eq!(config.retention_generations, 1); + } + + #[test] + fn verification_zero_opts_out() { + let config = ReplicaConfig::try_from(file_json(|json| { + json.verification_interval_secs = Some(0); + })) + .unwrap(); + assert_eq!(config.verification_interval, None); + } + + #[test] + fn checkpoint_interval_zero_rejected() { + let error = ReplicaConfig::try_from(file_json(|json| { + json.checkpoint_interval_secs = Some(0); + })) + .unwrap_err(); + assert!( + matches!(error, super::ReplicaConfigError::ZeroCheckpointInterval), + "expected ZeroCheckpointInterval, got {error:?}" + ); + } + + #[test] + fn snapshot_interval_zero_rejected() { + let error = ReplicaConfig::try_from(file_json(|json| { + json.snapshot_interval_secs = Some(0); + })) + .unwrap_err(); + assert!( + matches!(error, super::ReplicaConfigError::ZeroSnapshotInterval), + "expected ZeroSnapshotInterval, got {error:?}" + ); + } + + #[test] + fn max_transaction_bytes_defaults_to_decompression_bound() { + let config = ReplicaConfig::try_from(file_json(|_| {})).unwrap(); + assert_eq!( + config.max_transaction_bytes, + crate::segment::MAX_DECOMPRESSED_BYTES + ); + } + + #[test] + fn s3_valid_target_accepted() { + let config = + ReplicaConfig::try_from(s3_json("bucket", "AKIA", Some("https://r2.example.com"))) + .unwrap(); + assert!(matches!(config.target, ReplicationTarget::S3 { .. })); + } + + #[test] + fn s3_empty_bucket_rejected() { + let error = ReplicaConfig::try_from(s3_json(" ", "AKIA", None)).unwrap_err(); + assert!( + matches!(error, super::ReplicaConfigError::EmptyS3Bucket), + "expected EmptyS3Bucket, got {error:?}" + ); + } + + #[test] + fn s3_empty_access_key_rejected() { + let error = ReplicaConfig::try_from(s3_json("bucket", "", None)).unwrap_err(); + assert!( + matches!(error, super::ReplicaConfigError::EmptyS3AccessKeyId), + "expected EmptyS3AccessKeyId, got {error:?}" + ); + } + + #[test] + fn s3_whitespace_secret_rejected() { + let mut json = s3_json("bucket", "AKIA", None); + let ReplicationTarget::S3 { + secret_access_key, .. + } = &mut json.target + else { + panic!("s3_json builds an S3 target"); + }; + *secret_access_key = " ".parse().expect("whitespace passes Secret validation"); + let error = ReplicaConfig::try_from(json).unwrap_err(); + assert!( + matches!(error, super::ReplicaConfigError::EmptyS3SecretAccessKey), + "expected EmptyS3SecretAccessKey, got {error:?}" + ); + } + + #[test] + fn s3_malformed_endpoint_rejected() { + let error = + ReplicaConfig::try_from(s3_json("bucket", "AKIA", Some("not a url"))).unwrap_err(); + assert!( + matches!(error, super::ReplicaConfigError::InvalidS3Endpoint(_)), + "expected InvalidS3Endpoint, got {error:?}" + ); + } + + #[test] + fn file_target_builds_local_storage() { + let config = ReplicaConfig::try_from(file_json(|_| {})).unwrap(); + let storage = config.build_storage(); + assert!(matches!(storage, crate::storage::ReplicaStorage::Local(_))); + } +} diff --git a/plus/bencher_replica/src/lib.rs b/plus/bencher_replica/src/lib.rs new file mode 100644 index 000000000..03e26b11d --- /dev/null +++ b/plus/bencher_replica/src/lib.rs @@ -0,0 +1,86 @@ +#![cfg(feature = "plus")] + +//! # Bencher Replica +//! +//! In-process `SQLite` streaming replication: continuously ships WAL frames +//! to a local-filesystem or S3-compatible replica, takes periodic snapshots, +//! and restores the latest state at API server startup. Replaces Litestream. +//! +//! ## Core invariants +//! +//! - **I1 Ship-before-checkpoint**: a checkpoint is only issued after every +//! valid committed WAL frame has been durably uploaded to the replica. +//! - **I2 Sole checkpointer**: with replication configured, +//! `wal_autocheckpoint = 0` is set on every connection that can write. +//! - **I3 Restart safety**: `SQLite` restarts the WAL (new salts) only when a +//! writer finds it fully backfilled; given I1+I2, fully backfilled implies +//! fully shipped, so no unshipped frame is ever overwritten. +//! - **I4 DB-file writes = checkpoints**: the sync task is sequential, so the +//! main DB file is frozen during a snapshot copy; snapshots take no locks. +//! - **I5 Prime directive**: the `SQLite` write lock is only ever held for +//! O(WAL-tail) work, never O(database). +//! - **I6 Replica is the source of truth**: the local meta file is advisory; +//! any mismatch resolves to a new generation, never guessing. +//! +//! ## Operational assumptions +//! +//! - **One server per replica target.** There is no fencing: two servers +//! pointed at the same bucket/prefix will interleave generations and +//! segments, and the result is undefined. Each replica target must have +//! exactly one writer. +//! - **Verification needs transient headroom.** A restore-and-compare +//! verification restores a full copy of the database next to the live one, +//! so peak disk use is roughly 2x the database size for the duration of the +//! check. +//! - **Shipping pauses during the snapshot copy.** The sync task is +//! sequential, so while a snapshot body is copied and uploaded no WAL +//! segments ship; the recovery point objective widens to the copy duration +//! (bounded by the throttle) and then catches up. +//! - **S3 targets want an incomplete-multipart lifecycle rule.** Aborted or +//! crashed snapshot uploads leave orphaned multipart parts that this crate +//! does not garbage-collect; configure a bucket lifecycle rule to expire +//! incomplete multipart uploads. + +// Dev-dependency used by the integration tests in tests/, not the lib tests. +#[cfg(test)] +use futures as _; + +mod backoff; +mod checkpoint; +mod config; +mod local; +mod meta; +mod position; +mod replicator; +mod restore; +mod s3; +mod segment; +mod snapshot; +mod snapshot_meta; +mod storage; +mod sync; +mod verify; +mod wal; + +#[cfg(any(test, feature = "testing"))] +pub mod testing; + +pub use backoff::Backoff; +pub use checkpoint::{CheckpointError, CheckpointOutcome}; +pub use config::{DEFAULT_SHUTDOWN_SYNC_TIMEOUT, ReplicaConfig, ReplicaConfigError}; +pub use local::LocalStorage; +pub use meta::{MetaError, ReplicaMeta}; +pub use position::{GenerationId, Position, SegmentKey}; +pub use replicator::{ReplicaDb, Replicator, ReplicatorHandle}; +pub use restore::{RestoreError, RestoreOutcome, restore_if_missing}; +pub use s3::S3Storage; +pub use segment::{SEGMENT_MAX_BYTES, SegmentError, compress_segment, decompress_segment}; +pub use snapshot::{SnapshotError, SnapshotStatus}; +pub use snapshot_meta::{SnapshotMeta, SnapshotMetaError, WalBoundary}; +pub use storage::{MultipartUpload, ReplicaStorage, StorageError}; +pub use sync::{EngineState, SyncEngine, SyncError, SyncProgress}; +pub use verify::{VerifyError, VerifyReport, fingerprint_database, verify_against_replica}; +pub use wal::{ + CommittedChunk, FRAME_HEADER_SIZE, FrameHeader, WAL_HEADER_SIZE, WalError, WalHeader, + WalScanner, wal_checksum, +}; diff --git a/plus/bencher_replica/src/local.rs b/plus/bencher_replica/src/local.rs new file mode 100644 index 000000000..342bc2358 --- /dev/null +++ b/plus/bencher_replica/src/local.rs @@ -0,0 +1,852 @@ +//! Local-filesystem replica backend. +//! +//! Objects live under a root directory; keys map to relative paths. Writes +//! land in a `.partial-` sibling then rename into place (atomic +//! visibility), with the file fsynced before the rename (matching the +//! `rate_limiting.json` convention in `bencher_schema`). The parent directory +//! is fsynced after the rename, and each newly created parent directory is +//! fsynced as it is made, so an acknowledged `put` is durable across a power +//! loss (not merely atomically visible). A `put` that fails removes its own +//! partial file so a crash-looping writer cannot accumulate them. Because +//! partial files live inside the object tree, `list` filters out any file +//! whose name contains the partial infix; replica keys never contain it. + +use std::io; + +use bytes::Bytes; +use camino::{Utf8Path, Utf8PathBuf}; +use tokio::fs; +use tokio::io::AsyncWriteExt as _; +use uuid::Uuid; + +use crate::storage::StorageError; + +/// The infix separating a final object name from its temp-write suffix. +/// Files containing it are in-progress writes, never objects. +const PARTIAL_INFIX: &str = ".partial-"; + +/// Local filesystem backend rooted at a directory. +#[derive(Debug, Clone)] +pub struct LocalStorage { + root: Utf8PathBuf, +} + +#[derive(Debug, thiserror::Error)] +pub enum LocalError { + #[error("Failed to create directory ({path}): {error}")] + CreateDir { path: Utf8PathBuf, error: io::Error }, + #[error("Failed to read ({path}): {error}")] + Read { path: Utf8PathBuf, error: io::Error }, + #[error("Failed to write ({path}): {error}")] + Write { path: Utf8PathBuf, error: io::Error }, + #[error("Failed to rename ({from} -> {to}): {error}")] + Rename { + from: Utf8PathBuf, + to: Utf8PathBuf, + error: io::Error, + }, + #[error("Failed to remove ({path}): {error}")] + Remove { path: Utf8PathBuf, error: io::Error }, + #[error("Failed to list ({path}): {error}")] + List { path: Utf8PathBuf, error: io::Error }, + #[error("Non-UTF-8 path in replica directory: {0}")] + NonUtf8Path(std::path::PathBuf), +} + +impl LocalStorage { + #[must_use] + pub fn new(root: Utf8PathBuf) -> Self { + Self { root } + } + + #[must_use] + pub fn root(&self) -> &Utf8Path { + &self.root + } + + pub(crate) async fn put(&self, key: &str, bytes: Bytes) -> Result<(), StorageError> { + crate::storage::validate_key(key)?; + let final_path = self.root.join(key); + create_parent_dirs(&final_path).await?; + let partial_path = partial_path_for(&final_path); + // On any failure, remove the partial so a crash-looping writer cannot + // accumulate unbounded `.partial-` files; the original error wins. + match write_then_rename(&partial_path, &final_path, &bytes).await { + Ok(()) => Ok(()), + Err(error) => { + remove_partial(&partial_path).await; + Err(error.into()) + }, + } + } + + pub(crate) async fn get(&self, key: &str) -> Result { + crate::storage::validate_key(key)?; + let path = self.root.join(key); + match fs::read(&path).await { + Ok(bytes) => Ok(Bytes::from(bytes)), + Err(error) if error.kind() == io::ErrorKind::NotFound => Err(StorageError::NotFound { + key: key.to_owned(), + }), + Err(error) => Err(LocalError::Read { path, error }.into()), + } + } + + pub(crate) async fn get_stream( + &self, + key: &str, + ) -> Result, StorageError> { + crate::storage::validate_key(key)?; + let path = self.root.join(key); + match fs::File::open(&path).await { + Ok(file) => Ok(Box::new(file)), + Err(error) if error.kind() == io::ErrorKind::NotFound => Err(StorageError::NotFound { + key: key.to_owned(), + }), + Err(error) => Err(LocalError::Read { path, error }.into()), + } + } + + pub(crate) async fn list(&self, prefix: &str) -> Result, StorageError> { + // Start the walk at the deepest directory implied by the prefix so + // unrelated trees are not scanned; filter by full key prefix below + // (S3 semantics: the prefix may end mid-component). + let start_dir = match prefix.rsplit_once('/') { + Some((dir, _)) => self.root.join(dir), + None => self.root.clone(), + }; + // Iterative walk with an explicit stack: async fns cannot recurse + // without boxing. + let mut keys = Vec::new(); + let mut stack = vec![start_dir]; + while let Some(dir) = stack.pop() { + let mut entries = match fs::read_dir(&dir).await { + Ok(entries) => entries, + // A missing prefix directory is an empty result, not an error. + Err(error) if error.kind() == io::ErrorKind::NotFound => continue, + Err(error) => return Err(LocalError::List { path: dir, error }.into()), + }; + while let Some(entry) = + entries + .next_entry() + .await + .map_err(|error| LocalError::List { + path: dir.clone(), + error, + })? + { + let path = + Utf8PathBuf::from_path_buf(entry.path()).map_err(LocalError::NonUtf8Path)?; + let file_type = entry.file_type().await.map_err(|error| LocalError::List { + path: path.clone(), + error, + })?; + if file_type.is_dir() { + stack.push(path); + } else if let Ok(relative) = path.strip_prefix(&self.root) { + let key = relative_key(relative); + if key.starts_with(prefix) && !key.contains(PARTIAL_INFIX) { + keys.push(key); + } + } + } + } + keys.sort(); + Ok(keys) + } + + /// Best-effort reaper for `.partial-` fragments orphaned by a crash + /// (SIGKILL/OOM) mid-write: the error-path cleanup in [`Self::put`] and + /// the multipart finish never ran, so a potentially multi-GB fragment + /// lingers invisibly (listings filter the infix) forever. Safe at + /// startup: exactly one writer per target (see the crate operational + /// assumptions), so no in-flight write can race the sweep. Walk errors + /// are logged and swallowed; this never fails the caller. + pub(crate) async fn abort_incomplete_uploads(&self, log: &slog::Logger) { + let mut removed = 0u64; + let mut stack = vec![self.root.clone()]; + while let Some(dir) = stack.pop() { + let mut entries = match fs::read_dir(&dir).await { + Ok(entries) => entries, + Err(error) if error.kind() == io::ErrorKind::NotFound => continue, + Err(error) => { + slog::warn!(log, "Failed to read directory during partial-file sweep"; + "path" => dir.as_str(), "error" => %error); + continue; + }, + }; + loop { + let entry = match entries.next_entry().await { + Ok(Some(entry)) => entry, + Ok(None) => break, + Err(error) => { + slog::warn!(log, "Failed to read directory entry during partial-file sweep"; + "path" => dir.as_str(), "error" => %error); + break; + }, + }; + let Ok(path) = Utf8PathBuf::from_path_buf(entry.path()) else { + continue; + }; + let is_dir = entry + .file_type() + .await + .is_ok_and(|file_type| file_type.is_dir()); + if is_dir { + stack.push(path); + } else if path + .file_name() + .is_some_and(|name| name.contains(PARTIAL_INFIX)) + { + match fs::remove_file(&path).await { + Ok(()) => { + removed += 1; + slog::info!(log, "Removed crash-orphaned partial file"; + "path" => path.as_str()); + }, + Err(error) => { + slog::warn!(log, "Failed to remove crash-orphaned partial file"; + "path" => path.as_str(), "error" => %error); + }, + } + } + } + } + if removed > 0 { + slog::info!(log, "Partial-file sweep complete"; "removed" => removed); + } + } + + pub(crate) async fn list_dirs(&self, prefix: &str) -> Result, StorageError> { + let dir = if prefix.is_empty() { + self.root.clone() + } else { + self.root.join(prefix) + }; + let mut entries = match fs::read_dir(&dir).await { + Ok(entries) => entries, + // A missing prefix directory is an empty result, not an error. + Err(error) if error.kind() == io::ErrorKind::NotFound => return Ok(Vec::new()), + Err(error) => return Err(LocalError::List { path: dir, error }.into()), + }; + let mut dirs = Vec::new(); + while let Some(entry) = entries + .next_entry() + .await + .map_err(|error| LocalError::List { + path: dir.clone(), + error, + })? + { + let file_type = entry.file_type().await.map_err(|error| LocalError::List { + path: dir.clone(), + error, + })?; + if file_type.is_dir() { + let name = entry + .file_name() + .into_string() + .map_err(|name| LocalError::NonUtf8Path(name.into()))?; + dirs.push(name); + } + } + dirs.sort(); + Ok(dirs) + } + + pub(crate) async fn delete(&self, key: &str) -> Result<(), StorageError> { + crate::storage::validate_key(key)?; + let path = self.root.join(key); + match fs::remove_file(&path).await { + Ok(()) => Ok(()), + // Idempotent: deleting a missing key succeeds. + Err(error) if error.kind() == io::ErrorKind::NotFound => Ok(()), + Err(error) => Err(LocalError::Remove { path, error }.into()), + } + } + + pub(crate) async fn delete_prefix(&self, prefix: &str) -> Result<(), StorageError> { + // The local backend treats the prefix as a directory path; replica + // pruning always deletes directory-aligned prefixes (generations). + if prefix.is_empty() { + // Delete the root's children, never the root itself, so the + // backend stays usable afterwards. + return self.delete_root_children().await; + } + let path = self.root.join(prefix); + match fs::remove_dir_all(&path).await { + Ok(()) => Ok(()), + // A missing prefix is already deleted. + Err(error) if error.kind() == io::ErrorKind::NotFound => Ok(()), + Err(error) => Err(LocalError::Remove { path, error }.into()), + } + } + + pub(crate) async fn start_multipart(&self, key: &str) -> Result { + crate::storage::validate_key(key)?; + let final_path = self.root.join(key); + create_parent_dirs(&final_path).await?; + let partial_path = partial_path_for(&final_path); + let file = fs::File::create(&partial_path) + .await + .map_err(|error| LocalError::Write { + path: partial_path.clone(), + error, + })?; + Ok(LocalMultipart { + final_path, + partial_path, + file, + }) + } + + async fn delete_root_children(&self) -> Result<(), StorageError> { + let mut entries = match fs::read_dir(&self.root).await { + Ok(entries) => entries, + Err(error) if error.kind() == io::ErrorKind::NotFound => return Ok(()), + Err(error) => { + return Err(LocalError::List { + path: self.root.clone(), + error, + } + .into()); + }, + }; + while let Some(entry) = entries + .next_entry() + .await + .map_err(|error| LocalError::List { + path: self.root.clone(), + error, + })? + { + let path = Utf8PathBuf::from_path_buf(entry.path()).map_err(LocalError::NonUtf8Path)?; + let file_type = entry.file_type().await.map_err(|error| LocalError::List { + path: path.clone(), + error, + })?; + let removed = if file_type.is_dir() { + fs::remove_dir_all(&path).await + } else { + fs::remove_file(&path).await + }; + match removed { + Ok(()) => {}, + Err(error) if error.kind() == io::ErrorKind::NotFound => {}, + Err(error) => return Err(LocalError::Remove { path, error }.into()), + } + } + Ok(()) + } +} + +/// Streaming upload to the local backend: writes to `.partial-`, +/// fsyncs, then renames on finish. Dropping without finish leaves only the +/// partial file, never the final name; `list` never reports partial files. +pub struct LocalMultipart { + final_path: Utf8PathBuf, + partial_path: Utf8PathBuf, + file: fs::File, +} + +impl LocalMultipart { + pub(crate) async fn write_part(&mut self, bytes: Bytes) -> Result<(), StorageError> { + self.file + .write_all(&bytes) + .await + .map_err(|error| LocalError::Write { + path: self.partial_path.clone(), + error, + })?; + Ok(()) + } + + pub(crate) async fn finish(self) -> Result<(), StorageError> { + let Self { + final_path, + partial_path, + file, + } = self; + file.sync_all().await.map_err(|error| LocalError::Write { + path: partial_path.clone(), + error, + })?; + drop(file); + rename_into_place(&partial_path, &final_path).await?; + fsync_parent(&final_path).await?; + Ok(()) + } + + pub(crate) async fn abort(self) -> Result<(), StorageError> { + let Self { + final_path: _, + partial_path, + file, + } = self; + drop(file); + match fs::remove_file(&partial_path).await { + Ok(()) => Ok(()), + Err(error) if error.kind() == io::ErrorKind::NotFound => Ok(()), + Err(error) => Err(LocalError::Remove { + path: partial_path, + error, + } + .into()), + } + } +} + +/// The unique in-progress sibling for a final path: `.partial-`. +fn partial_path_for(final_path: &Utf8Path) -> Utf8PathBuf { + let mut name = final_path.to_string(); + name.push_str(PARTIAL_INFIX); + name.push_str(&Uuid::new_v4().simple().to_string()); + Utf8PathBuf::from(name) +} + +/// Create every missing parent directory of `path`, fsyncing the containing +/// directory after each new level appears. Object durability depends on the +/// parent directories surviving a crash, so the directory entries must be +/// flushed, not just the file content. +async fn create_parent_dirs(path: &Utf8Path) -> Result<(), LocalError> { + let Some(parent) = path.parent() else { + return Ok(()); + }; + // Walk up to the nearest existing ancestor, recording the missing levels. + let mut missing = Vec::new(); + let mut cursor = Some(parent); + while let Some(dir) = cursor { + match fs::try_exists(dir).await { + Ok(true) => break, + Ok(false) => { + missing.push(dir.to_owned()); + cursor = dir.parent(); + }, + Err(error) => { + return Err(LocalError::CreateDir { + path: dir.to_owned(), + error, + }); + }, + } + } + // Create shallow-to-deep, fsyncing each new directory's container so the + // new entry is durable. + for dir in missing.iter().rev() { + match fs::create_dir(dir).await { + Ok(()) => {}, + // A concurrent put may have created it first. + Err(error) if error.kind() == io::ErrorKind::AlreadyExists => {}, + Err(error) => { + return Err(LocalError::CreateDir { + path: dir.clone(), + error, + }); + }, + } + fsync_parent(dir).await?; + } + Ok(()) +} + +/// Write to the partial path, rename it into place, and fsync the parent +/// directory so the rename is durable. The single fallible region a failing +/// `put` guards with partial-file cleanup. +async fn write_then_rename( + partial: &Utf8Path, + final_path: &Utf8Path, + bytes: &[u8], +) -> Result<(), LocalError> { + write_file(partial, bytes).await?; + rename_into_place(partial, final_path).await?; + fsync_parent(final_path).await?; + Ok(()) +} + +/// Best-effort removal of a partial file on a `put` error path. A failure +/// here is ignored so the original error is preserved; the leftover partial +/// is inert either way (`list` filters the partial infix). +async fn remove_partial(partial: &Utf8Path) { + drop(fs::remove_file(partial).await); +} + +/// Fsync the parent directory of `path` so a rename or child creation within +/// it is durable (the directory metadata is only guaranteed on disk after the +/// containing directory is fsynced; see `restore::finalize` for the same +/// technique). +async fn fsync_parent(path: &Utf8Path) -> Result<(), LocalError> { + if let Some(parent) = path.parent() { + fsync_dir(parent).await?; + } + Ok(()) +} + +/// Open a directory and fsync it, flushing its entries to disk. +async fn fsync_dir(dir: &Utf8Path) -> Result<(), LocalError> { + let handle = fs::File::open(dir) + .await + .map_err(|error| LocalError::Read { + path: dir.to_owned(), + error, + })?; + handle.sync_all().await.map_err(|error| LocalError::Write { + path: dir.to_owned(), + error, + })?; + Ok(()) +} + +/// Write `bytes` to `path` and fsync so the content is durable before the +/// atomic rename (see `bencher_schema::context::rate_limiting` for the +/// convention). +async fn write_file(path: &Utf8Path, bytes: &[u8]) -> Result<(), LocalError> { + let mut file = fs::File::create(path) + .await + .map_err(|error| LocalError::Write { + path: path.to_owned(), + error, + })?; + file.write_all(bytes) + .await + .map_err(|error| LocalError::Write { + path: path.to_owned(), + error, + })?; + file.sync_all().await.map_err(|error| LocalError::Write { + path: path.to_owned(), + error, + })?; + Ok(()) +} + +async fn rename_into_place(from: &Utf8Path, to: &Utf8Path) -> Result<(), LocalError> { + fs::rename(from, to) + .await + .map_err(|error| LocalError::Rename { + from: from.to_owned(), + to: to.to_owned(), + error, + }) +} + +/// A storage key from a root-relative path: components joined with `/`. +fn relative_key(relative: &Utf8Path) -> String { + relative + .components() + .map(|component| component.as_str()) + .collect::>() + .join("/") +} + +#[cfg(test)] +mod tests { + use pretty_assertions::assert_eq; + + use super::*; + + fn storage(tmp: &tempfile::TempDir) -> LocalStorage { + let root = Utf8Path::from_path(tmp.path()) + .expect("tempdir path is UTF-8") + .to_path_buf(); + LocalStorage::new(root) + } + + /// Every file in the tree under `root`, as root-relative strings, + /// including partial files (unlike `list`). + fn all_files(root: &Utf8Path) -> Vec { + let mut files = Vec::new(); + let mut stack = vec![root.to_path_buf()]; + while let Some(dir) = stack.pop() { + for entry in std::fs::read_dir(&dir).expect("read_dir failed") { + let path = Utf8PathBuf::from_path_buf(entry.expect("entry failed").path()) + .expect("path is UTF-8"); + if path.is_dir() { + stack.push(path); + } else { + files.push(relative_key(path.strip_prefix(root).expect("under root"))); + } + } + } + files.sort(); + files + } + + #[tokio::test] + async fn put_creates_parents_and_leaves_no_partial() { + let tmp = tempfile::tempdir().expect("tempdir failed"); + let local = storage(&tmp); + local + .put("a/b/c.txt", Bytes::from_static(b"content")) + .await + .expect("put failed"); + assert_eq!( + all_files(local.root()), + vec!["a/b/c.txt".to_owned()], + "put must leave exactly the final file, no partial" + ); + } + + #[tokio::test] + async fn get_missing_maps_not_found() { + let tmp = tempfile::tempdir().expect("tempdir failed"); + let local = storage(&tmp); + let error = local.get("nope.txt").await.expect_err("get must fail"); + assert!( + matches!(&error, StorageError::NotFound { key } if key == "nope.txt"), + "expected NotFound, got: {error}" + ); + } + + #[tokio::test] + async fn list_filters_partial_files() { + let tmp = tempfile::tempdir().expect("tempdir failed"); + let local = storage(&tmp); + local + .put("wal/seg1", Bytes::from_static(b"1")) + .await + .expect("put failed"); + // Simulate a crashed mid-put: a stale partial file in the tree. + std::fs::write( + local.root().join("wal/seg2.partial-deadbeef"), + b"incomplete", + ) + .expect("write partial failed"); + let keys = local.list("").await.expect("list failed"); + assert_eq!( + keys, + vec!["wal/seg1".to_owned()], + "list must filter out partial files" + ); + } + + #[tokio::test] + async fn list_walks_deep_trees_iteratively() { + let tmp = tempfile::tempdir().expect("tempdir failed"); + let local = storage(&tmp); + let dirs = (0..64).map(|depth| format!("d{depth}")).collect::>(); + let key = format!("{}/leaf.txt", dirs.join("/")); + local + .put(&key, Bytes::from_static(b"deep")) + .await + .expect("put failed"); + let keys = local.list("").await.expect("list failed"); + assert_eq!(keys, vec![key], "deep tree must be fully listed"); + } + + #[tokio::test] + async fn list_mid_component_prefix_matches() { + let tmp = tempfile::tempdir().expect("tempdir failed"); + let local = storage(&tmp); + local + .put("generations/2026a/x", Bytes::from_static(b"1")) + .await + .expect("put failed"); + local + .put("generations/2027b/y", Bytes::from_static(b"2")) + .await + .expect("put failed"); + // S3 semantics: a prefix may end in the middle of a component. + let keys = local.list("generations/2026").await.expect("list failed"); + assert_eq!( + keys, + vec!["generations/2026a/x".to_owned()], + "mid-component prefix must match like S3" + ); + } + + #[tokio::test] + async fn delete_prefix_empty_prefix_keeps_root_usable() { + let tmp = tempfile::tempdir().expect("tempdir failed"); + let local = storage(&tmp); + local + .put("a/1.txt", Bytes::from_static(b"1")) + .await + .expect("put failed"); + local + .put("top.txt", Bytes::from_static(b"2")) + .await + .expect("put failed"); + local.delete_prefix("").await.expect("delete_prefix failed"); + assert_eq!( + local.list("").await.expect("list failed"), + Vec::::new(), + "delete_prefix of empty prefix must remove everything" + ); + // The root itself survives: a subsequent put works. + local + .put("again.txt", Bytes::from_static(b"3")) + .await + .expect("put after root wipe failed"); + } + + #[tokio::test] + async fn multipart_drop_leaves_only_partial() { + let tmp = tempfile::tempdir().expect("tempdir failed"); + let local = storage(&tmp); + let mut upload = local + .start_multipart("snap/db.zst") + .await + .expect("start failed"); + upload + .write_part(Bytes::from_static(b"chunk")) + .await + .expect("write failed"); + drop(upload); + let files = all_files(local.root()); + assert_eq!(files.len(), 1, "exactly one file expected: {files:?}"); + let only = files.first().expect("one file"); + assert!( + only.starts_with("snap/db.zst.partial-"), + "dropped upload must leave only the partial file, got {only}" + ); + } + + #[tokio::test] + async fn multipart_finish_renames_partial_away() { + let tmp = tempfile::tempdir().expect("tempdir failed"); + let local = storage(&tmp); + let mut upload = local + .start_multipart("snap/db.zst") + .await + .expect("start failed"); + upload + .write_part(Bytes::from_static(b"part one ")) + .await + .expect("write failed"); + upload + .write_part(Bytes::from_static(b"part two")) + .await + .expect("write failed"); + upload.finish().await.expect("finish failed"); + assert_eq!( + all_files(local.root()), + vec!["snap/db.zst".to_owned()], + "finish must leave exactly the final file" + ); + let got = local.get("snap/db.zst").await.expect("get failed"); + assert_eq!( + got.as_ref(), + b"part one part two".as_slice(), + "finished object must concatenate all parts" + ); + } + + #[tokio::test] + async fn multipart_abort_removes_partial() { + let tmp = tempfile::tempdir().expect("tempdir failed"); + let local = storage(&tmp); + let mut upload = local + .start_multipart("snap/db.zst") + .await + .expect("start failed"); + upload + .write_part(Bytes::from_static(b"chunk")) + .await + .expect("write failed"); + upload.abort().await.expect("abort failed"); + assert_eq!( + all_files(local.root()), + Vec::::new(), + "abort must remove the partial file" + ); + } + + #[tokio::test] + async fn sweep_removes_crash_orphaned_partials() { + let tmp = tempfile::tempdir().expect("tempdir failed"); + let local = storage(&tmp); + local + .put( + "generations/g1/wal/e0/seg.wal.zst", + Bytes::from_static(b"x"), + ) + .await + .expect("put"); + // A crash-orphaned partial: written directly, no cleanup path ran. + let orphan = local + .root() + .join("generations/g1/snapshot.db.zst.partial-deadbeef"); + std::fs::write(orphan.as_std_path(), b"torn").expect("plant orphan"); + let log = slog::Logger::root(slog::Discard, slog::o!()); + local.abort_incomplete_uploads(&log).await; + assert!( + !orphan.as_std_path().exists(), + "the sweep removes the crash-orphaned partial" + ); + assert_eq!( + local + .get("generations/g1/wal/e0/seg.wal.zst") + .await + .unwrap(), + Bytes::from_static(b"x"), + "real objects survive the sweep" + ); + } + + #[tokio::test] + async fn put_error_removes_partial_file() { + let tmp = tempfile::tempdir().expect("tempdir failed"); + let local = storage(&tmp); + // Pre-create the final key as a directory so the rename fails AFTER + // the partial file has already been written: this exercises the + // created-then-removed cleanup path, not just an early write failure. + std::fs::create_dir_all(local.root().join("k").as_std_path()) + .expect("plant directory at key path"); + let error = local + .put("k", Bytes::from_static(b"data")) + .await + .expect_err("put onto a directory key must fail"); + assert!( + matches!(error, StorageError::Local(_)), + "the rename failure surfaces as a local error: {error}" + ); + let partials: Vec = all_files(local.root()) + .into_iter() + .filter(|file| file.contains(PARTIAL_INFIX)) + .collect(); + assert!( + partials.is_empty(), + "a failed put must leave no partial file behind, found: {partials:?}" + ); + } + + #[tokio::test] + async fn object_methods_reject_escaping_keys() { + let tmp = tempfile::tempdir().expect("tempdir failed"); + let local = storage(&tmp); + for key in ["../escape", "/leading", "gen/../escape", "gen//doubled"] { + let error = local + .put(key, Bytes::from_static(b"x")) + .await + .expect_err("put must reject an escaping key"); + assert!( + matches!(&error, StorageError::InvalidKey { key: found, .. } if found == key), + "put({key}) must be InvalidKey, got: {error}" + ); + assert!( + matches!(local.get(key).await, Err(StorageError::InvalidKey { .. })), + "get({key}) must be InvalidKey" + ); + assert!( + matches!( + local.delete(key).await, + Err(StorageError::InvalidKey { .. }) + ), + "delete({key}) must be InvalidKey" + ); + assert!( + matches!( + local.start_multipart(key).await, + Err(StorageError::InvalidKey { .. }) + ), + "start_multipart({key}) must be InvalidKey" + ); + } + // A rejected put wrote nothing at all, not even a partial. + assert_eq!( + all_files(local.root()), + Vec::::new(), + "rejected keys must never touch the filesystem" + ); + } +} diff --git a/plus/bencher_replica/src/meta.rs b/plus/bencher_replica/src/meta.rs new file mode 100644 index 000000000..eb9ad0f57 --- /dev/null +++ b/plus/bencher_replica/src/meta.rs @@ -0,0 +1,309 @@ +//! Local advisory replication state: `.replica.json`. +//! +//! Written atomically (temp file + fsync + rename, matching the +//! `rate_limiting.json` convention) after generation creation, after each +//! segment ship, after each completed checkpoint, and at shutdown. +//! +//! The meta file is ADVISORY ONLY (invariant I6): it is never consulted by +//! restore and never trusted over the replica. It exists for exactly two +//! cases that a replica LIST plus the local WAL cannot resolve: +//! +//! 1. Crash after our checkpoint but before any new-epoch ship: local WAL +//! salts no longer match the replica's last epoch. Meta proving "epoch +//! fully shipped through checkpoint" allows resuming as epoch+1 instead +//! of re-snapshotting the whole database. +//! 2. Restored-old-volume detection: a stale meta disagrees with the replica +//! LIST, forcing a new generation instead of appending old-state frames +//! onto a newer lineage. +//! +//! Any mismatch between meta, local WAL, and replica resolves to "new +//! generation". + +use std::fs::{self, File}; +use std::io::{ErrorKind, Write as _}; + +use camino::{Utf8Path, Utf8PathBuf}; +use serde::{Deserialize, Serialize}; + +use crate::position::GenerationId; + +/// Current meta file schema version. +pub const META_VERSION: u32 = 1; + +/// Suffix appended to the full database file name to form the meta path. +const META_SUFFIX: &str = ".replica.json"; +/// Suffix appended to the meta path for the atomic-write temp sibling. +const PARTIAL_SUFFIX: &str = ".partial"; + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct ReplicaMeta { + pub version: u32, + pub generation: String, + pub epoch: u64, + pub salt1: u32, + pub salt2: u32, + /// Raw WAL byte offset shipped through (commit-aligned). + pub shipped_offset: u64, + /// True when every frame of `epoch` was shipped AND a checkpoint fully + /// backfilled it, so a subsequent WAL restart (new salts) is legitimate. + pub epoch_shipped_through_checkpoint: bool, + /// True when written while running in shadow mode (alongside + /// Litestream). Cutover to sole mode forces a new generation. + pub shadow: bool, +} + +#[derive(Debug, thiserror::Error)] +pub enum MetaError { + #[error("Failed to read replica meta ({path}): {error}")] + Read { + path: Utf8PathBuf, + error: std::io::Error, + }, + #[error("Failed to write replica meta ({path}): {error}")] + Write { + path: Utf8PathBuf, + error: std::io::Error, + }, + #[error("Failed to parse replica meta ({path}): {error}")] + Parse { + path: Utf8PathBuf, + error: serde_json::Error, + }, + #[error("Failed to serialize replica meta: {0}")] + Serialize(serde_json::Error), + #[error("Failed to remove replica meta ({path}): {error}")] + Remove { + path: Utf8PathBuf, + error: std::io::Error, + }, +} + +impl ReplicaMeta { + /// `.replica.json` next to the database file: the suffix is appended + /// to the full file name (`bencher.db` becomes `bencher.db.replica.json`), + /// never substituted for an existing extension. + #[must_use] + pub fn path_for_db(db_path: &Utf8Path) -> Utf8PathBuf { + Utf8PathBuf::from(format!("{db_path}{META_SUFFIX}")) + } + + /// Load the meta file. `Ok(None)` when missing or unparsable (the file is + /// advisory, so corrupt data is treated as absent per invariant I6); + /// `Err` only on I/O failures other than not-found. + pub fn load(db_path: &Utf8Path) -> Result, MetaError> { + let path = Self::path_for_db(db_path); + let json = match fs::read_to_string(&path) { + Ok(json) => json, + Err(error) if error.kind() == ErrorKind::NotFound => return Ok(None), + Err(error) => return Err(MetaError::Read { path, error }), + }; + // Unparsable or version-mismatched meta is treated as absent: the + // resume then takes the conservative path (a fresh snapshot) rather + // than trusting advisory fields whose semantics may have changed. A + // structurally-compatible future version must NOT be silently + // accepted; the version gate is what makes the format evolvable. + Ok(serde_json::from_str::(&json) + .ok() + .filter(|meta| meta.version == META_VERSION)) + } + + /// Atomically persist (temp sibling + fsync + rename), so a crash mid-way + /// leaves either the previous meta or the new one, never a torn file. + pub fn store(&self, db_path: &Utf8Path) -> Result<(), MetaError> { + let path = Self::path_for_db(db_path); + let partial_path = Utf8PathBuf::from(format!("{path}{PARTIAL_SUFFIX}")); + let json = serde_json::to_string(self).map_err(MetaError::Serialize)?; + // Write to a temp sibling and fsync it before the atomic rename, so + // the meta survives an abrupt VM stop rather than lingering in the + // page cache. + let write_err = |error| MetaError::Write { + path: partial_path.clone(), + error, + }; + let mut file = File::create(&partial_path).map_err(write_err)?; + file.write_all(json.as_bytes()).map_err(write_err)?; + file.sync_all().map_err(write_err)?; + drop(file); + fs::rename(&partial_path, &path).map_err(|error| MetaError::Write { + path: path.clone(), + error, + })?; + // Best-effort fsync of the parent directory so the rename itself is + // durable. The data file is already fsynced and the rename atomic, so + // a directory-sync failure must not fail the store. + if let Some(parent) = path.parent() + && let Ok(dir) = File::open(parent) + { + drop(dir.sync_all()); + } + Ok(()) + } + + /// Remove the meta file if present (used when the DB file is missing at + /// restore time). Idempotent. + pub fn remove(db_path: &Utf8Path) -> Result<(), MetaError> { + let path = Self::path_for_db(db_path); + match fs::remove_file(&path) { + Ok(()) => Ok(()), + Err(error) if error.kind() == ErrorKind::NotFound => Ok(()), + Err(error) => Err(MetaError::Remove { path, error }), + } + } + + #[must_use] + pub fn generation_id(&self) -> Option { + GenerationId::parse(&self.generation) + } +} + +#[cfg(test)] +mod tests { + use camino::{Utf8Path, Utf8PathBuf}; + use pretty_assertions::assert_eq; + + use super::{META_VERSION, MetaError, ReplicaMeta}; + use crate::position::GenerationId; + + fn test_db_path(tmp: &tempfile::TempDir) -> Utf8PathBuf { + Utf8Path::from_path(tmp.path()).unwrap().join("bencher.db") + } + + fn test_meta() -> ReplicaMeta { + ReplicaMeta { + version: META_VERSION, + generation: "20260710T145900Z-3f8a2c1d".to_owned(), + epoch: 3, + salt1: 0x9d2f_1c4a, + salt2: 0x8b3e_6f70, + shipped_offset: 524_320, + epoch_shipped_through_checkpoint: true, + shadow: false, + } + } + + #[test] + fn version_mismatch_loads_as_absent() { + let tmp = tempfile::tempdir().unwrap(); + let db_path = test_db_path(&tmp); + let mut meta = test_meta(); + meta.version = META_VERSION + 1; + meta.store(&db_path).unwrap(); + assert_eq!( + ReplicaMeta::load(&db_path).unwrap(), + None, + "a structurally-compatible future version must load as absent, \ + forcing the conservative resume path" + ); + } + + #[test] + fn serde_round_trips() { + let meta = test_meta(); + let json = serde_json::to_string(&meta).unwrap(); + let parsed: ReplicaMeta = serde_json::from_str(&json).unwrap(); + assert_eq!(parsed, meta); + } + + #[test] + fn path_for_db_appends_to_full_file_name() { + assert_eq!( + ReplicaMeta::path_for_db(Utf8Path::new("/var/lib/bencher/bencher.db")), + Utf8PathBuf::from("/var/lib/bencher/bencher.db.replica.json") + ); + } + + #[test] + fn store_then_load_round_trips() { + let tmp = tempfile::tempdir().unwrap(); + let db_path = test_db_path(&tmp); + let meta = test_meta(); + meta.store(&db_path).unwrap(); + assert_eq!(ReplicaMeta::load(&db_path).unwrap(), Some(meta)); + } + + #[test] + fn store_overwrites_and_leaves_no_temp_file() { + let tmp = tempfile::tempdir().unwrap(); + let db_path = test_db_path(&tmp); + let first = test_meta(); + first.store(&db_path).unwrap(); + let second = ReplicaMeta { + epoch: 4, + shipped_offset: 0, + epoch_shipped_through_checkpoint: false, + ..test_meta() + }; + second.store(&db_path).unwrap(); + assert_eq!(ReplicaMeta::load(&db_path).unwrap(), Some(second)); + // The atomic write leaves exactly the meta file behind: no temp + // sibling droppings. + let entries: Vec = std::fs::read_dir(tmp.path()) + .unwrap() + .map(|entry| entry.unwrap().file_name().to_string_lossy().into_owned()) + .collect(); + assert_eq!(entries, vec!["bencher.db.replica.json".to_owned()]); + } + + #[test] + fn load_missing_file_is_none() { + let tmp = tempfile::tempdir().unwrap(); + let db_path = test_db_path(&tmp); + assert_eq!(ReplicaMeta::load(&db_path).unwrap(), None); + } + + #[test] + fn load_corrupt_json_is_none() { + let tmp = tempfile::tempdir().unwrap(); + let db_path = test_db_path(&tmp); + for corrupt in ["", "not json", "{\"version\": \"wat\"}", "[1, 2, 3]"] { + std::fs::write(ReplicaMeta::path_for_db(&db_path), corrupt).unwrap(); + assert_eq!(ReplicaMeta::load(&db_path).unwrap(), None, "{corrupt:?}"); + } + } + + #[test] + fn load_io_error_is_error() { + let tmp = tempfile::tempdir().unwrap(); + let db_path = test_db_path(&tmp); + // A directory in place of the meta file: reading it is a real I/O + // error, not a not-found, and must surface as Err. + std::fs::create_dir(ReplicaMeta::path_for_db(&db_path)).unwrap(); + let err = ReplicaMeta::load(&db_path).unwrap_err(); + assert!(matches!(err, MetaError::Read { .. }), "{err}"); + } + + #[test] + fn remove_is_idempotent() { + let tmp = tempfile::tempdir().unwrap(); + let db_path = test_db_path(&tmp); + // Removing a missing file is fine. + ReplicaMeta::remove(&db_path).unwrap(); + test_meta().store(&db_path).unwrap(); + ReplicaMeta::remove(&db_path).unwrap(); + assert_eq!(ReplicaMeta::load(&db_path).unwrap(), None); + ReplicaMeta::remove(&db_path).unwrap(); + } + + #[test] + fn version_field_round_trips_meta_version() { + let tmp = tempfile::tempdir().unwrap(); + let db_path = test_db_path(&tmp); + test_meta().store(&db_path).unwrap(); + let loaded = ReplicaMeta::load(&db_path).unwrap().unwrap(); + assert_eq!(loaded.version, META_VERSION); + } + + #[test] + fn generation_id_parses_stored_generation() { + let meta = test_meta(); + assert_eq!( + meta.generation_id(), + GenerationId::parse("20260710T145900Z-3f8a2c1d") + ); + let bogus = ReplicaMeta { + generation: "not-a-generation".to_owned(), + ..test_meta() + }; + assert_eq!(bogus.generation_id(), None); + } +} diff --git a/plus/bencher_replica/src/position.rs b/plus/bencher_replica/src/position.rs new file mode 100644 index 000000000..c1982153d --- /dev/null +++ b/plus/bencher_replica/src/position.rs @@ -0,0 +1,674 @@ +//! Replica addressing: generations, epochs, and object key naming. +//! +//! Replica layout (all keys relative to the configured root/prefix): +//! +//! ```text +//! generations/ +//! 20260710T145900Z-3f8a2c1d/ # GenerationId; lexicographic order == age +//! snapshot.db.zst +//! snapshot.json # atomic "generation is valid" marker +//! wal/ +//! 0000000000-9d2f1c4a8b3e6f70/ # - +//! 00000000000000000000-00000000000524320.wal.zst +//! ``` +//! +//! Key-naming rules pinned here (heavily unit-tested, parsed by LIST-driven +//! resume and restore): +//! +//! - Generation IDs are `%Y%m%dT%H%M%SZ-<8 lowercase hex>`: lexicographic +//! order equals creation order (single process; random suffix breaks +//! same-second ties). +//! - Epoch directories embed the WAL salts so resume can compare the local +//! WAL header against the replica from a LIST alone. +//! - Segment names are `-.wal.zst` with zero-padded +//! decimal byte offsets: lexicographic order equals numeric order, and the +//! shipped position is derivable from the last key alone. +//! - The first segment of every epoch starts at offset 0 and therefore +//! contains the 32-byte WAL header: restore rebuilds a `-wal` file by +//! decompress-and-concatenate. + +use bencher_json::DateTime; + +/// Prefix under which all generations live. +pub const GENERATIONS_PREFIX: &str = "generations"; +/// Snapshot object file name within a generation. +pub const SNAPSHOT_FILE: &str = "snapshot.db.zst"; +/// Snapshot metadata (generation commit marker) file name. +pub const SNAPSHOT_META_FILE: &str = "snapshot.json"; +/// Directory under a generation holding the WAL epoch directories. +pub const WAL_DIR: &str = "wal"; + +/// A replica generation identifier: `%Y%m%dT%H%M%SZ-<8 hex>`. +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct GenerationId(String); + +impl GenerationId { + /// Build a generation ID from a timestamp and an explicit 32-bit suffix + /// (rendered as 8 lowercase hex digits). Deterministic for tests. + #[must_use] + pub fn new(created: DateTime, suffix: u32) -> Self { + let timestamp = created.into_inner().format(TIMESTAMP_FORMAT); + Self(format!("{timestamp}-{suffix:08x}")) + } + + /// Build a generation ID with a random suffix (the leading 32 random bits + /// of a v4 UUID). + #[must_use] + pub fn generate(created: DateTime) -> Self { + let (suffix, ..) = uuid::Uuid::new_v4().as_fields(); + Self::new(created, suffix) + } + + /// Parse a generation ID from a single path component; `None` when the + /// component does not match the expected shape. + #[must_use] + pub fn parse(component: &str) -> Option { + let (timestamp, suffix) = component.split_once('-')?; + validate_utc_second(timestamp)?; + parse_hex_u32(suffix)?; + Some(Self(component.to_owned())) + } + + #[must_use] + pub fn as_str(&self) -> &str { + &self.0 + } + + /// The lexicographically NEXT id with the same timestamp: suffix plus + /// one. `None` when the suffix is saturated (2^32 - 1). Used as the + /// monotonic fallback when the clock has rewound behind an observed + /// generation: new ids must always sort after the replica tip or restore + /// would silently pick the older lineage. + #[must_use] + pub fn successor(&self) -> Option { + let (timestamp, suffix) = self.0.split_once('-')?; + let next = parse_hex_u32(suffix)?.checked_add(1)?; + Some(Self(format!("{timestamp}-{next:08x}"))) + } +} + +/// In-memory replication position: the next unshipped byte of the WAL. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Position { + pub generation: GenerationId, + /// WAL salt cycle, numbered from 0 within the generation. + pub epoch: u64, + /// Header salts of this epoch. + pub salt: (u32, u32), + /// Next unshipped byte offset (commit-aligned; 0 means the WAL header + /// itself has not shipped yet). + pub offset: u64, + /// Running WAL checksum at `offset` (chain continuation seed). + pub checksum: (u32, u32), +} + +/// Parsed segment key components within a generation. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct SegmentKey { + pub epoch: u64, + pub salt: (u32, u32), + /// Raw WAL byte range `[start, end)` covered by the segment. + pub start: u64, + pub end: u64, +} + +/// `generations//` (trailing slash included). +#[must_use] +pub fn generation_prefix(generation: &GenerationId) -> String { + format!("{GENERATIONS_PREFIX}/{}/", generation.as_str()) +} + +/// `generations//snapshot.db.zst` +#[must_use] +pub fn snapshot_key(generation: &GenerationId) -> String { + format!("{}{SNAPSHOT_FILE}", generation_prefix(generation)) +} + +/// `generations//snapshot.json` +#[must_use] +pub fn snapshot_meta_key(generation: &GenerationId) -> String { + format!("{}{SNAPSHOT_META_FILE}", generation_prefix(generation)) +} + +/// `generations//wal/-/` (trailing slash +/// included). +/// +/// Epochs at or above `10^EPOCH_WIDTH` would render wider than the fixed +/// width, breaking lexicographic-equals-numeric ordering, and +/// `parse_epoch_dir` would then reject the directory on restore. Unreachable +/// in practice (the epoch resets per generation and increments once per WAL +/// restart), so it is a debug assertion rather than a runtime error. +#[must_use] +pub fn epoch_dir(generation: &GenerationId, epoch: u64, salt: (u32, u32)) -> String { + debug_assert!( + epoch < 10u64.pow(u32::try_from(EPOCH_WIDTH).unwrap_or(u32::MAX)), + "epoch {epoch} exceeds the fixed {EPOCH_WIDTH}-digit directory width" + ); + let (salt1, salt2) = salt; + format!( + "{}{WAL_DIR}/{epoch:010}-{salt1:08x}{salt2:08x}/", + generation_prefix(generation) + ) +} + +/// Full segment key: +/// `generations//wal//-.wal.zst` +#[must_use] +pub fn segment_key(generation: &GenerationId, segment: &SegmentKey) -> String { + let SegmentKey { + epoch, + salt, + start, + end, + } = *segment; + format!( + "{}{start:020}-{end:020}{SEGMENT_SUFFIX}", + epoch_dir(generation, epoch, salt) + ) +} + +/// Parse an epoch directory component (`-`); +/// `None` when malformed. +#[must_use] +pub fn parse_epoch_dir(component: &str) -> Option<(u64, (u32, u32))> { + let (epoch, hex) = component.split_once('-')?; + if epoch.len() != EPOCH_WIDTH || !epoch.bytes().all(|byte| byte.is_ascii_digit()) { + return None; + } + let (salt1, salt2) = hex.split_at_checked(SALT_HEX_WIDTH)?; + Some(( + epoch.parse().ok()?, + (parse_hex_u32(salt1)?, parse_hex_u32(salt2)?), + )) +} + +/// Parse a full segment key back into its generation and components; `None` +/// when any component is malformed (including empty or inverted ranges: a +/// real segment always covers at least one byte). +#[must_use] +pub fn parse_segment_key(key: &str) -> Option<(GenerationId, SegmentKey)> { + let mut parts = key.split('/'); + let prefix = parts.next()?; + let generation = parts.next()?; + let wal = parts.next()?; + let epoch_component = parts.next()?; + let file = parts.next()?; + if parts.next().is_some() || prefix != GENERATIONS_PREFIX || wal != WAL_DIR { + return None; + } + let generation = GenerationId::parse(generation)?; + let (epoch, salt) = parse_epoch_dir(epoch_component)?; + let (start, end) = parse_segment_file(file)?; + Some(( + generation, + SegmentKey { + epoch, + salt, + start, + end, + }, + )) +} + +/// `chrono` strftime format behind [`GenerationId`] timestamps. +const TIMESTAMP_FORMAT: &str = "%Y%m%dT%H%M%SZ"; +/// Rendered length of [`TIMESTAMP_FORMAT`]. +const TIMESTAMP_LEN: usize = 16; +/// Zero-padded decimal width of an epoch number in an epoch directory. +const EPOCH_WIDTH: usize = 10; +/// Lowercase hex width of a single WAL salt in an epoch directory. +const SALT_HEX_WIDTH: usize = 8; +/// Zero-padded decimal width of a segment byte offset (full `u64` range). +const OFFSET_WIDTH: usize = 20; +/// Segment object file name suffix. +const SEGMENT_SUFFIX: &str = ".wal.zst"; + +/// Validate a `%Y%m%dT%H%M%SZ` timestamp component; `Option<()>` purely for +/// `?` chaining. +fn validate_utc_second(timestamp: &str) -> Option<()> { + if timestamp.len() != TIMESTAMP_LEN { + return None; + } + let (year, rest) = split_digits(timestamp, 4)?; + let (month, rest) = split_digits(rest, 2)?; + let (day, rest) = split_digits(rest, 2)?; + let rest = rest.strip_prefix('T')?; + let (hour, rest) = split_digits(rest, 2)?; + let (minute, rest) = split_digits(rest, 2)?; + let (second, rest) = split_digits(rest, 2)?; + // An out-of-range month has zero days, so the day check rejects it too. + let valid = rest == "Z" + && (1..=days_in_month(year, month)).contains(&day) + && hour <= 23 + && minute <= 59 + && second <= 59; + valid.then_some(()) +} + +/// Split `len` leading ASCII digits off `component` and parse them. +fn split_digits(component: &str, len: usize) -> Option<(u32, &str)> { + let (digits, rest) = component.split_at_checked(len)?; + if !digits.bytes().all(|byte| byte.is_ascii_digit()) { + return None; + } + Some((digits.parse().ok()?, rest)) +} + +fn days_in_month(year: u32, month: u32) -> u32 { + match month { + 1 | 3 | 5 | 7 | 8 | 10 | 12 => 31, + 4 | 6 | 9 | 11 => 30, + 2 if is_leap_year(year) => 29, + 2 => 28, + _ => 0, + } +} + +fn is_leap_year(year: u32) -> bool { + year.is_multiple_of(4) && (!year.is_multiple_of(100) || year.is_multiple_of(400)) +} + +/// Parse exactly [`SALT_HEX_WIDTH`] lowercase hex digits (uppercase rejected +/// so every key has one canonical spelling). +fn parse_hex_u32(hex: &str) -> Option { + if hex.len() != SALT_HEX_WIDTH || !hex.bytes().all(is_lower_hex_digit) { + return None; + } + u32::from_str_radix(hex, 16).ok() +} + +fn is_lower_hex_digit(byte: u8) -> bool { + byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte) +} + +/// Parse a segment file name (`-.wal.zst`) into its +/// `[start, end)` range; requires `start < end`. +fn parse_segment_file(file: &str) -> Option<(u64, u64)> { + let range = file.strip_suffix(SEGMENT_SUFFIX)?; + let (start, rest) = range.split_at_checked(OFFSET_WIDTH)?; + let end = rest.strip_prefix('-')?; + let start = parse_padded_u64(start)?; + let end = parse_padded_u64(end)?; + (start < end).then_some((start, end)) +} + +/// Parse exactly [`OFFSET_WIDTH`] zero-padded decimal digits; values beyond +/// `u64::MAX` fail the inner parse. +fn parse_padded_u64(digits: &str) -> Option { + if digits.len() != OFFSET_WIDTH || !digits.bytes().all(|byte| byte.is_ascii_digit()) { + return None; + } + digits.parse().ok() +} + +#[cfg(test)] +mod tests { + use bencher_json::DateTime; + use pretty_assertions::assert_eq; + + use super::{ + GenerationId, SegmentKey, epoch_dir, generation_prefix, parse_epoch_dir, parse_segment_key, + segment_key, snapshot_key, snapshot_meta_key, + }; + + /// 2026-07-10T14:59:00Z + const CREATED_SECS: i64 = 1_783_695_540; + const SUFFIX: u32 = 0x3f8a_2c1d; + const SALT: (u32, u32) = (0x9d2f_1c4a, 0x8b3e_6f70); + + fn created() -> DateTime { + DateTime::try_from(CREATED_SECS).unwrap() + } + + fn generation() -> GenerationId { + GenerationId::new(created(), SUFFIX) + } + + #[test] + fn generation_id_new_renders_timestamp_and_suffix() { + assert_eq!(generation().as_str(), "20260710T145900Z-3f8a2c1d"); + // The suffix is zero-padded, lowercase hex. + assert_eq!( + GenerationId::new(created(), 0xa).as_str(), + "20260710T145900Z-0000000a" + ); + } + + #[test] + fn generation_id_parse_round_trips() { + let generation = generation(); + assert_eq!( + GenerationId::parse(generation.as_str()), + Some(generation.clone()) + ); + assert_eq!(GenerationId::parse(generation.as_str()), Some(generation)); + } + + #[test] + fn generation_id_generate_keeps_timestamp_and_parses() { + let generated = GenerationId::generate(created()); + assert!( + generated.as_str().starts_with("20260710T145900Z-"), + "{generated:?}" + ); + assert_eq!(GenerationId::parse(generated.as_str()), Some(generated)); + } + + #[test] + fn generation_id_parse_rejects_malformed() { + let malformed = [ + "", + // Missing separator or suffix + "20260710T145900Z", + "20260710T145900Z-", + // Wrong lengths + "20260710T145900Z-3f8a2c1", + "20260710T145900Z-3f8a2c1d1", + "2026710T145900Z-3f8a2c1d", + "202607100T145900Z-3f8a2c1d", + // Bad timestamps + "20261310T145900Z-3f8a2c1d", + "20260732T145900Z-3f8a2c1d", + "20260700T145900Z-3f8a2c1d", + "20260710T245900Z-3f8a2c1d", + "20260710T146000Z-3f8a2c1d", + "20260710T145960Z-3f8a2c1d", + "20260710X145900Z-3f8a2c1d", + "20260710T145900X-3f8a2c1d", + "20260710t145900Z-3f8a2c1d", + "20260710T145900z-3f8a2c1d", + "2026a710T145900Z-3f8a2c1d", + // Non-hex or uppercase-hex suffixes + "20260710T145900Z-3f8a2c1g", + "20260710T145900Z-3F8A2C1D", + "20260710T145900Z-3f8a-c1d", + "20260710T145900Z-+f8a2c1d", + // Trailing garbage + "20260710T145900Z-3f8a2c1d/", + " 20260710T145900Z-3f8a2c1d", + ]; + for component in malformed { + assert_eq!(GenerationId::parse(component), None, "{component:?}"); + } + } + + #[test] + fn generation_id_parse_handles_leap_years() { + // 2024 is a leap year; 2023 is not; 1900 is not (century rule); + // 2000 is (400-year rule). + assert!(GenerationId::parse("20240229T000000Z-00000000").is_some()); + assert_eq!(GenerationId::parse("20230229T000000Z-00000000"), None); + assert_eq!(GenerationId::parse("19000229T000000Z-00000000"), None); + assert!(GenerationId::parse("20000229T000000Z-00000000").is_some()); + } + + #[test] + fn generation_id_orders_by_timestamp_then_suffix() { + let earlier = DateTime::try_from(CREATED_SECS).unwrap(); + let later = DateTime::try_from(CREATED_SECS + 1).unwrap(); + // Timestamp dominates even when the earlier suffix is numerically larger. + assert!(GenerationId::new(earlier, u32::MAX) < GenerationId::new(later, 0)); + // Same second: suffix order is numeric because the hex is fixed-width. + assert!(GenerationId::new(earlier, 0x1) < GenerationId::new(earlier, 0xa)); + assert!(GenerationId::new(earlier, 0xa) < GenerationId::new(earlier, 0x10)); + assert!(GenerationId::new(earlier, 0x10) < GenerationId::new(earlier, u32::MAX)); + } + + #[test] + fn generation_id_successor_increments_suffix_and_saturates() { + // A normal suffix increments by one, carrying across the hex-width + // boundary, and the result sorts strictly after the original. + let base = GenerationId::new(created(), 0x0000_000f); + let next = base + .successor() + .expect("successor of a non-saturated suffix"); + assert_eq!(next.as_str(), "20260710T145900Z-00000010"); + assert!(next > base, "the successor sorts after the original"); + // The maximum suffix has no successor (would overflow the 8-hex field). + assert_eq!(GenerationId::new(created(), u32::MAX).successor(), None); + // One below the maximum still has a successor: the field saturates. + let penultimate = GenerationId::new(created(), u32::MAX - 1); + assert_eq!( + penultimate.successor().map(|id| id.as_str().to_owned()), + Some("20260710T145900Z-ffffffff".to_owned()) + ); + } + + #[test] + fn generation_and_snapshot_keys_render_exactly() { + let generation = generation(); + assert_eq!( + generation_prefix(&generation), + "generations/20260710T145900Z-3f8a2c1d/" + ); + assert_eq!( + snapshot_key(&generation), + "generations/20260710T145900Z-3f8a2c1d/snapshot.db.zst" + ); + assert_eq!( + snapshot_meta_key(&generation), + "generations/20260710T145900Z-3f8a2c1d/snapshot.json" + ); + } + + #[test] + fn epoch_dir_renders_zero_padded_salts() { + assert_eq!( + epoch_dir(&generation(), 0, SALT), + "generations/20260710T145900Z-3f8a2c1d/wal/0000000000-9d2f1c4a8b3e6f70/" + ); + assert_eq!( + epoch_dir(&generation(), 42, (0xa, 0xb)), + "generations/20260710T145900Z-3f8a2c1d/wal/0000000042-0000000a0000000b/" + ); + } + + #[test] + fn segment_key_renders_zero_padded_offsets() { + let segment = SegmentKey { + epoch: 0, + salt: SALT, + start: 0, + end: 524_320, + }; + assert_eq!( + segment_key(&generation(), &segment), + "generations/20260710T145900Z-3f8a2c1d/wal/0000000000-9d2f1c4a8b3e6f70/00000000000000000000-00000000000000524320.wal.zst" + ); + } + + #[test] + fn segment_key_is_idempotent_and_collision_free() { + let generation = generation(); + let segment = SegmentKey { + epoch: 1, + salt: SALT, + start: 32, + end: 4128, + }; + // Same range yields the same key (idempotent re-upload target). + assert_eq!( + segment_key(&generation, &segment), + segment_key(&generation, &segment) + ); + // Any differing component yields a different key. + let variants = [ + SegmentKey { + start: 4128, + end: 8224, + ..segment + }, + SegmentKey { + end: 8224, + ..segment + }, + SegmentKey { + epoch: 2, + ..segment + }, + SegmentKey { + salt: (0x1, 0x2), + ..segment + }, + ]; + for variant in variants { + assert_ne!( + segment_key(&generation, &segment), + segment_key(&generation, &variant), + "{variant:?}" + ); + } + } + + #[test] + fn lexicographic_order_equals_numeric_order() { + let generation = generation(); + // Epoch directories across decimal-width boundaries. + let epoch_pairs = [ + (9u64, 10u64), + (99, 100), + (999_999_999, 1_000_000_000), + (9_999_999_998, 9_999_999_999), + ]; + for (lo, hi) in epoch_pairs { + let lo_dir = epoch_dir(&generation, lo, SALT); + let hi_dir = epoch_dir(&generation, hi, SALT); + assert!(lo_dir < hi_dir, "{lo_dir} vs {hi_dir}"); + } + // Segment offsets across decimal-width boundaries up to u64::MAX. + let offset_pairs = [(9u64, 10u64), (99, 100), (u64::MAX - 2, u64::MAX - 1)]; + for (lo, hi) in offset_pairs { + let lo_key = segment_key( + &generation, + &SegmentKey { + epoch: 0, + salt: SALT, + start: lo, + end: lo + 1, + }, + ); + let hi_key = segment_key( + &generation, + &SegmentKey { + epoch: 0, + salt: SALT, + start: hi, + end: hi + 1, + }, + ); + assert!(lo_key < hi_key, "{lo_key} vs {hi_key}"); + } + // Generation IDs across second/day/month/year rollovers. + let time_pairs = [ + ("20261231T235959Z", "20270101T000000Z"), + ("20270101T000000Z", "20270101T000001Z"), + ("20270101T000001Z", "20270102T000000Z"), + ("20270102T000000Z", "20270201T000000Z"), + ]; + for (lo, hi) in time_pairs { + let lo_id = GenerationId::parse(&format!("{lo}-ffffffff")).unwrap(); + let hi_id = GenerationId::parse(&format!("{hi}-00000000")).unwrap(); + assert!(lo_id < hi_id, "{lo_id:?} vs {hi_id:?}"); + } + } + + #[test] + fn parse_epoch_dir_round_trips() { + for (epoch, salt) in [ + (0u64, (0u32, 0u32)), + (42, SALT), + (9_999_999_999, (u32::MAX, u32::MAX)), + ] { + let component = format!("{epoch:010}-{:08x}{:08x}", salt.0, salt.1); + assert_eq!( + parse_epoch_dir(&component), + Some((epoch, salt)), + "{component}" + ); + } + } + + #[test] + fn parse_epoch_dir_rejects_malformed() { + let malformed = [ + "", + // Wrong widths + "000000000-9d2f1c4a8b3e6f70", + "00000000000-9d2f1c4a8b3e6f70", + "0000000000-9d2f1c4a8b3e6f7", + "0000000000-9d2f1c4a8b3e6f700", + // Non-digit epoch, non-hex or uppercase salts + "000000000a-9d2f1c4a8b3e6f70", + "0000000000-9D2F1C4A8B3E6F70", + "0000000000-9d2f1c4a8b3e6g70", + // Missing separator; trailing slash + "00000000009d2f1c4a8b3e6f70", + "0000000000-9d2f1c4a8b3e6f70/", + ]; + for component in malformed { + assert_eq!(parse_epoch_dir(component), None, "{component:?}"); + } + } + + #[test] + fn parse_segment_key_round_trips() { + let generation = generation(); + let segments = [ + SegmentKey { + epoch: 0, + salt: SALT, + start: 0, + end: 32, + }, + SegmentKey { + epoch: 7, + salt: (0xa, 0xb), + start: 32, + end: 524_320, + }, + SegmentKey { + epoch: 9_999_999_999, + salt: (u32::MAX, u32::MAX), + start: u64::MAX - 1, + end: u64::MAX, + }, + ]; + for segment in segments { + let key = segment_key(&generation, &segment); + assert_eq!( + parse_segment_key(&key), + Some((generation.clone(), segment)), + "{key}" + ); + } + } + + #[test] + fn parse_segment_key_rejects_malformed() { + let malformed = [ + "", + // Wrong prefix or missing wal dir + "generation/20260710T145900Z-3f8a2c1d/wal/0000000000-9d2f1c4a8b3e6f70/00000000000000000000-00000000000000000032.wal.zst", + "generations/20260710T145900Z-3f8a2c1d/0000000000-9d2f1c4a8b3e6f70/00000000000000000000-00000000000000000032.wal.zst", + // Malformed generation id + "generations/20260710T145900Z-3F8A2C1D/wal/0000000000-9d2f1c4a8b3e6f70/00000000000000000000-00000000000000000032.wal.zst", + // Malformed epoch dir + "generations/20260710T145900Z-3f8a2c1d/wal/000000000-9d2f1c4a8b3e6f70/00000000000000000000-00000000000000000032.wal.zst", + // Bad extension + "generations/20260710T145900Z-3f8a2c1d/wal/0000000000-9d2f1c4a8b3e6f70/00000000000000000000-00000000000000000032.wal", + // Unpadded, non-numeric, or overflowing offsets + "generations/20260710T145900Z-3f8a2c1d/wal/0000000000-9d2f1c4a8b3e6f70/0-32.wal.zst", + "generations/20260710T145900Z-3f8a2c1d/wal/0000000000-9d2f1c4a8b3e6f70/000000000000000000zz-00000000000000000032.wal.zst", + "generations/20260710T145900Z-3f8a2c1d/wal/0000000000-9d2f1c4a8b3e6f70/00000000000000000000-99999999999999999999.wal.zst", + // Empty or inverted ranges + "generations/20260710T145900Z-3f8a2c1d/wal/0000000000-9d2f1c4a8b3e6f70/00000000000000000032-00000000000000000032.wal.zst", + "generations/20260710T145900Z-3f8a2c1d/wal/0000000000-9d2f1c4a8b3e6f70/00000000000000000032-00000000000000000000.wal.zst", + // Extra path components + "generations/20260710T145900Z-3f8a2c1d/wal/0000000000-9d2f1c4a8b3e6f70/x/00000000000000000000-00000000000000000032.wal.zst", + "generations/20260710T145900Z-3f8a2c1d/wal/0000000000-9d2f1c4a8b3e6f70/00000000000000000000-00000000000000000032.wal.zst/x", + ]; + for key in malformed { + assert_eq!(parse_segment_key(key), None, "{key:?}"); + } + } +} diff --git a/plus/bencher_replica/src/replicator.rs b/plus/bencher_replica/src/replicator.rs new file mode 100644 index 000000000..42c4e5e99 --- /dev/null +++ b/plus/bencher_replica/src/replicator.rs @@ -0,0 +1,246 @@ +//! Public entry points: start the replicator task, shut it down with a +//! final ship, and expose the fatal-error future for the server's race. +//! +//! The production shell is deliberately trivial: one tokio task ticking +//! [`SyncEngine::sync_once`] every `sync_interval`. Engine construction +//! happens INSIDE the spawned task so [`Replicator::start`] is synchronous +//! and fast; a retryable (storage) construction failure is retried with +//! capped backoff (an unreachable replica at boot must never be mistaken +//! for an empty one), while any other construction failure surfaces through +//! [`ReplicatorHandle::wait_fatal`]. Per-tick storage errors are NOT fatal: +//! the engine backs off internally and the WAL is the buffer. + +use std::sync::Arc; +use std::time::Duration; + +use bencher_json::Clock; +use bencher_json::system::config::ReplicationTarget; +use camino::Utf8PathBuf; +use slog::Logger; +use tokio::sync::{mpsc, oneshot}; +use tokio::time::timeout; + +use crate::backoff::Backoff; +use crate::config::ReplicaConfig; +use crate::sync::{SyncEngine, SyncError}; + +/// The only thing the replicator needs from the application, generic over +/// the writer connection type so this crate never depends on diesel or the +/// schema: the engine only ever HOLDS the mutex guard, never uses `C`. +pub struct ReplicaDb { + pub db_path: Utf8PathBuf, + /// The app's single-writer mutex; the checkpoint critical section + /// acquires it so app writers queue on tokio instead of burning their + /// `SQLite` `busy_timeout`. + pub writer: Arc>, + /// The app's configured busy timeout, for observability parity. + pub busy_timeout_ms: u32, +} + +// Manual impl: `#[derive(Clone)]` would wrongly require `C: Clone`. +impl Clone for ReplicaDb { + fn clone(&self) -> Self { + Self { + db_path: self.db_path.clone(), + writer: Arc::clone(&self.writer), + busy_timeout_ms: self.busy_timeout_ms, + } + } +} + +/// Starts the single replication task. +pub struct Replicator; + +impl Replicator { + /// Spawn the replication task and return its handle. Construction and + /// resume run inside the task; see the module docs for the failure + /// contract. + pub fn start( + log: Logger, + config: ReplicaConfig, + db: ReplicaDb, + clock: Clock, + shadow: bool, + ) -> ReplicatorHandle { + let (commands, command_rx) = mpsc::channel(1); + let (fatal_tx, fatal_rx) = oneshot::channel(); + drop(tokio::spawn(run_replicator( + log, config, db, clock, shadow, command_rx, fatal_tx, + ))); + ReplicatorHandle { + commands, + fatal: Some(fatal_rx), + } + } +} + +/// Handle to the running replication task. +pub struct ReplicatorHandle { + commands: mpsc::Sender, + fatal: Option>, +} + +impl ReplicatorHandle { + /// Ship the remaining WAL tail (`deadline`-bounded) and stop the task. + /// After a COMPLETE drain in sole mode a final checkpoint runs (after the + /// deadline, unbounded) to seal the epoch so the next boot resumes in + /// place rather than re-snapshotting; see [`SyncEngine::final_sync`]. On + /// the deadline the un-shipped tail simply stays in the local WAL and the + /// next boot resumes by salt match (lag, never loss). + /// + /// PRECONDITION: application writers must already be quiesced (the server + /// drained) before calling this. The final checkpoint acquires the app + /// writer mutex with no deadline of its own, and this call awaits the + /// task's completion unboundedly; a writer still holding the mutex would + /// stall shutdown. The production caller orders `shutdown(server)` before + /// `replica_handle.shutdown(deadline)` for exactly this reason. + pub async fn shutdown(self, deadline: Duration) -> Result<(), SyncError> { + let Self { commands, fatal: _ } = self; + let (done, done_rx) = oneshot::channel(); + if commands + .send(Command::Shutdown { deadline, done }) + .await + .is_err() + { + return Err(SyncError::TaskExited); + } + match done_rx.await { + Ok(result) => result, + Err(_closed) => Err(SyncError::TaskExited), + } + } + + /// Resolves only if the replication task DIED: a non-retryable + /// construction failure, a fatal tick error, or a panic. In normal + /// operation (including across storage outages, which back off + /// internally) this future NEVER resolves; it exists for the server's + /// shutdown race. + pub async fn wait_fatal(&mut self) -> Result<(), SyncError> { + // Take the receiver inside the async body, on first POLL rather than + // at call time: constructing this future and dropping it unpolled + // must not disarm fatal detection for a later call. + let Some(receiver) = self.fatal.take() else { + // Already consumed by an earlier call: nothing further can + // ever be reported. + return std::future::pending().await; + }; + match receiver.await { + Ok(error) => Err(error), + // The sender dropped without a fatal report while this + // handle still exists: the task panicked. + Err(_closed) => Err(SyncError::TaskExited), + } + } +} + +enum Command { + Shutdown { + deadline: Duration, + done: oneshot::Sender>, + }, +} + +/// The task body: construct (with storage-retry), then tick until shutdown. +async fn run_replicator( + log: Logger, + config: ReplicaConfig, + db: ReplicaDb, + clock: Clock, + shadow: bool, + mut commands: mpsc::Receiver, + fatal: oneshot::Sender, +) { + // Endpoint overrides are for S3-compatible stores on trusted networks; + // still, plaintext transit deserves a loud line in the boot log. SigV4 + // never transmits the secret, so only the replica DATA is exposed. + if let ReplicationTarget::S3 { + endpoint: Some(endpoint), + .. + } = &config.target + && endpoint + .get(..7) + .is_some_and(|scheme| scheme.eq_ignore_ascii_case("http://")) + { + slog::warn!(log, "Replica S3 endpoint uses plaintext http; replica data transits unencrypted"; + "endpoint" => endpoint.as_str()); + } + // A zero interval would busy-spin the tick loop. + let tick = if config.sync_interval.is_zero() { + Duration::from_secs(1) + } else { + config.sync_interval + }; + let mut construction_backoff = Backoff::default(); + let mut engine = loop { + match SyncEngine::new( + log.clone(), + config.clone(), + db.clone(), + clock.clone(), + shadow, + ) + .await + { + Ok(engine) => break engine, + Err(error) => { + // Every failed construction attempt is an adverse event: a + // retryable failure (e.g. bad S3 credentials) otherwise retries + // forever with only WARN logs and no metric to alert on. + #[cfg(feature = "otel")] + bencher_otel::ApiMeter::increment(bencher_otel::ApiCounter::ReplicaInitFailed); + if error.is_retryable() { + let delay = construction_backoff.next_delay(); + slog::warn!(log, "Replica unreachable at startup; retrying resume"; + "error" => %error, "retry_in_secs" => delay.as_secs()); + match timeout(delay, commands.recv()).await { + Ok(Some(Command::Shutdown { done, .. })) => { + // Nothing was running yet; nothing to ship. + drop(done.send(Ok(()))); + return; + }, + Ok(None) => return, + Err(_elapsed) => {}, + } + } else { + slog::error!(log, "Replicator failed to start"; "error" => %error); + drop(fatal.send(error)); + return; + } + }, + } + }; + // Best-effort reclamation of multipart uploads orphaned by an earlier + // crash mid-snapshot: each orphan otherwise accrues storage cost until a + // bucket lifecycle rule (if any) reaps it. Runs once per boot, inside + // this task so server startup is never delayed, and never fails the + // replicator. + engine.storage().abort_incomplete_uploads(&log).await; + loop { + match timeout(tick, commands.recv()).await { + Ok(Some(Command::Shutdown { deadline, done })) => { + let result = engine.final_sync(deadline).await; + drop(done.send(result)); + return; + }, + Ok(None) => { + slog::warn!( + log, + "Replicator handle dropped without shutdown; replication stops" + ); + return; + }, + Err(_elapsed) => match engine.sync_once().await { + Ok(progress) => { + if let Some(error) = &progress.error { + slog::warn!(log, "Sync tick failed; backing off"; "error" => %error); + } + }, + Err(error) => { + slog::error!(log, "Replicator died"; "error" => %error); + drop(fatal.send(error)); + return; + }, + }, + } + } +} diff --git a/plus/bencher_replica/src/restore.rs b/plus/bencher_replica/src/restore.rs new file mode 100644 index 000000000..bc6ad3068 --- /dev/null +++ b/plus/bencher_replica/src/restore.rs @@ -0,0 +1,1373 @@ +//! Startup-only restore: rebuild the database from the latest valid +//! generation (snapshot + ordered WAL segment replay). +//! +//! Runs BEFORE any application `SQLite` connection exists, in the same slot as +//! the old `litestream restore -if-replica-exists -if-db-not-exists`. +//! +//! Failure policy, per the design's disaster-recovery stance: +//! +//! - A broken WAL lineage (missing/corrupt segment, epoch gap) is a SOFT +//! stop: replay ends at the last fully-valid epoch, the offending key is +//! logged loudly, and the server boots on the older CONSISTENT state (an +//! older commit boundary beats refusing to boot). +//! - A corrupt snapshot (checksum mismatch) or a restored file that fails +//! `quick_check` is a HARD error: never boot on a corrupt database, and +//! never leave a partial file behind. + +use std::future::Future; +use std::io::ErrorKind; +use std::pin::Pin; +use std::task::{Context, Poll}; +use std::time::Instant; + +use async_compression::tokio::bufread::ZstdDecoder; +use camino::{Utf8Path, Utf8PathBuf}; +use sha2::Digest as _; +use slog::Logger; +use tokio::io::{AsyncRead, AsyncReadExt as _, AsyncWriteExt as _, BufReader, ReadBuf}; +use tokio::task::spawn_blocking; + +use crate::backoff::Backoff; +use crate::config::ReplicaConfig; +use crate::meta::{META_VERSION, MetaError, ReplicaMeta}; +use crate::position::{ + GENERATIONS_PREFIX, GenerationId, Position, SegmentKey, WAL_DIR, generation_prefix, + parse_segment_key, segment_key, snapshot_key, snapshot_meta_key, +}; +use crate::segment::decompress_segment_with_cap; +use crate::snapshot_meta::SnapshotMeta; +use crate::storage::{ReplicaStorage, StorageError}; +use crate::wal::{WalError, WalScanner}; + +/// Outcome of the startup restore handshake. +#[derive(Debug)] +pub enum RestoreOutcome { + /// The database file already exists: nothing to do + /// (`-if-db-not-exists` semantics). + Skipped, + /// No valid generation exists on the replica: fresh server + /// (`-if-replica-exists` semantics). + NoReplica, + /// The database was rebuilt from the replica. + Restored { + generation: GenerationId, + /// Restored database size in bytes. + db_bytes: u64, + /// Number of WAL segments replayed on top of the snapshot. + segments: u64, + }, +} + +#[derive(Debug, thiserror::Error)] +pub enum RestoreError { + #[error("Replica storage: {0}")] + Storage(#[from] StorageError), + #[error("Replica meta: {0}")] + Meta(#[from] MetaError), + #[error("Segment: {0}")] + Segment(#[from] crate::segment::SegmentError), + #[error("Failed to write restore file ({path}): {error}")] + Io { + path: Utf8PathBuf, + error: std::io::Error, + }, + #[error("Failed to download snapshot for generation {}: {}", .generation.as_str(), .error)] + SnapshotDownload { + generation: GenerationId, + error: std::io::Error, + }, + #[error( + "Snapshot checksum mismatch for generation {}: expected {expected}, computed {computed}", + .generation.as_str() + )] + SnapshotChecksum { + generation: GenerationId, + expected: String, + computed: String, + }, + #[error( + "Snapshot body for generation {} exceeds its recorded size of {db_bytes} bytes", + .generation.as_str() + )] + SnapshotTooLarge { + generation: GenerationId, + db_bytes: u64, + }, + #[error("Failed to apply WAL epoch {epoch}: {error}")] + Apply { epoch: u64, error: rusqlite::Error }, + #[error( + "WAL epoch {epoch} application incomplete: expected {expected_frames} frames, checkpoint reported busy {busy}, log {in_log}, checkpointed {checkpointed}" + )] + ApplyIncomplete { + epoch: u64, + expected_frames: u64, + busy: i64, + in_log: i64, + checkpointed: i64, + }, + #[error("Failed to run quick_check on the restored database: {0}")] + QuickCheckRun(rusqlite::Error), + #[error("Restored database failed quick_check: {0}")] + QuickCheck(String), + #[error("Restore task panicked: {0}")] + Join(tokio::task::JoinError), +} + +/// Bounded attempts for each transient storage read in the restore path +/// (initial try plus retries). Restore runs once at startup, so a small +/// bound rides out a blip without stalling boot indefinitely. +const RESTORE_ATTEMPTS: u32 = 3; + +/// Adverse-event counter for every restore decision that drops replica data +/// (falling back to an older generation, truncating a broken lineage, or +/// discarding a short boundary epoch): the restore still "succeeds", so +/// without this a silent data regression is indistinguishable from a clean +/// restore in metrics. +fn count_soft_stop() { + #[cfg(feature = "otel")] + bencher_otel::ApiMeter::increment(bencher_otel::ApiCounter::ReplicaRestoreSoftStop); +} + +/// Restore the latest replica state into `db_path` if (and only if) the +/// database file does not exist. Removes any stale `.replica.json` when +/// the database is missing, and writes a fresh one on success so the +/// replicator resumes cleanly. +pub async fn restore_if_missing( + log: &Logger, + config: &ReplicaConfig, + db_path: &Utf8Path, +) -> Result { + let exists = db_path + .as_std_path() + .try_exists() + .map_err(|error| RestoreError::Io { + path: db_path.to_owned(), + error, + })?; + if exists { + slog::info!(log, "Database file exists; skipping restore"; + "db_path" => db_path.as_str()); + return Ok(RestoreOutcome::Skipped); + } + // The database is missing, so any meta file is stale advisory data from + // a previous life of this volume: remove it before anything else. + ReplicaMeta::remove(db_path)?; + let storage = config.build_storage(); + let start = Instant::now(); + let Some(restored) = restore_db(log, &storage, db_path, None).await? else { + slog::info!(log, "No valid replica generation found; starting fresh"; + "db_path" => db_path.as_str()); + return Ok(RestoreOutcome::NoReplica); + }; + // Fresh advisory meta: after a restore the local WAL is empty, so the + // first server write creates a fresh WAL with new salts. Marking the + // restored epoch as fully shipped through a checkpoint lets the + // replicator's resume take the meta-verified epoch+1 path instead of + // forcing a whole-database re-snapshot. + let meta = ReplicaMeta { + version: META_VERSION, + generation: restored.generation.as_str().to_owned(), + epoch: restored.epoch, + salt1: restored.salt.0, + salt2: restored.salt.1, + shipped_offset: restored.shipped_offset, + epoch_shipped_through_checkpoint: true, + shadow: false, + }; + // Best-effort: the database is already restored and valid, and the meta + // is advisory (without it the replicator's resume takes the conservative + // re-snapshot path). Failing boot here over an advisory write would turn + // a healthy restore into an outage. + if let Err(error) = meta.store(db_path) { + slog::warn!(log, "Failed to write advisory replica meta after restore; \ + the next resume will re-snapshot"; + "db_path" => db_path.as_str(), "error" => %error); + } + #[cfg(feature = "otel")] + bencher_otel::ApiMeter::increment(bencher_otel::ApiCounter::ReplicaRestore); + slog::info!(log, "Restored database from replica"; + "db_path" => db_path.as_str(), + "generation" => restored.generation.as_str(), + "db_bytes" => restored.db_bytes, + "segments" => restored.segments, + "duration_ms" => u64::try_from(start.elapsed().as_millis()).unwrap_or(u64::MAX)); + Ok(RestoreOutcome::Restored { + generation: restored.generation, + db_bytes: restored.db_bytes, + segments: restored.segments, + }) +} + +/// Restore into an explicit `scratch_db` path, unconditionally (no +/// database-exists check, no meta handling). Used by verification. +/// +/// With `up_to = Some(position)`, only the position's generation is +/// considered and replay covers epochs `< position.epoch` plus, within +/// `position.epoch`, segments with `end <= position.offset` (positions are +/// commit-aligned, so segment boundaries line up exactly). +/// +/// `Ok(None)` means no valid generation was found. +pub(crate) async fn restore_to( + log: &Logger, + storage: &ReplicaStorage, + scratch_db: &Utf8Path, + up_to: Option<&Position>, +) -> Result, RestoreError> { + // Verification restores log through the caller's logger so that the root + // cause of a verification failure (a vanished segment, a soft-stopped + // epoch) is observable; divergence itself is reported through the verify + // result. + restore_db(log, storage, scratch_db, up_to).await +} + +/// A completed restore, before any advisory meta is written. +pub(crate) struct RestoredDb { + pub generation: GenerationId, + /// Restored database file size in bytes. + pub db_bytes: u64, + /// WAL segments replayed on top of the snapshot. + pub segments: u64, + /// Last epoch whose state the restored file reflects (the snapshot's + /// `wal_boundary.epoch` when no segments were applied). + pub epoch: u64, + /// Salts of that epoch: from its directory name, or from the snapshot + /// boundary when no segments were applied. + pub salt: (u32, u32), + /// Raw WAL offset restored through within `epoch` (0 = snapshot only). + pub shipped_offset: u64, +} + +/// Scratch-file cleanup wrapper around [`restore_db_inner`]: stale leftovers +/// from a previously crashed restore are removed up front (a stale +/// `.restore-wal` would otherwise be replayed into the fresh snapshot), and +/// nothing partial survives a hard error. +async fn restore_db( + log: &Logger, + storage: &ReplicaStorage, + db_path: &Utf8Path, + up_to: Option<&Position>, +) -> Result, RestoreError> { + let paths = RestorePaths::new(db_path); + paths.clean_scratch(log).await; + let result = restore_db_inner(log, storage, &paths, up_to).await; + if result.is_err() { + paths.clean_scratch(log).await; + } + result +} + +async fn restore_db_inner( + log: &Logger, + storage: &ReplicaStorage, + paths: &RestorePaths, + up_to: Option<&Position>, +) -> Result, RestoreError> { + let want = up_to.map(|position| &position.generation); + for generation in candidate_generations(log, storage, want).await? { + let Some(snapshot) = load_snapshot_meta(log, storage, &generation).await? else { + continue; + }; + // A missing snapshot BODY is the partial-prune state (marker present, + // body already deleted), not corruption: fall through to the + // next-older restorable generation instead of refusing to boot. + // Checksum mismatches and every other download error still HARD-fail, + // because silently regressing to older data on real corruption would + // be an operator's decision, not ours. + match download_snapshot_retrying(log, storage, &generation, &snapshot, paths).await { + Ok(()) => {}, + Err(RestoreError::Storage(StorageError::NotFound { key })) => { + slog::warn!(log, "Snapshot body missing (partially pruned generation); trying an older generation"; + "generation" => generation.as_str(), "key" => key); + count_soft_stop(); + // A failed download may have left a partial scratch file. + paths.clean_scratch(log).await; + continue; + }, + Err(error) => return Err(error), + } + let restored = finish_restore(log, storage, generation, snapshot, paths, up_to).await?; + return Ok(Some(restored)); + } + Ok(None) +} + +/// After the snapshot body has downloaded and verified: replay the WAL +/// segments on top of it, `quick_check` the result, and finalize it into +/// place. +async fn finish_restore( + log: &Logger, + storage: &ReplicaStorage, + generation: GenerationId, + snapshot: SnapshotMeta, + paths: &RestorePaths, + up_to: Option<&Position>, +) -> Result { + let segments = list_segments(log, storage, &generation).await?; + let boundary_epoch = snapshot.wal_boundary.epoch; + let limit = up_to.map(|position| (position.epoch, position.offset)); + let (mut plans, violation) = plan_epochs(segments, boundary_epoch, limit); + if let Some(offender) = violation { + slog::error!(log, "WAL segment lineage is broken; replaying only the leading valid epochs (best-effort disaster recovery)"; + "generation" => generation.as_str(), + "offending_key" => segment_key(&generation, &offender)); + count_soft_stop(); + } + // `wal_boundary.offset` is the mandatory-replay threshold: at least the + // committed extent the snapshot body captured (overstated is safe, never + // understated). Replaying only a PREFIX of the boundary epoch that stops + // below it could regress pages the snapshot holds at a newer state, so + // the boundary epoch is all-or-nothing up to the threshold: when the + // available segments fall short, restore the snapshot alone (a consistent + // committed state by itself). Overstating merely forces that fallback. + if let Some(first) = plans.first() + && first.epoch == boundary_epoch + // The threshold is meaningful only for the SAME salt cycle it was + // measured in. A rebound epoch 0 (different salts) means the WAL + // restarted before anything shipped, which proves the whole old + // cycle was backfilled into the backup: no frame of the new cycle + // is mandatory. + && first.salt == (snapshot.wal_boundary.salt1, snapshot.wal_boundary.salt2) + { + let available = first.segments.last().map_or(0, |segment| segment.end); + if available < snapshot.wal_boundary.offset { + slog::error!(log, "Boundary epoch segments end below the snapshot's mandatory-replay offset; restoring the snapshot alone"; + "generation" => generation.as_str(), + "available" => available, + "mandatory" => snapshot.wal_boundary.offset); + count_soft_stop(); + plans.clear(); + } + } + let applied = apply_epochs(log, storage, &generation, &plans, paths).await?; + quick_check(paths.restore.clone()).await?; + let db_bytes = tokio::fs::metadata(&paths.restore) + .await + .map_err(|error| RestoreError::Io { + path: paths.restore.clone(), + error, + })? + .len(); + finalize(log, paths).await?; + let (epoch, salt, shipped_offset) = applied.last.unwrap_or(( + boundary_epoch, + (snapshot.wal_boundary.salt1, snapshot.wal_boundary.salt2), + 0, + )); + Ok(RestoredDb { + generation, + db_bytes, + segments: applied.segments, + epoch, + salt, + shipped_offset, + }) +} + +/// The scratch and target paths of one restore. +struct RestorePaths { + /// Final destination (``). + target: Utf8PathBuf, + /// Scratch database (`.restore`). + restore: Utf8PathBuf, + /// Scratch WAL (`.restore-wal`), rebuilt and replayed per epoch. + wal: Utf8PathBuf, + /// Scratch shared-memory sibling (`.restore-shm`). + shm: Utf8PathBuf, +} + +impl RestorePaths { + fn new(db_path: &Utf8Path) -> Self { + let restore = Utf8PathBuf::from(format!("{db_path}.restore")); + Self { + target: db_path.to_owned(), + wal: Utf8PathBuf::from(format!("{restore}-wal")), + shm: Utf8PathBuf::from(format!("{restore}-shm")), + restore, + } + } + + async fn clean_scratch(&self, log: &Logger) { + remove_quiet(log, &self.restore).await; + self.clean_scratch_wal(log).await; + } + + async fn clean_scratch_wal(&self, log: &Logger) { + remove_quiet(log, &self.wal).await; + remove_quiet(log, &self.shm).await; + } +} + +/// The restorable generation directories, newest first. Unparsable directory +/// names are skipped. `want` restricts the search to one exact generation +/// (verification restores to a recorded position). +async fn candidate_generations( + log: &Logger, + storage: &ReplicaStorage, + want: Option<&GenerationId>, +) -> Result, RestoreError> { + let components = with_retry(log, Backoff::default(), "list generations", || { + storage.list_dirs(GENERATIONS_PREFIX) + }) + .await?; + let mut generations = Vec::new(); + for component in components.iter().rev() { + let Some(generation) = GenerationId::parse(component) else { + slog::warn!(log, "Skipping unparsable generation directory"; + "component" => component.as_str()); + continue; + }; + if let Some(want) = want + && &generation != want + { + continue; + } + generations.push(generation); + } + Ok(generations) +} + +/// Load and parse a generation's `snapshot.json` commit marker. `Ok(None)` +/// when the marker is absent (crashed mid-snapshot) or unparsable: such a +/// generation is invisible to restore. +async fn load_snapshot_meta( + log: &Logger, + storage: &ReplicaStorage, + generation: &GenerationId, +) -> Result, RestoreError> { + let key = snapshot_meta_key(generation); + let bytes = match with_retry(log, Backoff::default(), "get snapshot.json", || { + storage.get(&key) + }) + .await + { + Ok(bytes) => bytes, + Err(StorageError::NotFound { .. }) => { + slog::warn!(log, "Generation has no snapshot.json (crashed mid-snapshot); skipping"; + "generation" => generation.as_str()); + return Ok(None); + }, + Err(error) => return Err(error.into()), + }; + match SnapshotMeta::from_bytes(&bytes) { + Ok(snapshot) => Ok(Some(snapshot)), + Err(error) => { + slog::error!(log, "Skipping generation with unparsable snapshot.json"; + "key" => key, "error" => %error); + count_soft_stop(); + Ok(None) + }, + } +} + +/// Retry a storage read up to [`RESTORE_ATTEMPTS`] times, waiting `backoff` +/// between attempts, retrying only TRANSIENT backend/network failures. +/// Definitive answers (`NotFound`, `InvalidKey`) return on the first attempt: +/// they are not going to change, and callers depend on `NotFound` to fall +/// through or soft-stop. The `backoff` is passed in (rather than built here) +/// so tests can drive it with zero delay. +async fn with_retry( + log: &Logger, + mut backoff: Backoff, + what: &str, + mut op: F, +) -> Result +where + F: FnMut() -> Fut, + Fut: Future>, +{ + let mut attempt = 1u32; + loop { + match op().await { + Ok(value) => return Ok(value), + Err(error) => { + if attempt >= RESTORE_ATTEMPTS || !is_retryable(&error) { + return Err(error); + } + let delay = backoff.next_delay(); + slog::warn!(log, "Transient replica storage error during restore; retrying"; + "operation" => what, "attempt" => attempt, "error" => %error, + "delay_ms" => u64::try_from(delay.as_millis()).unwrap_or(u64::MAX)); + tokio::time::sleep(delay).await; + attempt = attempt.saturating_add(1); + }, + } + } +} + +/// Whether a storage error is worth retrying. `NotFound` and `InvalidKey` are +/// definitive; backend/network failures (and injected faults, in tests) are +/// transient. +fn is_retryable(error: &StorageError) -> bool { + match error { + StorageError::NotFound { .. } | StorageError::InvalidKey { .. } => false, + StorageError::Local(_) | StorageError::S3(_) => true, + #[cfg(any(test, feature = "testing"))] + StorageError::Injected { .. } => true, + } +} + +/// Retry a whole snapshot-body download on TRANSIENT failures (a mid-stream +/// read reset or a transient backend/network error), up to +/// [`RESTORE_ATTEMPTS`] times. The stream cannot be resumed mid-way, so the +/// entire streamed download is re-run from scratch each attempt. Validation +/// verdicts (checksum mismatch, oversize) and a missing body (`NotFound`, +/// which the caller falls through on) are definitive and returned at once. +async fn download_snapshot_retrying( + log: &Logger, + storage: &ReplicaStorage, + generation: &GenerationId, + snapshot: &SnapshotMeta, + paths: &RestorePaths, +) -> Result<(), RestoreError> { + let mut backoff = Backoff::default(); + let mut attempt = 1u32; + loop { + match download_snapshot(storage, generation, snapshot, paths).await { + Ok(()) => return Ok(()), + Err(error) => { + // `SnapshotDownload` wraps the streamed read+decode io::Error, + // where a mid-stream network reset is indistinguishable from a + // deterministic zstd corruption error, so BOTH retry (a + // corrupt object wastes the remaining attempts, then + // hard-fails). A checksum mismatch is computed from fully + // downloaded bytes, is definitively corruption, and never + // retries. + let retryable = matches!(&error, RestoreError::SnapshotDownload { .. }) + || matches!(&error, RestoreError::Storage(storage_error) if is_retryable(storage_error)); + if attempt >= RESTORE_ATTEMPTS || !retryable { + return Err(error); + } + slog::warn!(log, "Transient error downloading the snapshot body; retrying"; + "generation" => generation.as_str(), "attempt" => attempt, "error" => %error); + // Discard the partial scratch before the next attempt. + paths.clean_scratch(log).await; + let delay = backoff.next_delay(); + tokio::time::sleep(delay).await; + attempt = attempt.saturating_add(1); + }, + } + } +} + +/// Stream the snapshot object through a zstd decoder into the scratch +/// database file, hashing the COMPRESSED bytes and verifying them against +/// `snapshot.json`. The file is fsynced before the hash verdict. +async fn download_snapshot( + storage: &ReplicaStorage, + generation: &GenerationId, + snapshot: &SnapshotMeta, + paths: &RestorePaths, +) -> Result<(), RestoreError> { + let stream = storage.get_stream(&snapshot_key(generation)).await?; + let mut decoder = ZstdDecoder::new(BufReader::new(HashingReader::new(stream))); + let io_err = |error| RestoreError::Io { + path: paths.restore.clone(), + error, + }; + let download_err = |error| RestoreError::SnapshotDownload { + generation: generation.clone(), + error, + }; + let mut file = tokio::fs::File::create(&paths.restore) + .await + .map_err(io_err)?; + // Bound the decompressed write at the snapshot's recorded size (plus one + // byte, to detect an overrun) so a corrupt or hostile object cannot force + // an unbounded write to disk before the checksum is even consulted. + let limit = snapshot.db_bytes.saturating_add(1); + let written = tokio::io::copy(&mut (&mut decoder).take(limit), &mut file) + .await + .map_err(download_err)?; + if written > snapshot.db_bytes { + return Err(RestoreError::SnapshotTooLarge { + generation: generation.clone(), + db_bytes: snapshot.db_bytes, + }); + } + file.sync_all().await.map_err(io_err)?; + drop(file); + // Drain anything after the zstd frame so the hash covers the whole + // object (trailing garbage must fail the checksum, not slip past it). + let mut reader = decoder.into_inner(); + tokio::io::copy(&mut reader, &mut tokio::io::sink()) + .await + .map_err(download_err)?; + let computed = reader.into_inner().finalize_hex(); + if !computed.eq_ignore_ascii_case(&snapshot.sha256) { + return Err(RestoreError::SnapshotChecksum { + generation: generation.clone(), + expected: snapshot.sha256.clone(), + computed, + }); + } + Ok(()) +} + +/// List and parse every segment key of the generation; foreign keys under +/// the WAL prefix are skipped with a warning. +async fn list_segments( + log: &Logger, + storage: &ReplicaStorage, + generation: &GenerationId, +) -> Result, RestoreError> { + let prefix = format!("{}{WAL_DIR}/", generation_prefix(generation)); + let keys = with_retry(log, Backoff::default(), "list segments", || { + storage.list(&prefix) + }) + .await?; + let mut segments = Vec::with_capacity(keys.len()); + for key in keys { + match parse_segment_key(&key) { + Some((parsed, segment)) if &parsed == generation => segments.push(segment), + Some(_) | None => { + slog::warn!(log, "Skipping foreign object under the WAL prefix"; "key" => key); + }, + } + } + Ok(segments) +} + +/// One epoch's replay plan: its segments in offset order. +struct EpochPlan { + epoch: u64, + salt: (u32, u32), + segments: Vec, +} + +/// Group segments into per-epoch replay plans and validate the lineage: +/// epochs must be contiguous starting exactly at `boundary_epoch`; within an +/// epoch the first segment must start at 0 (it carries the WAL header) and +/// each segment must start where the previous one ended. +/// +/// Returns the plans for the leading run of fully-valid epochs plus the +/// first offending segment, if any (an epoch with any violation is discarded +/// entirely: replay stops at the last FULLY-valid epoch). +/// +/// With `up_to = Some((epoch, offset))`, segments beyond that commit-aligned +/// position are excluded up front; their absence is not a violation. +fn plan_epochs( + mut segments: Vec, + boundary_epoch: u64, + up_to: Option<(u64, u64)>, +) -> (Vec, Option) { + if let Some((epoch_limit, offset_limit)) = up_to { + segments.retain(|segment| { + segment.epoch < epoch_limit + || (segment.epoch == epoch_limit && segment.end <= offset_limit) + }); + } + segments.sort_by_key(|segment| (segment.epoch, segment.start, segment.end)); + let mut groups: Vec = Vec::new(); + for segment in segments { + match groups.last_mut() { + Some(group) if group.epoch == segment.epoch => group.segments.push(segment), + Some(_) | None => groups.push(EpochPlan { + epoch: segment.epoch, + salt: segment.salt, + segments: vec![segment], + }), + } + } + let mut plans = Vec::with_capacity(groups.len()); + let mut expected_epoch = boundary_epoch; + for group in groups { + if group.epoch != expected_epoch { + return (plans, group.segments.first().copied()); + } + let mut expected_start = 0u64; + for segment in &group.segments { + if segment.start != expected_start || segment.salt != group.salt { + return (plans, Some(*segment)); + } + expected_start = segment.end; + } + plans.push(group); + expected_epoch = expected_epoch.saturating_add(1); + } + (plans, None) +} + +/// Segments and end position actually applied. +struct AppliedEpochs { + segments: u64, + /// `(epoch, salt, end offset)` of the last fully applied epoch. + last: Option<(u64, (u32, u32), u64)>, +} + +/// Apply the planned epochs in order. +/// +/// Failure policy, chosen so a restored database is NEVER a torn mixture: +/// +/// - Replica DATA problems (vanished, corrupt, wrongly sized, or +/// chain-broken segments) are detected BEFORE any database mutation +/// (assembly checks plus a full checksum-chain pre-validation of the +/// assembled epoch WAL) and soft-stop the replay: the state produced by +/// the previous epochs is kept, loudly (an older CONSISTENT state beats +/// refusing to boot when the data is simply gone). +/// - LOCAL failures during application (`SQLite` errors, or a checkpoint +/// that consumed fewer frames than the pre-validated WAL contains) are +/// HARD errors: the application may have partially backfilled pages, so +/// the scratch database is discarded and the whole restore fails, +/// retryable once the local condition clears. +/// - Storage failures other than `NotFound` are hard errors too: an +/// unreachable replica is not a broken lineage. +async fn apply_epochs( + log: &Logger, + storage: &ReplicaStorage, + generation: &GenerationId, + plans: &[EpochPlan], + paths: &RestorePaths, +) -> Result { + let mut applied = AppliedEpochs { + segments: 0, + last: None, + }; + for plan in plans { + if !build_epoch_wal(log, storage, generation, plan, paths).await? { + paths.clean_scratch_wal(log).await; + break; + } + // Pre-validate the assembled WAL with our own scanner: the salts, + // the full cumulative checksum chain, and that every byte belongs + // to a committed frame. A shipped-then-forked lineage (or any + // corruption the per-segment checks missed) is caught HERE, before + // `SQLite` recovery could silently truncate at the break and let + // later epochs replay on top of the wrong base. + let Some(expected_frames) = validate_epoch_wal(paths.wal.clone()).await? else { + slog::error!(log, "Assembled epoch WAL fails checksum-chain validation; stopping at the previous epoch"; + "generation" => generation.as_str(), + "epoch" => plan.epoch); + paths.clean_scratch_wal(log).await; + break; + }; + let (busy, in_log, checkpointed) = + match checkpoint_scratch_wal(paths.restore.clone()).await? { + Ok(row) => row, + Err(error) => { + // The checkpoint may have partially backfilled: torn state. + paths.clean_scratch(log).await; + return Err(RestoreError::Apply { + epoch: plan.epoch, + error, + }); + }, + }; + let expected = i64::try_from(expected_frames).unwrap_or(i64::MAX); + if busy != 0 || in_log != checkpointed || checkpointed != expected { + // Pre-validation proved the WAL is fully committed and intact, + // so a shortfall here is local trouble, not missing data. + paths.clean_scratch(log).await; + return Err(RestoreError::ApplyIncomplete { + epoch: plan.epoch, + expected_frames, + busy, + in_log, + checkpointed, + }); + } + paths.clean_scratch_wal(log).await; + applied.segments = applied + .segments + .saturating_add(u64::try_from(plan.segments.len()).unwrap_or(u64::MAX)); + if let Some(last) = plan.segments.last() { + applied.last = Some((plan.epoch, plan.salt, last.end)); + } + } + Ok(applied) +} + +/// Validate an assembled epoch WAL end to end with the crate's own scanner +/// (async wrapper around [`validate_epoch_wal_blocking`]). +async fn validate_epoch_wal(wal_path: Utf8PathBuf) -> Result, RestoreError> { + spawn_blocking(move || validate_epoch_wal_blocking(&wal_path)) + .await + .map_err(RestoreError::Join)? +} + +/// Validate an assembled epoch WAL: the header parses, the cumulative +/// checksum chain is intact, and EVERY byte past the header belongs to a +/// committed frame (assembled epochs are commit-aligned by construction). +/// Returns the frame count, or `None` when the file does not fully validate. +/// +/// This file was just written and fsynced by this process, so a read or seek +/// FAILURE is a hard local I/O error (propagated), never a reason to silently +/// boot an older epoch; only invalid WAL CONTENT is the soft stop that yields +/// `None` (see [`epoch_wal_validation_stop`]). +fn validate_epoch_wal_blocking(wal_path: &Utf8Path) -> Result, RestoreError> { + let io_err = |error| RestoreError::Io { + path: wal_path.to_owned(), + error, + }; + let file_len = std::fs::metadata(wal_path).map_err(io_err)?.len(); + let file = std::fs::File::open(wal_path).map_err(io_err)?; + let mut scanner = match WalScanner::open(file) { + Ok(Some(scanner)) => scanner, + // A file shorter than a full header is an empty/truncated WAL. + Ok(None) => return Ok(None), + Err(error) => return epoch_wal_validation_stop(error, wal_path), + }; + let frame_size = scanner.header().frame_size(); + loop { + // A content stop returns Ok(None) from next_committed; only a real + // read/seek failure surfaces as Err, and that must hard-fail. + match scanner.next_committed(u64::MAX) { + Ok(Some(_chunk)) => {}, + Ok(None) => break, + Err(error) => return epoch_wal_validation_stop(error, wal_path), + } + } + if scanner.offset() != file_len { + return Ok(None); + } + let body = file_len.saturating_sub(crate::wal::WAL_HEADER_SIZE); + if body % frame_size != 0 { + return Ok(None); + } + // Exact by the modulus check above; checked_div sidesteps the + // integer_division lint without a bare `/`. + Ok(body.checked_div(frame_size)) +} + +/// Classify a [`WalError`] raised while validating an epoch WAL this process +/// just assembled and fsynced. A read or seek failure is a HARD local I/O +/// error (per this module's failure policy); every other variant means the +/// assembled WAL CONTENT does not validate, which is a soft stop (`Ok(None)`: +/// replay ends at the previous epoch). +fn epoch_wal_validation_stop( + error: WalError, + wal_path: &Utf8Path, +) -> Result, RestoreError> { + match error { + WalError::Read(error) | WalError::Seek(error) => Err(RestoreError::Io { + path: wal_path.to_owned(), + error, + }), + WalError::TruncatedHeader(_) + | WalError::BadMagic(_) + | WalError::UnsupportedFormat(_) + | WalError::InvalidPageSize(_) + | WalError::HeaderChecksum { .. } + | WalError::MisalignedOffset { .. } + | WalError::TransactionTooLarge { .. } => Ok(None), + } +} + +/// Decompress-and-concatenate one epoch's segments into the scratch `-wal` +/// file (the first segment carries the 32-byte WAL header, so the +/// concatenation IS a valid WAL file). Returns `Ok(false)` on a soft stop: +/// a vanished, corrupt, or wrongly-sized segment. +async fn build_epoch_wal( + log: &Logger, + storage: &ReplicaStorage, + generation: &GenerationId, + plan: &EpochPlan, + paths: &RestorePaths, +) -> Result { + let io_err = |error| RestoreError::Io { + path: paths.wal.clone(), + error, + }; + let mut wal_file = tokio::fs::File::create(&paths.wal).await.map_err(io_err)?; + for segment in &plan.segments { + let key = segment_key(generation, segment); + let compressed = match with_retry(log, Backoff::default(), "get segment", || { + storage.get(&key) + }) + .await + { + Ok(bytes) => bytes, + Err(StorageError::NotFound { .. }) => { + slog::error!(log, "WAL segment vanished during restore; stopping at the previous epoch"; + "key" => key); + return Ok(false); + }, + Err(error) => return Err(error.into()), + }; + let expected = segment.end.saturating_sub(segment.start); + // The key encodes the exact raw size, so decompress with that as the + // cap (plus one byte, to catch an overrun): a corrupt or hostile + // segment fails immediately instead of inflating toward the generic + // multi-gigabyte ceiling before the length check below rejects it. + let raw = spawn_blocking(move || { + decompress_segment_with_cap(&compressed, expected.saturating_add(1)) + }) + .await + .map_err(RestoreError::Join)?; + let raw = match raw { + Ok(raw) => raw, + Err(error) => { + let segment_error = RestoreError::Segment(error); + slog::error!(log, "Corrupt WAL segment; stopping at the previous epoch"; + "key" => key, "error" => %segment_error); + return Ok(false); + }, + }; + if u64::try_from(raw.len()).unwrap_or(u64::MAX) != expected { + slog::error!(log, "WAL segment length does not match its key; stopping at the previous epoch"; + "key" => key, + "expected_bytes" => expected, + "actual_bytes" => raw.len()); + return Ok(false); + } + wal_file.write_all(&raw).await.map_err(io_err)?; + } + wal_file.sync_all().await.map_err(io_err)?; + Ok(true) +} + +/// Replay the assembled epoch WAL through `SQLite`'s own recovery: +/// `PRAGMA wal_checkpoint(FULL)` validates the salts and the cumulative +/// checksum chain and applies exactly the committed frames (zero +/// hand-written page application). FULL (not TRUNCATE) so the result row +/// `(busy, log, checkpointed)` reports the REAL frame counts (TRUNCATE +/// zeroes them after truncating), letting the caller verify that EVERY +/// frame of the pre-validated WAL was consumed; the scratch WAL file is +/// deleted by the caller afterward. The inner `Result` carries `SQLite` +/// failures, the outer one is task failure. +async fn checkpoint_scratch_wal( + restore_path: Utf8PathBuf, +) -> Result, RestoreError> { + spawn_blocking(move || { + let conn = rusqlite::Connection::open(&restore_path)?; + let row: (i64, i64, i64) = conn.query_row("PRAGMA wal_checkpoint(FULL)", [], |row| { + Ok((row.get(0)?, row.get(1)?, row.get(2)?)) + })?; + conn.close().map_err(|(_conn, error)| error)?; + Ok(row) + }) + .await + .map_err(RestoreError::Join) +} + +/// `PRAGMA quick_check` on the restored file: anything but "ok" is a HARD +/// error (never boot on a corrupt database; contrast with the soft stop for +/// missing tail data, which yields an older but CONSISTENT state). +async fn quick_check(restore_path: Utf8PathBuf) -> Result<(), RestoreError> { + let report = spawn_blocking(move || -> Result, rusqlite::Error> { + let conn = rusqlite::Connection::open(&restore_path)?; + let mut statement = conn.prepare("PRAGMA quick_check")?; + let rows = statement.query_map([], |row| row.get(0))?; + rows.collect() + }) + .await + .map_err(RestoreError::Join)? + .map_err(RestoreError::QuickCheckRun)?; + if report == ["ok"] { + Ok(()) + } else { + Err(RestoreError::QuickCheck(report.join("; "))) + } +} + +/// Atomically move the verified scratch database into place. Stale `-wal` / +/// `-shm` siblings of the TARGET are removed first: `SQLite` would replay a +/// leftover WAL into the freshly restored file. +async fn finalize(log: &Logger, paths: &RestorePaths) -> Result<(), RestoreError> { + remove_quiet(log, Utf8Path::new(&format!("{}-wal", paths.target))).await; + remove_quiet(log, Utf8Path::new(&format!("{}-shm", paths.target))).await; + tokio::fs::rename(&paths.restore, &paths.target) + .await + .map_err(|error| RestoreError::Io { + path: paths.target.clone(), + error, + })?; + // Best-effort directory fsync so the rename survives an abrupt stop; + // the file itself is already fsynced. + if let Some(parent) = paths.target.parent() + && let Ok(dir) = std::fs::File::open(parent.as_std_path()) + { + drop(dir.sync_all()); + } + Ok(()) +} + +/// Remove a file, treating "already gone" as success; other failures are +/// logged and swallowed (cleanup is best-effort). +async fn remove_quiet(log: &Logger, path: &Utf8Path) { + match tokio::fs::remove_file(path).await { + Ok(()) => {}, + Err(error) if error.kind() == ErrorKind::NotFound => {}, + Err(error) => { + slog::warn!(log, "Failed to remove restore scratch file"; + "path" => path.as_str(), "error" => %error); + }, + } +} + +/// Wraps the raw snapshot byte stream, hashing every COMPRESSED byte read +/// through it so the object hash can be verified against `snapshot.json` +/// without a second pass. +struct HashingReader { + inner: R, + hasher: sha2::Sha256, +} + +impl HashingReader { + fn new(inner: R) -> Self { + Self { + inner, + hasher: sha2::Sha256::new(), + } + } + + /// Hex digest of everything read so far. + fn finalize_hex(self) -> String { + hex::encode(self.hasher.finalize()) + } +} + +impl AsyncRead for HashingReader { + fn poll_read( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll> { + let this = self.get_mut(); + let already_filled = buf.filled().len(); + let poll = Pin::new(&mut this.inner).poll_read(cx, buf); + if let Poll::Ready(Ok(())) = &poll { + this.hasher + .update(buf.filled().get(already_filled..).unwrap_or(&[])); + } + poll + } +} + +#[cfg(test)] +mod tests { + use camino::Utf8Path; + use pretty_assertions::assert_eq; + + use super::{ + RESTORE_ATTEMPTS, RestoreError, epoch_wal_validation_stop, is_retryable, plan_epochs, + validate_epoch_wal_blocking, with_retry, + }; + use crate::backoff::Backoff; + use crate::position::SegmentKey; + use crate::storage::StorageError; + use crate::wal::WalError; + + const SALT: (u32, u32) = (0x9d2f_1c4a, 0x8b3e_6f70); + const OTHER_SALT: (u32, u32) = (0x0101_0101, 0x0202_0202); + + fn seg(epoch: u64, start: u64, end: u64) -> SegmentKey { + seg_with_salt(epoch, SALT, start, end) + } + + fn seg_with_salt(epoch: u64, salt: (u32, u32), start: u64, end: u64) -> SegmentKey { + SegmentKey { + epoch, + salt, + start, + end, + } + } + + fn plan_shape(plans: &[super::EpochPlan]) -> Vec<(u64, usize)> { + plans + .iter() + .map(|plan| (plan.epoch, plan.segments.len())) + .collect() + } + + #[test] + fn plan_contiguous_epochs_all_valid() { + // Unsorted input on purpose: the planner sorts. + let segments = vec![ + seg(1, 0, 100), + seg(0, 32, 64), + seg(0, 0, 32), + seg(1, 100, 250), + seg(0, 64, 200), + ]; + let (plans, violation) = plan_epochs(segments, 0, None); + assert_eq!(violation, None, "a contiguous lineage has no violation"); + assert_eq!(plan_shape(&plans), vec![(0, 3), (1, 2)]); + } + + #[test] + fn plan_stops_at_offset_gap() { + let segments = vec![seg(0, 0, 32), seg(0, 64, 128), seg(1, 0, 32)]; + let (plans, violation) = plan_epochs(segments, 0, None); + assert_eq!( + violation, + Some(seg(0, 64, 128)), + "the first non-contiguous segment is the offender" + ); + assert_eq!( + plan_shape(&plans), + Vec::<(u64, usize)>::new(), + "an epoch with a gap is discarded entirely, along with everything after it" + ); + } + + #[test] + fn plan_requires_first_segment_at_zero() { + let segments = vec![seg(0, 32, 128)]; + let (plans, violation) = plan_epochs(segments, 0, None); + assert_eq!(violation, Some(seg(0, 32, 128))); + assert_eq!(plan_shape(&plans), Vec::<(u64, usize)>::new()); + } + + #[test] + fn plan_requires_epochs_to_start_at_boundary() { + // The snapshot boundary is epoch 0, but only epoch 1 was shipped. + let segments = vec![seg(1, 0, 32)]; + let (plans, violation) = plan_epochs(segments, 0, None); + assert_eq!(violation, Some(seg(1, 0, 32))); + assert_eq!(plan_shape(&plans), Vec::<(u64, usize)>::new()); + } + + #[test] + fn plan_stops_at_epoch_number_gap() { + let segments = vec![seg(0, 0, 32), seg(2, 0, 32)]; + let (plans, violation) = plan_epochs(segments, 0, None); + assert_eq!(violation, Some(seg(2, 0, 32)), "epoch 1 is missing"); + assert_eq!(plan_shape(&plans), vec![(0, 1)], "epoch 0 still replays"); + } + + #[test] + fn plan_stops_at_salt_mismatch_within_epoch() { + // Two epoch dirs claiming the same epoch number with different + // salts collapse into one group; the second lineage is an offender. + let segments = vec![seg(0, 0, 32), seg_with_salt(0, OTHER_SALT, 32, 64)]; + let (plans, violation) = plan_epochs(segments, 0, None); + assert_eq!(violation, Some(seg_with_salt(0, OTHER_SALT, 32, 64))); + assert_eq!(plan_shape(&plans), Vec::<(u64, usize)>::new()); + } + + #[test] + fn plan_discards_epoch_with_duplicate_salt_lineages() { + // The salt-collision case: epoch 1 exists as TWO complete, individually + // well-formed lineages under different salts (two epoch directories with + // the same epoch number). Lineage A (SALT) chains 0..40..90 and lineage + // B (OTHER_SALT) chains 0..50..120; each is valid on its own. The sort + // key is (epoch, start, end) with salt EXCLUDED, so both collapse into + // one epoch-1 group whose salt is the first sorted segment's. The second + // lineage's leading segment then fails the group-salt / offset-chain + // check, discarding the WHOLE epoch. Restore soft-stops at epoch 0: it + // never tears the two lineages together and never silently applies + // either one. This is safe but pessimistic (a complete valid lineage is + // dropped), and is the reason the resume path must never create the + // collision in the first place (see `meta_matches` in `sync.rs`). + let segments = vec![ + seg(0, 0, 32), + seg(1, 0, 40), + seg(1, 40, 90), + seg_with_salt(1, OTHER_SALT, 0, 50), + seg_with_salt(1, OTHER_SALT, 50, 120), + ]; + let (plans, violation) = plan_epochs(segments, 0, None); + assert_eq!( + plan_shape(&plans), + vec![(0, 1)], + "epoch 0 applies; the duplicate-epoch collision discards epoch 1 entirely" + ); + assert_eq!( + violation, + Some(seg_with_salt(1, OTHER_SALT, 0, 50)), + "the offender is the second lineage's leading segment (start 0 under a foreign salt)" + ); + } + + #[test] + fn plan_up_to_filters_epochs_and_offsets() { + let segments = vec![ + seg(0, 0, 100), + seg(1, 0, 40), + seg(1, 40, 90), + seg(1, 90, 150), + seg(2, 0, 32), + ]; + // Stop mid-epoch-1 at the commit-aligned offset 90. + let (plans, violation) = plan_epochs(segments.clone(), 0, Some((1, 90))); + assert_eq!(violation, None, "excluded segments are not violations"); + assert_eq!(plan_shape(&plans), vec![(0, 1), (1, 2)]); + // An offset before any segment of the epoch keeps only prior epochs. + let (plans, violation) = plan_epochs(segments, 0, Some((1, 0))); + assert_eq!(violation, None); + assert_eq!(plan_shape(&plans), vec![(0, 1)]); + } + + #[test] + fn epoch_wal_validation_io_error_hard_fails_content_soft_stops() { + let path = Utf8Path::new("/does/not/matter.wal"); + // A read or seek failure on the file this process just wrote and + // fsynced is a HARD error: it must never masquerade as a broken chain + // and silently boot an older epoch. + for error in [ + WalError::Read(std::io::Error::other("boom")), + WalError::Seek(std::io::Error::other("boom")), + ] { + assert!( + matches!( + epoch_wal_validation_stop(error, path), + Err(RestoreError::Io { .. }) + ), + "an I/O error must hard-fail" + ); + } + // Invalid WAL CONTENT is a soft stop: replay ends at the previous + // epoch (Ok(None)). + for error in [ + WalError::TruncatedHeader(5), + WalError::BadMagic(0xdead_beef), + WalError::UnsupportedFormat(1), + WalError::InvalidPageSize(3), + WalError::HeaderChecksum { + stored: (1, 2), + computed: (3, 4), + }, + WalError::MisalignedOffset { + offset: 33, + page_size: 4096, + }, + WalError::TransactionTooLarge { + bytes: 1 << 30, + max_bytes: 1 << 20, + }, + ] { + assert!( + matches!(epoch_wal_validation_stop(error, path), Ok(None)), + "invalid content is a soft stop" + ); + } + } + + #[test] + fn validate_epoch_wal_blocking_hard_fails_on_read_error() { + // A directory where the WAL file should be: metadata and open succeed + // on Unix, but the header read fails (EISDIR) -> WalError::Read -> a + // HARD RestoreError::Io, never a silent Ok(None) soft stop. Even if a + // platform instead fails at open, the result is still a hard error, + // which is the property under test. + let tmp = tempfile::tempdir().unwrap(); + let dir = Utf8Path::from_path(tmp.path()).unwrap().join("epoch.wal"); + std::fs::create_dir(&dir).unwrap(); + assert!( + matches!( + validate_epoch_wal_blocking(&dir), + Err(RestoreError::Io { .. }) + ), + "a read failure on the assembled WAL must hard-fail" + ); + } + + fn discard_logger() -> slog::Logger { + slog::Logger::root(slog::Discard, slog::o!()) + } + + /// A retryable (transient) storage error for the retry tests. + fn injected() -> StorageError { + StorageError::Injected { + op: "get", + key: "k".to_owned(), + } + } + + /// Zero-delay backoff so the retry tests never touch the wall clock. + fn no_delay() -> Backoff { + Backoff::new(std::time::Duration::ZERO, std::time::Duration::ZERO) + } + + #[test] + fn is_retryable_splits_definitive_from_transient() { + assert!( + !is_retryable(&StorageError::NotFound { + key: "k".to_owned() + }), + "NotFound is definitive" + ); + assert!( + !is_retryable(&StorageError::InvalidKey { + key: "k".to_owned(), + reason: "leading slash", + }), + "InvalidKey is a programming error, not transient" + ); + assert!( + is_retryable(&injected()), + "an injected backend fault is transient" + ); + } + + #[tokio::test] + async fn with_retry_recovers_from_a_transient_failure() { + use std::sync::Arc; + use std::sync::atomic::{AtomicU32, Ordering}; + + let calls = Arc::new(AtomicU32::new(0)); + let result = with_retry(&discard_logger(), no_delay(), "get", || { + let calls = Arc::clone(&calls); + async move { + let seen = calls.fetch_add(1, Ordering::SeqCst) + 1; + if seen < 2 { + Err(injected()) + } else { + Ok::(42) + } + } + }) + .await; + assert_eq!(result.unwrap(), 42, "the retry succeeds"); + assert_eq!( + calls.load(Ordering::SeqCst), + 2, + "exactly one retry after the transient failure" + ); + } + + #[tokio::test] + async fn with_retry_gives_up_after_the_attempt_bound() { + use std::sync::Arc; + use std::sync::atomic::{AtomicU32, Ordering}; + + let calls = Arc::new(AtomicU32::new(0)); + let result: Result = + with_retry(&discard_logger(), no_delay(), "get", || { + let calls = Arc::clone(&calls); + async move { + calls.fetch_add(1, Ordering::SeqCst); + Err(injected()) + } + }) + .await; + assert!( + matches!(result, Err(StorageError::Injected { .. })), + "a persistent transient error still fails after the bound" + ); + assert_eq!( + calls.load(Ordering::SeqCst), + RESTORE_ATTEMPTS, + "tried exactly RESTORE_ATTEMPTS times" + ); + } + + #[tokio::test] + async fn with_retry_does_not_retry_definitive_errors() { + use std::sync::Arc; + use std::sync::atomic::{AtomicU32, Ordering}; + + let calls = Arc::new(AtomicU32::new(0)); + let result: Result = + with_retry(&discard_logger(), no_delay(), "get", || { + let calls = Arc::clone(&calls); + async move { + calls.fetch_add(1, Ordering::SeqCst); + Err(StorageError::NotFound { + key: "k".to_owned(), + }) + } + }) + .await; + assert!( + matches!(result, Err(StorageError::NotFound { .. })), + "NotFound propagates unchanged" + ); + assert_eq!( + calls.load(Ordering::SeqCst), + 1, + "a definitive NotFound is tried exactly once" + ); + } +} diff --git a/plus/bencher_replica/src/s3.rs b/plus/bencher_replica/src/s3.rs new file mode 100644 index 000000000..23101be25 --- /dev/null +++ b/plus/bencher_replica/src/s3.rs @@ -0,0 +1,1073 @@ +//! S3-compatible replica backend (AWS S3, or R2/MinIO via `endpoint`). +//! +//! Construction mirrors the existing patterns in +//! `lib/bencher_schema/src/context/database.rs` (static credentials via +//! `aws_credential_types::Credentials::new`) with an optional endpoint +//! override for S3-compatible services (which also forces path-style +//! addressing). +//! +//! Error variants wrap the concrete `SdkError` types (never +//! `String`, never `Box`), one variant per operation. Native +//! missing-object errors (`NoSuchKey`) are mapped to +//! `StorageError::NotFound` by the callers in this module, not surfaced as +//! `S3Error`. + +use aws_credential_types::provider::SharedCredentialsProvider; +use aws_sdk_s3::config::Region; +use aws_sdk_s3::error::SdkError; +use aws_sdk_s3::operation::abort_multipart_upload::AbortMultipartUploadError; +use aws_sdk_s3::operation::complete_multipart_upload::CompleteMultipartUploadError; +use aws_sdk_s3::operation::create_multipart_upload::CreateMultipartUploadError; +use aws_sdk_s3::operation::delete_object::DeleteObjectError; +use aws_sdk_s3::operation::delete_objects::DeleteObjectsError; +use aws_sdk_s3::operation::get_object::GetObjectError; +use aws_sdk_s3::operation::list_multipart_uploads::ListMultipartUploadsError; +use aws_sdk_s3::operation::list_objects_v2::ListObjectsV2Error; +use aws_sdk_s3::operation::list_objects_v2::builders::ListObjectsV2FluentBuilder; +use aws_sdk_s3::operation::put_object::PutObjectError; +use aws_sdk_s3::operation::upload_part::UploadPartError; +use aws_sdk_s3::primitives::ByteStream; +use aws_sdk_s3::types::{CompletedMultipartUpload, CompletedPart, Delete, ObjectIdentifier}; +use bytes::Bytes; +use slog::Logger; + +use crate::storage::StorageError; + +/// The S3 minimum part size for all but the last part of a multipart +/// upload: writes are buffered until at least this many bytes are pending. +const MIN_PART_BYTES: usize = 5 * 1024 * 1024; + +/// The S3 maximum part number for a multipart upload. +const MAX_PART_NUMBER: i32 = 10_000; + +/// The S3 maximum number of objects per `DeleteObjects` request. +const DELETE_BATCH_SIZE: usize = 1_000; + +/// The region assumed when none is configured (endpoint-only targets like +/// `MinIO` still require a signing region). +const DEFAULT_REGION: &str = "us-east-1"; + +/// S3-compatible backend: bucket plus optional key prefix. +pub struct S3Storage { + client: aws_sdk_s3::Client, + bucket: String, + /// Key prefix (no leading or trailing slash); empty for bucket root. + prefix: String, + /// Test-only `ListObjectsV2` page-size cap, so the continuation-token + /// pagination loop can be exercised against a small object count. + #[cfg(feature = "testing")] + max_keys: Option, +} + +#[derive(Debug, thiserror::Error)] +pub enum S3Error { + #[error("Failed to put object ({key}): {error}")] + Put { + key: String, + error: SdkError, + }, + #[error("Failed to get object ({key}): {error}")] + Get { + key: String, + error: SdkError, + }, + #[error("Failed to read object body ({key}): {error}")] + Body { + key: String, + error: aws_sdk_s3::primitives::ByteStreamError, + }, + #[error("Failed to list objects ({prefix}): {error}")] + List { + prefix: String, + error: SdkError, + }, + #[error( + "Truncated listing without a continuation token ({prefix}); the listing would be silently short" + )] + TruncatedList { prefix: String }, + #[error("Failed to list multipart uploads ({prefix}): {error}")] + ListMultipart { + prefix: String, + error: SdkError, + }, + #[error( + "Truncated multipart-upload listing without a continuation marker ({prefix}); the sweep would be silently short" + )] + TruncatedMultipartList { prefix: String }, + #[error("Failed to delete object ({key}): {error}")] + Delete { + key: String, + error: SdkError, + }, + #[error("Failed to batch delete objects ({prefix}): {error}")] + DeleteBatch { + prefix: String, + error: SdkError, + }, + #[error("Failed to build batch delete request ({prefix}): {error}")] + DeleteBatchBuild { + prefix: String, + error: aws_sdk_s3::error::BuildError, + }, + #[error("Batch delete reported a per-key failure ({prefix}, key {key}): {code} {message}")] + DeleteBatchKey { + prefix: String, + key: String, + code: String, + message: String, + }, + #[error("Failed to create multipart upload ({key}): {error}")] + CreateMultipart { + key: String, + error: SdkError, + }, + #[error("Failed to upload part {part} ({key}): {error}")] + UploadPart { + key: String, + part: i32, + error: SdkError, + }, + #[error("Failed to complete multipart upload ({key}): {error}")] + CompleteMultipart { + key: String, + error: SdkError, + }, + #[error("Failed to abort multipart upload ({key}): {error}")] + AbortMultipart { + key: String, + error: SdkError, + }, + #[error("Multipart upload part count overflow ({key})")] + PartCountOverflow { key: String }, + #[error("Multipart upload id missing from response ({key})")] + MissingUploadId { key: String }, + #[error("ETag missing from upload part {part} response ({key})")] + MissingETag { key: String, part: i32 }, +} + +impl S3Storage { + /// Build a client from static credentials with an optional endpoint + /// override (R2/MinIO) and optional region (defaults to `us-east-1` when + /// none is given, e.g. endpoint-only targets). + #[must_use] + pub fn new( + bucket: String, + prefix: Option, + endpoint: Option, + region: Option, + access_key_id: String, + secret_access_key: &str, + ) -> Self { + let credentials = aws_credential_types::Credentials::new( + access_key_id, + secret_access_key, + None, + None, + "bencher_replica", + ); + let region = Region::new(region.unwrap_or_else(|| DEFAULT_REGION.to_owned())); + let mut config = aws_sdk_s3::Config::builder() + .credentials_provider(SharedCredentialsProvider::new(credentials)) + .region(region); + if let Some(endpoint) = endpoint { + // S3-compatible services (MinIO, R2) are addressed by URL, not + // by virtual-hosted bucket subdomains. + config = config.endpoint_url(endpoint).force_path_style(true); + } + let client = aws_sdk_s3::Client::from_conf(config.build()); + Self { + client, + bucket, + prefix: normalize_prefix(prefix), + #[cfg(feature = "testing")] + max_keys: None, + } + } + + /// Test-only: cap the `ListObjectsV2` page size so a modest object count + /// still spans several continuation-token round-trips. + #[cfg(feature = "testing")] + pub fn set_max_keys(&mut self, max_keys: i32) { + self.max_keys = Some(max_keys); + } + + pub(crate) async fn put(&self, key: &str, bytes: Bytes) -> Result<(), StorageError> { + crate::storage::validate_key(key)?; + let object_key = self.object_key(key); + self.client + .put_object() + .bucket(&self.bucket) + .key(&object_key) + .body(ByteStream::from(bytes)) + .send() + .await + .map_err(|error| S3Error::Put { + key: object_key, + error, + })?; + Ok(()) + } + + pub(crate) async fn get(&self, key: &str) -> Result { + crate::storage::validate_key(key)?; + let object_key = self.object_key(key); + let response = self + .client + .get_object() + .bucket(&self.bucket) + .key(&object_key) + .send() + .await + .map_err(|error| map_get_error(key, &object_key, error))?; + let aggregated = response + .body + .collect() + .await + .map_err(|error| S3Error::Body { + key: object_key, + error, + })?; + Ok(aggregated.into_bytes()) + } + + pub(crate) async fn get_stream( + &self, + key: &str, + ) -> Result, StorageError> { + crate::storage::validate_key(key)?; + let object_key = self.object_key(key); + let response = self + .client + .get_object() + .bucket(&self.bucket) + .key(&object_key) + .send() + .await + .map_err(|error| map_get_error(key, &object_key, error))?; + Ok(Box::new(response.body.into_async_read())) + } + + pub(crate) async fn list(&self, prefix: &str) -> Result, StorageError> { + let full_prefix = self.full_prefix(prefix); + let root = self.root_prefix(); + let mut keys = Vec::new(); + let mut continuation: Option = None; + loop { + let mut request = self + .client + .list_objects_v2() + .bucket(&self.bucket) + .prefix(&full_prefix); + if let Some(token) = &continuation { + request = request.continuation_token(token); + } + request = self.with_page_size(request); + let response = request.send().await.map_err(|error| S3Error::List { + prefix: full_prefix.clone(), + error, + })?; + for object in response.contents() { + if let Some(key) = object.key() + && let Some(stripped) = key.strip_prefix(root.as_str()) + { + keys.push(stripped.to_owned()); + } + } + match next_list_page(response.is_truncated(), response.next_continuation_token()) { + ListPage::Continue(token) => continuation = Some(token), + ListPage::Done => break, + ListPage::Truncated => { + return Err(S3Error::TruncatedList { + prefix: full_prefix, + } + .into()); + }, + } + } + keys.sort(); + Ok(keys) + } + + pub(crate) async fn list_dirs(&self, prefix: &str) -> Result, StorageError> { + let full_prefix = self.full_prefix(&normalize_dir_prefix(prefix)); + let mut dirs = Vec::new(); + let mut continuation: Option = None; + loop { + let mut request = self + .client + .list_objects_v2() + .bucket(&self.bucket) + .prefix(&full_prefix) + .delimiter("/"); + if let Some(token) = &continuation { + request = request.continuation_token(token); + } + request = self.with_page_size(request); + let response = request.send().await.map_err(|error| S3Error::List { + prefix: full_prefix.clone(), + error, + })?; + for common_prefix in response.common_prefixes() { + // A common prefix is `/`; return the + // bare component. + if let Some(prefixed) = common_prefix.prefix() + && let Some(component) = prefixed + .strip_prefix(full_prefix.as_str()) + .and_then(|component| component.strip_suffix('/')) + && !component.is_empty() + { + dirs.push(component.to_owned()); + } + } + match next_list_page(response.is_truncated(), response.next_continuation_token()) { + ListPage::Continue(token) => continuation = Some(token), + ListPage::Done => break, + ListPage::Truncated => { + return Err(S3Error::TruncatedList { + prefix: full_prefix, + } + .into()); + }, + } + } + dirs.sort(); + Ok(dirs) + } + + pub(crate) async fn delete(&self, key: &str) -> Result<(), StorageError> { + crate::storage::validate_key(key)?; + let object_key = self.object_key(key); + // S3 DeleteObject on a missing key succeeds, giving idempotency + // naturally. + self.client + .delete_object() + .bucket(&self.bucket) + .key(&object_key) + .send() + .await + .map_err(|error| S3Error::Delete { + key: object_key, + error, + })?; + Ok(()) + } + + pub(crate) async fn delete_prefix(&self, prefix: &str) -> Result<(), StorageError> { + let keys = self.list(prefix).await?; + for chunk in keys.chunks(DELETE_BATCH_SIZE) { + let mut objects = Vec::with_capacity(chunk.len()); + for key in chunk { + let identifier = ObjectIdentifier::builder() + .key(self.object_key(key)) + .build() + .map_err(|error| S3Error::DeleteBatchBuild { + prefix: prefix.to_owned(), + error, + })?; + objects.push(identifier); + } + let delete = Delete::builder() + .set_objects(Some(objects)) + .build() + .map_err(|error| S3Error::DeleteBatchBuild { + prefix: prefix.to_owned(), + error, + })?; + let response = self + .client + .delete_objects() + .bucket(&self.bucket) + .delete(delete) + .send() + .await + .map_err(|error| S3Error::DeleteBatch { + prefix: self.full_prefix(prefix), + error, + })?; + // S3 reports per-key failures inside an HTTP 200 response; a + // swallowed one would make prune claim success while objects + // (possibly a snapshot body without its snapshot.json) remain. + if let Some(first) = response.errors().first() { + return Err(S3Error::DeleteBatchKey { + prefix: self.full_prefix(prefix), + key: first.key().unwrap_or_default().to_owned(), + code: first.code().unwrap_or_default().to_owned(), + message: first.message().unwrap_or_default().to_owned(), + } + .into()); + } + } + Ok(()) + } + + pub(crate) async fn start_multipart(&self, key: &str) -> Result { + crate::storage::validate_key(key)?; + let object_key = self.object_key(key); + let response = self + .client + .create_multipart_upload() + .bucket(&self.bucket) + .key(&object_key) + .send() + .await + .map_err(|error| S3Error::CreateMultipart { + key: object_key.clone(), + error, + })?; + let upload_id = response + .upload_id() + .ok_or_else(|| S3Error::MissingUploadId { + key: object_key.clone(), + })? + .to_owned(); + Ok(S3Multipart { + client: self.client.clone(), + bucket: self.bucket.clone(), + key: object_key, + upload_id, + part_number: 0, + buffer: Vec::new(), + completed: Vec::new(), + }) + } + + /// Best-effort sweep of crash-orphaned incomplete multipart uploads under + /// the configured prefix. A process killed (SIGKILL/OOM) between + /// `create_multipart_upload` and completion loses the upload id (held only + /// in memory), and the incomplete upload then accrues storage cost until + /// aborted. This reconciles by listing every incomplete upload scoped to + /// the prefix and aborting it. Per-upload failures are logged and + /// swallowed; a listing failure stops the sweep with a log line. Never + /// fails the caller. NOT wired into the sync engine here: the orchestrator + /// invokes it at startup. + /// + /// Refuses to run without a configured path prefix: with an empty prefix + /// the listing is BUCKET-WIDE, and a shared bucket (database backups, OCI + /// storage, another tenant) would have its in-flight multipart uploads + /// aborted by every replica boot. Rely on a bucket lifecycle rule instead + /// in that layout. + pub(crate) async fn abort_incomplete_uploads(&self, log: &Logger) { + if self.prefix.is_empty() { + slog::warn!( + log, + "No S3 path prefix configured; skipping the incomplete multipart upload sweep \ + (a bucket-wide sweep could abort other systems' uploads). \ + Configure a bucket lifecycle rule to reap crash-orphaned uploads." + ); + return; + } + let uploads = match self.list_incomplete_uploads().await { + Ok(uploads) => uploads, + Err(error) => { + slog::warn!(log, "Failed to list incomplete multipart uploads; sweep skipped"; + "error" => %error); + return; + }, + }; + let found = uploads.len(); + let mut aborted = 0usize; + for (key, upload_id) in uploads { + match abort_multipart(&self.client, &self.bucket, &key, &upload_id).await { + Ok(()) => { + aborted += 1; + slog::info!(log, "Aborted crash-orphaned multipart upload"; + "key" => key.as_str(), "upload_id" => upload_id.as_str()); + }, + Err(error) => { + slog::warn!(log, "Failed to abort incomplete multipart upload; continuing"; + "key" => key.as_str(), "upload_id" => upload_id.as_str(), "error" => %error); + }, + } + } + slog::info!(log, "Incomplete multipart upload sweep complete"; + "found" => found, "aborted" => aborted); + } + + /// Every incomplete multipart upload under the configured prefix as + /// `(object_key, upload_id)` pairs, paginated with the key-marker / + /// upload-id-marker cursor and the same truncation guard as [`Self::list`] + /// (a truncated page with no continuation marker is an error, never a + /// silently short listing). + async fn list_incomplete_uploads(&self) -> Result, S3Error> { + let prefix = self.root_prefix(); + let mut uploads = Vec::new(); + let mut key_marker: Option = None; + let mut upload_id_marker: Option = None; + loop { + // Deliberately NO server-side `prefix` parameter: prefix + // filtering on `ListMultipartUploads` is unreliable across + // S3-compatible servers (MinIO matches the prefix only against + // exact object keys, so a directory prefix returns nothing). + // List everything and filter client-side; the caller-enforced + // non-empty prefix keeps the ABORT scoped to our uploads, and + // incomplete uploads are a small set. + let mut request = self.client.list_multipart_uploads().bucket(&self.bucket); + if let Some(marker) = &key_marker { + request = request.key_marker(marker); + } + if let Some(marker) = &upload_id_marker { + request = request.upload_id_marker(marker); + } + let response = request + .send() + .await + .map_err(|error| S3Error::ListMultipart { + prefix: prefix.clone(), + error, + })?; + for upload in response.uploads() { + if let (Some(key), Some(upload_id)) = (upload.key(), upload.upload_id()) + && key.starts_with(&prefix) + { + uploads.push((key.to_owned(), upload_id.to_owned())); + } + } + match next_upload_page( + response.is_truncated(), + response.next_key_marker(), + response.next_upload_id_marker(), + ) { + UploadPage::Continue { + key_marker: next_key, + upload_id_marker: next_upload_id, + } => { + key_marker = Some(next_key); + upload_id_marker = next_upload_id; + }, + UploadPage::Done => break, + UploadPage::Truncated => { + return Err(S3Error::TruncatedMultipartList { prefix }); + }, + } + } + Ok(uploads) + } + + /// Test-only: how many incomplete multipart uploads exist under the + /// configured prefix (a listing failure reports 0). + #[cfg(feature = "testing")] + pub async fn incomplete_upload_count(&self) -> usize { + self.list_incomplete_uploads() + .await + .map(|uploads| uploads.len()) + .unwrap_or_default() + } + + /// The full object key for a caller-relative key. + fn object_key(&self, key: &str) -> String { + if self.prefix.is_empty() { + key.to_owned() + } else { + format!("{}/{key}", self.prefix) + } + } + + /// The configured root as a listing prefix (`""` or `"/"`); also + /// what gets stripped from listed keys. + fn root_prefix(&self) -> String { + if self.prefix.is_empty() { + String::new() + } else { + format!("{}/", self.prefix) + } + } + + /// A caller-relative listing prefix scoped under the configured root. + fn full_prefix(&self, prefix: &str) -> String { + format!("{}{prefix}", self.root_prefix()) + } + + /// Apply the test-only page-size cap to a `ListObjectsV2` request; a no-op + /// in production builds. + fn with_page_size(&self, request: ListObjectsV2FluentBuilder) -> ListObjectsV2FluentBuilder { + #[cfg(feature = "testing")] + if let Some(max_keys) = self.max_keys { + return request.max_keys(max_keys); + } + request + } +} + +/// How a paginated `ListObjectsV2` loop proceeds from one response. +enum ListPage { + /// More results remain; continue with this continuation token. + Continue(String), + /// The listing is complete. + Done, + /// The response was truncated but carried no usable continuation token: + /// continuing is impossible and stopping would silently drop keys. The + /// caller must raise an error naming the prefix. + Truncated, +} + +/// Decide how a paginated listing continues. `is_truncated == Some(true)` with +/// a non-empty token continues; not truncated is done; truncated with no +/// usable token is [`ListPage::Truncated`] (an error, never a short listing, +/// because the engine treats LIST as the source of truth: a missing key would +/// masquerade as a deleted object). +fn next_list_page(is_truncated: Option, next_token: Option<&str>) -> ListPage { + if is_truncated == Some(true) { + match next_token { + Some(token) if !token.is_empty() => ListPage::Continue(token.to_owned()), + _ => ListPage::Truncated, + } + } else { + ListPage::Done + } +} + +/// How a paginated `ListMultipartUploads` loop proceeds from one response. The +/// cursor is a `(key_marker, upload_id_marker)` pair; the key marker is the +/// primary continuation token and the upload-id marker is only meaningful +/// alongside it. +enum UploadPage { + /// More results remain; continue with these markers. + Continue { + key_marker: String, + upload_id_marker: Option, + }, + /// The listing is complete. + Done, + /// Truncated but with no usable key marker: continuing is impossible and + /// stopping would silently drop uploads. The caller must raise an error. + Truncated, +} + +/// Decide how a `ListMultipartUploads` listing continues, applying the same +/// truncation rigor as [`next_list_page`]: truncated with a key marker +/// continues; not truncated is done; truncated with no usable key marker is +/// [`UploadPage::Truncated`]. +fn next_upload_page( + is_truncated: Option, + next_key_marker: Option<&str>, + next_upload_id_marker: Option<&str>, +) -> UploadPage { + if is_truncated == Some(true) { + match next_key_marker { + Some(key_marker) if !key_marker.is_empty() => UploadPage::Continue { + key_marker: key_marker.to_owned(), + upload_id_marker: next_upload_id_marker + .filter(|marker| !marker.is_empty()) + .map(ToOwned::to_owned), + }, + _ => UploadPage::Truncated, + } + } else { + UploadPage::Done + } +} + +/// Streaming upload to S3 via multipart. Buffers writes internally until the +/// 5 MiB minimum part size is reached, then uploads parts sequentially. +/// Dropping without [`Self::finish`] leaves only an uncompleted multipart +/// upload: the object never becomes visible under its key. +pub struct S3Multipart { + client: aws_sdk_s3::Client, + bucket: String, + key: String, + upload_id: String, + part_number: i32, + buffer: Vec, + completed: Vec, +} + +impl S3Multipart { + pub(crate) async fn write_part(&mut self, bytes: Bytes) -> Result<(), StorageError> { + self.buffer.extend_from_slice(&bytes); + if self.buffer.len() >= MIN_PART_BYTES { + self.upload_buffered_part().await?; + } + Ok(()) + } + + pub(crate) async fn finish(self) -> Result<(), StorageError> { + if self.completed.is_empty() && self.buffer.is_empty() { + // S3 rejects CompleteMultipartUpload with zero parts: abort the + // upload and store the zero-byte object with a plain PUT. + let Self { + client, + bucket, + key, + upload_id, + .. + } = self; + abort_multipart(&client, &bucket, &key, &upload_id).await?; + client + .put_object() + .bucket(&bucket) + .key(&key) + .body(ByteStream::from_static(&[])) + .send() + .await + .map_err(|error| S3Error::Put { key, error })?; + return Ok(()); + } + // Capture the abort coordinates before the fallible tail consumes + // self, so ANY failure inside it (a final part upload, a part-count + // overflow, or the completion call) aborts the upload instead of + // leaking an uncompleted multipart upload that accrues storage cost + // forever. + let client = self.client.clone(); + let bucket = self.bucket.clone(); + let key = self.key.clone(); + let upload_id = self.upload_id.clone(); + if let Err(error) = self.complete().await { + // Best-effort abort; the original error is always returned and the + // abort's own failure is swallowed. + drop(abort_multipart(&client, &bucket, &key, &upload_id).await); + return Err(error); + } + Ok(()) + } + + /// The fallible completion tail: upload any final buffered part, then + /// complete the multipart upload. Kept as a single guarded region so a + /// future error path added here cannot forget the abort in + /// [`Self::finish`]. + async fn complete(mut self) -> Result<(), StorageError> { + if !self.buffer.is_empty() { + // The final part may be smaller than the 5 MiB minimum. + self.upload_buffered_part().await?; + } + let Self { + client, + bucket, + key, + upload_id, + completed, + .. + } = self; + let parts = CompletedMultipartUpload::builder() + .set_parts(Some(completed)) + .build(); + client + .complete_multipart_upload() + .bucket(&bucket) + .key(&key) + .upload_id(&upload_id) + .multipart_upload(parts) + .send() + .await + .map_err(|error| S3Error::CompleteMultipart { key, error })?; + Ok(()) + } + + pub(crate) async fn abort(self) -> Result<(), StorageError> { + let Self { + client, + bucket, + key, + upload_id, + .. + } = self; + abort_multipart(&client, &bucket, &key, &upload_id).await?; + Ok(()) + } + + /// Upload the entire pending buffer as the next part. + async fn upload_buffered_part(&mut self) -> Result<(), StorageError> { + let part_number = self + .part_number + .checked_add(1) + .filter(|part_number| *part_number <= MAX_PART_NUMBER) + .ok_or_else(|| S3Error::PartCountOverflow { + key: self.key.clone(), + })?; + let body = std::mem::take(&mut self.buffer); + let response = self + .client + .upload_part() + .bucket(&self.bucket) + .key(&self.key) + .upload_id(&self.upload_id) + .part_number(part_number) + .body(ByteStream::from(body)) + .send() + .await + .map_err(|error| S3Error::UploadPart { + key: self.key.clone(), + part: part_number, + error, + })?; + let e_tag = response + .e_tag() + .ok_or_else(|| S3Error::MissingETag { + key: self.key.clone(), + part: part_number, + })? + .to_owned(); + self.part_number = part_number; + self.completed.push( + CompletedPart::builder() + .part_number(part_number) + .e_tag(e_tag) + .build(), + ); + Ok(()) + } +} + +/// Map a `GetObject` failure: a missing key becomes +/// [`StorageError::NotFound`] (with the caller-relative key), anything else +/// wraps the original error. +fn map_get_error(key: &str, object_key: &str, error: SdkError) -> StorageError { + if error + .as_service_error() + .is_some_and(GetObjectError::is_no_such_key) + { + StorageError::NotFound { + key: key.to_owned(), + } + } else { + S3Error::Get { + key: object_key.to_owned(), + error, + } + .into() + } +} + +async fn abort_multipart( + client: &aws_sdk_s3::Client, + bucket: &str, + key: &str, + upload_id: &str, +) -> Result<(), S3Error> { + client + .abort_multipart_upload() + .bucket(bucket) + .key(key) + .upload_id(upload_id) + .send() + .await + .map_err(|error| S3Error::AbortMultipart { + key: key.to_owned(), + error, + })?; + Ok(()) +} + +/// Normalize a configured prefix: no leading/trailing/doubled slashes; +/// `None` and `"/"` become the empty (bucket-root) prefix. +fn normalize_prefix(prefix: Option) -> String { + prefix + .unwrap_or_default() + .split('/') + .filter(|component| !component.is_empty()) + .collect::>() + .join("/") +} + +/// Directory listings require a trailing slash on non-empty prefixes so S3 +/// common prefixes align on whole components. +fn normalize_dir_prefix(prefix: &str) -> String { + if prefix.is_empty() || prefix.ends_with('/') { + prefix.to_owned() + } else { + format!("{prefix}/") + } +} + +#[cfg(test)] +mod tests { + use pretty_assertions::assert_eq; + + use super::*; + + fn storage(prefix: Option<&str>, endpoint: Option<&str>) -> S3Storage { + S3Storage::new( + "test-bucket".to_owned(), + prefix.map(ToOwned::to_owned), + endpoint.map(ToOwned::to_owned), + None, + "test-access-key".to_owned(), + "test-secret-key", + ) + } + + #[test] + fn prefix_is_normalized() { + assert_eq!(storage(None, None).prefix, "", "None prefix"); + assert_eq!(storage(Some(""), None).prefix, "", "empty prefix"); + assert_eq!(storage(Some("/"), None).prefix, "", "slash-only prefix"); + assert_eq!( + storage(Some("/replica//prod/"), None).prefix, + "replica/prod", + "leading, trailing, and doubled slashes are removed" + ); + } + + #[test] + fn object_key_joins_with_single_slash() { + let scoped = storage(Some("replica"), None); + assert_eq!( + scoped.object_key("generations/g1/snapshot.json"), + "replica/generations/g1/snapshot.json", + "scoped key" + ); + let root = storage(None, None); + assert_eq!( + root.object_key("generations/g1/snapshot.json"), + "generations/g1/snapshot.json", + "root key has no leading slash" + ); + } + + #[test] + fn full_prefix_scopes_listings_to_the_root() { + let scoped = storage(Some("replica"), None); + // An empty caller prefix must still scope to `replica/`, never + // matching sibling prefixes like `replica2/`. + assert_eq!(scoped.full_prefix(""), "replica/", "empty caller prefix"); + assert_eq!( + scoped.full_prefix("generations/"), + "replica/generations/", + "nested caller prefix" + ); + let root = storage(None, None); + assert_eq!(root.full_prefix(""), "", "root empty prefix"); + assert_eq!( + root.full_prefix("generations/"), + "generations/", + "root nested prefix" + ); + } + + #[test] + fn root_prefix_strips_listed_keys() { + let scoped = storage(Some("replica"), None); + let listed = "replica/generations/g1/snapshot.json"; + assert_eq!( + listed.strip_prefix(scoped.root_prefix().as_str()), + Some("generations/g1/snapshot.json"), + "listed keys are returned caller-relative" + ); + } + + #[test] + fn dir_prefix_gets_trailing_slash() { + assert_eq!(normalize_dir_prefix(""), "", "empty stays empty"); + assert_eq!( + normalize_dir_prefix("generations"), + "generations/", + "missing slash is added" + ); + assert_eq!( + normalize_dir_prefix("generations/"), + "generations/", + "existing slash is kept" + ); + } + + #[test] + fn endpoint_override_builds_a_client() { + // MinIO/R2-style construction: endpoint plus default region; the + // client must build without panicking. + let scoped = storage(Some("replica"), Some("http://localhost:9000")); + assert_eq!(scoped.bucket, "test-bucket", "bucket is stored"); + } + + #[test] + fn next_list_page_decides_continue_done_or_error() { + assert!( + matches!( + next_list_page(Some(true), Some("token")), + ListPage::Continue(token) if token == "token" + ), + "truncated with a token must continue" + ); + assert!( + matches!(next_list_page(Some(false), None), ListPage::Done), + "not truncated must be done" + ); + assert!( + matches!(next_list_page(None, None), ListPage::Done), + "absent truncation flag must be done" + ); + assert!( + matches!(next_list_page(Some(false), Some("token")), ListPage::Done), + "a stray token without truncation is still done" + ); + assert!( + matches!(next_list_page(Some(true), None), ListPage::Truncated), + "truncated with no token must be an error" + ); + assert!( + matches!(next_list_page(Some(true), Some("")), ListPage::Truncated), + "truncated with an empty token must be an error" + ); + } + + #[test] + fn next_upload_page_decides_continue_done_or_error() { + assert!( + matches!( + next_upload_page(Some(true), Some("key"), Some("uid")), + UploadPage::Continue { key_marker, upload_id_marker } + if key_marker == "key" && upload_id_marker.as_deref() == Some("uid") + ), + "truncated with both markers must continue with both" + ); + assert!( + matches!( + next_upload_page(Some(true), Some("key"), None), + UploadPage::Continue { key_marker, upload_id_marker } + if key_marker == "key" && upload_id_marker.is_none() + ), + "truncated with only a key marker continues without an upload-id marker" + ); + assert!( + matches!(next_upload_page(Some(false), None, None), UploadPage::Done), + "not truncated must be done" + ); + assert!( + matches!(next_upload_page(None, None, None), UploadPage::Done), + "absent truncation flag must be done" + ); + assert!( + matches!( + next_upload_page(Some(true), None, Some("uid")), + UploadPage::Truncated + ), + "truncated with no key marker must be an error" + ); + assert!( + matches!( + next_upload_page(Some(true), Some(""), None), + UploadPage::Truncated + ), + "truncated with an empty key marker must be an error" + ); + } + + #[tokio::test] + async fn object_methods_reject_escaping_keys() { + // Validation happens before any network call, so a fake client is + // fine: an escaping key never reaches S3. + let s3 = storage(Some("replica"), None); + for key in ["../escape", "/leading", "gen/../escape", "gen//doubled"] { + assert!( + matches!( + s3.put(key, Bytes::from_static(b"x")).await, + Err(StorageError::InvalidKey { .. }) + ), + "put({key}) must be InvalidKey before any request" + ); + assert!( + matches!(s3.get(key).await, Err(StorageError::InvalidKey { .. })), + "get({key}) must be InvalidKey" + ); + assert!( + matches!(s3.delete(key).await, Err(StorageError::InvalidKey { .. })), + "delete({key}) must be InvalidKey" + ); + assert!( + matches!( + s3.start_multipart(key).await, + Err(StorageError::InvalidKey { .. }) + ), + "start_multipart({key}) must be InvalidKey" + ); + } + } +} diff --git a/plus/bencher_replica/src/segment.rs b/plus/bencher_replica/src/segment.rs new file mode 100644 index 000000000..a9edd9ae5 --- /dev/null +++ b/plus/bencher_replica/src/segment.rs @@ -0,0 +1,189 @@ +//! WAL segment compression. +//! +//! A segment is a run of raw WAL bytes `[start, end)` that always ends on a +//! commit-frame boundary (commit alignment is enforced upstream by +//! [`crate::wal::WalScanner::next_committed`]). Segments are stored +//! zstd-compressed with zstd's built-in content checksum enabled, so a +//! corrupted object fails loudly at decode time rather than yielding garbage +//! frames. + +use std::io::{Read as _, Write as _}; + +/// Maximum raw (uncompressed) WAL bytes per segment. A single transaction +/// larger than this is shipped as one oversized segment: commit atomicity +/// beats the cap. +pub const SEGMENT_MAX_BYTES: u64 = 8 * 1024 * 1024; + +/// zstd compression level for segments and snapshots. +pub const ZSTD_LEVEL: i32 = 3; + +/// Hard cap on decompressed segment output. Segments are bounded well below +/// this (see [`SEGMENT_MAX_BYTES`]; a single oversized transaction is the only +/// exception), so anything larger is a corrupt or hostile object +/// (zstd-bomb), not data: refuse instead of allocating without bound. +pub const MAX_DECOMPRESSED_BYTES: u64 = 4 * 1024 * 1024 * 1024; + +#[derive(Debug, thiserror::Error)] +pub enum SegmentError { + #[error("Failed to compress segment: {0}")] + Compress(std::io::Error), + #[error("Failed to decompress segment: {0}")] + Decompress(std::io::Error), + #[error("Decompressed segment exceeds the {max_bytes} byte cap")] + TooLarge { max_bytes: u64 }, +} + +/// Compress raw WAL bytes (zstd, content checksum on). +pub fn compress_segment(raw: &[u8]) -> Result, SegmentError> { + let mut encoder = + zstd::stream::Encoder::new(Vec::new(), ZSTD_LEVEL).map_err(SegmentError::Compress)?; + encoder + .include_checksum(true) + .map_err(SegmentError::Compress)?; + encoder.write_all(raw).map_err(SegmentError::Compress)?; + encoder.finish().map_err(SegmentError::Compress) +} + +/// Decompress a segment object back to raw WAL bytes, verifying the zstd +/// content checksum and refusing output beyond [`MAX_DECOMPRESSED_BYTES`]. +pub fn decompress_segment(compressed: &[u8]) -> Result, SegmentError> { + decompress_segment_with_cap(compressed, MAX_DECOMPRESSED_BYTES) +} + +/// Decompress with an explicit output cap; reading stops one byte past the +/// cap, so a zstd bomb never allocates more than `cap + 1` bytes. Restore +/// passes the segment's exact expected size so an oversized or hostile +/// object is rejected up front rather than after a multi-gigabyte allocation. +pub(crate) fn decompress_segment_with_cap( + compressed: &[u8], + cap: u64, +) -> Result, SegmentError> { + let decoder = + zstd::stream::Decoder::with_buffer(compressed).map_err(SegmentError::Decompress)?; + let mut raw = Vec::new(); + decoder + .take(cap.saturating_add(1)) + .read_to_end(&mut raw) + .map_err(SegmentError::Decompress)?; + if u64::try_from(raw.len()).unwrap_or(u64::MAX) > cap { + return Err(SegmentError::TooLarge { max_bytes: cap }); + } + Ok(raw) +} + +#[cfg(test)] +mod tests { + use pretty_assertions::assert_eq; + + use super::{ + MAX_DECOMPRESSED_BYTES, SegmentError, compress_segment, decompress_segment, + decompress_segment_with_cap, + }; + + /// Deterministic pseudo-random bytes (xorshift64): incompressible enough + /// that the corruption tests have a real compressed body to flip bits in. + fn noise_bytes(len: usize) -> Vec { + let mut state = 0x9e37_79b9_7f4a_7c15u64; + let mut bytes = Vec::with_capacity(len); + for _ in 0..len { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + bytes.push(u8::try_from(state >> 56).unwrap()); + } + bytes + } + + #[test] + fn zstd_roundtrip_identity() { + let raw = noise_bytes(64 * 1024); + let compressed = compress_segment(&raw).unwrap(); + let decompressed = decompress_segment(&compressed).unwrap(); + assert_eq!(decompressed, raw); + } + + #[test] + fn compressed_frame_has_content_checksum_flag() { + let compressed = compress_segment(b"bencher replica segment").unwrap(); + // zstd frame magic number, little-endian on the wire. + assert_eq!(&compressed[..4], &[0x28, 0xb5, 0x2f, 0xfd]); + // Frame_Header_Descriptor bit 2 is the Content_Checksum_flag. + assert_eq!(compressed[4] & 0b100, 0b100); + } + + #[test] + fn zstd_corrupt_payload_detected() { + let raw = noise_bytes(8 * 1024); + let mut compressed = compress_segment(&raw).unwrap(); + assert!(compressed.len() > 1024); + // Flip a byte well past the frame header, in the compressed body. + compressed[512] ^= 0xff; + decompress_segment(&compressed).unwrap_err(); + } + + #[test] + fn zstd_corrupt_checksum_detected() { + let raw = noise_bytes(4 * 1024); + let mut compressed = compress_segment(&raw).unwrap(); + // The content checksum is the final 4 bytes of the frame; corrupting + // it proves the checksum is both present and verified on decode. + let last = compressed.len() - 1; + compressed[last] ^= 0xff; + decompress_segment(&compressed).unwrap_err(); + } + + #[test] + fn empty_input_roundtrip() { + let compressed = compress_segment(&[]).unwrap(); + let decompressed = decompress_segment(&compressed).unwrap(); + assert_eq!(decompressed, Vec::::new()); + } + + #[test] + fn compression_actually_compresses() { + let raw = b"bencher replica wal frame ".repeat(40_330); + assert!(raw.len() > 1024 * 1024); + let compressed = compress_segment(&raw).unwrap(); + // Repetitive input must shrink by orders of magnitude + // (multiplication instead of division to keep clippy quiet). + assert!(compressed.len() * 100 < raw.len()); + } + + #[test] + fn decompress_rejects_output_beyond_cap() { + let raw = vec![0u8; 1000]; + let compressed = compress_segment(&raw).unwrap(); + let err = decompress_segment_with_cap(&compressed, 100).unwrap_err(); + assert!(matches!(err, SegmentError::TooLarge { max_bytes: 100 })); + // At or below the cap decompresses fine. + let decompressed = decompress_segment_with_cap(&compressed, 1000).unwrap(); + assert_eq!(decompressed, raw); + assert!(MAX_DECOMPRESSED_BYTES >= u64::from(u32::MAX)); + } + + #[test] + fn decompress_garbage_is_error() { + let err = decompress_segment(b"definitely not a zstd frame").unwrap_err(); + assert!(matches!(err, SegmentError::Decompress(_))); + } + + #[test] + fn exact_size_cap_rejects_oversize_without_full_allocation() { + // Restore knows each segment's exact raw size from its key and passes + // `expected + 1` as the cap, so a corrupt object that decompresses + // larger fails immediately instead of forcing the generic 4 GiB + // allocation. A highly compressible body stands in for a zstd bomb. + let expected = 4096u64; + let raw = vec![0u8; usize::try_from(expected).unwrap() + 8192]; + let compressed = compress_segment(&raw).unwrap(); + let err = decompress_segment_with_cap(&compressed, expected + 1).unwrap_err(); + assert!(matches!(err, SegmentError::TooLarge { max_bytes } if max_bytes == expected + 1)); + // A body of exactly the expected size decompresses fine at that cap. + let exact = vec![0u8; usize::try_from(expected).unwrap()]; + let compressed = compress_segment(&exact).unwrap(); + assert_eq!( + decompress_segment_with_cap(&compressed, expected + 1).unwrap(), + exact + ); + } +} diff --git a/plus/bencher_replica/src/snapshot.rs b/plus/bencher_replica/src/snapshot.rs new file mode 100644 index 000000000..0ef4cefac --- /dev/null +++ b/plus/bencher_replica/src/snapshot.rs @@ -0,0 +1,514 @@ +//! Lock-free, incremental generation snapshots. +//! +//! The snapshot body comes from a SINGLE-STEP `SQLite` online backup into a +//! local scratch file: one `sqlite3_backup_step(-1)` call copies the entire +//! committed state through the pager under one read transaction, so the +//! result is transactionally consistent NO MATTER WHO checkpoints +//! concurrently (a raw file copy would be torn by any external backfill, +//! which matters in shadow mode where Litestream checkpoints freely, and +//! whenever an operator tool touches the database). Writers are never +//! blocked: the backup holds only a read transaction, and a single step +//! never observes writer restarts. +//! +//! The upload is step-driven: each [`copy_step`] call processes at most the +//! throttle budget (`snapshot_throttle_mib * sync_interval`) of the SCRATCH +//! file in 1 MiB reads through a streaming zstd encoder into a multipart +//! upload, so the engine stays responsive between steps. The state machine +//! itself ([`SnapshotJob`]) is advanced by `SyncEngine::snapshot_step`: +//! +//! `ShipTail` (drain committed frames so the OLD generation is complete) +//! -> `CreateGeneration` (new id, online backup into the scratch file, +//! record boundary salts + the mandatory-replay offset, open the multipart +//! upload) -> `Copying` (budgeted steps over the scratch) -> `Finalize` +//! (finish the upload, re-check the boundary salts, PUT `snapshot.json` +//! LAST as the atomic commit marker, reset the position to the new +//! generation's epoch 0, delete the scratch, prune). +//! +//! ## The epoch-0 lazy binding rule +//! +//! The new generation's epoch 0 must bind to the salt cycle from which the +//! first post-snapshot ship actually happens. If the WAL restarts before +//! any epoch-0 segment ships, that is only possible when the old cycle was +//! fully backfilled and hence fully contained in the snapshot just taken, +//! so epoch 0 is REBOUND to the new cycle instead of leaving an empty +//! epoch. This preserves the restore invariant that epochs `0..N` are +//! contiguous and non-empty. Concretely: (a) boundary salts are recorded at +//! `CreateGeneration`; (b) the body is copied and uploaded; (c) BEFORE +//! `snapshot.json` is uploaded the WAL header is re-read and, since nothing +//! has shipped into the new generation yet by construction, a changed salt +//! rebinds the boundary; (d) `snapshot.json` is uploaded LAST. Any restart +//! AFTER that is handled by the ship path's offset-0 rebind, which keeps +//! the position's epoch number and adopts the new salts. + +use std::fs::File; +use std::io::{ErrorKind, Read as _, Seek as _, SeekFrom, Write as _}; + +use camino::{Utf8Path, Utf8PathBuf}; +use sha2::Sha256; + +use crate::position::GenerationId; +use crate::segment::ZSTD_LEVEL; +use crate::snapshot_meta::SnapshotMetaError; +use crate::storage::MultipartUpload; + +/// One mebibyte: the copy read granularity. +pub(crate) const MIB: u64 = 1024 * 1024; + +/// Incomplete generations (no `snapshot.json`) older than this are crashed +/// snapshots and get pruned. +pub(crate) const STALE_INCOMPLETE_SECS: i64 = 86_400; + +/// The streaming snapshot encoder: compressed output accumulates in the +/// inner buffer and is drained into multipart parts after each step. +pub(crate) type SnapshotEncoder = zstd::stream::Encoder<'static, Vec>; + +/// The status `SyncEngine::snapshot_step` reports to its driver. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum SnapshotStatus { + /// More steps remain; call again. + InProgress, + /// `snapshot.json` is uploaded and the position now points at the new + /// generation's epoch 0. + Finished, +} + +#[derive(Debug, thiserror::Error)] +pub enum SnapshotError { + #[error("Failed to back up the database into the snapshot scratch ({path}): {error}")] + Backup { + path: Utf8PathBuf, + error: rusqlite::Error, + }, + #[error("Failed to prepare the snapshot scratch file ({path}): {error}")] + Scratch { + path: Utf8PathBuf, + error: std::io::Error, + }, + #[error("Failed to read the snapshot scratch file ({path}): {error}")] + DbRead { + path: Utf8PathBuf, + error: std::io::Error, + }, + #[error("Failed to compress the snapshot stream: {0}")] + Compress(std::io::Error), + #[error("Snapshot backup did not complete in a single step: {outcome:?}")] + BackupIncomplete { + outcome: rusqlite::backup::StepResult, + }, + #[error("Snapshot meta: {0}")] + Meta(#[from] SnapshotMetaError), +} + +/// The snapshot state machine, held by the engine between steps. +pub(crate) enum SnapshotJob { + /// Drain committed WAL frames into the old generation first, so it + /// remains a complete restore point under retention. + ShipTail, + /// Fix the new generation id, run the online backup into the scratch + /// file, record the boundary; open the multipart upload. + CreateGeneration, + /// Budgeted upload steps over the scratch file through the zstd encoder. + Copying(Box), + /// Finish the upload, verify the boundary, PUT `snapshot.json` LAST. + Finalize(Box), +} + +/// In-flight copy state. +pub(crate) struct CopyJob { + pub generation: GenerationId, + /// WAL salts recorded at `CreateGeneration` (may rebind at finalize). + pub boundary_salt: (u32, u32), + /// The mandatory-replay threshold: the committed WAL extent read AFTER + /// the backup. It is always at least as large as the extent the backup + /// body captured (a frame committed after the backup only OVERSTATES it, + /// the safe direction), and restore requires the boundary epoch's + /// segments to reach at least this far before replaying it. Overstating + /// merely makes restore fall back to the snapshot alone; understating + /// (a value below the backup extent) could let a prefix replay that + /// regresses pages the snapshot already holds, so the read must never + /// move before the backup. + pub boundary_offset: u64, + pub upload: MultipartUpload, + /// SHA-256 over the COMPRESSED stream, part by part. + pub hasher: Sha256, + pub encoder: SnapshotEncoder, + /// The backup scratch file being uploaded. + pub scratch: Utf8PathBuf, + /// Scratch file length in bytes (the uncompressed snapshot size). + pub db_len: u64, + /// Next unuploaded byte offset in the scratch file. + pub offset: u64, + /// Database page size (informational, stored in `snapshot.json`). + pub page_size: u32, + /// LIVE database fingerprint right after the backup; re-checked at + /// finalize to detect external checkpoints mid-snapshot (sole mode). + pub live_fingerprint: (u64, [u8; 4]), +} + +/// Everything needed to commit the generation after the copy. +pub(crate) struct FinalizeJob { + pub generation: GenerationId, + pub boundary_salt: (u32, u32), + pub boundary_offset: u64, + pub upload: MultipartUpload, + /// Hex SHA-256 of the compressed snapshot object. + pub sha256: String, + /// Uncompressed database size in bytes. + pub db_bytes: u64, + pub page_size: u32, + /// The scratch file to delete once the generation commits. + pub scratch: Utf8PathBuf, + /// LIVE database fingerprint right after the backup. + pub live_fingerprint: (u64, [u8; 4]), +} + +/// The result of one budgeted copy step. +pub(crate) enum CopyStepResult { + /// More database bytes remain. + Continue { + encoder: SnapshotEncoder, + hasher: Sha256, + /// Compressed bytes drained this step (may be empty when zstd is + /// still buffering). + compressed: Vec, + /// Next uncopied byte offset. + offset: u64, + }, + /// The whole file was copied and the zstd frame is finished. + Done { + hasher: Sha256, + /// The final compressed tail (frame epilogue included). + compressed: Vec, + }, +} + +/// A fresh streaming encoder for one snapshot body (content checksum on, so +/// a corrupted object fails loudly at restore). +pub(crate) fn new_snapshot_encoder() -> Result { + let mut encoder = + zstd::stream::Encoder::new(Vec::new(), ZSTD_LEVEL).map_err(SnapshotError::Compress)?; + encoder + .include_checksum(true) + .map_err(SnapshotError::Compress)?; + Ok(encoder) +} + +/// Run a SINGLE-STEP online backup of the database into `scratch`: one +/// `sqlite3_backup_step(-1)` copies the whole committed state through the +/// pager under one read transaction. Writers proceed concurrently (a read +/// lock only), no restart can occur within the single step, and the result +/// is transactionally consistent regardless of concurrent checkpoints. +/// Runs inside `spawn_blocking`. +pub(crate) fn backup_to_scratch( + db_path: &Utf8Path, + scratch: &Utf8Path, +) -> Result<(), SnapshotError> { + // A stale scratch from a crashed snapshot must not leak into the copy. + match std::fs::remove_file(scratch) { + Ok(()) => {}, + Err(error) if error.kind() == ErrorKind::NotFound => {}, + Err(error) => { + return Err(SnapshotError::Scratch { + path: scratch.to_owned(), + error, + }); + }, + } + let backup_err = |error| SnapshotError::Backup { + path: scratch.to_owned(), + error, + }; + let src = rusqlite::Connection::open(db_path).map_err(backup_err)?; + // The backup source is a reader, but invariant I2 applies to every + // connection defensively. + let _pages: i64 = src + .query_row("PRAGMA wal_autocheckpoint = 0", [], |row| row.get(0)) + .map_err(backup_err)?; + let mut dst = rusqlite::Connection::open(scratch).map_err(backup_err)?; + let backup = rusqlite::backup::Backup::new(&src, &mut dst).map_err(backup_err)?; + let state = backup.step(-1).map_err(backup_err)?; + // A single `step(-1)` copies the whole database in one pass, so anything + // but `Done` (legitimately `Busy`/`Locked`) is surfaced honestly rather + // than disguised as a rusqlite row-count error. + require_backup_done(state)?; + drop(backup); + drop(dst); + Ok(()) +} + +/// A single-step online backup either finishes (`Done`) or is reporting a +/// live contention outcome (`Busy`/`Locked`, or a `More` that a one-pass +/// step should never return). The non-`Done` cases become +/// [`SnapshotError::BackupIncomplete`]. +fn require_backup_done(outcome: rusqlite::backup::StepResult) -> Result<(), SnapshotError> { + use rusqlite::backup::StepResult; + + match outcome { + StepResult::Done => Ok(()), + // `StepResult` is `#[non_exhaustive]`, so the trailing `_` is required + // even with every present variant listed. + StepResult::More | StepResult::Busy | StepResult::Locked | _ => { + Err(SnapshotError::BackupIncomplete { outcome }) + }, + } +} + +/// A cheap fingerprint of the LIVE database file (length plus header change +/// counter, bytes 24..28). In sole mode nothing legitimate mutates the file +/// during a snapshot (the engine is sequential and is the only +/// checkpointer), so any change between backup and finalize proves an +/// external checkpointer that may have buried post-backup frames: the +/// finalize step then schedules a follow-up generation to recapture them. +pub(crate) fn live_db_fingerprint(db_path: &Utf8Path) -> Result<(u64, [u8; 4]), SnapshotError> { + let read_err = |error| SnapshotError::DbRead { + path: db_path.to_owned(), + error, + }; + let mut file = File::open(db_path).map_err(read_err)?; + let len = file.metadata().map_err(read_err)?.len(); + let mut header = [0u8; 28]; + let mut change_counter = [0u8; 4]; + match file.read_exact(&mut header) { + Ok(()) => { + if let Some(counter) = header.get(24..28) { + change_counter.copy_from_slice(counter); + } + }, + Err(error) if error.kind() == ErrorKind::UnexpectedEof => {}, + Err(error) => return Err(read_err(error)), + } + Ok((len, change_counter)) +} + +/// Read the scratch file length plus the database page size (header bytes +/// 16..18 big-endian; the stored value 1 encodes 65536). A file shorter +/// than the 100-byte header yields a zero page size; the length still +/// bounds the upload. +pub(crate) fn read_scratch_info(scratch: &Utf8Path) -> Result<(u64, u32), SnapshotError> { + let read_err = |error| SnapshotError::DbRead { + path: scratch.to_owned(), + error, + }; + let mut file = File::open(scratch).map_err(read_err)?; + let db_len = file.metadata().map_err(read_err)?.len(); + let mut header = [0u8; 28]; + let mut page_size = 0u32; + match file.read_exact(&mut header) { + Ok(()) => { + let high = header.get(16).copied().unwrap_or(0); + let low = header.get(17).copied().unwrap_or(0); + let raw = (u32::from(high) << 8) | u32::from(low); + page_size = if raw == 1 { 0x0001_0000 } else { raw }; + }, + Err(error) if error.kind() == ErrorKind::UnexpectedEof => {}, + Err(error) => return Err(read_err(error)), + } + Ok((db_len, page_size)) +} + +/// Copy up to `budget` bytes of the database file (1 MiB reads) into the +/// encoder, then drain whatever compressed output accumulated. Runs inside +/// `spawn_blocking`. +pub(crate) fn copy_step( + db_path: &Utf8Path, + mut encoder: SnapshotEncoder, + mut hasher: Sha256, + offset: u64, + db_len: u64, + budget: u64, +) -> Result { + use sha2::Digest as _; + + let read_err = |error| SnapshotError::DbRead { + path: db_path.to_owned(), + error, + }; + let mut cursor = offset; + let mut remaining = db_len.saturating_sub(offset).min(budget.max(MIB)); + if remaining > 0 { + let mut file = File::open(db_path).map_err(read_err)?; + file.seek(SeekFrom::Start(offset)).map_err(read_err)?; + let mut buffer = vec![0u8; usize::try_from(remaining.min(MIB)).unwrap_or(1)]; + while remaining > 0 { + let want = usize::try_from(remaining.min(MIB)).unwrap_or(1); + let Some(chunk) = buffer.get_mut(..want) else { + break; + }; + file.read_exact(chunk).map_err(read_err)?; + encoder.write_all(chunk).map_err(SnapshotError::Compress)?; + cursor = cursor.saturating_add(u64::try_from(want).unwrap_or(0)); + remaining = remaining.saturating_sub(u64::try_from(want).unwrap_or(remaining)); + } + } + if cursor >= db_len { + let compressed = encoder.finish().map_err(SnapshotError::Compress)?; + hasher.update(&compressed); + Ok(CopyStepResult::Done { hasher, compressed }) + } else { + let compressed = std::mem::take(encoder.get_mut()); + hasher.update(&compressed); + Ok(CopyStepResult::Continue { + encoder, + hasher, + compressed, + offset: cursor, + }) + } +} + +#[cfg(test)] +mod tests { + use camino::Utf8Path; + use pretty_assertions::assert_eq; + use sha2::{Digest as _, Sha256}; + + use super::{ + CopyStepResult, MIB, SnapshotError, backup_to_scratch, copy_step, new_snapshot_encoder, + read_scratch_info, require_backup_done, + }; + use crate::segment::decompress_segment; + + fn tempdir_path(dir: &tempfile::TempDir) -> &Utf8Path { + Utf8Path::from_path(dir.path()).unwrap() + } + + /// Deterministic pseudo-random bytes (xorshift64), incompressible + /// enough to exercise multi-part output. + fn noise_bytes(len: usize) -> Vec { + let mut state = 0x1234_5678_9abc_def0u64; + let mut bytes = Vec::with_capacity(len); + for _ in 0..len { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + bytes.push(u8::try_from(state >> 56).unwrap()); + } + bytes + } + + /// Drive `copy_step` to completion with the given budget; returns the + /// concatenated compressed parts, the digest, and the step count. + fn copy_all(db_path: &Utf8Path, db_len: u64, budget: u64) -> (Vec, String, u32) { + let mut encoder = new_snapshot_encoder().unwrap(); + let mut hasher = Sha256::new(); + let mut offset = 0u64; + let mut compressed_all = Vec::new(); + let mut steps = 0u32; + loop { + steps += 1; + assert!(steps < 1000, "copy must terminate"); + match copy_step(db_path, encoder, hasher, offset, db_len, budget).unwrap() { + CopyStepResult::Continue { + encoder: next_encoder, + hasher: next_hasher, + compressed, + offset: next_offset, + } => { + assert!(next_offset > offset, "each step advances"); + compressed_all.extend_from_slice(&compressed); + encoder = next_encoder; + hasher = next_hasher; + offset = next_offset; + }, + CopyStepResult::Done { hasher, compressed } => { + compressed_all.extend_from_slice(&compressed); + return (compressed_all, hex::encode(hasher.finalize()), steps); + }, + } + } + } + + #[test] + fn copy_step_round_trips_and_hashes_compressed_stream() { + let tmp = tempfile::tempdir().unwrap(); + let db_path = tempdir_path(&tmp).join("fake.db"); + let body = noise_bytes(usize::try_from(3 * MIB + 500).unwrap()); + std::fs::write(&db_path, &body).unwrap(); + let db_len = u64::try_from(body.len()).unwrap(); + + let (compressed, sha256, steps) = copy_all(&db_path, db_len, MIB); + assert!(steps > 3, "a 3.5 MiB copy at 1 MiB budget takes >3 steps"); + assert_eq!( + decompress_segment(&compressed).unwrap(), + body, + "the concatenated parts form one valid zstd frame" + ); + assert_eq!( + sha256, + hex::encode(Sha256::digest(&compressed)), + "the incremental hash covers exactly the compressed bytes" + ); + + // A one-shot copy (huge budget) produces the identical content. + let (_, sha_oneshot, steps) = copy_all(&db_path, db_len, u64::MAX); + assert_eq!(steps, 1, "one budgeted step suffices"); + drop(sha_oneshot); + } + + #[test] + fn copy_step_empty_file_finishes_immediately() { + let tmp = tempfile::tempdir().unwrap(); + let db_path = tempdir_path(&tmp).join("empty.db"); + std::fs::write(&db_path, b"").unwrap(); + let (compressed, _, steps) = copy_all(&db_path, 0, MIB); + assert_eq!(steps, 1, "an empty file is one step"); + assert_eq!( + decompress_segment(&compressed).unwrap(), + Vec::::new(), + "an empty frame decodes to nothing" + ); + } + + #[test] + fn backup_to_scratch_copies_committed_state_including_wal() { + let tmp = tempfile::tempdir().unwrap(); + let db_path = tempdir_path(&tmp).join("source.db"); + let scratch = tempdir_path(&tmp).join("source.db.snapshot-scratch"); + let conn = rusqlite::Connection::open(&db_path).unwrap(); + conn.execute_batch( + "PRAGMA journal_mode = WAL; + PRAGMA wal_autocheckpoint = 0; + CREATE TABLE t (id INTEGER PRIMARY KEY, data TEXT); + INSERT INTO t (data) VALUES ('committed'), ('in the wal');", + ) + .unwrap(); + // The rows above live only in the WAL (no checkpoint ran): the + // backup must still include them, unlike a raw file copy. + backup_to_scratch(&db_path, &scratch).unwrap(); + let copy = rusqlite::Connection::open(&scratch).unwrap(); + let rows: i64 = copy + .query_row("SELECT count(*) FROM t", [], |row| row.get(0)) + .unwrap(); + assert_eq!(rows, 2, "WAL content is part of the backup"); + + // A stale scratch is replaced, not appended to. + backup_to_scratch(&db_path, &scratch).unwrap(); + let (db_len, page_size) = read_scratch_info(&scratch).unwrap(); + assert!(db_len > 0, "scratch has content"); + assert_eq!(page_size, 4096, "page size parsed from the header"); + } + + #[test] + fn require_backup_done_rejects_incomplete_outcomes() { + use rusqlite::backup::StepResult; + + require_backup_done(StepResult::Done).unwrap(); + for outcome in [StepResult::Busy, StepResult::Locked, StepResult::More] { + let err = require_backup_done(outcome).unwrap_err(); + assert!( + matches!(err, SnapshotError::BackupIncomplete { outcome: reported } if reported == outcome), + "a non-Done step is reported honestly, got {err}" + ); + } + } + + #[test] + fn read_scratch_info_short_file_yields_zero_page_size() { + let tmp = tempfile::tempdir().unwrap(); + let db_path = tempdir_path(&tmp).join("short.db"); + std::fs::write(&db_path, b"tiny").unwrap(); + let (db_len, page_size) = read_scratch_info(&db_path).unwrap(); + assert_eq!(page_size, 0, "no header, no page size"); + assert_eq!(db_len, 4, "length still bounds the upload"); + } +} diff --git a/plus/bencher_replica/src/snapshot_meta.rs b/plus/bencher_replica/src/snapshot_meta.rs new file mode 100644 index 000000000..49cc08f80 --- /dev/null +++ b/plus/bencher_replica/src/snapshot_meta.rs @@ -0,0 +1,164 @@ +//! The `snapshot.json` object: a generation's atomic commit marker. +//! +//! Uploaded LAST, after the snapshot body (`snapshot.db.zst`) has fully +//! landed: its presence is what makes a generation visible to restore. +//! Generations without a `snapshot.json` are invisible and eventually +//! pruned. + +use serde::{Deserialize, Serialize}; + +/// Current snapshot meta schema version. +pub const SNAPSHOT_META_VERSION: u32 = 1; + +#[derive(Debug, thiserror::Error)] +pub enum SnapshotMetaError { + #[error("Failed to serialize snapshot meta: {0}")] + Serialize(serde_json::Error), + #[error("Failed to parse snapshot meta: {0}")] + Parse(serde_json::Error), + #[error("Snapshot meta version {found} is not the supported version {SNAPSHOT_META_VERSION}")] + Version { found: u32 }, +} + +/// Generation commit marker and restore metadata. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct SnapshotMeta { + pub version: u32, + /// RFC 3339 creation timestamp (informational). + pub created: String, + /// Uncompressed database size in bytes at snapshot time. + pub db_bytes: u64, + /// Database page size. + pub page_size: u32, + /// Hex SHA-256 of the COMPRESSED snapshot object, verified on restore. + pub sha256: String, + /// The WAL epoch boundary: restore replays all segments of epochs + /// `>= wal_boundary.epoch` (always 0 by construction; kept explicit). + pub wal_boundary: WalBoundary, +} + +/// The WAL salt cycle current when the snapshot began. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub struct WalBoundary { + pub epoch: u64, + pub salt1: u32, + pub salt2: u32, + /// The mandatory-replay threshold for the boundary epoch, a raw WAL byte + /// offset at least as large as the committed extent the snapshot body + /// captured (it is read AFTER the backup, so a frame committed after the + /// backup only OVERSTATES it, the safe direction). Restore replays the + /// boundary epoch only when its available segments reach at least this + /// far; below it, the snapshot alone (a consistent committed state) is + /// used. Overstating merely forces that snapshot-only fallback; + /// understating (a value below the backup extent) could let a PREFIX of + /// the boundary epoch replay and regress pages the snapshot already + /// holds at a newer state, so this must never be measured before the + /// backup. + /// Deliberately NOT `#[serde(default)]`: defaulting a missing field to 0 + /// would silently disable this guard (0 means "replay any prefix"), so a + /// marker without it must fail to parse, which makes its generation + /// unrestorable rather than wrongly restorable. + pub offset: u64, +} + +impl SnapshotMeta { + /// Serialize for upload. + pub fn to_bytes(&self) -> Result, SnapshotMetaError> { + serde_json::to_vec_pretty(self).map_err(SnapshotMetaError::Serialize) + } + + /// Parse a downloaded `snapshot.json`. + /// + /// A version other than [`SNAPSHOT_META_VERSION`] is rejected even when + /// the fields parse: the marker's values feed the decompression cap and + /// the mandatory-replay boundary, so a structurally-compatible marker + /// whose field SEMANTICS changed must make its generation unrestorable + /// (skipped, falling back to an older generation) rather than misread. + /// Unknown ADDITIVE fields within the same version remain tolerated for + /// forward compatibility. + pub fn from_bytes(bytes: &[u8]) -> Result { + let meta: Self = serde_json::from_slice(bytes).map_err(SnapshotMetaError::Parse)?; + if meta.version != SNAPSHOT_META_VERSION { + return Err(SnapshotMetaError::Version { + found: meta.version, + }); + } + Ok(meta) + } +} + +#[cfg(test)] +mod tests { + use pretty_assertions::assert_eq; + + use super::{SNAPSHOT_META_VERSION, SnapshotMeta, WalBoundary}; + + fn test_meta() -> SnapshotMeta { + SnapshotMeta { + version: SNAPSHOT_META_VERSION, + created: "2026-07-10T14:59:00Z".to_owned(), + db_bytes: 4_350_000_000, + page_size: 4096, + sha256: "ab".repeat(32), + wal_boundary: WalBoundary { + epoch: 0, + salt1: 0x9d2f_1c4a, + salt2: 0x8b3e_6f70, + offset: 0, + }, + } + } + + #[test] + fn round_trips() { + let meta = test_meta(); + let bytes = meta.to_bytes().unwrap(); + assert_eq!(SnapshotMeta::from_bytes(&bytes).unwrap(), meta); + } + + #[test] + fn rejects_garbage() { + SnapshotMeta::from_bytes(b"not json").unwrap_err(); + } + + #[test] + fn version_mismatch_fails_to_parse() { + let mut meta = test_meta(); + meta.version = SNAPSHOT_META_VERSION + 1; + let bytes = meta.to_bytes().unwrap(); + let error = SnapshotMeta::from_bytes(&bytes) + .expect_err("a structurally-compatible future version must be rejected"); + assert!( + matches!( + error, + super::SnapshotMetaError::Version { found } if found == SNAPSHOT_META_VERSION + 1 + ), + "expected a version error, got: {error:?}" + ); + } + + #[test] + fn missing_boundary_offset_fails_to_parse() { + // The offset is the anti-regression guard; a marker without it must + // be unparseable (generation skipped), never defaulted to 0, which + // would mean "replay any prefix". + let mut value: serde_json::Value = + serde_json::from_slice(&test_meta().to_bytes().unwrap()).unwrap(); + let boundary = value + .get_mut("wal_boundary") + .and_then(serde_json::Value::as_object_mut) + .expect("wal_boundary object"); + boundary.remove("offset").expect("offset present"); + let bytes = serde_json::to_vec(&value).unwrap(); + SnapshotMeta::from_bytes(&bytes).expect_err("a marker without an offset must not parse"); + } + + #[test] + fn unknown_fields_are_ignored_for_forward_compat() { + let mut value: serde_json::Value = + serde_json::from_slice(&test_meta().to_bytes().unwrap()).unwrap(); + value["future_field"] = serde_json::json!("ignored"); + let bytes = serde_json::to_vec(&value).unwrap(); + assert_eq!(SnapshotMeta::from_bytes(&bytes).unwrap(), test_meta()); + } +} diff --git a/plus/bencher_replica/src/storage.rs b/plus/bencher_replica/src/storage.rs new file mode 100644 index 000000000..c19198e97 --- /dev/null +++ b/plus/bencher_replica/src/storage.rs @@ -0,0 +1,459 @@ +//! Replica storage backends: enum dispatch over local filesystem and +//! S3-compatible object storage (plus a fault-injection wrapper for tests). +//! +//! Semantics pinned here and enforced by the contract test suite +//! (`tests/storage_contract.rs`, run against every backend): +//! +//! - Keys are `/`-separated relative paths with no leading slash, no `.` or +//! `..` components, and no empty components. The object methods +//! (`put`, `get`, `get_stream`, `delete`, `start_multipart`) reject anything +//! else with [`StorageError::InvalidKey`] before touching the backend, so a +//! crate-public caller cannot escape the storage root. (`list`, `list_dirs`, +//! and `delete_prefix` take a prefix, which may legitimately end +//! mid-component, and so are not validated this way.) +//! - `put` is atomically visible AND durable when it returns `Ok`: a reader +//! never observes a partial object under its final key, and an acknowledged +//! write survives a power loss (crate invariant I1 gates checkpoints on +//! frames being durably uploaded). Local: the temp file is fsynced, renamed +//! into place, and the parent directory is fsynced so the rename itself is +//! durable; newly created parent directories are fsynced as they are made. +//! S3: durability is the 200 response to the single PUT (or, for a streaming +//! upload, to `CompleteMultipartUpload`). +//! - `get` of a missing key returns [`StorageError::NotFound`], never empty +//! bytes. +//! - `list` returns full keys, sorted lexicographically, with pagination +//! handled internally. A backend failure is an `Err`, NEVER an empty list: +//! conflating the two would make an unreachable replica look empty and +//! trigger a spurious new generation. (S3: a truncated page that carries no +//! continuation token is an `Err`, never a silently short listing, since the +//! engine treats LIST as the source of truth per invariant I6.) +//! - `delete` is idempotent: deleting a missing key succeeds. +//! +//! ## Pinned backend divergences +//! +//! These behaviors differ between backends by design. The engine only ever +//! uses directory-aligned prefixes and never relies on the divergent cases, +//! but they are pinned here (and in the backend-specific sections of the +//! contract suite) so they cannot change silently: +//! +//! - `delete_prefix` interprets its argument differently. S3 deletes every +//! object whose key has the raw string as a prefix, so `delete_prefix("gen")` +//! would also remove `generations/...`. The local backend treats the prefix +//! as a directory path and calls `remove_dir_all`, so a prefix that ends +//! mid-component (`"gen"` when only `generations/` exists) matches no +//! directory and no-ops. Always pass a directory-aligned prefix. +//! - Local `list` errors when a prefix names a path whose component is a +//! regular file (e.g. `list("a/")` when `a` is an object); S3 returns an +//! empty listing. +//! - Local `get` of a key that is actually a directory errors (`EISDIR`, +//! surfaced as [`StorageError::Local`]); S3 has no directories and returns +//! [`StorageError::NotFound`]. + +use bytes::Bytes; + +use crate::local::{LocalError, LocalStorage}; +use crate::s3::{S3Error, S3Storage}; + +/// A replica storage backend. +pub enum ReplicaStorage { + Local(LocalStorage), + S3(Box), + /// Fault-injection wrapper around another backend; test-only. + #[cfg(any(test, feature = "testing"))] + Flaky(Box), +} + +#[derive(Debug, thiserror::Error)] +pub enum StorageError { + /// The requested key does not exist. Every backend maps its native + /// missing-object error to this variant (never to an empty result). + #[error("Object not found: {key}")] + NotFound { key: String }, + /// An object key was not a valid `/`-separated relative path (see + /// [`validate_key`]); rejected before touching the backend so a key can + /// never escape the storage root. + #[error("Invalid object key ({key}): {reason}")] + InvalidKey { key: String, reason: &'static str }, + /// Boxed to keep the common `Ok`/`NotFound` paths small. + #[error("Local replica storage: {0}")] + Local(#[source] Box), + /// Boxed to keep the common `Ok`/`NotFound` paths small. + #[error("S3 replica storage: {0}")] + S3(#[source] Box), + /// Injected by the fault-injection test backend. + #[cfg(any(test, feature = "testing"))] + #[error("Injected fault: {op} {key}")] + Injected { op: &'static str, key: String }, +} + +impl From for StorageError { + fn from(error: LocalError) -> Self { + Self::Local(Box::new(error)) + } +} + +impl From for StorageError { + fn from(error: S3Error) -> Self { + Self::S3(Box::new(error)) + } +} + +/// Reject an object key that is not a `/`-separated relative path: a leading +/// slash (which would replace the local root via `Utf8PathBuf::join`), a `.` +/// or `..` component (which would escape the root), or an empty component +/// (a leading, trailing, or doubled slash). Both backends call this before +/// any object operation so the contract holds regardless of who supplied the +/// key. +pub(crate) fn validate_key(key: &str) -> Result<(), StorageError> { + let reject = |reason| { + Err(StorageError::InvalidKey { + key: key.to_owned(), + reason, + }) + }; + if key.starts_with('/') { + return reject("leading slash"); + } + for component in key.split('/') { + match component { + "" => return reject("empty path component"), + "." | ".." => return reject("relative path component"), + _ => {}, + } + } + Ok(()) +} + +/// Validate a listing/deletion prefix. Looser than [`validate_key`]: a prefix +/// may be empty (the whole store), may carry a trailing slash, and may end +/// mid-component. Still rejected are dot components and a leading slash: the +/// local backend maps prefixes onto the filesystem, where `..` in a +/// destructive path like `delete_prefix` would escape the storage root. +pub(crate) fn validate_prefix(prefix: &str) -> Result<(), StorageError> { + if prefix.is_empty() { + return Ok(()); + } + let reject = |reason| { + Err(StorageError::InvalidKey { + key: prefix.to_owned(), + reason, + }) + }; + if prefix.starts_with('/') { + return reject("leading slash"); + } + let mut components = prefix.split('/').peekable(); + while let Some(component) = components.next() { + match component { + // A single trailing empty component is the trailing slash. + "" if components.peek().is_none() => {}, + "" => return reject("empty path component"), + "." | ".." => return reject("relative path component"), + _ => {}, + } + } + Ok(()) +} + +impl ReplicaStorage { + /// Store a whole object atomically under `key`. + pub async fn put(&self, key: &str, bytes: Bytes) -> Result<(), StorageError> { + match self { + Self::Local(local) => local.put(key, bytes).await, + Self::S3(s3) => s3.put(key, bytes).await, + #[cfg(any(test, feature = "testing"))] + Self::Flaky(flaky) => flaky.put(key, bytes).await, + } + } + + /// Fetch a whole object. Missing keys are [`StorageError::NotFound`]. + pub async fn get(&self, key: &str) -> Result { + match self { + Self::Local(local) => local.get(key).await, + Self::S3(s3) => s3.get(key).await, + #[cfg(any(test, feature = "testing"))] + Self::Flaky(flaky) => flaky.get(key).await, + } + } + + /// Fetch an object as a byte stream (used for multi-GB snapshot + /// downloads during restore). + pub async fn get_stream( + &self, + key: &str, + ) -> Result, StorageError> { + match self { + Self::Local(local) => local.get_stream(key).await, + Self::S3(s3) => s3.get_stream(key).await, + #[cfg(any(test, feature = "testing"))] + Self::Flaky(flaky) => flaky.get_stream(key).await, + } + } + + /// List all keys under `prefix` (recursive), sorted lexicographically. + /// Pagination is handled internally. Errors are errors, never `vec![]`. + pub async fn list(&self, prefix: &str) -> Result, StorageError> { + validate_prefix(prefix)?; + match self { + Self::Local(local) => local.list(prefix).await, + Self::S3(s3) => s3.list(prefix).await, + #[cfg(any(test, feature = "testing"))] + Self::Flaky(flaky) => flaky.list(prefix).await, + } + } + + /// List the immediate "subdirectory" components under `prefix` (S3: + /// delimiter `/` common prefixes; local: child directories), sorted + /// lexicographically. Returned values are the bare path components, not + /// full keys. + pub async fn list_dirs(&self, prefix: &str) -> Result, StorageError> { + validate_prefix(prefix)?; + match self { + Self::Local(local) => local.list_dirs(prefix).await, + Self::S3(s3) => s3.list_dirs(prefix).await, + #[cfg(any(test, feature = "testing"))] + Self::Flaky(flaky) => flaky.list_dirs(prefix).await, + } + } + + /// Delete one object. Idempotent: deleting a missing key is `Ok`. + pub async fn delete(&self, key: &str) -> Result<(), StorageError> { + match self { + Self::Local(local) => local.delete(key).await, + Self::S3(s3) => s3.delete(key).await, + #[cfg(any(test, feature = "testing"))] + Self::Flaky(flaky) => flaky.delete(key).await, + } + } + + /// Delete every object under `prefix` (generation pruning). + pub async fn delete_prefix(&self, prefix: &str) -> Result<(), StorageError> { + validate_prefix(prefix)?; + match self { + Self::Local(local) => local.delete_prefix(prefix).await, + Self::S3(s3) => s3.delete_prefix(prefix).await, + #[cfg(any(test, feature = "testing"))] + Self::Flaky(flaky) => flaky.delete_prefix(prefix).await, + } + } + + /// Best-effort sweep of crash-orphaned incomplete uploads: S3 aborts + /// uncompleted multipart uploads (which accrue storage cost until + /// aborted), and the local backend removes orphaned `.partial-` write + /// fragments. The fault-injection backend has nothing to reclaim. + /// Never fails the caller. + pub async fn abort_incomplete_uploads(&self, log: &slog::Logger) { + match self { + Self::S3(s3) => s3.abort_incomplete_uploads(log).await, + Self::Local(local) => local.abort_incomplete_uploads(log).await, + #[cfg(any(test, feature = "testing"))] + Self::Flaky(_) => {}, + } + } + + /// Begin a streaming upload for a large object (snapshots). Parts are + /// buffered/uploaded as written; the object becomes visible under `key` + /// only when [`MultipartUpload::finish`] succeeds. + pub async fn start_multipart(&self, key: &str) -> Result { + match self { + Self::Local(local) => Ok(MultipartUpload::Local(local.start_multipart(key).await?)), + Self::S3(s3) => Ok(MultipartUpload::S3(Box::new( + s3.start_multipart(key).await?, + ))), + #[cfg(any(test, feature = "testing"))] + Self::Flaky(flaky) => Ok(MultipartUpload::Flaky(Box::new( + flaky.start_multipart(key).await?, + ))), + } + } +} + +/// An in-progress streaming upload. Dropping without [`Self::finish`] must +/// never leave a visible object under the final key (local: temp file; S3: +/// uncompleted multipart upload). +pub enum MultipartUpload { + Local(crate::local::LocalMultipart), + S3(Box), + #[cfg(any(test, feature = "testing"))] + Flaky(Box), +} + +impl MultipartUpload { + /// Append a part. Parts may be any size; the S3 backend buffers + /// internally to satisfy the 5 MiB minimum-part rule. + pub async fn write_part(&mut self, bytes: Bytes) -> Result<(), StorageError> { + match self { + Self::Local(local) => local.write_part(bytes).await, + Self::S3(s3) => s3.write_part(bytes).await, + #[cfg(any(test, feature = "testing"))] + Self::Flaky(flaky) => flaky.write_part(bytes).await, + } + } + + /// Complete the upload, making the object visible under its final key. + pub async fn finish(self) -> Result<(), StorageError> { + match self { + Self::Local(local) => local.finish().await, + Self::S3(s3) => s3.finish().await, + #[cfg(any(test, feature = "testing"))] + Self::Flaky(flaky) => flaky.finish().await, + } + } + + /// Abort the upload, releasing any partial state. + pub async fn abort(self) -> Result<(), StorageError> { + match self { + Self::Local(local) => local.abort().await, + Self::S3(s3) => s3.abort().await, + #[cfg(any(test, feature = "testing"))] + Self::Flaky(flaky) => flaky.abort().await, + } + } +} + +#[cfg(test)] +mod tests { + use camino::Utf8Path; + use pretty_assertions::assert_eq; + + use super::*; + + #[test] + fn not_found_display_names_key() { + let error = StorageError::NotFound { + key: "generations/g1/snapshot.json".to_owned(), + }; + assert_eq!( + error.to_string(), + "Object not found: generations/g1/snapshot.json", + "NotFound display must name the key" + ); + } + + #[test] + fn validate_key_accepts_relative_paths() { + for key in [ + "snapshot.json", + "generations/g1/snapshot.json", + "generations/20260710T145900Z-3f8a2c1d/wal/0000000000-9d2f1c4a8b3e6f70/x.wal.zst", + ] { + validate_key(key).unwrap_or_else(|error| panic!("{key} must be valid: {error}")); + } + } + + #[test] + fn validate_prefix_accepts_prefix_shapes() { + for prefix in [ + "", + "generations/", + "generations/g1", + "generations/g1/wal/epoch00", + "generations/20260710T145900Z-3f8a2c1d/wal/", + ] { + validate_prefix(prefix) + .unwrap_or_else(|error| panic!("{prefix:?} must be valid: {error}")); + } + } + + #[test] + fn validate_prefix_rejects_escaping_shapes() { + let cases = [ + ("/leading", "leading slash"), + ("/", "leading slash"), + ("..", "relative path component"), + ("../escape", "relative path component"), + ("gen/../escape", "relative path component"), + ("gen/./here", "relative path component"), + ("gen/..", "relative path component"), + ("gen//doubled", "empty path component"), + ("gen//", "empty path component"), + ]; + for (prefix, reason) in cases { + match validate_prefix(prefix) { + Err(StorageError::InvalidKey { + key: found, + reason: found_reason, + }) => { + assert_eq!(found, prefix, "InvalidKey names the wrong prefix"); + assert_eq!(found_reason, reason, "wrong reason for {prefix}"); + }, + other => panic!("{prefix} must be rejected, got: {other:?}"), + } + } + } + + #[test] + fn validate_key_rejects_escaping_shapes() { + let cases = [ + ("/leading", "leading slash"), + ("/", "leading slash"), + ("..", "relative path component"), + ("../escape", "relative path component"), + ("gen/../escape", "relative path component"), + ("gen/./here", "relative path component"), + ("", "empty path component"), + ("gen//doubled", "empty path component"), + ("trailing/", "empty path component"), + ]; + for (key, reason) in cases { + match validate_key(key) { + Err(StorageError::InvalidKey { + key: found, + reason: found_reason, + }) => { + assert_eq!(found, key, "InvalidKey names the wrong key"); + assert_eq!(found_reason, reason, "wrong reason for {key}"); + }, + other => panic!("{key} must be rejected, got: {other:?}"), + } + } + } + + #[test] + fn injected_display_names_op_and_key() { + let error = StorageError::Injected { + op: "list", + key: "generations/".to_owned(), + }; + assert_eq!( + error.to_string(), + "Injected fault: list generations/", + "Injected display must name the op and key" + ); + } + + /// Smoke test that the enum dispatch delegates to the wrapped backend; + /// full behavior is covered by `tests/storage_contract.rs`. + #[tokio::test] + async fn local_variant_dispatches() { + let tmp = tempfile::tempdir().expect("tempdir failed"); + let root = Utf8Path::from_path(tmp.path()) + .expect("tempdir path is UTF-8") + .to_path_buf(); + let storage = ReplicaStorage::Local(LocalStorage::new(root)); + storage + .put("dispatch/key", Bytes::from_static(b"value")) + .await + .expect("put failed"); + let got = storage.get("dispatch/key").await.expect("get failed"); + assert_eq!(got.as_ref(), b"value".as_slice(), "dispatch roundtrip"); + let keys = storage.list("").await.expect("list failed"); + assert_eq!(keys, vec!["dispatch/key".to_owned()], "dispatch list"); + let mut upload = storage + .start_multipart("dispatch/multi") + .await + .expect("start failed"); + upload + .write_part(Bytes::from_static(b"m")) + .await + .expect("write failed"); + upload.finish().await.expect("finish failed"); + let got = storage.get("dispatch/multi").await.expect("get failed"); + assert_eq!(got.as_ref(), b"m".as_slice(), "dispatch multipart"); + storage.delete("dispatch/key").await.expect("delete failed"); + storage + .delete_prefix("dispatch/") + .await + .expect("delete_prefix failed"); + } +} diff --git a/plus/bencher_replica/src/sync.rs b/plus/bencher_replica/src/sync.rs new file mode 100644 index 000000000..fe857b3ef --- /dev/null +++ b/plus/bencher_replica/src/sync.rs @@ -0,0 +1,2342 @@ +//! The sync engine: a single sequential task that ships WAL frames, +//! checkpoints, snapshots, and prunes. +//! +//! Step-driven core: [`SyncEngine::sync_once`] is one production tick, and +//! the finer steps ([`SyncEngine::ship_once`], [`SyncEngine::checkpoint_once`], +//! [`SyncEngine::snapshot_step`], [`SyncEngine::prune_once`]) are public so +//! tests drive every transition deterministically; the production loop in +//! `replicator.rs` is a trivial tick shell. All scheduling decisions go +//! through the injected [`bencher_json::Clock`]. +//! +//! Blocking work (rusqlite, WAL file scans, zstd, meta writes) runs under +//! `tokio::task::spawn_blocking` with owned data; the engine is safe on +//! current-thread runtimes. +//! +//! ## Position resume (five-step decision table) +//! +//! At construction the engine reconciles three sources: the replica LIST +//! (the source of truth, invariant I6), the local WAL header, and the +//! advisory meta file. +//! +//! 1. LIST the latest valid generation (has `snapshot.json`) and its last +//! segment, giving the replica tip `(epoch_r, salt_r, end_r)`. No +//! generation: schedule the first snapshot (state `PendingSnapshot`). +//! 2. Read the local WAL header (missing or short: step 5). +//! 3. Local salts equal `salt_r`: rescan the local chain from offset 0 and +//! require it valid through exactly `end_r` (recovering the running +//! checksum there); resume at `(epoch_r, end_r)`. A local WAL shorter +//! than `end_r` (replica ahead: `synchronous=NORMAL` rewind after a +//! crash) or a broken chain is DIVERGENCE: new generation. +//! 4. Salt mismatch: if the meta file matches the replica exactly +//! (generation, `epoch == epoch_r`, `shipped_offset == end_r`, shipped +//! through a completed checkpoint, same shadow mode), the old epoch is +//! provably complete: resume as `epoch_r + 1` at offset 0 with the LOCAL +//! header's salts (or rebind an empty `epoch_r` in place when +//! `end_r == 0`). Anything else: new generation. +//! 5. Local WAL absent or shorter than a header: with the same meta proof, +//! wait in `AwaitingEpoch` and bind the salts when the first frames +//! appear; otherwise new generation. +//! +//! In shadow mode a resume divergence also just starts a new generation +//! (the shadow replica is disposable). + +use std::fs::File; +use std::io::{ErrorKind, Read as _}; +use std::sync::Arc; +use std::time::Duration; + +use bencher_json::Clock; +use bytes::Bytes; +use camino::{Utf8Path, Utf8PathBuf}; +use sha2::{Digest as _, Sha256}; +use slog::Logger; +use tokio::task::spawn_blocking; + +use crate::backoff::Backoff; +use crate::checkpoint::{ + CheckpointConns, CheckpointError, CheckpointOutcome, PinOutcome, checkpoint_locked, pin_locked, +}; +use crate::config::ReplicaConfig; +use crate::meta::{META_VERSION, MetaError, ReplicaMeta}; +use crate::position::{ + GENERATIONS_PREFIX, GenerationId, Position, SegmentKey, WAL_DIR, generation_prefix, + parse_segment_key, segment_key, snapshot_key, snapshot_meta_key, +}; +use crate::replicator::ReplicaDb; +use crate::segment::{SEGMENT_MAX_BYTES, SegmentError, compress_segment}; +use crate::snapshot::{ + CopyJob, CopyStepResult, FinalizeJob, MIB, STALE_INCOMPLETE_SECS, SnapshotError, SnapshotJob, + SnapshotStatus, backup_to_scratch, copy_step, live_db_fingerprint, new_snapshot_encoder, + read_scratch_info, +}; +use crate::snapshot_meta::{SNAPSHOT_META_VERSION, SnapshotMeta, SnapshotMetaError, WalBoundary}; +use crate::storage::{MultipartUpload, ReplicaStorage, StorageError}; +use crate::verify::{VerifyError, VerifyReport, fingerprint_database, verify_against_replica}; +use crate::wal::{ + CommittedChunk, WAL_HEADER_SIZE, WalError, WalHeader, WalScanner, parse_wal_header, +}; + +/// The single-task replication engine, generic over the app writer +/// connection type: the engine only ever HOLDS the mutex guard (to freeze +/// app writers during the checkpoint critical section), never uses `C`. +pub struct SyncEngine { + log: Logger, + config: ReplicaConfig, + db: ReplicaDb, + clock: Clock, + shadow: bool, + storage: ReplicaStorage, + /// Current shipping position; `None` when no lineage is usable. + position: Option, + /// Salts-unbound resume: `(generation, epoch)` awaiting first frames. + awaiting: Option<(GenerationId, u64)>, + /// In-flight snapshot state machine. + snapshot: Option, + /// A new-generation snapshot has been requested (divergence, tripwire + /// abort, or explicit trigger). + pending_new_generation: bool, + /// The greatest generation id ever observed (replica tip at resume or + /// our own snapshots); new generation ids must sort after it so restore + /// always picks the newest lineage, even across a divergence. + generation_floor: Option, + /// Whether the current epoch is fully shipped AND fully backfilled by a + /// completed checkpoint (mirrors the advisory meta flag). Cleared by + /// every segment ship; set by [`CheckpointOutcome::Completed`]. + epoch_checkpointed: bool, + backoff: Backoff, + /// Clock second before which ticks are no-ops (storage error backoff). + backoff_until: Option, + /// Lazily opened checkpoint connections, reused across checkpoints. + conns: Option, + /// Clock second of the last Completed or Partial checkpoint (Partial + /// counts for retry pacing). Initialized to construction time. + last_checkpoint_secs: i64, + /// Clock second the current generation was created or resumed. + generation_birth_secs: i64, + /// Clock second of the last completed verification (pass or fail). + /// Initialized to construction time, so the first verification runs one + /// interval after startup. + last_verify_secs: i64, + /// Clock second before which verification is not retried after it could + /// not complete (execution error, or a Busy/Unshipped pin). Paces + /// verification independently of the WAL-ship backoff so a persistently + /// broken verify never throttles shipping. + verify_retry_until: Option, + /// The oversized-transaction drain state machine (see [`OversizedDrain`]). + oversized_drain: OversizedDrain, + /// Clock second before which an oversized-transaction divergence is not + /// re-triggered (a broken drain, e.g. a permanently pinned reader or a + /// shadow replica awaiting Litestream, must not churn snapshots tick after + /// tick). + oversized_retrigger_until: Option, +} + +/// How the engine is recovering from an oversized COMMITTED transaction that +/// cannot ship (see [`SyncEngine::diverge_oversized`]). +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum OversizedDrain { + /// Not draining. + Idle, + /// A re-snapshot is scheduled to capture the transaction; its finalize + /// pins epoch 0 at the boundary and advances to `AwaitingRestart`. + Pending, + /// Epoch 0 is pinned at the boundary awaiting the WAL restart that a + /// completed checkpoint allows: the ship path ships nothing, and the next + /// restart rebinds epoch 0 to the fresh cycle (offset 0) rather than + /// advancing to epoch+1 (which would leave epoch 0 empty and break + /// restore's epoch contiguity). + AwaitingRestart, +} + +#[derive(Debug, thiserror::Error)] +pub enum SyncError { + #[error("Replica storage: {0}")] + Storage(#[from] StorageError), + #[error("WAL: {0}")] + Wal(#[from] WalError), + #[error("Failed to read local WAL ({path}): {error}")] + WalIo { + path: Utf8PathBuf, + error: std::io::Error, + }, + #[error("Segment: {0}")] + Segment(#[from] SegmentError), + #[error("Replica meta: {0}")] + Meta(#[from] MetaError), + #[error("Checkpoint: {0}")] + Checkpoint(#[from] CheckpointError), + #[error("Snapshot: {0}")] + Snapshot(#[from] SnapshotError), + #[error("Snapshot meta: {0}")] + SnapshotMeta(#[from] SnapshotMetaError), + #[error("Verification: {0}")] + Verify(#[from] VerifyError), + #[error( + "A single transaction wrote {bytes} WAL bytes, beyond the restorable segment bound of {max_bytes}; refusing to ship it (split the write, or it would poison every restore of this generation)" + )] + TransactionTooLarge { bytes: u64, max_bytes: u64 }, + #[error( + "An external checkpoint restarted the WAL during a snapshot; the generation was aborted (no snapshot.json) and a fresh one scheduled" + )] + SnapshotBoundaryDiverged, + #[error("Sync task panicked: {0}")] + Join(#[from] tokio::task::JoinError), + #[error("The replicator task exited unexpectedly")] + TaskExited, +} + +impl SyncError { + /// Whether the error is a transient storage failure that the caller + /// should retry with backoff (as opposed to a fatal local failure). + #[must_use] + pub fn is_retryable(&self) -> bool { + matches!(self, Self::Storage(_)) + } + + /// Whether the error is structurally unrecoverable ("poison"): retrying + /// the identical operation can never succeed, and because a failed ship + /// blocks every downstream step (checkpoints, snapshots), a poison error + /// silently retried every tick would let the local WAL grow without bound. + /// Poison errors are propagated out of [`SyncEngine::sync_once`] so the + /// replicator's fatal channel fires and the operator is alerted. + /// + /// An oversized transaction qualifies as a LAST-RESORT safety net: the + /// ship path intercepts it first (see [`Self::is_transaction_too_large`]) + /// and converts it into a recoverable re-snapshot, so in normal operation + /// it never reaches here. Should it ever surface from an unconverted path, + /// a loud fatal exit still beats an unbounded retry loop. + #[must_use] + pub fn is_poison(&self) -> bool { + self.is_transaction_too_large() + } + + /// Whether the error is an oversized transaction (raw WAL bytes since the + /// last commit beyond the configured bound), from either the WAL scan or + /// the ship-time chunk check. The ship path converts a COMMITTED oversized + /// transaction into a re-snapshot: it is already durable, so it cannot be + /// un-committed and "split", but the snapshot captures its state without + /// shipping the unshippable WAL. + #[must_use] + pub fn is_transaction_too_large(&self) -> bool { + matches!( + self, + Self::TransactionTooLarge { .. } | Self::Wal(WalError::TransactionTooLarge { .. }) + ) + } +} + +/// What one [`SyncEngine::sync_once`] tick did. +#[derive(Debug, Default)] +pub struct SyncProgress { + /// Segments shipped this tick. + pub shipped_segments: u64, + /// Checkpoint outcome, when one was attempted. + pub checkpoint: Option, + /// Snapshot step status, when one ran. + pub snapshot: Option, + /// Verification report, when a scheduled verification completed. + pub verify: Option, + /// The tick was skipped because a backoff delay is still pending. + pub backing_off: bool, + /// The error that armed a backoff this tick (storage or snapshot); the + /// tick itself still returns `Ok` so the task keeps running. + pub error: Option, +} + +/// The engine's observable state. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum EngineState { + /// Shipping from a bound position. + Streaming, + /// A known generation and epoch await their first frames to bind salts. + AwaitingEpoch, + /// No usable lineage: a new-generation snapshot must run first. + PendingSnapshot, + /// A snapshot is in progress. + Snapshotting, +} + +/// Upper bound on segments shipped per ship pass (approximately 256 MiB of +/// raw WAL; a single oversized transaction ships whole, so one pass can +/// exceed that). A continuously committing writer would otherwise keep +/// feeding the scan loop forever and starve the tick (checkpoints, snapshots, +/// shutdown); the next tick simply resumes where this one stopped. +const MAX_SEGMENTS_PER_PASS: u64 = 32; + +/// Minimum age of the current generation before a failed verification may +/// force another one (rate limit: repeated verify failures must not churn +/// whole-database snapshots back to back). +const VERIFY_RETRIGGER_MIN_SECS: i64 = 6 * 60 * 60; + +/// Retry spacing after a verification that could not complete (an execution +/// error, or a Busy/Unshipped pin under sustained writes). Modest, so a +/// failed pin does not re-drain the WAL tail on every tick, yet far below the +/// verification interval, so a transient failure heals well within one cycle. +const VERIFY_RETRY_SECS: i64 = 60; + +/// Minimum spacing between oversized-transaction re-snapshots. The drain +/// normally breaks the loop (the position pins at the boundary, so the ship +/// path stops re-hitting the oversized transaction), but if the drain cannot +/// complete (a permanently pinned reader, or a shadow replica whose WAL only +/// Litestream can restart), this bounds how often a fresh generation is minted +/// while the condition persists. +const OVERSIZED_RETRIGGER_MIN_SECS: i64 = 60 * 60; + +/// The replica's tip, derived from a LIST alone (invariant I6). +struct ReplicaTip { + generation: GenerationId, + epoch: u64, + salt: (u32, u32), + /// Raw WAL byte offset shipped through in `epoch` (0: nothing shipped). + end: u64, + /// The tip epoch's last segment, for content verification at resume. + last_segment: Option, +} + +/// The local WAL header, read without touching `SQLite`. +enum WalHeaderState { + /// No WAL file, or shorter than a full header. + Missing, + /// A header exists but does not parse (e.g. torn concurrent write). + Unreadable(WalError), + Present { + header: WalHeader, + /// The exact on-disk header bytes (shipped verbatim in the first + /// segment of every epoch). + raw: [u8; 32], + }, +} + +impl SyncEngine { + /// Build the engine and resolve the resume position against the replica + /// (see the module docs for the decision table). + /// + /// Storage errors here are retryable: an unreachable replica at boot + /// must NEVER be conflated with an empty one (a new generation is only + /// ever created from a successful LIST). + pub async fn new( + log: Logger, + config: ReplicaConfig, + db: ReplicaDb, + clock: Clock, + shadow: bool, + ) -> Result { + let storage = config.build_storage(); + Self::new_with_storage(log, config, db, clock, shadow, storage).await + } + + /// [`Self::new`] with an explicit storage backend. Production callers + /// use [`Self::new`]; tests inject `ReplicaStorage::Flaky` here. + pub async fn new_with_storage( + log: Logger, + config: ReplicaConfig, + db: ReplicaDb, + clock: Clock, + shadow: bool, + storage: ReplicaStorage, + ) -> Result { + let now = clock.timestamp(); + let mut engine = Self { + log, + config, + db, + clock, + shadow, + storage, + position: None, + awaiting: None, + snapshot: None, + pending_new_generation: false, + generation_floor: None, + epoch_checkpointed: false, + backoff: Backoff::default(), + backoff_until: None, + conns: None, + last_checkpoint_secs: now, + generation_birth_secs: now, + last_verify_secs: now, + verify_retry_until: None, + oversized_drain: OversizedDrain::Idle, + oversized_retrigger_until: None, + }; + engine.resume().await?; + Ok(engine) + } + + /// One full production tick: honor backoff, advance a snapshot if one + /// is active (or required), otherwise ship, checkpoint when due, and + /// start a snapshot when due. Transient errors (storage and the like) arm + /// a capped exponential backoff and are reported in the progress instead + /// of failing the task; a POISON error ([`SyncError::is_poison`]) is + /// returned as `Err` so the replicator's fatal channel fires (retrying it + /// forever would only let the WAL grow without bound). + pub async fn sync_once(&mut self) -> Result { + let mut progress = SyncProgress::default(); + let now = self.now_secs(); + if let Some(until) = self.backoff_until { + if now < until { + progress.backing_off = true; + return Ok(progress); + } + self.backoff_until = None; + } + + // Snapshotting, or no usable lineage: the snapshot machine runs. + if self.snapshot.is_some() || (self.position.is_none() && self.awaiting.is_none()) { + match self.snapshot_step().await { + Ok(status) => { + progress.snapshot = Some(status); + self.backoff.reset(); + }, + Err(error) if error.is_poison() => return Err(error), + Err(error) => { + self.arm_backoff(now); + progress.error = Some(error); + }, + } + return Ok(progress); + } + + // Streaming (or awaiting the first frames of a known epoch). + match self.ship_once().await { + Ok(segments) => progress.shipped_segments = segments, + Err(error) => { + #[cfg(feature = "otel")] + bencher_otel::ApiMeter::increment(bencher_otel::ApiCounter::ReplicaShipFailed); + if error.is_poison() { + return Err(error); + } + self.arm_backoff(now); + progress.error = Some(error); + return Ok(progress); + }, + } + if self.checkpoint_due(now).await { + match self.checkpoint_once().await { + Ok(outcome) => progress.checkpoint = Some(outcome), + Err(error) => { + #[cfg(feature = "otel")] + bencher_otel::ApiMeter::increment( + bencher_otel::ApiCounter::ReplicaCheckpointFailed, + ); + if error.is_poison() { + return Err(error); + } + self.arm_backoff(now); + progress.error = Some(error); + return Ok(progress); + }, + } + } + if self.verify_due(now) { + match self.verify_once().await { + // A completed verification (Pass or Fail); `verify_once` + // already advanced `last_verify_secs`. + Ok(Some(report)) => progress.verify = Some(report), + // Could not pin (Busy/Unshipped) or not ready: space out the + // retry so a sustained-write pin failure does not re-run the + // full drain on every tick. + Ok(None) => { + self.verify_retry_until = Some(now.saturating_add(VERIFY_RETRY_SECS)); + }, + // A verification EXECUTION error (restore/fingerprint/storage) + // must NOT arm the global ship backoff: a persistently broken + // verify would otherwise throttle WAL shipping forever. Pace + // it on its own retry clock and keep the tick going. Poison + // errors still escalate: `verify_once` runs its own ship + // drain, so a structurally-unrecoverable ship error can + // surface here first and must reach the fatal channel like + // the ship/checkpoint arms above. + Err(error) => { + #[cfg(feature = "otel")] + bencher_otel::ApiMeter::increment( + bencher_otel::ApiCounter::ReplicaVerifyFailed, + ); + if error.is_poison() { + return Err(error); + } + slog::error!(self.log, "Replica verification could not complete"; + "error" => %error); + self.verify_retry_until = Some(now.saturating_add(VERIFY_RETRY_SECS)); + progress.error = Some(error); + }, + } + } + if self.snapshot_due(now) { + self.snapshot = Some(SnapshotJob::ShipTail); + } + self.backoff.reset(); + self.backoff_until = None; + Ok(progress) + } + + /// Ship every complete committed transaction currently in the WAL, + /// handling salt transitions first. Returns the number of segments + /// shipped. A divergence (illegitimate salt change) schedules a new + /// generation and returns 0. No-op while a snapshot is in flight (the + /// snapshot machine runs its own ship pass). + pub async fn ship_once(&mut self) -> Result { + if self.snapshot.is_some() { + return Ok(0); + } + self.ship_pass().await + } + + /// Run the checkpoint critical section (see `checkpoint.rs`). The app + /// writer mutex is held for the WHOLE call, so app writers queue on the + /// tokio mutex and never burn their `busy_timeout`. + pub async fn checkpoint_once(&mut self) -> Result { + if self.shadow { + return Ok(CheckpointOutcome::SkippedShadow); + } + let Some(position) = self.position.clone() else { + // Nothing is provably shipped without a bound position. + return Ok(CheckpointOutcome::SkippedUnshipped); + }; + let guard = Arc::clone(&self.db.writer).lock_owned().await; + let mut conns = if let Some(conns) = self.conns.take() { + conns + } else { + let db_path = self.db.db_path.clone(); + spawn_blocking(move || CheckpointConns::open(&db_path)) + .await + .map_err(SyncError::Join)? + .map_err(|error| SyncError::Checkpoint(error.into()))? + }; + let wal_path = self.wal_path(); + let max_transaction_bytes = self.config.max_transaction_bytes; + #[cfg(feature = "otel")] + let critical_started = std::time::Instant::now(); + let (conns, outcome) = spawn_blocking(move || { + let outcome = + checkpoint_locked(&mut conns, &wal_path, &position, max_transaction_bytes); + (conns, outcome) + }) + .await + .map_err(SyncError::Join)?; + #[cfg(feature = "otel")] + bencher_otel::ApiMeter::record( + bencher_otel::ApiHistogram::ReplicaCriticalSectionDuration, + critical_started.elapsed().as_secs_f64(), + ); + let outcome = match outcome { + // Only reuse the connections after a clean section: an error + // may have left the lock connection inside a transaction, and + // dropping it rolls back and releases everything. + Ok(outcome) => { + self.conns = Some(conns); + outcome + }, + Err(error) => return Err(SyncError::Checkpoint(error)), + }; + if matches!(outcome, CheckpointOutcome::Completed) { + // The salt change at the next WAL restart is now legitimate. + self.epoch_checkpointed = true; + if let Some(position) = self.position.clone() { + // Crash window (churn, never loss): a crash AFTER the + // checkpoint completes but BEFORE this meta store persists the + // `epoch_shipped_through_checkpoint` flag leaves the meta + // saying the epoch was NOT sealed. If the WAL then restarts + // before the next boot, resume cannot prove the clean epoch + // transition and forces a spurious full re-snapshot. This is + // wasteful churn, not data loss (every frame was shipped + // before the checkpoint, invariant I1), so it is tolerated + // rather than guarded with a heavier commit protocol. + self.store_meta_for(&position).await?; + } + } + drop(guard); + if matches!( + outcome, + CheckpointOutcome::Completed | CheckpointOutcome::Partial + ) { + // Partial counts for pacing: retry next interval, not sooner. + self.last_checkpoint_secs = self.now_secs(); + } + #[cfg(feature = "otel")] + bencher_otel::ApiMeter::increment(match outcome { + CheckpointOutcome::Completed => bencher_otel::ApiCounter::ReplicaCheckpoint, + CheckpointOutcome::Partial => bencher_otel::ApiCounter::ReplicaCheckpointPartial, + // Unshipped skips are the checkpoint-liveness alert signal: a + // sustained-write workload can starve checkpoints (see + // `checkpoint_due`), so they get their own counter. + CheckpointOutcome::SkippedUnshipped => { + bencher_otel::ApiCounter::ReplicaCheckpointSkippedUnshipped + }, + CheckpointOutcome::SkippedBusy | CheckpointOutcome::SkippedShadow => { + bencher_otel::ApiCounter::ReplicaCheckpointSkipped + }, + }); + slog::debug!(self.log, "Checkpoint attempted"; "outcome" => format!("{outcome:?}")); + Ok(outcome) + } + + /// Run one restore-and-compare verification: prove the replica + /// reproduces the source database at the shipped position (the shadow + /// burn-in check and the post-cutover replacement for Litestream's + /// `validation`). + /// + /// Choreography: drain the WAL tail with ship passes (no locks), then + /// briefly gate writers (app mutex + `BEGIN IMMEDIATE` in sole mode) to + /// pin a read snapshot at exactly the shipped position, release the + /// gate, and compare fingerprints off-lock: the multi-second fingerprint + /// and restore never block writes (invariant I5). + /// + /// Returns `Ok(None)` when verification cannot run right now (no bound + /// position, a snapshot in flight, a busy stray writer, or frames that + /// slipped in after the drain): `last_verify_secs` is left untouched, and + /// the caller (`sync_once`) paces the retry so a persistent Busy/Unshipped + /// pin does not re-drain every tick. On `Fail`, a new generation is + /// triggered unless the current one is younger than six hours (rate limit). + pub async fn verify_once(&mut self) -> Result, SyncError> { + if self.snapshot.is_some() { + return Ok(None); + } + // Drain the committed tail so the pin lands at the true tip. Bounded: + // a continuously committing writer is caught by the gate re-check. + for _pass in 0u8..8 { + if self.ship_once().await? == 0 { + break; + } + } + let Some(position) = self.position.clone() else { + return Ok(None); + }; + let guard = Arc::clone(&self.db.writer).lock_owned().await; + let mut conns = if let Some(conns) = self.conns.take() { + conns + } else { + let db_path = self.db.db_path.clone(); + spawn_blocking(move || CheckpointConns::open(&db_path)) + .await + .map_err(SyncError::Join)? + .map_err(|error| SyncError::Checkpoint(error.into()))? + }; + let db_path = self.db.db_path.clone(); + let wal_path = self.wal_path(); + // The pin holds the write lock even in shadow mode: without it a + // commit can slip between the unshipped-tail scan and the snapshot + // pin, producing a spurious verification failure that would churn a + // whole new generation. Litestream lock contention just yields Busy, + // which retries next tick. + let pin_position = position.clone(); + let max_transaction_bytes = self.config.max_transaction_bytes; + let (conns, pin) = spawn_blocking(move || { + let pin = pin_locked( + &mut conns, + &db_path, + &wal_path, + &pin_position, + max_transaction_bytes, + ); + (conns, pin) + }) + .await + .map_err(SyncError::Join)?; + let pin = match pin { + Ok(pin) => { + self.conns = Some(conns); + pin + }, + Err(error) => { + drop(guard); + return Err(SyncError::Checkpoint(error)); + }, + }; + drop(guard); + let pinned = match pin { + PinOutcome::Pinned(pinned) => pinned, + PinOutcome::Busy | PinOutcome::Unshipped => return Ok(None), + }; + + // Off-lock from here: writers proceed while we fingerprint, restore, + // and compare. + let fingerprint = spawn_blocking(move || { + let fingerprint = fingerprint_database(&pinned); + // Dropping the connection releases the pinned read snapshot. + drop(pinned); + fingerprint + }) + .await + .map_err(SyncError::Join)? + .map_err(SyncError::Verify)?; + + let scratch = self.verify_scratch_dir(); + let result = + verify_against_replica(&self.log, &self.storage, &position, &fingerprint, &scratch) + .await; + // Reclaim the DB-sized scratch copy off the async executor, on BOTH + // the success and error paths: a `?` on the result above would leak a + // full database copy until the next run, and `remove_dir_all` blocks. + let cleanup = scratch.clone(); + drop(spawn_blocking(move || std::fs::remove_dir_all(&cleanup)).await); + let report = result.map_err(SyncError::Verify)?; + self.record_verify_outcome(&position, &report); + Ok(Some(report)) + } + + /// Log and meter one COMPLETED verification, and retrigger a fresh + /// generation on divergence (outside the retrigger window). + fn record_verify_outcome(&mut self, position: &Position, report: &VerifyReport) { + let now = self.now_secs(); + self.last_verify_secs = now; + match report { + VerifyReport::Pass => { + #[cfg(feature = "otel")] + bencher_otel::ApiMeter::increment(bencher_otel::ApiCounter::ReplicaVerifyPass); + slog::info!(self.log, "Replica verification passed"; + "generation" => position.generation.as_str(), + "epoch" => position.epoch, + "offset" => position.offset, + ); + }, + VerifyReport::Fail { detail } => { + #[cfg(feature = "otel")] + bencher_otel::ApiMeter::increment( + bencher_otel::ApiCounter::ReplicaVerifyDivergence, + ); + slog::error!(self.log, "Replica verification FAILED"; + "generation" => position.generation.as_str(), + "epoch" => position.epoch, + "offset" => position.offset, + "detail" => detail, + ); + if now.saturating_sub(self.generation_birth_secs) >= VERIFY_RETRIGGER_MIN_SECS { + self.pending_new_generation = true; + } else { + slog::warn!( + self.log, + "Verification failure within the retrigger window: not forcing another generation yet" + ); + } + }, + } + } + + /// Whether a scheduled verification is due. + fn verify_due(&self, now: i64) -> bool { + let Some(interval) = self.config.verification_interval else { + return false; + }; + if self.snapshot.is_some() || self.position.is_none() { + return false; + } + // Respect the independent verify retry pacing after a run that could + // not complete, so a broken verify does not re-run every tick. + if let Some(until) = self.verify_retry_until + && now < until + { + return false; + } + let interval = i64::try_from(interval.as_secs()).unwrap_or(i64::MAX); + now.saturating_sub(self.last_verify_secs) >= interval + } + + /// Scratch directory for verification restores, next to the database. + fn verify_scratch_dir(&self) -> Utf8PathBuf { + let parent = self + .db + .db_path + .parent() + .map_or_else(|| Utf8PathBuf::from("."), Utf8Path::to_path_buf); + parent.join(".replica-verify") + } + + /// Advance the snapshot state machine by one bounded unit of work, + /// starting a new generation snapshot when none is active. On error the + /// in-flight upload is aborted and a retrigger is scheduled; the caller + /// arms the backoff. + pub async fn snapshot_step(&mut self) -> Result { + let job = self.snapshot.take().unwrap_or(SnapshotJob::ShipTail); + match self.run_snapshot_step(job).await { + Ok(Some(next)) => { + self.snapshot = Some(next); + Ok(SnapshotStatus::InProgress) + }, + Ok(None) => { + #[cfg(feature = "otel")] + bencher_otel::ApiMeter::increment(bencher_otel::ApiCounter::ReplicaSnapshot); + Ok(SnapshotStatus::Finished) + }, + Err(error) => { + #[cfg(feature = "otel")] + bencher_otel::ApiMeter::increment(bencher_otel::ApiCounter::ReplicaSnapshotFailed); + // Retrigger: the WAL keeps buffering, nothing is lost. + self.pending_new_generation = true; + Err(error) + }, + } + } + + /// Delete generations beyond retention plus stale incomplete ones. + /// Whole-generation deletes only; the current (and any in-flight) + /// generation is never pruned. Failures are safe to retry at the next + /// snapshot. + pub async fn prune_once(&mut self) -> Result<(), SyncError> { + let components = self.storage.list_dirs(GENERATIONS_PREFIX).await?; + let mut complete = Vec::new(); + let mut incomplete = Vec::new(); + for component in &components { + let Some(generation) = GenerationId::parse(component) else { + slog::warn!(self.log, "Skipping foreign directory under generations/"; + "component" => component.as_str()); + continue; + }; + match self.storage.get(&snapshot_meta_key(&generation)).await { + Ok(_bytes) => complete.push(generation), + Err(StorageError::NotFound { .. }) => incomplete.push(generation), + Err(error) => return Err(error.into()), + } + } + let protected = self.protected_generations(); + // `list_dirs` is sorted, so `complete` is oldest first. + let retention = usize::try_from(self.config.retention_generations).unwrap_or(usize::MAX); + let prune_count = complete.len().saturating_sub(retention.max(1)); + for generation in complete.iter().take(prune_count) { + if protected.contains(generation) { + continue; + } + slog::info!(self.log, "Pruning generation beyond retention"; + "generation" => generation.as_str()); + // Delete the `snapshot.json` marker FIRST, as a single explicit + // delete, before the unordered `delete_prefix` batch. Batch + // deletion order is unspecified, so a crash mid-prune could + // otherwise remove `snapshot.db.zst` while leaving the marker, + // producing a marker-without-body generation that resume/restore + // would trust. Marker-first leaves only an invisible, markerless + // generation instead (pruned later as stale-incomplete). + self.storage.delete(&snapshot_meta_key(generation)).await?; + self.storage + .delete_prefix(&generation_prefix(generation)) + .await?; + } + // Incomplete generations older than the cutoff are crashed + // snapshots; lexicographic order equals chronological order. + let cutoff_secs = self.now_secs().saturating_sub(STALE_INCOMPLETE_SECS); + let Ok(cutoff_time) = bencher_json::DateTime::try_from(cutoff_secs) else { + return Ok(()); + }; + let cutoff = GenerationId::new(cutoff_time, 0); + for generation in &incomplete { + if protected.contains(generation) || *generation >= cutoff { + continue; + } + slog::info!(self.log, "Pruning stale incomplete generation (crashed snapshot)"; + "generation" => generation.as_str()); + self.storage + .delete_prefix(&generation_prefix(generation)) + .await?; + } + Ok(()) + } + + /// Shutdown ship of the remaining committed tail, `deadline`-bounded. + /// + /// The `deadline` bounds ONLY the ship loop. After a COMPLETE drain in + /// sole mode a final checkpoint then runs (unbounded, after the deadline): + /// it seals the epoch through a completed checkpoint so the next boot can + /// trust the WAL-less state and resume in place instead of churning a + /// whole new generation on every graceful restart (see the checkpoint + /// rationale below). The checkpoint is crash-safe if it is truncated: it + /// only ever backfills already-shipped frames (invariant I1), so an + /// interrupted checkpoint costs a re-snapshot, never data. + /// + /// On the deadline (an incomplete drain) or a storage failure NO + /// checkpoint runs and the remaining frames simply stay in the local WAL: + /// nothing is lost, only lagged, and the next boot resumes by salt match. + pub async fn final_sync(&mut self, deadline: Duration) -> Result<(), SyncError> { + // An in-flight snapshot would make ship_once a silent no-op (and its + // half generation is invisible to restore anyway): abort it so the + // tail ships into the CURRENT lineage. The next process re-snapshots + // when due. + if let Some(job) = self.snapshot.take() { + slog::info!( + self.log, + "Aborting the in-flight snapshot for shutdown; it will be retaken after restart" + ); + match job { + SnapshotJob::ShipTail | SnapshotJob::CreateGeneration => {}, + SnapshotJob::Copying(copy) => self.abort_upload(copy.upload).await, + SnapshotJob::Finalize(finalize) => self.abort_upload(finalize.upload).await, + } + self.pending_new_generation = true; + } + let drained = tokio::time::timeout(deadline, async { + loop { + match self.ship_once().await { + Ok(0) => return true, + Ok(_segments) => {}, + Err(error) => { + slog::warn!(self.log, + "Final sync ship failed; remaining WAL frames stay local (no data loss)"; + "error" => %error); + return false; + }, + } + } + }) + .await; + match drained { + Ok(true) => { + slog::info!(self.log, "Final sync complete: WAL tail fully shipped"); + // Checkpoint AFTER the full ship (never before: invariant + // I1), so the meta records a completed checkpoint. This is + // load-bearing: when the last in-process connection closes, + // `SQLite` checkpoints and DELETES the WAL file on its own; + // without the meta proof the next boot could not trust the + // WAL-less state and would churn a whole new generation on + // every graceful restart. + if !self.shadow { + match self.checkpoint_once().await { + Ok(outcome) => { + slog::info!(self.log, "Shutdown checkpoint"; + "outcome" => format!("{outcome:?}")); + }, + Err(error) => { + slog::warn!(self.log, "Shutdown checkpoint failed; the next boot may re-snapshot"; + "error" => %error); + }, + } + } + }, + Ok(false) => {}, + Err(_elapsed) => slog::warn!( + self.log, + "Final sync deadline reached; remaining WAL frames stay local (no data loss: the next boot resumes by salt match)" + ), + } + Ok(()) + } + + /// Request a new-generation snapshot at the next opportunity. + pub fn trigger_snapshot(&mut self) { + self.pending_new_generation = true; + } + + /// Current shipping position, when one is bound. + #[must_use] + pub fn position(&self) -> Option<&Position> { + self.position.as_ref() + } + + /// The generation currently being written to (or awaited). + #[must_use] + pub fn generation(&self) -> Option<&GenerationId> { + self.position + .as_ref() + .map(|position| &position.generation) + .or_else(|| self.awaiting.as_ref().map(|(generation, _)| generation)) + } + + /// The engine's observable state. + #[must_use] + pub fn state(&self) -> EngineState { + if self.snapshot.is_some() { + EngineState::Snapshotting + } else if self.position.is_some() { + EngineState::Streaming + } else if self.awaiting.is_some() { + EngineState::AwaitingEpoch + } else { + EngineState::PendingSnapshot + } + } + + /// The storage backend (tests reach the fault-injection wrapper here). + #[must_use] + pub fn storage(&self) -> &ReplicaStorage { + &self.storage + } + + fn now_secs(&self) -> i64 { + self.clock.timestamp() + } + + fn wal_path(&self) -> Utf8PathBuf { + Utf8PathBuf::from(format!("{}-wal", self.db.db_path)) + } + + /// Arm the capped exponential backoff after a failed tick. + fn arm_backoff(&mut self, now: i64) { + let delay = i64::try_from(self.backoff.next_delay().as_secs()).unwrap_or(i64::MAX); + self.backoff_until = Some(now.saturating_add(delay)); + } + + /// Resolve the resume position (see the module docs). + async fn resume(&mut self) -> Result<(), SyncError> { + let Some(tip) = self.replica_tip().await? else { + slog::info!( + self.log, + "Replica has no valid generation; scheduling the first snapshot" + ); + self.pending_new_generation = true; + return Ok(()); + }; + self.generation_floor = Some(tip.generation.clone()); + let wal_path = self.wal_path(); + let header_state = spawn_blocking(move || read_wal_header_state(&wal_path)) + .await + .map_err(SyncError::Join)??; + let db_path = self.db.db_path.clone(); + let meta = spawn_blocking(move || ReplicaMeta::load(&db_path)) + .await + .map_err(SyncError::Join)??; + match header_state { + WalHeaderState::Present { header, .. } => { + self.resume_with_header(tip, header, meta.as_ref()).await + }, + WalHeaderState::Missing => { + self.resume_without_wal(tip, meta.as_ref()); + Ok(()) + }, + WalHeaderState::Unreadable(error) => { + slog::warn!(self.log, "Local WAL header unreadable at resume; treating as absent"; + "error" => %error); + self.resume_without_wal(tip, meta.as_ref()); + Ok(()) + }, + } + } + + /// Resume steps 3 and 4: a local WAL header exists. + async fn resume_with_header( + &mut self, + tip: ReplicaTip, + header: WalHeader, + meta: Option<&ReplicaMeta>, + ) -> Result<(), SyncError> { + // A shadow<->sole config flip invalidates the lineage regardless of + // how well the WAL lines up: cutover forces one clean generation. + if let Some(meta) = meta + && meta.shadow != self.shadow + { + self.divergence( + &tip, + "shadow mode changed since the last run; forcing a clean generation", + ); + return Ok(()); + } + if header.salt == tip.salt { + // Step 3: same salt cycle; verify the local chain reaches the + // replica end and recover the running checksum there. + let checksum = if tip.end == 0 { + Some((0, 0)) + } else { + let wal_path = self.wal_path(); + let end = tip.end; + spawn_blocking(move || checksum_at_offset(&wal_path, end)) + .await + .map_err(SyncError::Join)?? + }; + let Some(checksum) = checksum else { + self.divergence( + &tip, + "local WAL chain does not reach the replica end (replica ahead or broken chain)", + ); + return Ok(()); + }; + // CONTENT proof, not just length: with `synchronous = NORMAL` a + // power loss can rewind the local WAL below the shipped offset + // (committed frames lived only in the page cache); if writers + // then re-extend it past `tip.end` with DIFFERENT transactions, + // the chain reaches `tip.end` again but the content forks from + // what the replica stored. The cumulative WAL checksum at + // `tip.end` fingerprints the whole prefix, so comparing it + // against the replica tip segment's stored value detects the + // fork (invariant I6: never guess). + if tip.end > 0 && !self.tip_content_matches(&tip, &header, checksum).await? { + self.divergence( + &tip, + "local WAL content at the shipped offset differs from the replica (post-crash rewind fork)", + ); + return Ok(()); + } + self.epoch_checkpointed = meta_matches(meta, &tip, self.shadow); + slog::info!(self.log, "Resuming by salt match"; + "generation" => tip.generation.as_str(), + "epoch" => tip.epoch, "offset" => tip.end); + self.position = Some(Position { + generation: tip.generation, + epoch: tip.epoch, + salt: tip.salt, + offset: tip.end, + checksum, + }); + return Ok(()); + } + // Step 4: salt mismatch; only an exact meta proof avoids a + // re-snapshot. The proof must include salt CONTINUITY: `SQLite` + // increments salt1 by exactly one per WAL restart, so a larger jump + // means an external writer buried at least one whole cycle + // (checkpointed then overwrote commits we never shipped) while the + // replicator was down. + let salt_continuous = meta.is_some_and(|meta| header.salt.0 == meta.salt1.wrapping_add(1)); + if meta_matches(meta, &tip, self.shadow) && salt_continuous { + let (epoch, note) = if tip.end == 0 { + // Nothing ever shipped into the tip epoch: rebind it in + // place so epochs stay contiguous and non-empty. + (tip.epoch, "rebinding empty epoch to the local WAL salts") + } else { + ( + tip.epoch.saturating_add(1), + "meta-verified resume as the next epoch", + ) + }; + slog::info!(self.log, "Resuming after WAL restart"; "note" => note, + "generation" => tip.generation.as_str(), "epoch" => epoch); + self.epoch_checkpointed = false; + self.position = Some(Position { + generation: tip.generation, + epoch, + salt: header.salt, + offset: 0, + checksum: (0, 0), + }); + } else { + self.divergence(&tip, "local WAL salts do not match the replica and the meta cannot prove a clean, salt-continuous epoch transition"); + } + Ok(()) + } + + /// Whether the local WAL's running checksum at the replica tip offset + /// matches the checksum stored in the replica tip segment's final frame + /// header. Downloads and decompresses one segment (bounded by the + /// segment size cap); runs only at resume. + /// + /// Conservative: any anomaly (segment gone, undecodable, or too short) + /// reports a mismatch, so resume diverges instead of guessing. + async fn tip_content_matches( + &mut self, + tip: &ReplicaTip, + header: &WalHeader, + local_checksum: (u32, u32), + ) -> Result { + let Some(segment) = &tip.last_segment else { + // `tip.end > 0` without a segment cannot happen (the end comes + // from the segment key); treat it as a mismatch if it does. + return Ok(false); + }; + let key = segment_key(&tip.generation, segment); + let compressed = match self.storage.get(&key).await { + Ok(bytes) => bytes, + Err(StorageError::NotFound { .. }) => { + slog::warn!(self.log, "Replica tip segment vanished during resume"; "key" => &key); + return Ok(false); + }, + Err(error) => return Err(error.into()), + }; + let raw = match spawn_blocking(move || crate::segment::decompress_segment(&compressed)) + .await + .map_err(SyncError::Join)? + { + Ok(raw) => raw, + Err(error) => { + slog::warn!(self.log, "Replica tip segment undecodable during resume"; + "key" => &key, "error" => %error); + return Ok(false); + }, + }; + let Some(stored) = segment_tail_checksum(&raw, header.frame_size()) else { + slog::warn!(self.log, "Replica tip segment too short during resume"; "key" => &key); + return Ok(false); + }; + Ok(stored == local_checksum) + } + + /// Resume step 5: no usable local WAL. + /// + /// Documented limitation: a brand-new WAL file gets RANDOM salts (salt1 + /// increments only on the restart of an EXISTING WAL), so salt + /// continuity cannot be checked here. An external writer that buries a + /// cycle AND removes the WAL file while the replicator is down is + /// undetectable at resume; the periodic restore-and-compare + /// verification is the backstop for that case. + fn resume_without_wal(&mut self, tip: ReplicaTip, meta: Option<&ReplicaMeta>) { + if meta_matches(meta, &tip, self.shadow) { + let epoch = if tip.end == 0 { + tip.epoch + } else { + tip.epoch.saturating_add(1) + }; + slog::info!(self.log, "Local WAL absent; awaiting first frames to bind the epoch"; + "generation" => tip.generation.as_str(), "epoch" => epoch); + self.awaiting = Some((tip.generation, epoch)); + } else { + self.divergence( + &tip, + "local WAL absent and the meta cannot prove a full ship through checkpoint", + ); + } + } + + /// Record a divergence: the replica is the source of truth (I6), so any + /// unprovable mismatch resolves to a new generation, never guessing. + fn divergence(&mut self, tip: &ReplicaTip, reason: &'static str) { + if self.shadow { + // The shadow replica is disposable; boundary losses are + // expected while Litestream owns checkpoints. + slog::info!(self.log, "Shadow replica divergence; starting a new generation"; + "reason" => reason, "generation" => tip.generation.as_str()); + } else { + slog::error!(self.log, "Replica divergence; starting a new generation"; + "reason" => reason, "generation" => tip.generation.as_str()); + } + #[cfg(feature = "otel")] + bencher_otel::ApiMeter::increment(bencher_otel::ApiCounter::ReplicaDivergence); + self.position = None; + self.awaiting = None; + self.pending_new_generation = true; + } + + /// Locate the replica tip from a LIST alone: the latest generation with + /// a `snapshot.json`, and its last segment (or the snapshot boundary + /// when no segment has shipped). + async fn replica_tip(&mut self) -> Result, SyncError> { + let components = self.storage.list_dirs(GENERATIONS_PREFIX).await?; + for component in components.iter().rev() { + let Some(generation) = GenerationId::parse(component) else { + slog::warn!(self.log, "Skipping foreign directory under generations/"; + "component" => component.as_str()); + continue; + }; + let bytes = match self.storage.get(&snapshot_meta_key(&generation)).await { + Ok(bytes) => bytes, + Err(StorageError::NotFound { .. }) => continue, + Err(error) => return Err(error.into()), + }; + let Ok(snapshot) = SnapshotMeta::from_bytes(&bytes) else { + slog::error!(self.log, "Skipping generation with unparsable snapshot.json"; + "generation" => generation.as_str()); + continue; + }; + let wal_prefix = format!("{}{WAL_DIR}/", generation_prefix(&generation)); + let keys = self.storage.list(&wal_prefix).await?; + // Keys sort lexicographically == (epoch, offset) numerically, + // so the last parseable key is the tip. + let mut tip_segment = None; + for key in &keys { + if let Some((parsed, segment)) = parse_segment_key(key) + && parsed == generation + { + tip_segment = Some(segment); + } + } + return Ok(Some(match tip_segment { + Some(segment) => ReplicaTip { + generation, + epoch: segment.epoch, + salt: segment.salt, + end: segment.end, + last_segment: Some(segment), + }, + None => ReplicaTip { + generation, + epoch: snapshot.wal_boundary.epoch, + salt: (snapshot.wal_boundary.salt1, snapshot.wal_boundary.salt2), + end: 0, + last_segment: None, + }, + })); + } + Ok(None) + } + + /// One full ship pass over the local WAL (used by both `ship_once` and + /// the snapshot's `ShipTail` phase). + async fn ship_pass(&mut self) -> Result { + let wal_path = self.wal_path(); + let header_state = { + let wal_path = wal_path.clone(); + spawn_blocking(move || read_wal_header_state(&wal_path)) + .await + .map_err(SyncError::Join)?? + }; + let (header, raw_header) = match header_state { + WalHeaderState::Present { header, raw } => (header, raw), + WalHeaderState::Missing => return Ok(0), + WalHeaderState::Unreadable(error) => { + // Possibly a torn concurrent header write; retry next tick. + slog::warn!(self.log, "Local WAL header unreadable; skipping this ship pass"; + "error" => %error); + return Ok(0); + }, + }; + // Bind an awaited epoch to the first observed salt cycle. + if let Some((generation, epoch)) = self.awaiting.take() { + slog::info!(self.log, "Binding awaited epoch to the local WAL salts"; + "generation" => generation.as_str(), "epoch" => epoch); + self.epoch_checkpointed = false; + self.position = Some(Position { + generation, + epoch, + salt: header.salt, + offset: 0, + checksum: (0, 0), + }); + } + let Some(mut position) = self.position.clone() else { + return Ok(0); + }; + if header.salt != position.salt && !self.transition_epoch(&mut position, &header) { + return Ok(0); + } + // An oversized-drain generation pins epoch 0 at the boundary until the + // WAL restarts (the restart is handled by `transition_epoch` above, + // which rebinds and clears this state). Until then, ship nothing: the + // below-boundary frames are already in the snapshot, and any new + // committed frames above the boundary cannot ship into epoch 0 (its + // first segment must start at 0). They are backfilled by the draining + // checkpoint and captured by the next snapshot. + if self.oversized_drain == OversizedDrain::AwaitingRestart { + return Ok(0); + } + // Discard-scan the committed extent from the current position FIRST + // (page bytes discarded), then retention-read only up to that extent. + // Without this, every ship poll re-buffered any large open + // transaction past the last commit just to discover there was nothing + // new to ship (a multi-GiB re-allocation each tick). Bounded per pass: + // the scan stops after roughly a full pass worth of committed data, + // and the next pass resumes where this one left off. + let (scan_offset, scan_checksum) = if position.offset == 0 { + (WAL_HEADER_SIZE, header.checksum) + } else { + (position.offset, position.checksum) + }; + let scan_budget = MAX_SEGMENTS_PER_PASS.saturating_mul(SEGMENT_MAX_BYTES); + let max_txn_bytes = self.config.max_transaction_bytes; + let extent = { + let wal_path = wal_path.clone(); + let scanned = spawn_blocking(move || { + scan_committed_end( + &wal_path, + header, + scan_offset, + scan_checksum, + scan_budget, + max_txn_bytes, + ) + }) + .await + .map_err(SyncError::Join)?; + match scanned { + Ok(extent) => extent, + // A committed transaction beyond the bound cannot ship (restore + // would reject the oversized segment) and, being durable, cannot + // be split retroactively. Capture it with a re-snapshot and + // drain the WAL instead of wedging the ship path. + Err(error) if error.is_transaction_too_large() => { + self.diverge_oversized(); + return Ok(0); + }, + Err(error) => return Err(error), + } + }; + if extent <= scan_offset { + return Ok(0); + } + self.ship_up_to_extent(&wal_path, header, &raw_header, &mut position, extent) + .await + } + + /// Retention-read and ship committed segments from the current position up + /// to `extent` (the boundary the discard scan found), at most + /// [`MAX_SEGMENTS_PER_PASS`]. Each read is capped at the extent so the + /// final chunk stops exactly on the last committed frame and never buffers + /// the uncommitted tail beyond it. + async fn ship_up_to_extent( + &mut self, + wal_path: &Utf8Path, + header: WalHeader, + raw_header: &[u8; 32], + position: &mut Position, + extent: u64, + ) -> Result { + let mut shipped = 0u64; + while shipped < MAX_SEGMENTS_PER_PASS { + let (offset, checksum) = if position.offset == 0 { + (WAL_HEADER_SIZE, header.checksum) + } else { + (position.offset, position.checksum) + }; + if offset >= extent { + break; + } + let max_bytes = SEGMENT_MAX_BYTES.min(extent.saturating_sub(offset)); + let chunk = { + let wal_path = wal_path.to_owned(); + spawn_blocking(move || { + scan_next_chunk(&wal_path, header, offset, checksum, max_bytes) + }) + .await + .map_err(SyncError::Join)?? + }; + let Some(chunk) = chunk else { break }; + // Defensive: the discard scan already intercepts an oversized + // transaction before the retaining read buffers it, but should one + // reach here, treat it the same (re-snapshot, never a hard error). + if let Err(error) = self.ship_chunk(chunk, position, raw_header).await { + if error.is_transaction_too_large() { + self.diverge_oversized(); + return Ok(shipped); + } + return Err(error); + } + shipped = shipped.saturating_add(1); + } + Ok(shipped) + } + + /// Ship one committed chunk: refuse an oversized transaction, compress the + /// segment (prepending the 32-byte WAL header for the first segment of an + /// epoch, so restore can rebuild a `-wal` file by decompress-and- + /// concatenate), upload it, advance `position`, and persist the meta. + async fn ship_chunk( + &mut self, + chunk: CommittedChunk, + position: &mut Position, + raw_header: &[u8; 32], + ) -> Result<(), SyncError> { + // A single transaction is always shipped whole (commit atomicity beats + // the size cap), but a transaction beyond `max_transaction_bytes` + // (defaulting to the restore decompression bound) cannot ship: restore + // would reject the oversized segment. Signal it; `ship_pass` converts + // it into an oversized-transaction re-snapshot rather than shipping + // anything corrupt. The discard scan normally intercepts it first, so + // this is the defensive backstop. + let chunk_bytes = u64::try_from(chunk.bytes.len()).unwrap_or(u64::MAX); + let max_transaction_bytes = self.config.max_transaction_bytes; + if chunk_bytes.saturating_add(WAL_HEADER_SIZE) > max_transaction_bytes { + return Err(SyncError::TransactionTooLarge { + bytes: chunk_bytes, + max_bytes: max_transaction_bytes, + }); + } + let start = if position.offset == 0 { + 0 + } else { + chunk.start_offset + }; + let key = segment_key( + &position.generation, + &SegmentKey { + epoch: position.epoch, + salt: position.salt, + start, + end: chunk.end_offset, + }, + ); + let CommittedChunk { + end_offset, + checksum_at_end, + bytes, + .. + } = chunk; + let raw = if start == 0 { + let mut with_header = raw_header.to_vec(); + with_header.extend_from_slice(&bytes); + with_header + } else { + bytes + }; + let compressed = spawn_blocking(move || compress_segment(&raw)) + .await + .map_err(SyncError::Join)??; + self.storage.put(&key, Bytes::from(compressed)).await?; + #[cfg(feature = "otel")] + bencher_otel::ApiMeter::increment(bencher_otel::ApiCounter::ReplicaSegmentShip); + position.offset = end_offset; + position.checksum = checksum_at_end; + self.epoch_checkpointed = false; + self.position = Some(position.clone()); + self.store_meta_for(position).await?; + slog::debug!(self.log, "Shipped WAL segment"; "key" => &key); + Ok(()) + } + + /// Handle a WAL salt change before shipping. Returns whether shipping + /// may proceed; `false` means a divergence was recorded. + fn transition_epoch(&mut self, position: &mut Position, header: &WalHeader) -> bool { + if self.oversized_drain == OversizedDrain::AwaitingRestart { + // The oversized-drain generation pinned epoch 0 at the boundary + // (offset > 0) waiting for exactly this restart: the checkpoint + // (ours in sole mode, Litestream's in shadow mode) backfilled the + // WAL below the boundary and a writer has now restarted it. REBIND + // epoch 0 to the fresh cycle at offset 0 (rather than advancing to + // epoch+1, which would leave epoch 0 with no segments and break + // restore's epoch contiguity); the oversized transaction lives on + // only in the snapshot body now, never a shipped segment. + slog::info!(self.log, "Oversized-drain complete: rebinding epoch 0 to the restarted WAL cycle"; + "epoch" => position.epoch); + position.salt = header.salt; + position.offset = 0; + position.checksum = (0, 0); + self.oversized_drain = OversizedDrain::Idle; + self.epoch_checkpointed = false; + self.position = Some(position.clone()); + return true; + } + if position.offset == 0 { + // Nothing shipped in this epoch yet: rebind it to the new + // cycle (the epoch-0 lazy binding rule; also covers any epoch + // start). The vanished cycle held nothing beyond what a full + // backfill already captured. + if self.shadow { + slog::warn!(self.log, "Shadow: rebinding epoch to a new WAL cycle (boundary unverified)"; + "epoch" => position.epoch); + } else { + slog::info!(self.log, "Rebinding unshipped epoch to the new WAL salts"; + "epoch" => position.epoch); + } + position.salt = header.salt; + position.checksum = (0, 0); + self.epoch_checkpointed = false; + self.position = Some(position.clone()); + true + } else if self.shadow { + if !self.epoch_checkpointed { + // Commits between the last shadow tick and a + // Litestream-driven restart can be lost from the shadow + // replica; the daily verify catches any divergence. + slog::warn!(self.log, "Shadow: WAL restarted without our checkpoint; epoch boundary unverified"; + "epoch" => position.epoch); + #[cfg(feature = "otel")] + bencher_otel::ApiMeter::increment( + bencher_otel::ApiCounter::ReplicaShadowUnverifiedBoundary, + ); + } + position.epoch = position.epoch.saturating_add(1); + position.salt = header.salt; + position.offset = 0; + position.checksum = (0, 0); + self.epoch_checkpointed = false; + self.position = Some(position.clone()); + true + } else if self.epoch_checkpointed { + // `SQLite` increments salt1 by exactly one per WAL restart + // (`walRestartHdr`). A larger jump proves at least one WHOLE + // cycle was written and buried (checkpointed then overwritten) + // by an external writer: its commits never shipped, so the + // lineage is broken (invariant I6: never guess). + // + // Residual limitation (same as the resume path, see + // `resume_without_wal`): this only catches jumps GREATER than one. + // A SINGLE external checkpoint-plus-restart increments salt1 by + // exactly one, indistinguishable from a legitimate restart of our + // own fully-shipped, fully-backfilled WAL, so if it buried frames + // committed after our last shipped frame those are silently lost + // from the replica. The periodic restore-and-compare verification + // is the backstop for that case. + if header.salt.0 != position.salt.0.wrapping_add(1) { + self.transition_divergence( + position, + "WAL restarted more than once since the last shipped frame (buried cycle)", + ); + return false; + } + slog::info!(self.log, "WAL restarted after a completed checkpoint; starting the next epoch"; + "epoch" => position.epoch.saturating_add(1)); + position.epoch = position.epoch.saturating_add(1); + position.salt = header.salt; + position.offset = 0; + position.checksum = (0, 0); + self.epoch_checkpointed = false; + self.position = Some(position.clone()); + true + } else { + self.transition_divergence( + position, + "WAL restarted with unshipped frames (salt change without our checkpoint)", + ); + false + } + } + + /// Record a mid-stream divergence discovered during an epoch transition. + fn transition_divergence(&mut self, position: &Position, reason: &'static str) { + slog::error!(self.log, "Replica divergence; starting a new generation"; + "reason" => reason, + "epoch" => position.epoch, "shipped_offset" => position.offset); + #[cfg(feature = "otel")] + bencher_otel::ApiMeter::increment(bencher_otel::ApiCounter::ReplicaDivergence); + self.position = None; + self.awaiting = None; + self.pending_new_generation = true; + } + + /// Handle an oversized COMMITTED transaction found in the ship path. + /// + /// The transaction is already durable in the source, so "split the write" + /// cannot apply retroactively and shipping it is impossible (restore + /// rejects a segment beyond the decompression bound). Force a whole- + /// database re-snapshot to capture its state, then drain the WAL below the + /// boundary via a checkpoint (the snapshot IS the shipped state for those + /// frames): [`Self::epoch0_bind_offset`] pins the new position at the + /// boundary ([`OversizedDrain::AwaitingRestart`]) so the ship path stops + /// re-hitting the transaction and a normal checkpoint can restart the WAL. + /// + /// Rate-limited: the drain normally breaks the loop, but a drain that + /// cannot complete (a pinned reader, or a shadow replica awaiting + /// Litestream) must not mint a fresh generation every tick. + fn diverge_oversized(&mut self) { + let now = self.now_secs(); + if let Some(until) = self.oversized_retrigger_until + && now < until + { + return; + } + slog::error!(self.log, + "A committed transaction exceeds the shippable WAL bound; forcing a re-snapshot to capture it and draining the WAL below the boundary (split the write to avoid recurrence)"; + "max_transaction_bytes" => self.config.max_transaction_bytes); + #[cfg(feature = "otel")] + bencher_otel::ApiMeter::increment(bencher_otel::ApiCounter::ReplicaDivergence); + self.position = None; + self.awaiting = None; + self.pending_new_generation = true; + self.oversized_drain = OversizedDrain::Pending; + self.oversized_retrigger_until = Some(now.saturating_add(OVERSIZED_RETRIGGER_MIN_SECS)); + } + + /// Where a freshly finalized generation binds epoch 0, as + /// `(offset, checksum)`. Normally offset 0 (re-ship the WAL from the + /// start). For an oversized-transaction re-snapshot ([`OversizedDrain::Pending`]) + /// it pins epoch 0 at the boundary and advances to + /// [`OversizedDrain::AwaitingRestart`] so the ship path stops re-hitting the + /// unshippable transaction and a normal checkpoint can backfill and restart + /// the WAL. A stale drain state from any other finalize is reset to + /// [`OversizedDrain::Idle`]. + async fn epoch0_bind_offset( + &mut self, + generation: &GenerationId, + boundary_offset: u64, + ) -> Result<(u64, (u32, u32)), SyncError> { + if self.oversized_drain != OversizedDrain::Pending { + self.oversized_drain = OversizedDrain::Idle; + return Ok((0, (0, 0))); + } + self.oversized_drain = OversizedDrain::Idle; + let recovered = if boundary_offset == 0 { + Some((0, 0)) + } else { + let wal_path = self.wal_path(); + spawn_blocking(move || checksum_at_offset(&wal_path, boundary_offset)) + .await + .map_err(SyncError::Join)?? + }; + if let Some(checksum) = recovered { + self.oversized_drain = OversizedDrain::AwaitingRestart; + slog::info!(self.log, "Oversized-drain: pinning epoch 0 at the boundary to drain the WAL"; + "generation" => generation.as_str(), "boundary" => boundary_offset); + Ok((boundary_offset, checksum)) + } else { + // The local WAL no longer reaches the boundary (a concurrent + // change); fall back to a plain epoch-0 bind. The oversized + // transaction is re-detected and the drain retried (rate-limited). + slog::warn!(self.log, "Oversized-drain boundary checksum unavailable; deferring the drain"; + "generation" => generation.as_str()); + Ok((0, (0, 0))) + } + } + + /// Whether a checkpoint should run this tick. + /// + /// Checkpoint liveness is not guaranteed under a continuously committing + /// writer: each attempt re-verifies the tail under `BEGIN IMMEDIATE` and + /// defers (`SkippedUnshipped`) whenever a fresh commit has landed past the + /// shipped position (invariant I1), so sustained write traffic can starve + /// checkpoints and let the WAL grow. The `ReplicaCheckpointSkippedUnshipped` + /// counter is the alert signal for that condition; a persistently + /// unshipped-skipped checkpoint means shipping is not keeping up. + async fn checkpoint_due(&mut self, now: i64) -> bool { + if self.shadow || self.position.is_none() { + return false; + } + let interval = i64::try_from(self.config.checkpoint_interval.as_secs()).unwrap_or(i64::MAX); + if now.saturating_sub(self.last_checkpoint_secs) < interval { + return false; + } + let wal_path = self.wal_path(); + let pages = spawn_blocking(move || count_wal_pages(&wal_path)) + .await + .unwrap_or(0); + pages >= u64::from(self.config.min_checkpoint_pages) + } + + /// Whether a snapshot should start this tick. + fn snapshot_due(&self, now: i64) -> bool { + if self.snapshot.is_some() { + return false; + } + if self.pending_new_generation { + return true; + } + let interval = i64::try_from(self.config.snapshot_interval.as_secs()).unwrap_or(i64::MAX); + self.position.is_some() && now.saturating_sub(self.generation_birth_secs) >= interval + } + + /// Advance the snapshot machine one step; `Ok(None)` means finished. + async fn run_snapshot_step( + &mut self, + job: SnapshotJob, + ) -> Result, SyncError> { + match job { + SnapshotJob::ShipTail => { + // Drain committed frames into the old generation before + // fixing the boundary, so it stays a complete restore point + // under retention; skipped when no lineage is usable (fresh + // replica or divergence). Only a capped pass (a genuine + // backlog) repeats: commits racing in DURING the drain are + // covered by the new generation, which re-ships the whole + // WAL from offset 0, so chasing them would never terminate + // under a busy writer. + if self.position.is_some() { + let shipped = self.ship_pass().await?; + if shipped >= MAX_SEGMENTS_PER_PASS { + return Ok(Some(SnapshotJob::ShipTail)); + } + } + Ok(Some(SnapshotJob::CreateGeneration)) + }, + SnapshotJob::CreateGeneration => self.snapshot_create().await.map(Some), + SnapshotJob::Copying(copy) => self.snapshot_copy(copy).await, + SnapshotJob::Finalize(finalize) => { + self.snapshot_finalize(finalize).await?; + Ok(None) + }, + } + } + + /// A fresh generation id strictly greater than the current one, so + /// lexicographic order keeps matching creation order even when a + /// retrigger lands within the same clock second (restore always picks + /// the greatest valid generation). Bumps the timestamp by up to a + /// minute before giving up (a rewound clock beyond that is logged and + /// tolerated: restore would prefer the older lineage until the clock + /// catches up). + fn next_generation_id(&self) -> GenerationId { + let floor = self.generation().max(self.generation_floor.as_ref()); + let now = self.now_secs(); + for bump in 0..60 { + let Ok(created) = bencher_json::DateTime::try_from(now.saturating_add(bump)) else { + continue; + }; + let candidate = GenerationId::generate(created); + if floor.is_none_or(|floor| candidate > *floor) { + return candidate; + } + } + // The clock has rewound more than a minute behind the observed tip: + // fall back to the floor's own SUCCESSOR (same timestamp, suffix + // plus one), which is guaranteed to sort after it. A below-floor id + // would make restore silently pick the older lineage forever. + slog::warn!( + self.log, + "Clock is behind the newest generation; minting its successor instead" + ); + if let Some(successor) = floor.and_then(GenerationId::successor) { + return successor; + } + GenerationId::generate(self.clock.now()) + } + + /// Fix the generation id, run the single-step online backup into the + /// scratch file, record the boundary; open the multipart upload. + async fn snapshot_create(&mut self) -> Result { + #[cfg(feature = "otel")] + bencher_otel::ApiMeter::increment(bencher_otel::ApiCounter::ReplicaGenerationCreate); + let generation = self.next_generation_id(); + let db_path = self.db.db_path.clone(); + let scratch = self.snapshot_scratch_path(); + let (db_len, page_size, live_fingerprint) = { + let scratch = scratch.clone(); + spawn_blocking(move || { + backup_to_scratch(&db_path, &scratch)?; + let (db_len, page_size) = read_scratch_info(&scratch)?; + let live_fingerprint = live_db_fingerprint(&db_path)?; + Ok::<_, SnapshotError>((db_len, page_size, live_fingerprint)) + }) + .await + .map_err(SyncError::Join)?? + }; + // The boundary, read AFTER the backup: the current cycle's salts + // plus its valid committed length. The length is the mandatory- + // replay threshold: every backup-contained boundary-cycle frame + // lies below it (frames committed after the backup only OVERSTATE + // the threshold, which is the safe direction; they are replayed + // once shipped, like everything else). + let (boundary_salt, boundary_offset) = { + let wal_path = self.wal_path(); + spawn_blocking(move || wal_committed_extent(&wal_path)) + .await + .map_err(SyncError::Join)?? + .map_or(((0, 0), 0), |extent| (extent.salt, extent.committed_end)) + }; + let upload = self + .storage + .start_multipart(&snapshot_key(&generation)) + .await?; + let encoder = new_snapshot_encoder()?; + slog::info!(self.log, "Snapshot started"; + "generation" => generation.as_str(), "db_bytes" => db_len); + Ok(SnapshotJob::Copying(Box::new(CopyJob { + generation, + boundary_salt, + boundary_offset, + upload, + hasher: Sha256::new(), + encoder, + scratch, + db_len, + offset: 0, + page_size, + live_fingerprint, + }))) + } + + /// One budgeted upload step over the backup scratch file. + async fn snapshot_copy( + &mut self, + copy: Box, + ) -> Result, SyncError> { + let CopyJob { + generation, + boundary_salt, + boundary_offset, + mut upload, + hasher, + encoder, + scratch, + db_len, + offset, + page_size, + live_fingerprint, + } = *copy; + let budget = self.copy_budget(); + let step = { + let scratch = scratch.clone(); + spawn_blocking(move || copy_step(&scratch, encoder, hasher, offset, db_len, budget)) + .await + .map_err(SyncError::Join)? + }; + let step = match step { + Ok(step) => step, + Err(error) => { + slog::error!(self.log, "Snapshot aborted"; + "generation" => generation.as_str(), "error" => %error); + self.abort_upload(upload).await; + self.remove_scratch(&scratch).await; + return Err(error.into()); + }, + }; + match step { + CopyStepResult::Continue { + encoder, + hasher, + compressed, + offset, + } => { + if !compressed.is_empty() + && let Err(error) = upload.write_part(Bytes::from(compressed)).await + { + self.abort_upload(upload).await; + self.remove_scratch(&scratch).await; + return Err(error.into()); + } + Ok(Some(SnapshotJob::Copying(Box::new(CopyJob { + generation, + boundary_salt, + boundary_offset, + upload, + hasher, + encoder, + scratch, + db_len, + offset, + page_size, + live_fingerprint, + })))) + }, + CopyStepResult::Done { hasher, compressed } => { + if !compressed.is_empty() + && let Err(error) = upload.write_part(Bytes::from(compressed)).await + { + self.abort_upload(upload).await; + self.remove_scratch(&scratch).await; + return Err(error.into()); + } + Ok(Some(SnapshotJob::Finalize(Box::new(FinalizeJob { + generation, + boundary_salt, + boundary_offset, + upload, + sha256: hex::encode(hasher.finalize()), + db_bytes: db_len, + page_size, + scratch, + live_fingerprint, + })))) + }, + } + } + + /// Commit the generation: finish the body upload, re-check the boundary + /// salts, PUT `snapshot.json` LAST, reset the position, and prune. A salt + /// change since `CreateGeneration` is the epoch-0 lazy binding rule in + /// shadow mode (rebind) but an external-checkpointer divergence in sole + /// mode (abort with [`SyncError::SnapshotBoundaryDiverged`], no marker). + /// Validate (or rebind) the snapshot boundary against the CURRENT WAL + /// salt at finalize time. Nothing has shipped into the new generation yet + /// by construction (the engine is sequential), so a salt change since + /// `CreateGeneration` means the WAL restarted between the backup read point + /// and here; the engine's own checkpoints are suppressed while a snapshot + /// job is active, so an engine-own path cannot cause this. + /// + /// Shadow mode: Litestream owns checkpoints and legitimately restarts the + /// WAL. Keep committing (the shadow replica is disposable) but flag the + /// UNVERIFIED boundary and rebind to the new cycle at offset 0, mirroring + /// the ship-path rebind's observability. + /// + /// Sole mode: the restart is an EXTERNAL checkpointer that may have + /// backfilled (then buried) frames committed after the backup read point + /// which the replica never saw. Rebinding would commit a marker whose + /// restore replays new-cycle segments onto a snapshot body missing those + /// frames (a torn restore). ABORT instead: no snapshot.json is written + /// (the uploaded body stays invisible and is pruned as stale) and a + /// fresh generation is scheduled; the WAL keeps buffering, so nothing is + /// lost, only re-snapshotted. + async fn finalize_boundary( + &mut self, + generation: &GenerationId, + boundary_salt: (u32, u32), + boundary_offset: u64, + ) -> Result<((u32, u32), u64), SyncError> { + let Some(salt) = self.current_wal_salt().await? else { + return Ok((boundary_salt, boundary_offset)); + }; + if salt == boundary_salt { + return Ok((boundary_salt, boundary_offset)); + } + if self.shadow { + slog::warn!(self.log, + "Shadow: rebinding snapshot boundary to a restarted WAL cycle (boundary unverified; restore from this generation may be inconsistent if the buried cycle held unshipped frames)"; + "generation" => generation.as_str()); + #[cfg(feature = "otel")] + bencher_otel::ApiMeter::increment( + bencher_otel::ApiCounter::ReplicaShadowUnverifiedBoundary, + ); + Ok((salt, 0)) + } else { + slog::error!(self.log, + "External checkpoint restarted the WAL during the snapshot; aborting the generation (no snapshot.json) and scheduling a fresh one"; + "generation" => generation.as_str(), + "boundary_salt" => format!("{boundary_salt:?}"), + "current_salt" => format!("{salt:?}")); + #[cfg(feature = "otel")] + bencher_otel::ApiMeter::increment(bencher_otel::ApiCounter::ReplicaDivergence); + Err(SyncError::SnapshotBoundaryDiverged) + } + } + + async fn snapshot_finalize(&mut self, finalize: Box) -> Result<(), SyncError> { + let FinalizeJob { + generation, + boundary_salt, + boundary_offset, + upload, + sha256, + db_bytes, + page_size, + scratch, + live_fingerprint, + } = *finalize; + if let Err(error) = upload.finish().await { + self.remove_scratch(&scratch).await; + return Err(error.into()); + } + self.remove_scratch(&scratch).await; + let (boundary_salt, boundary_offset) = self + .finalize_boundary(&generation, boundary_salt, boundary_offset) + .await?; + let snapshot_meta = SnapshotMeta { + version: SNAPSHOT_META_VERSION, + created: self.clock.now().into_inner().to_rfc3339(), + db_bytes, + page_size, + sha256, + wal_boundary: WalBoundary { + epoch: 0, + salt1: boundary_salt.0, + salt2: boundary_salt.1, + offset: boundary_offset, + }, + }; + let marker = snapshot_meta.to_bytes().map_err(SyncError::SnapshotMeta)?; + self.storage + .put(&snapshot_meta_key(&generation), Bytes::from(marker)) + .await?; + slog::info!(self.log, "Snapshot complete"; + "generation" => generation.as_str(), "db_bytes" => db_bytes); + // A fresh generation normally binds epoch 0 at offset 0; an oversized- + // transaction re-snapshot instead pins it at the boundary to drain the + // WAL (see `epoch0_bind_offset`). + let (offset, checksum) = self + .epoch0_bind_offset(&generation, boundary_offset) + .await?; + let position = Position { + generation, + epoch: 0, + salt: boundary_salt, + offset, + checksum, + }; + self.position = Some(position.clone()); + self.generation_floor = Some(position.generation.clone()); + self.awaiting = None; + self.epoch_checkpointed = false; + self.pending_new_generation = false; + self.generation_birth_secs = self.now_secs(); + // The meta file is advisory: log a failure instead of failing the + // completed snapshot. + if let Err(error) = self.store_meta_for(&position).await { + slog::warn!(self.log, "Failed to store replica meta after snapshot"; + "error" => %error); + } + if let Err(error) = self.prune_once().await { + slog::warn!(self.log, "Prune failed; retrying at the next snapshot"; + "error" => %error); + } + // Sole mode: nothing legitimate mutates the LIVE database file + // during a snapshot, so a moved fingerprint proves an external + // checkpointer that may have buried post-backup frames the replica + // never saw. The committed snapshot stays (it is consistent); a + // follow-up generation recaptures whatever was buried. Shadow mode + // skips this (Litestream checkpoints constantly; the daily + // verification is the backstop there). + if !self.shadow { + let db_path = self.db.db_path.clone(); + let current = spawn_blocking(move || live_db_fingerprint(&db_path)) + .await + .map_err(SyncError::Join)??; + if current != live_fingerprint { + slog::error!(self.log, "External checkpoint detected during the snapshot; scheduling a follow-up generation"; + "recorded" => format!("{live_fingerprint:?}"), + "current" => format!("{current:?}")); + #[cfg(feature = "otel")] + bencher_otel::ApiMeter::increment(bencher_otel::ApiCounter::ReplicaDivergence); + self.pending_new_generation = true; + } + } + Ok(()) + } + + /// The snapshot backup scratch file, next to the database (same + /// volume). Deleted on completion, on abort, and before reuse. + fn snapshot_scratch_path(&self) -> Utf8PathBuf { + Utf8PathBuf::from(format!("{}.snapshot-scratch", self.db.db_path)) + } + + /// Best-effort scratch removal (a leftover is replaced by the next + /// snapshot's backup). + async fn remove_scratch(&mut self, scratch: &Utf8Path) { + let scratch = scratch.to_owned(); + let removed = spawn_blocking(move || match std::fs::remove_file(&scratch) { + Ok(()) => Ok(()), + Err(error) if error.kind() == ErrorKind::NotFound => Ok(()), + Err(error) => Err(error), + }) + .await; + match removed { + Ok(Ok(())) => {}, + Ok(Err(error)) => { + slog::warn!(self.log, "Failed to remove the snapshot scratch file"; + "error" => %error); + }, + Err(error) => { + slog::warn!(self.log, "Scratch removal task panicked"; "error" => %error); + }, + } + } + + /// Per-step copy budget: `snapshot_throttle_mib * sync_interval`, at + /// least 1 MiB. + fn copy_budget(&self) -> u64 { + u64::from(self.config.snapshot_throttle_mib) + .saturating_mul(MIB) + .saturating_mul(self.config.sync_interval.as_secs().max(1)) + .max(MIB) + } + + /// Best-effort multipart abort (a failed abort only leaves an invisible + /// partial upload behind). + async fn abort_upload(&mut self, upload: MultipartUpload) { + if let Err(error) = upload.abort().await { + slog::warn!(self.log, "Failed to abort snapshot upload"; "error" => %error); + } + } + + /// The current local WAL header salts, if a parseable header exists. + async fn current_wal_salt(&mut self) -> Result, SyncError> { + let wal_path = self.wal_path(); + let state = spawn_blocking(move || read_wal_header_state(&wal_path)) + .await + .map_err(SyncError::Join)??; + if let WalHeaderState::Present { header, .. } = state { + Ok(Some(header.salt)) + } else { + Ok(None) + } + } + + /// Generations that must never be pruned: the one being written to and + /// the one an in-flight snapshot is building. + fn protected_generations(&self) -> Vec { + let mut protected = Vec::new(); + if let Some(generation) = self.generation() { + protected.push(generation.clone()); + } + match &self.snapshot { + Some(SnapshotJob::Copying(copy)) => protected.push(copy.generation.clone()), + Some(SnapshotJob::Finalize(finalize)) => protected.push(finalize.generation.clone()), + Some(SnapshotJob::ShipTail | SnapshotJob::CreateGeneration) | None => {}, + } + protected + } + + /// Persist the advisory meta file (temp + fsync + rename) after a state + /// change. + async fn store_meta_for(&mut self, position: &Position) -> Result<(), SyncError> { + let meta = ReplicaMeta { + version: META_VERSION, + generation: position.generation.as_str().to_owned(), + epoch: position.epoch, + salt1: position.salt.0, + salt2: position.salt.1, + shipped_offset: position.offset, + epoch_shipped_through_checkpoint: self.epoch_checkpointed, + shadow: self.shadow, + }; + let db_path = self.db.db_path.clone(); + spawn_blocking(move || meta.store(&db_path)) + .await + .map_err(SyncError::Join)??; + Ok(()) + } +} + +/// The cumulative checksum stored in the final frame header of a raw WAL +/// segment: frame-header bytes 16..24 of the last frame. +fn segment_tail_checksum(raw: &[u8], frame_size: u64) -> Option<(u32, u32)> { + let frame_size = usize::try_from(frame_size).ok()?; + let frame_start = raw.len().checked_sub(frame_size)?; + let frame_header: &[u8; 24] = raw + .get(frame_start..frame_start.checked_add(24)?)? + .try_into() + .ok()?; + Some(crate::wal::FrameHeader::parse(frame_header).checksum) +} + +/// Whether the advisory meta exactly matches the replica tip: the proof +/// required for an epoch+1 (or awaiting) resume instead of a re-snapshot. +/// +/// The `meta.epoch == tip.epoch` clause is what keeps a same-epoch salt +/// collision (two epoch directories with the same number but different salts, +/// which `plan_epochs` would soft-stop on forever) from ever being created. +/// The tip epoch is the highest epoch with a shipped segment, so an epoch+1 +/// resume only ever extends at `tip.epoch + 1`, a number no directory holds +/// yet. If the restored meta records a LOWER epoch than the tip (e.g. a restore +/// soft-stopped at epoch N below a corrupt-but-present epoch N+1), this clause +/// fails and resume diverges to a fresh generation rather than binding new +/// salts onto the already-occupied epoch N+1. See the pinned regression tests +/// `plan_discards_epoch_with_duplicate_salt_lineages` (restore.rs) and +/// `resume_after_soft_stop_below_corrupt_epoch_forces_new_generation`. +fn meta_matches(meta: Option<&ReplicaMeta>, tip: &ReplicaTip, shadow: bool) -> bool { + meta.is_some_and(|meta| { + meta.generation == tip.generation.as_str() + && meta.epoch == tip.epoch + && meta.shipped_offset == tip.end + && meta.epoch_shipped_through_checkpoint + && meta.shadow == shadow + }) +} + +/// Read and parse the local WAL header without opening `SQLite`. +fn read_wal_header_state(wal_path: &Utf8Path) -> Result { + let read_err = |error| SyncError::WalIo { + path: wal_path.to_owned(), + error, + }; + let mut file = match File::open(wal_path) { + Ok(file) => file, + Err(error) if error.kind() == ErrorKind::NotFound => return Ok(WalHeaderState::Missing), + Err(error) => return Err(read_err(error)), + }; + let mut raw = [0u8; 32]; + match file.read_exact(&mut raw) { + Ok(()) => {}, + Err(error) if error.kind() == ErrorKind::UnexpectedEof => { + return Ok(WalHeaderState::Missing); + }, + Err(error) => return Err(read_err(error)), + } + match parse_wal_header(&raw) { + Ok(header) => Ok(WalHeaderState::Present { header, raw }), + Err(error) => Ok(WalHeaderState::Unreadable(error)), + } +} + +/// Scan the next committed chunk from `offset` (up to `max_bytes`), +/// re-verifying that the WAL header still carries the expected salts (a +/// concurrent restart otherwise invalidates the resume seed). `Ok(None)` +/// means nothing shippable now. +fn scan_next_chunk( + wal_path: &Utf8Path, + expected: WalHeader, + offset: u64, + checksum: (u32, u32), + max_bytes: u64, +) -> Result, SyncError> { + match open_resumed_scanner(wal_path, expected, offset, checksum)? { + Some(mut scanner) => Ok(scanner.next_committed(max_bytes)?), + None => Ok(None), + } +} + +/// Discard-scan from `offset` to the committed extent (page bytes discarded), +/// re-verifying the header salts. Returns the absolute WAL offset just past +/// the last commit within `max_bytes` (equal to `offset` when nothing new is +/// committed, the WAL is missing, or the salts have changed). `max_txn_bytes` +/// bounds the bytes since the last commit, so an oversized open transaction +/// errors instead of scanning without bound. +fn scan_committed_end( + wal_path: &Utf8Path, + expected: WalHeader, + offset: u64, + checksum: (u32, u32), + max_bytes: u64, + max_txn_bytes: u64, +) -> Result { + match open_resumed_scanner(wal_path, expected, offset, checksum)? { + Some(mut scanner) => Ok(scanner.scan_committed_extent(max_bytes, max_txn_bytes)?), + None => Ok(offset), + } +} + +/// Open the WAL, re-check the header salts against `expected`, and resume a +/// scanner at `(offset, checksum)`. `Ok(None)` when the WAL is missing, +/// header-short, unparsable, or salt-changed (all handled by the caller as +/// "nothing to scan from here"). +fn open_resumed_scanner( + wal_path: &Utf8Path, + expected: WalHeader, + offset: u64, + checksum: (u32, u32), +) -> Result>, SyncError> { + let read_err = |error| SyncError::WalIo { + path: wal_path.to_owned(), + error, + }; + let mut file = match File::open(wal_path) { + Ok(file) => file, + Err(error) if error.kind() == ErrorKind::NotFound => return Ok(None), + Err(error) => return Err(read_err(error)), + }; + let mut raw = [0u8; 32]; + match file.read_exact(&mut raw) { + Ok(()) => {}, + Err(error) if error.kind() == ErrorKind::UnexpectedEof => return Ok(None), + Err(error) => return Err(read_err(error)), + } + let Ok(header) = parse_wal_header(&raw) else { + return Ok(None); + }; + if header.salt != expected.salt { + return Ok(None); + } + Ok(Some(WalScanner::resume(file, header, offset, checksum)?)) +} + +/// The live WAL's header salts plus the byte offset just past its last +/// committed frame. +struct WalExtent { + salt: (u32, u32), + committed_end: u64, +} + +/// Measure the live WAL's [`WalExtent`] (`None` when the WAL is missing or +/// shorter than a header). Used for the snapshot boundary. +fn wal_committed_extent(wal_path: &Utf8Path) -> Result, SyncError> { + let file = match File::open(wal_path) { + Ok(file) => file, + Err(error) if error.kind() == ErrorKind::NotFound => return Ok(None), + Err(error) => { + return Err(SyncError::WalIo { + path: wal_path.to_owned(), + error, + }); + }, + }; + let Some(mut scanner) = WalScanner::open(file)? else { + return Ok(None); + }; + let salt = scanner.header().salt; + // Discard-scan the whole committed WAL (offsets only): materializing every + // frame just to learn where committed data ends was a needless multi-GiB + // allocation. Errors PROPAGATE via `?`: swallowing a mid-scan read/seek + // error and returning the extent measured so far would UNDERSTATE the + // mandatory-replay boundary, and restore would then replay a truncated WAL + // prefix on top of a newer snapshot body and regress pages. Overstating + // the boundary is the only safe direction (see the boundary note in + // `snapshot_create`), so a scan error must fail the snapshot, not shorten + // the boundary. No transaction cap here: the boundary is purely the + // committed length, so an oversized uncommitted tail (which the ship path + // is what refuses) must not spuriously fail the snapshot. + let committed_end = scanner.scan_committed_extent(u64::MAX, u64::MAX)?; + Ok(Some(WalExtent { + salt, + committed_end, + })) +} + +/// Rescan the local WAL chain from offset 0 and recover the running +/// checksum at exactly `target` (a commit boundary the replica claims). +/// `Ok(None)` when the valid chain never reaches `target` exactly. +/// +/// Discard-scan: only offsets and the running checksum are needed, so page +/// bytes are validated and dropped instead of buffering every committed +/// transaction below `target` on the heap. +fn checksum_at_offset(wal_path: &Utf8Path, target: u64) -> Result, SyncError> { + let file = File::open(wal_path).map_err(|error| SyncError::WalIo { + path: wal_path.to_owned(), + error, + })?; + let Some(mut scanner) = WalScanner::open(file)? else { + return Ok(None); + }; + // Stop at the first commit whose end reaches `target` (no transaction cap: + // the checksum must be recoverable regardless of transaction sizes below + // the target). + let max_bytes = target.saturating_sub(WAL_HEADER_SIZE); + scanner.scan_committed_extent(max_bytes, u64::MAX)?; + if scanner.offset() == target { + Ok(Some(scanner.checksum())) + } else { + // The valid chain never lands exactly on `target` (short, or the next + // commit overshoots it). + Ok(None) + } +} + +/// Committed-or-not page count of the WAL file (checkpoint due-ness input); +/// best effort, 0 on any error. +fn count_wal_pages(wal_path: &Utf8Path) -> u64 { + let Ok(metadata) = std::fs::metadata(wal_path) else { + return 0; + }; + let Ok(WalHeaderState::Present { header, .. }) = read_wal_header_state(wal_path) else { + return 0; + }; + metadata + .len() + .saturating_sub(WAL_HEADER_SIZE) + .checked_div(header.frame_size()) + .unwrap_or(0) +} diff --git a/plus/bencher_replica/src/testing/compare.rs b/plus/bencher_replica/src/testing/compare.rs new file mode 100644 index 000000000..ab1fd19b9 --- /dev/null +++ b/plus/bencher_replica/src/testing/compare.rs @@ -0,0 +1,394 @@ +//! Equivalence oracles for replicated databases. +//! +//! The primary oracle is LOGICAL comparison, not whole-file byte compare: +//! DB header bytes 24..28 (file change counter) and 92..100 (version-valid- +//! for + `SQLite` version) legitimately differ between independently +//! checkpointed copies. +//! +//! Both oracles collect a list of human-readable differences and assert it +//! empty, so a failure prints every context-bearing line at once (and the +//! comparison logic stays a plain, panic-free function). + +use camino::Utf8Path; +use rusqlite::{Connection, OpenFlags}; + +use crate::verify::{ + query_serialized_rows, quote_ident, schema_rows, user_table_names, user_version, +}; + +/// Assert that `restored` is logically equivalent to `source`: +/// +/// 1. `PRAGMA integrity_check` == "ok" on `restored`; +/// 2. `sqlite_master` rows equal (`ORDER BY type, name`); +/// 3. every user table's rows equal (`SELECT * ORDER BY rowid`, values +/// serialized with type tags; `WITHOUT ROWID` tables are out of scope per +/// the workload module docs); +/// 4. `PRAGMA user_version` equal. +/// +/// Panics with a readable diff naming the first differing table. +pub fn assert_replica_equivalent(source: &Utf8Path, restored: &Utf8Path) { + let differences = replica_differences(source, restored); + assert!( + differences.is_empty(), + "restored database {restored} is NOT equivalent to source {source}:\n{}", + differences.join("\n") + ); +} + +/// Strict page-level comparison, masking DB header offsets 24..28 and +/// 92..100. Requires equal file lengths. Used only by the frame-apply +/// unit-test family where both sides are fully controlled. +pub fn assert_pages_equal(a: &Utf8Path, b: &Utf8Path) { + let differences = page_differences(a, b); + assert!( + differences.is_empty(), + "database pages of {a} and {b} are NOT equal:\n{}", + differences.join("\n") + ); +} + +/// Collect every logical difference between `source` and `restored`; empty +/// means equivalent. Row-level comparison stops at the first differing +/// table. +fn replica_differences(source: &Utf8Path, restored: &Utf8Path) -> Vec { + let mut differences = Vec::new(); + let Some(source_conn) = open_db(source, &mut differences) else { + return differences; + }; + let Some(restored_conn) = open_db(restored, &mut differences) else { + return differences; + }; + + match integrity_check(&restored_conn) { + Ok(report) if report == ["ok"] => {}, + Ok(report) => differences.push(format!( + "integrity_check failed on restored {restored}: {}", + report.join("; ") + )), + Err(error) => differences.push(format!( + "integrity_check could not run on restored {restored}: {error}" + )), + } + + push_list_difference( + "schema (sqlite_master ORDER BY type, name)", + schema_rows(&source_conn), + schema_rows(&restored_conn), + &mut differences, + ); + + match user_table_names(&source_conn) { + Ok(tables) => { + for table in tables { + let before = differences.len(); + push_list_difference( + &format!("table {table} (SELECT * ORDER BY rowid)"), + table_rows(&source_conn, &table), + table_rows(&restored_conn, &table), + &mut differences, + ); + if differences.len() > before { + // Name the FIRST differing table; later tables would + // only bury it. + break; + } + } + }, + Err(error) => differences.push(format!("failed to list source tables: {error}")), + } + + let source_version = user_version(&source_conn); + let restored_version = user_version(&restored_conn); + match (source_version, restored_version) { + (Ok(source_version), Ok(restored_version)) => { + if source_version != restored_version { + differences.push(format!( + "user_version differs: source {source_version} vs restored {restored_version}" + )); + } + }, + (Err(error), Ok(_) | Err(_)) | (Ok(_), Err(error)) => { + differences.push(format!("user_version query failed: {error}")); + }, + } + differences +} + +/// Open a database for comparison. `READ_WRITE` (not `READ_ONLY`) because +/// opening a database with a live WAL may need to create the `-shm` file; +/// no CREATE flag, so a missing file is a difference, not an empty database. +fn open_db(path: &Utf8Path, differences: &mut Vec) -> Option { + match Connection::open_with_flags( + path, + OpenFlags::SQLITE_OPEN_READ_WRITE | OpenFlags::SQLITE_OPEN_NO_MUTEX, + ) { + Ok(conn) => Some(conn), + Err(error) => { + differences.push(format!("failed to open {path}: {error}")); + None + }, + } +} + +fn integrity_check(conn: &Connection) -> Result, rusqlite::Error> { + let mut statement = conn.prepare("PRAGMA integrity_check")?; + let rows = statement.query_map([], |row| row.get(0))?; + rows.collect() +} + +fn table_rows(conn: &Connection, table: &str) -> Result, rusqlite::Error> { + query_serialized_rows( + conn, + &format!("SELECT * FROM {} ORDER BY rowid", quote_ident(table)), + ) +} + +/// Record the first difference between two serialized row lists (or the +/// query failure that prevented the comparison). +fn push_list_difference( + kind: &str, + source: Result, rusqlite::Error>, + restored: Result, rusqlite::Error>, + differences: &mut Vec, +) { + let source = match source { + Ok(rows) => rows, + Err(error) => { + differences.push(format!("{kind}: source query failed: {error}")); + return; + }, + }; + let restored = match restored { + Ok(rows) => rows, + Err(error) => { + differences.push(format!("{kind}: restored query failed: {error}")); + return; + }, + }; + if source == restored { + return; + } + for (index, (source_row, restored_row)) in source.iter().zip(&restored).enumerate() { + if source_row != restored_row { + differences.push(format!( + "{kind}: first difference at row {index}:\n source: {source_row}\n restored: {restored_row}" + )); + return; + } + } + // Identical shared prefix: the row counts must differ. + let shared = source.len().min(restored.len()); + let first_unmatched = source + .get(shared) + .or_else(|| restored.get(shared)) + .cloned() + .unwrap_or_default(); + differences.push(format!( + "{kind}: row counts differ (source {} vs restored {}); first unmatched row: {first_unmatched}", + source.len(), + restored.len() + )); +} + +/// Collect page-level differences after masking the mutable header ranges; +/// empty means byte-equal. Reports only the first differing page. +fn page_differences(a: &Utf8Path, b: &Utf8Path) -> Vec { + let mut differences = Vec::new(); + let a_bytes = read_masked(a, &mut differences); + let b_bytes = read_masked(b, &mut differences); + if !differences.is_empty() { + return differences; + } + if a_bytes.len() != b_bytes.len() { + differences.push(format!( + "file lengths differ: {a} is {} bytes, {b} is {} bytes", + a_bytes.len(), + b_bytes.len() + )); + return differences; + } + let page_size = header_page_size(&a_bytes); + for (page, (page_a, page_b)) in a_bytes + .chunks(page_size) + .zip(b_bytes.chunks(page_size)) + .enumerate() + { + if page_a != page_b { + let offset = page_a + .iter() + .zip(page_b) + .position(|(byte_a, byte_b)| byte_a != byte_b) + .unwrap_or_default(); + differences.push(format!( + "page {page} differs at byte offset {offset} (page size {page_size})" + )); + return differences; + } + } + differences +} + +/// Read a whole database file with the legitimately-mutable header ranges +/// zeroed: 24..28 (file change counter) and 92..100 (version-valid-for +/// change counter + `SQLITE_VERSION_NUMBER` of the last writer). +fn read_masked(path: &Utf8Path, differences: &mut Vec) -> Vec { + match std::fs::read(path) { + Ok(mut bytes) => { + for range in [24..28, 92..100] { + if let Some(masked) = bytes.get_mut(range) { + masked.fill(0); + } + } + bytes + }, + Err(error) => { + differences.push(format!("failed to read {path}: {error}")); + Vec::new() + }, + } +} + +/// Database page size from header bytes 16..18 (big-endian; the value 1 +/// encodes 65536). Falls back to whole-file comparison when the header is +/// absent or nonsensical. +fn header_page_size(bytes: &[u8]) -> usize { + let high = bytes.get(16).copied().unwrap_or(0); + let low = bytes.get(17).copied().unwrap_or(0); + let raw = (usize::from(high) << 8) | usize::from(low); + match raw { + 1 => 0x0001_0000, + 0 => bytes.len().max(1), + page_size => page_size, + } +} + +#[cfg(test)] +mod tests { + use camino::{Utf8Path, Utf8PathBuf}; + + use super::{assert_pages_equal, assert_replica_equivalent}; + + fn tempdir_path(dir: &tempfile::TempDir) -> &Utf8Path { + Utf8Path::from_path(dir.path()).unwrap() + } + + fn create_db(path: &Utf8Path, statements: &str) { + let conn = rusqlite::Connection::open(path).unwrap(); + conn.execute_batch(statements).unwrap(); + } + + const SEED: &str = "CREATE TABLE t (id INTEGER PRIMARY KEY, data TEXT); + CREATE TABLE mixed (id INTEGER PRIMARY KEY, n INTEGER, r REAL, b BLOB); + INSERT INTO t (data) VALUES ('alpha'), ('beta'); + INSERT INTO mixed (n, r, b) VALUES (1, 2.5, x'0102'), (NULL, NULL, NULL);"; + + fn twin_dbs(tmp: &tempfile::TempDir) -> (Utf8PathBuf, Utf8PathBuf) { + let source = tempdir_path(tmp).join("source.db"); + let restored = tempdir_path(tmp).join("restored.db"); + create_db(&source, SEED); + create_db(&restored, SEED); + (source, restored) + } + + #[test] + fn equivalent_databases_pass() { + let tmp = tempfile::tempdir().unwrap(); + let (source, restored) = twin_dbs(&tmp); + assert_replica_equivalent(&source, &restored); + } + + #[test] + #[should_panic(expected = "table t")] + fn row_difference_names_the_table() { + let tmp = tempfile::tempdir().unwrap(); + let (source, restored) = twin_dbs(&tmp); + let conn = rusqlite::Connection::open(&restored).unwrap(); + conn.execute("UPDATE t SET data = 'gamma' WHERE id = 2", []) + .unwrap(); + drop(conn); + assert_replica_equivalent(&source, &restored); + } + + #[test] + #[should_panic(expected = "row counts differ")] + fn row_count_difference_detected() { + let tmp = tempfile::tempdir().unwrap(); + let (source, restored) = twin_dbs(&tmp); + let conn = rusqlite::Connection::open(&restored).unwrap(); + conn.execute("INSERT INTO t (data) VALUES ('extra')", []) + .unwrap(); + drop(conn); + assert_replica_equivalent(&source, &restored); + } + + #[test] + #[should_panic(expected = "schema")] + fn schema_difference_detected() { + let tmp = tempfile::tempdir().unwrap(); + let (source, restored) = twin_dbs(&tmp); + let conn = rusqlite::Connection::open(&restored).unwrap(); + conn.execute_batch("CREATE INDEX t_data ON t (data)") + .unwrap(); + drop(conn); + assert_replica_equivalent(&source, &restored); + } + + #[test] + #[should_panic(expected = "user_version")] + fn user_version_difference_detected() { + let tmp = tempfile::tempdir().unwrap(); + let (source, restored) = twin_dbs(&tmp); + let conn = rusqlite::Connection::open(&restored).unwrap(); + conn.pragma_update(None, "user_version", 9).unwrap(); + drop(conn); + assert_replica_equivalent(&source, &restored); + } + + #[test] + #[should_panic(expected = "failed to open")] + fn missing_database_is_not_equivalent() { + let tmp = tempfile::tempdir().unwrap(); + let (source, _) = twin_dbs(&tmp); + assert_replica_equivalent(&source, &tempdir_path(&tmp).join("missing.db")); + } + + #[test] + fn pages_equal_masks_header_counters() { + let tmp = tempfile::tempdir().unwrap(); + let (source, _) = twin_dbs(&tmp); + let copy = tempdir_path(&tmp).join("copy.db"); + let mut bytes = std::fs::read(&source).unwrap(); + // Differing change counters and version-valid-for stamps are + // legitimate between independently checkpointed copies. + for offset in (24..28).chain(92..100) { + bytes[offset] ^= 0xff; + } + std::fs::write(©, &bytes).unwrap(); + assert_pages_equal(&source, ©); + } + + #[test] + #[should_panic(expected = "page 1 differs")] + fn page_content_difference_detected() { + let tmp = tempfile::tempdir().unwrap(); + let (source, _) = twin_dbs(&tmp); + let copy = tempdir_path(&tmp).join("copy.db"); + let mut bytes = std::fs::read(&source).unwrap(); + let page_size = usize::from(bytes[16]) << 8 | usize::from(bytes[17]); + bytes[page_size + 100] ^= 0xff; + std::fs::write(©, &bytes).unwrap(); + assert_pages_equal(&source, ©); + } + + #[test] + #[should_panic(expected = "file lengths differ")] + fn page_length_mismatch_detected() { + let tmp = tempfile::tempdir().unwrap(); + let (source, _) = twin_dbs(&tmp); + let copy = tempdir_path(&tmp).join("copy.db"); + let mut bytes = std::fs::read(&source).unwrap(); + bytes.truncate(bytes.len() - 1); + std::fs::write(©, &bytes).unwrap(); + assert_pages_equal(&source, ©); + } +} diff --git a/plus/bencher_replica/src/testing/flaky.rs b/plus/bencher_replica/src/testing/flaky.rs new file mode 100644 index 000000000..05385c04d --- /dev/null +++ b/plus/bencher_replica/src/testing/flaky.rs @@ -0,0 +1,817 @@ +//! Fault-injection storage wrapper: wraps any [`ReplicaStorage`] with a +//! scripted [`FailurePlan`] and records an operation journal for assertions +//! (retry counts, no-double-ship, ordering). +//! +//! Every operation, passthrough or injected, is journaled with its key and +//! outcome. A fired rule injects either BEFORE the wrapped call (the backend +//! is never touched) or, in ambiguous-success mode, AFTER it (the backend +//! applies the operation but the caller still sees an injected error: the +//! lost-200 case that exercises retry idempotency). The plan and journal are +//! shared behind `Arc` so in-flight [`FlakyMultipart`] uploads consult the +//! same live plan. + +use std::future::Future; +use std::sync::{Arc, Mutex, PoisonError}; + +use bytes::Bytes; + +use crate::storage::{MultipartUpload, ReplicaStorage, StorageError}; + +/// Fault-injection wrapper. Construct with [`FlakyStorage::new`], mutate the +/// live plan via [`FlakyStorage::set_plan`]/[`FlakyStorage::heal`], and read +/// the journal via [`FlakyStorage::journal`]. +pub struct FlakyStorage { + inner: Box, + shared: Shared, +} + +impl FlakyStorage { + #[must_use] + pub fn new(inner: ReplicaStorage, plan: FailurePlan) -> Self { + Self { + inner: Box::new(inner), + shared: Shared { + plan: Arc::new(Mutex::new(plan)), + journal: Arc::new(Mutex::new(Vec::new())), + }, + } + } + + /// Replace the live failure plan (e.g. heal after an outage). + pub fn set_plan(&self, plan: FailurePlan) { + *self.shared.lock_plan() = plan; + } + + /// Remove all failure rules: every subsequent operation passes through. + pub fn heal(&self) { + self.set_plan(FailurePlan::new()); + } + + /// Snapshot of the operation journal. + #[must_use] + pub fn journal(&self) -> Vec<(Op, OpOutcome)> { + self.shared.lock_journal().clone() + } + + // Every method boxes the inner future (`ReplicaStorage` methods can + // recurse back into this wrapper, so the cycle must be broken with a heap + // allocation) and routes it through `Shared::guard`, which applies the + // plan (before/after/none) and journals exactly once. + pub(crate) async fn put(&self, key: &str, bytes: Bytes) -> Result<(), StorageError> { + self.shared + .guard(OpKind::Put, key, Box::pin(self.inner.put(key, bytes))) + .await + } + + pub(crate) async fn get(&self, key: &str) -> Result { + self.shared + .guard(OpKind::Get, key, Box::pin(self.inner.get(key))) + .await + } + + pub(crate) async fn get_stream( + &self, + key: &str, + ) -> Result, StorageError> { + self.shared + .guard(OpKind::GetStream, key, Box::pin(self.inner.get_stream(key))) + .await + } + + pub(crate) async fn list(&self, prefix: &str) -> Result, StorageError> { + self.shared + .guard(OpKind::List, prefix, Box::pin(self.inner.list(prefix))) + .await + } + + pub(crate) async fn list_dirs(&self, prefix: &str) -> Result, StorageError> { + self.shared + .guard( + OpKind::ListDirs, + prefix, + Box::pin(self.inner.list_dirs(prefix)), + ) + .await + } + + pub(crate) async fn delete(&self, key: &str) -> Result<(), StorageError> { + self.shared + .guard(OpKind::Delete, key, Box::pin(self.inner.delete(key))) + .await + } + + pub(crate) async fn delete_prefix(&self, prefix: &str) -> Result<(), StorageError> { + self.shared + .guard( + OpKind::DeletePrefix, + prefix, + Box::pin(self.inner.delete_prefix(prefix)), + ) + .await + } + + pub(crate) async fn start_multipart(&self, key: &str) -> Result { + let inner = self + .shared + .guard( + OpKind::MultipartStart, + key, + Box::pin(self.inner.start_multipart(key)), + ) + .await?; + Ok(FlakyMultipart { + inner: Box::new(inner), + key: key.to_owned(), + shared: self.shared.clone(), + }) + } +} + +/// Multipart wrapper that consults the parent plan on every operation. An +/// injected `finish` drops the inner upload unfinished, so the object never +/// becomes visible (matching a real failure). +pub struct FlakyMultipart { + inner: Box, + key: String, + shared: Shared, +} + +impl FlakyMultipart { + pub(crate) async fn write_part(&mut self, bytes: Bytes) -> Result<(), StorageError> { + self.shared + .guard( + OpKind::MultipartWrite, + &self.key, + Box::pin(self.inner.write_part(bytes)), + ) + .await + } + + pub(crate) async fn finish(self) -> Result<(), StorageError> { + let Self { inner, key, shared } = self; + shared + .guard(OpKind::MultipartFinish, &key, Box::pin(inner.finish())) + .await + } + + pub(crate) async fn abort(self) -> Result<(), StorageError> { + let Self { inner, key, shared } = self; + shared + .guard(OpKind::MultipartAbort, &key, Box::pin(inner.abort())) + .await + } +} + +/// The plan and journal handles shared between the storage wrapper and its +/// in-flight multipart uploads. +#[derive(Clone)] +struct Shared { + plan: Arc>, + journal: Arc>>, +} + +impl Shared { + /// Run one wrapped operation under the live plan, journaling exactly once. + /// + /// - No rule fires: run the backend op and journal its real outcome. + /// - A rule fires [`Timing::Before`]: skip the backend, journal `Injected`, + /// return the injected error (the op had no effect). + /// - A rule fires [`Timing::After`]: run the backend op, discard its + /// result, journal `InjectedAfter`, return the injected error (the + /// backend applied the write but the caller sees failure). + async fn guard(&self, kind: OpKind, key: &str, op: Fut) -> Result + where + Fut: Future>, + { + // Drop the plan lock before awaiting: the guard is not held across the + // backend call. + let timing = self.lock_plan().evaluate(kind, key); + match timing { + Some(Timing::Before) => { + self.record(kind, key, OpOutcome::Injected); + Err(injected(kind, key)) + }, + Some(Timing::After) => { + drop(op.await); + self.record(kind, key, OpOutcome::InjectedAfter); + Err(injected(kind, key)) + }, + None => { + let result = op.await; + self.record_result(kind, key, &result); + result + }, + } + } + + fn record_result(&self, kind: OpKind, key: &str, result: &Result) { + let outcome = match result { + Ok(_) => OpOutcome::Ok, + Err(error) => OpOutcome::Err(error.to_string()), + }; + self.record(kind, key, outcome); + } + + fn record(&self, kind: OpKind, key: &str, outcome: OpOutcome) { + self.lock_journal().push(( + Op { + kind, + key: key.to_owned(), + }, + outcome, + )); + } + + fn lock_plan(&self) -> std::sync::MutexGuard<'_, FailurePlan> { + self.plan.lock().unwrap_or_else(PoisonError::into_inner) + } + + fn lock_journal(&self) -> std::sync::MutexGuard<'_, Vec<(Op, OpOutcome)>> { + self.journal.lock().unwrap_or_else(PoisonError::into_inner) + } +} + +/// The injected error for one operation. +fn injected(kind: OpKind, key: &str) -> StorageError { + StorageError::Injected { + op: kind.as_str(), + key: key.to_owned(), + } +} + +/// When a fired rule injects its failure relative to the wrapped backend call. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum Timing { + /// Fail before touching the backend: the operation has no effect. + Before, + /// Perform the operation, then return an injected error: the backend + /// applied the write but the caller sees failure (ambiguous success). + After, +} + +/// A scripted failure plan. Rules are evaluated per operation in the order +/// they were added; the first matching rule fires (and, for bounded rules, +/// is consumed). +#[derive(Debug, Default)] +pub struct FailurePlan { + rules: Vec, +} + +impl FailurePlan { + #[must_use] + pub fn new() -> Self { + Self::default() + } + + /// Fail the `n`th (1-indexed) operation of `kind` observed from now on. + #[must_use] + pub fn fail_nth(mut self, kind: OpKind, n: u64) -> Self { + self.rules.push(Rule::Nth { + kind, + n, + seen: 0, + timing: Timing::Before, + }); + self + } + + /// Fail the next `times` operations whose key contains `key_contains` + /// (any kind when `kind` is `None`). + #[must_use] + pub fn fail_matching(mut self, kind: Option, key_contains: &str, times: u64) -> Self { + self.rules.push(Rule::Matching { + kind, + key_contains: key_contains.to_owned(), + times, + timing: Timing::Before, + }); + self + } + + /// Fail every operation of `kind` until healed via + /// [`FlakyStorage::heal`] or [`FlakyStorage::set_plan`]. + #[must_use] + pub fn fail_all(mut self, kind: OpKind) -> Self { + self.rules.push(Rule::All { + kind, + timing: Timing::Before, + }); + self + } + + /// Make the most recently added rule fire in ambiguous-success mode: the + /// backend operation runs and THEN an injected error is returned (the + /// lost-200 case that exercises retry idempotency). The operation is + /// journaled as [`OpOutcome::InjectedAfter`]. A no-op when no rule has been + /// added yet. + #[must_use] + pub fn after(mut self) -> Self { + if let Some(rule) = self.rules.last_mut() { + rule.set_timing(Timing::After); + } + self + } + + /// Evaluate the plan for one operation: the first matching rule fires, + /// bounded rules are consumed once exhausted. Returns the fired rule's + /// [`Timing`] (or `None` when no rule fires). Rules after the firing rule + /// are untouched (their counters do not advance). + fn evaluate(&mut self, kind: OpKind, key: &str) -> Option { + let mut fired_index = None; + for (index, rule) in self.rules.iter_mut().enumerate() { + let fired = match rule { + Rule::Nth { + kind: rule_kind, + n, + seen, + .. + } => { + if *rule_kind == kind { + *seen += 1; + *seen == *n + } else { + false + } + }, + Rule::Matching { + kind: rule_kind, + key_contains, + times, + .. + } => { + if rule_kind.is_none_or(|rule_kind| rule_kind == kind) + && key.contains(key_contains.as_str()) + && *times > 0 + { + *times -= 1; + true + } else { + false + } + }, + Rule::All { + kind: rule_kind, .. + } => *rule_kind == kind, + }; + if fired { + fired_index = Some(index); + break; + } + } + let index = fired_index?; + let timing = self.rules.get(index).map(Rule::timing)?; + let consumed = match self.rules.get(index) { + Some(Rule::Nth { .. }) => true, + Some(Rule::Matching { times, .. }) => *times == 0, + Some(Rule::All { .. }) | None => false, + }; + if consumed { + self.rules.remove(index); + } + Some(timing) + } +} + +#[derive(Debug)] +enum Rule { + /// Fail the `n`th (1-indexed) operation of `kind` observed from now on. + Nth { + kind: OpKind, + n: u64, + seen: u64, + timing: Timing, + }, + /// Fail the next `times` operations whose key contains `key_contains` + /// (any kind when `kind` is `None`). + Matching { + kind: Option, + key_contains: String, + times: u64, + timing: Timing, + }, + /// Fail every operation of `kind` until the rule is removed via + /// [`FlakyStorage::heal`]. + All { kind: OpKind, timing: Timing }, +} + +impl Rule { + fn timing(&self) -> Timing { + match self { + Self::Nth { timing, .. } | Self::Matching { timing, .. } | Self::All { timing, .. } => { + *timing + }, + } + } + + fn set_timing(&mut self, timing: Timing) { + match self { + Self::Nth { timing: slot, .. } + | Self::Matching { timing: slot, .. } + | Self::All { timing: slot, .. } => *slot = timing, + } + } +} + +/// The kind of storage operation, for matching in failure plans and the +/// journal. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum OpKind { + Put, + Get, + GetStream, + List, + ListDirs, + Delete, + DeletePrefix, + MultipartStart, + MultipartWrite, + MultipartFinish, + MultipartAbort, +} + +impl OpKind { + /// The stable name used in [`StorageError::Injected`] messages. + #[must_use] + pub const fn as_str(self) -> &'static str { + match self { + Self::Put => "put", + Self::Get => "get", + Self::GetStream => "get_stream", + Self::List => "list", + Self::ListDirs => "list_dirs", + Self::Delete => "delete", + Self::DeletePrefix => "delete_prefix", + Self::MultipartStart => "multipart_start", + Self::MultipartWrite => "multipart_write", + Self::MultipartFinish => "multipart_finish", + Self::MultipartAbort => "multipart_abort", + } + } +} + +/// One recorded storage operation. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Op { + pub kind: OpKind, + pub key: String, +} + +/// The outcome recorded in the journal for one operation. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum OpOutcome { + Ok, + Injected, + /// The backend applied the operation but an injected error was returned + /// afterward (ambiguous success / lost-200). + InjectedAfter, + Err(String), +} + +#[cfg(test)] +mod tests { + use camino::Utf8Path; + use pretty_assertions::assert_eq; + + use crate::local::LocalStorage; + + use super::*; + + fn flaky(tmp: &tempfile::TempDir, plan: FailurePlan) -> FlakyStorage { + let root = Utf8Path::from_path(tmp.path()) + .expect("tempdir path is UTF-8") + .to_path_buf(); + FlakyStorage::new(ReplicaStorage::Local(LocalStorage::new(root)), plan) + } + + fn op(kind: OpKind, key: &str) -> Op { + Op { + kind, + key: key.to_owned(), + } + } + + fn assert_injected(result: Result<(), StorageError>, expected_op: &str, context: &str) { + match result { + Err(StorageError::Injected { op, .. }) => { + assert_eq!(op, expected_op, "{context}: wrong injected op"); + }, + Err(error) => panic!("{context}: unexpected error: {error}"), + Ok(()) => panic!("{context}: expected an injected failure"), + } + } + + #[tokio::test] + async fn passthrough_records_ok_in_journal() { + let tmp = tempfile::tempdir().expect("tempdir failed"); + let storage = flaky(&tmp, FailurePlan::new()); + storage + .put("wal/seg1", Bytes::from_static(b"1")) + .await + .expect("put failed"); + let got = storage.get("wal/seg1").await.expect("get failed"); + assert_eq!(got.as_ref(), b"1".as_slice(), "passthrough value"); + assert_eq!( + storage.journal(), + vec![ + (op(OpKind::Put, "wal/seg1"), OpOutcome::Ok), + (op(OpKind::Get, "wal/seg1"), OpOutcome::Ok), + ], + "journal must record passthrough operations" + ); + } + + #[tokio::test] + async fn passthrough_records_backend_error() { + let tmp = tempfile::tempdir().expect("tempdir failed"); + let storage = flaky(&tmp, FailurePlan::new()); + let error = storage.get("missing").await.expect_err("get must fail"); + assert!( + matches!(&error, StorageError::NotFound { key } if key == "missing"), + "NotFound must pass through unchanged, got: {error}" + ); + assert_eq!( + storage.journal(), + vec![( + op(OpKind::Get, "missing"), + OpOutcome::Err("Object not found: missing".to_owned()) + )], + "journal must record backend errors" + ); + } + + #[tokio::test] + async fn fail_nth_fires_once_then_consumed() { + let tmp = tempfile::tempdir().expect("tempdir failed"); + let storage = flaky(&tmp, FailurePlan::new().fail_nth(OpKind::Put, 2)); + storage + .put("a", Bytes::from_static(b"1")) + .await + .expect("first put must pass"); + assert_injected( + storage.put("b", Bytes::from_static(b"2")).await, + "put", + "second put", + ); + storage + .put("c", Bytes::from_static(b"3")) + .await + .expect("third put must pass: the rule is consumed"); + // The injected put never reached the backend. + let error = storage.get("b").await.expect_err("b must not exist"); + assert!( + matches!(error, StorageError::NotFound { .. }), + "injected put must not write" + ); + } + + #[tokio::test] + async fn fail_matching_bounded_times_consumed() { + let tmp = tempfile::tempdir().expect("tempdir failed"); + let storage = flaky(&tmp, FailurePlan::new().fail_matching(None, "wal/", 2)); + assert_injected( + storage.put("wal/1", Bytes::from_static(b"1")).await, + "put", + "first wal put", + ); + assert_injected( + storage.put("wal/2", Bytes::from_static(b"2")).await, + "put", + "second wal put", + ); + storage + .put("wal/3", Bytes::from_static(b"3")) + .await + .expect("third wal put must pass: rule exhausted"); + storage + .put("other", Bytes::from_static(b"4")) + .await + .expect("non-matching key must pass"); + } + + #[tokio::test] + async fn fail_matching_filters_by_kind() { + let tmp = tempfile::tempdir().expect("tempdir failed"); + let storage = flaky( + &tmp, + FailurePlan::new().fail_matching(Some(OpKind::Get), "x", 1), + ); + storage + .put("x", Bytes::from_static(b"1")) + .await + .expect("put must pass: rule only matches Get"); + let error = storage.get("x").await.expect_err("get must be injected"); + assert!( + matches!(error, StorageError::Injected { op: "get", .. }), + "expected injected get" + ); + let got = storage.get("x").await.expect("second get must pass"); + assert_eq!(got.as_ref(), b"1".as_slice(), "value after heal"); + } + + #[tokio::test] + async fn fail_all_persists_until_heal() { + let tmp = tempfile::tempdir().expect("tempdir failed"); + let storage = flaky(&tmp, FailurePlan::new().fail_all(OpKind::List)); + storage + .put("k", Bytes::from_static(b"1")) + .await + .expect("put must pass"); + for attempt in 0..3 { + let result = storage.list("").await; + assert!( + matches!(result, Err(StorageError::Injected { op: "list", .. })), + "list attempt {attempt} must stay injected" + ); + } + storage.heal(); + let keys = storage.list("").await.expect("healed list must pass"); + assert_eq!(keys, vec!["k".to_owned()], "healed list content"); + } + + #[tokio::test] + async fn set_plan_replaces_live_plan() { + let tmp = tempfile::tempdir().expect("tempdir failed"); + let storage = flaky(&tmp, FailurePlan::new()); + storage + .put("k", Bytes::from_static(b"1")) + .await + .expect("put must pass"); + storage.set_plan(FailurePlan::new().fail_all(OpKind::Get)); + let error = storage.get("k").await.expect_err("get must be injected"); + assert!( + matches!(error, StorageError::Injected { op: "get", .. }), + "expected injected get after set_plan" + ); + storage.set_plan(FailurePlan::new()); + let got = storage.get("k").await.expect("get after reset must pass"); + assert_eq!(got.as_ref(), b"1".as_slice(), "value after reset"); + } + + #[tokio::test] + async fn first_matching_rule_wins_and_later_counters_do_not_advance() { + let tmp = tempfile::tempdir().expect("tempdir failed"); + // Rule 1 matches only key "k" once; rule 2 fails the first Put it + // observes. If evaluation stopped at rule 1 for the first put, rule + // 2 must fire on the SECOND put (its counter did not advance). + let storage = flaky( + &tmp, + FailurePlan::new() + .fail_matching(None, "k", 1) + .fail_nth(OpKind::Put, 1), + ); + assert_injected( + storage.put("k", Bytes::from_static(b"1")).await, + "put", + "first put (rule 1)", + ); + assert_injected( + storage.put("z", Bytes::from_static(b"2")).await, + "put", + "second put (rule 2)", + ); + storage + .put("w", Bytes::from_static(b"3")) + .await + .expect("third put must pass: both rules consumed"); + } + + #[tokio::test] + async fn multipart_write_consults_parent_plan() { + let tmp = tempfile::tempdir().expect("tempdir failed"); + let storage = flaky(&tmp, FailurePlan::new().fail_nth(OpKind::MultipartWrite, 2)); + let mut upload = storage + .start_multipart("snap/db.zst") + .await + .expect("start failed"); + upload + .write_part(Bytes::from_static(b"one")) + .await + .expect("first write must pass"); + assert_injected( + upload.write_part(Bytes::from_static(b"two")).await, + "multipart_write", + "second write", + ); + upload + .write_part(Bytes::from_static(b"three")) + .await + .expect("third write must pass"); + upload.abort().await.expect("abort failed"); + assert_eq!( + storage.journal(), + vec![ + (op(OpKind::MultipartStart, "snap/db.zst"), OpOutcome::Ok), + (op(OpKind::MultipartWrite, "snap/db.zst"), OpOutcome::Ok), + ( + op(OpKind::MultipartWrite, "snap/db.zst"), + OpOutcome::Injected + ), + (op(OpKind::MultipartWrite, "snap/db.zst"), OpOutcome::Ok), + (op(OpKind::MultipartAbort, "snap/db.zst"), OpOutcome::Ok), + ], + "multipart journal must record every operation with its outcome" + ); + } + + #[tokio::test] + async fn multipart_finish_injection_leaves_object_invisible() { + let tmp = tempfile::tempdir().expect("tempdir failed"); + let storage = flaky( + &tmp, + FailurePlan::new().fail_nth(OpKind::MultipartFinish, 1), + ); + let mut upload = storage + .start_multipart("snap/db.zst") + .await + .expect("start failed"); + upload + .write_part(Bytes::from_static(b"data")) + .await + .expect("write failed"); + assert_injected(upload.finish().await, "multipart_finish", "finish"); + let error = storage + .get("snap/db.zst") + .await + .expect_err("object must not exist"); + assert!( + matches!(error, StorageError::NotFound { .. }), + "injected finish must leave the object invisible" + ); + let keys = storage.list("").await.expect("list failed"); + assert_eq!(keys, Vec::::new(), "no objects after failed finish"); + } + + #[tokio::test] + async fn after_mode_applies_operation_then_reports_failure() { + let tmp = tempfile::tempdir().expect("tempdir failed"); + let storage = flaky(&tmp, FailurePlan::new().fail_nth(OpKind::Put, 1).after()); + // The put reports failure to the caller... + assert_injected( + storage.put("k", Bytes::from_static(b"v")).await, + "put", + "ambiguous put", + ); + // ...but the backend actually applied it: a clean get finds the object. + let got = storage + .get("k") + .await + .expect("after-mode put must have landed server-side"); + assert_eq!(got.as_ref(), b"v".as_slice(), "after-mode put lands"); + assert_eq!( + storage.journal(), + vec![ + (op(OpKind::Put, "k"), OpOutcome::InjectedAfter), + (op(OpKind::Get, "k"), OpOutcome::Ok), + ], + "an after-mode injection journals distinctly as InjectedAfter" + ); + } + + #[tokio::test] + async fn paired_operations_have_distinct_op_kinds() { + let tmp = tempfile::tempdir().expect("tempdir failed"); + // A plan that fails `List` must NOT catch `list_dirs`: distinct kinds. + let storage = flaky(&tmp, FailurePlan::new().fail_all(OpKind::List)); + storage + .put("gen/g/x", Bytes::from_static(b"1")) + .await + .expect("put must pass"); + assert!( + matches!( + storage.list("").await, + Err(StorageError::Injected { op: "list", .. }) + ), + "list is failed" + ); + let dirs = storage + .list_dirs("") + .await + .expect("list_dirs is a distinct kind and passes through"); + assert_eq!(dirs, vec!["gen".to_owned()], "list_dirs result"); + storage + .get_stream("gen/g/x") + .await + .map(drop) + .expect("get_stream is a distinct kind"); + storage + .delete_prefix("gen/") + .await + .expect("delete_prefix is a distinct kind"); + let kinds: Vec = storage + .journal() + .into_iter() + .map(|(op, _)| op.kind) + .collect(); + assert_eq!( + kinds, + vec![ + OpKind::Put, + OpKind::List, + OpKind::ListDirs, + OpKind::GetStream, + OpKind::DeletePrefix, + ], + "each paired operation journals under its own kind" + ); + } +} diff --git a/plus/bencher_replica/src/testing/mod.rs b/plus/bencher_replica/src/testing/mod.rs new file mode 100644 index 000000000..3ef898571 --- /dev/null +++ b/plus/bencher_replica/src/testing/mod.rs @@ -0,0 +1,23 @@ +//! Test-only infrastructure: fixtures, fault injection, workload generation, +//! and equivalence oracles. +//! +//! Compiled under `#[cfg(any(test, feature = "testing"))]`. The `testing` +//! feature exists so downstream integration tests (e.g. `lib/api_server`) +//! can drive the same machinery. + +mod compare; +mod flaky; +mod probe; +mod synthetic_wal; +mod wal_fixture; +mod workload; + +pub use compare::{assert_pages_equal, assert_replica_equivalent}; +pub use flaky::{FailurePlan, FlakyMultipart, FlakyStorage, Op, OpKind, OpOutcome}; +pub use probe::{ProbeResult, WriteProbe}; +pub use synthetic_wal::SyntheticWal; +pub use wal_fixture::{CheckpointMode, FixtureError, WalFixture}; +pub use workload::{ + AppliedOp, WorkloadEnv, WorkloadError, WorkloadOp, WorkloadOpError, WorkloadRunner, + generate_workload, run_workload, +}; diff --git a/plus/bencher_replica/src/testing/probe.rs b/plus/bencher_replica/src/testing/probe.rs new file mode 100644 index 000000000..83a40d5e2 --- /dev/null +++ b/plus/bencher_replica/src/testing/probe.rs @@ -0,0 +1,96 @@ +//! Concurrent-writer probe: measures whether (and for how long) a real +//! second connection is blocked while the replicator holds locks. +//! +//! Each write runs on a fresh rusqlite connection inside `spawn_blocking`, +//! with a configurable `busy_timeout` and `wal_autocheckpoint = 0` (the +//! probe is a writer, so invariant I2 applies to it too). The probe writes +//! into its own `probe_writes` table, created on first use. + +use std::future::Future; +use std::task::Poll; +use std::time::{Duration, Instant}; + +use camino::{Utf8Path, Utf8PathBuf}; +use tokio::task::spawn_blocking; + +/// Result of one probe write attempt. +#[derive(Debug)] +pub struct ProbeResult { + /// The write's outcome; `Err` carries the `SQLite` error (e.g. busy). + pub result: Result<(), rusqlite::Error>, + /// Wall-clock duration the write took (observed block time). + pub blocked: Duration, +} + +/// Spawns blocking write attempts against a database on dedicated +/// connections with a configurable `busy_timeout`. +pub struct WriteProbe { + db_path: Utf8PathBuf, + busy_timeout: Duration, +} + +impl WriteProbe { + #[must_use] + pub fn new(db_path: &Utf8Path, busy_timeout: Duration) -> Self { + Self { + db_path: db_path.to_owned(), + busy_timeout, + } + } + + /// Perform one INSERT on a fresh connection, reporting the outcome and + /// observed block duration. + pub async fn write_once(&self) -> ProbeResult { + let db_path = self.db_path.clone(); + let busy_timeout = self.busy_timeout; + let outcome = spawn_blocking(move || { + let started = Instant::now(); + let result = probe_write(&db_path, busy_timeout); + ProbeResult { + result, + blocked: started.elapsed(), + } + }) + .await; + match outcome { + Ok(probe) => probe, + // Surface a task panic as a failed write rather than panicking + // the caller: the hammer keeps collecting. + Err(_join_error) => ProbeResult { + result: Err(rusqlite::Error::InvalidQuery), + blocked: Duration::ZERO, + }, + } + } + + /// Hammer writes continuously until `stop` resolves, returning all + /// probe results. `stop` is polled between writes, so the final write + /// always runs to completion. + pub async fn hammer(&self, stop: F) -> Vec + where + F: Future + Send, + { + let mut results = Vec::new(); + let mut stop = std::pin::pin!(stop); + loop { + let stopped = + std::future::poll_fn(|cx| Poll::Ready(stop.as_mut().poll(cx).is_ready())).await; + if stopped { + return results; + } + results.push(self.write_once().await); + } + } +} + +/// One synchronous probe write on a fresh connection. +fn probe_write(db_path: &Utf8Path, busy_timeout: Duration) -> Result<(), rusqlite::Error> { + let conn = rusqlite::Connection::open(db_path)?; + conn.busy_timeout(busy_timeout)?; + // The probe is a writer: it must never checkpoint behind the + // replicator's back (invariant I2). + let _pages: i64 = conn.query_row("PRAGMA wal_autocheckpoint = 0", [], |row| row.get(0))?; + conn.execute_batch("CREATE TABLE IF NOT EXISTS probe_writes(id INTEGER PRIMARY KEY, at TEXT)")?; + conn.execute("INSERT INTO probe_writes (at) VALUES ('probe')", [])?; + Ok(()) +} diff --git a/plus/bencher_replica/src/testing/synthetic_wal.rs b/plus/bencher_replica/src/testing/synthetic_wal.rs new file mode 100644 index 000000000..58b994381 --- /dev/null +++ b/plus/bencher_replica/src/testing/synthetic_wal.rs @@ -0,0 +1,296 @@ +//! Raw WAL bytes built from scratch, with an independent implementation of +//! the `SQLite` WAL checksum chain (double-entry bookkeeping against +//! [`crate::wal::wal_checksum`]). +//! +//! Covers inputs rusqlite on a little-endian host can never produce: +//! big-endian checksum order, stale salts, bit flips, torn files. + +/// Builder for synthetic WAL byte strings. +pub struct SyntheticWal { + bytes: Vec, + page_size: u32, + big_endian: bool, + salt: (u32, u32), + running: (u32, u32), +} + +impl SyntheticWal { + /// Start a WAL with the given page size and checksum byte order + /// (`big_endian` true selects magic `0x377f0683`). + #[must_use] + pub fn new(page_size: u32, big_endian: bool, salt: (u32, u32)) -> Self { + // Independent literals on purpose: reusing crate::wal constants here + // would defeat the double-entry bookkeeping. + let magic = if big_endian { 0x377f_0683 } else { 0x377f_0682 }; + let mut bytes = Vec::with_capacity(32); + push_be_u32(&mut bytes, magic); + push_be_u32(&mut bytes, 3_007_000); + push_be_u32(&mut bytes, page_size); + push_be_u32(&mut bytes, 0); // checkpoint sequence number + push_be_u32(&mut bytes, salt.0); + push_be_u32(&mut bytes, salt.1); + // Header self-checksum over bytes 0..24, seeded with (0, 0) + let running = fold_checksum(big_endian, (0, 0), &bytes); + push_be_u32(&mut bytes, running.0); + push_be_u32(&mut bytes, running.1); + Self { + bytes, + page_size, + big_endian, + salt, + running, + } + } + + /// Append a non-commit frame for `page_no` with the checksum chain + /// threaded automatically. + #[must_use] + pub fn frame(self, page_no: u32, page: &[u8]) -> Self { + let salt = self.salt; + self.push_frame(page_no, page, 0, salt) + } + + /// Append a commit frame (`db_size` pages). + #[must_use] + pub fn commit_frame(self, page_no: u32, page: &[u8], db_size: u32) -> Self { + let salt = self.salt; + self.push_frame(page_no, page, db_size, salt) + } + + /// Append a frame carrying explicit (stale) salts, simulating leftovers + /// from before a WAL restart. + /// + /// The checksum chain is still threaded normally, so the mismatched salts + /// are the frame's only defect. + #[must_use] + pub fn frame_with_salts(self, page_no: u32, page: &[u8], salt: (u32, u32)) -> Self { + self.push_frame(page_no, page, 0, salt) + } + + /// Flip one bit at `byte_offset` (corruption injection). + #[must_use] + pub fn flip_bit(mut self, byte_offset: usize, bit: u8) -> Self { + assert!(bit < 8, "bit index must be in 0..8, got {bit}"); + assert!( + byte_offset < self.bytes.len(), + "flip_bit offset {byte_offset} out of range for {} bytes", + self.bytes.len() + ); + if let Some(byte) = self.bytes.get_mut(byte_offset) { + *byte ^= 1u8 << bit; + } + self + } + + /// Truncate the byte string to `len` (torn header / mid-frame EOF). + /// A `len` beyond the current size leaves the bytes unchanged. + #[must_use] + pub fn truncate_at(mut self, len: usize) -> Self { + self.bytes.truncate(len); + self + } + + /// The finished WAL bytes. + #[must_use] + pub fn bytes(self) -> Vec { + self.bytes + } + + fn push_frame(mut self, page_no: u32, page: &[u8], db_size: u32, salt: (u32, u32)) -> Self { + assert_eq!( + page.len(), + self.page_size as usize, + "page payload must be exactly page_size bytes" + ); + // The frame checksum covers frame-header bytes 0..8 plus the page, + // seeded from the previous frame (or the WAL header) + let mut head = Vec::with_capacity(24); + push_be_u32(&mut head, page_no); + push_be_u32(&mut head, db_size); + let checksum = fold_checksum(self.big_endian, self.running, &head); + let checksum = fold_checksum(self.big_endian, checksum, page); + push_be_u32(&mut head, salt.0); + push_be_u32(&mut head, salt.1); + push_be_u32(&mut head, checksum.0); + push_be_u32(&mut head, checksum.1); + self.bytes.extend_from_slice(&head); + self.bytes.extend_from_slice(page); + self.running = checksum; + self + } +} + +/// Append a `u32` in big-endian byte order (all WAL integer fields, most +/// significant byte first). Each extracted byte is masked to 8 bits, so the +/// cast is exact; the assert surfaces a logic error loudly instead of +/// silently writing a zero byte if that masking is ever broken. +fn push_be_u32(bytes: &mut Vec, value: u32) { + for shift in [24u32, 16, 8, 0] { + let byte = (value >> shift) & 0xff; + let narrow = byte as u8; + // Masked to 8 bits, so the round-trip is exact; assert loudly rather + // than silently truncating if that masking is ever broken. + assert_eq!(u32::from(narrow), byte, "WAL byte {byte:#x} must fit in u8"); + bytes.push(narrow); + } +} + +/// Independent implementation of `SQLite`'s cumulative WAL checksum: the input +/// is folded as pairs of 32-bit words in the given byte order. +fn fold_checksum(big_endian: bool, seed: (u32, u32), data: &[u8]) -> (u32, u32) { + assert!( + data.len().is_multiple_of(8), + "checksum input must be a multiple of 8 bytes, got {}", + data.len() + ); + let mut sums = seed; + for pair in data.chunks_exact(8) { + let (first, second) = pair.split_at(4); + let x0 = word(big_endian, first); + let x1 = word(big_endian, second); + sums.0 = sums.0.wrapping_add(x0).wrapping_add(sums.1); + sums.1 = sums.1.wrapping_add(x1).wrapping_add(sums.0); + } + sums +} + +/// Read one 32-bit checksum word in the given byte order. +fn word(big_endian: bool, bytes: &[u8]) -> u32 { + let mut buf = [0u8; 4]; + buf.copy_from_slice(bytes); + if big_endian { + (u32::from(buf[0]) << 24) + | (u32::from(buf[1]) << 16) + | (u32::from(buf[2]) << 8) + | u32::from(buf[3]) + } else { + u32::from(buf[0]) + | (u32::from(buf[1]) << 8) + | (u32::from(buf[2]) << 16) + | (u32::from(buf[3]) << 24) + } +} + +#[cfg(test)] +mod tests { + use pretty_assertions::assert_eq; + + use super::*; + use crate::wal::{ + FrameHeader, WAL_FORMAT, WAL_MAGIC_BE, WAL_MAGIC_LE, parse_wal_header, wal_checksum, + }; + + const PAGE: u32 = 512; + const SALT: (u32, u32) = (0xdead_beef, 0xcafe_f00d); + + fn page_of(fill: u8) -> Vec { + vec![fill; usize::try_from(PAGE).unwrap()] + } + + #[test] + fn header_fields_roundtrip_both_orders() { + for (big_endian, magic) in [(false, WAL_MAGIC_LE), (true, WAL_MAGIC_BE)] { + let bytes = SyntheticWal::new(PAGE, big_endian, SALT).bytes(); + assert_eq!(bytes.len(), 32, "header only"); + let header = parse_wal_header(&bytes).unwrap(); + assert_eq!(header.magic, magic); + assert_eq!(header.format, WAL_FORMAT); + assert_eq!(header.page_size, PAGE); + assert_eq!(header.checkpoint_seq, 0); + assert_eq!(header.salt, SALT); + // The self-checksum verifies under the crate implementation + let computed = wal_checksum(big_endian, (0, 0), &bytes[..24]); + assert_eq!(header.checksum, computed); + } + } + + #[test] + fn frame_checksums_match_crate_checksum() { + // Double-entry bookkeeping: the builder's independent checksum chain + // must agree with crate::wal::wal_checksum frame by frame. + for big_endian in [false, true] { + let bytes = SyntheticWal::new(PAGE, big_endian, SALT) + .frame(1, &page_of(0x11)) + .commit_frame(2, &page_of(0x22), 2) + .bytes(); + let header = parse_wal_header(&bytes).unwrap(); + let frame_size = 24 + usize::try_from(PAGE).unwrap(); + let mut running = header.checksum; + for frame in bytes[32..].chunks_exact(frame_size) { + let mut head = [0u8; 24]; + head.copy_from_slice(&frame[..24]); + let frame_header = FrameHeader::parse(&head); + assert_eq!(frame_header.salt, SALT, "frames carry the header salts"); + let after_head = wal_checksum(big_endian, running, &frame[..8]); + let computed = wal_checksum(big_endian, after_head, &frame[24..]); + assert_eq!( + frame_header.checksum, computed, + "stored frame checksum matches the crate chain (big_endian: {big_endian})" + ); + running = computed; + } + } + } + + #[test] + fn commit_frame_records_db_size() { + let bytes = SyntheticWal::new(PAGE, false, SALT) + .frame(1, &page_of(0x11)) + .commit_frame(2, &page_of(0x22), 9) + .bytes(); + let frame_size = 24 + usize::try_from(PAGE).unwrap(); + let mut head = [0u8; 24]; + head.copy_from_slice(&bytes[32..32 + 24]); + let first = FrameHeader::parse(&head); + assert_eq!(first.page_no, 1); + assert_eq!(first.db_size, 0, "non-commit frame"); + assert!(!first.is_commit()); + head.copy_from_slice(&bytes[32 + frame_size..32 + frame_size + 24]); + let second = FrameHeader::parse(&head); + assert_eq!(second.page_no, 2); + assert_eq!(second.db_size, 9, "commit frame carries the db size"); + assert!(second.is_commit()); + } + + #[test] + fn frame_with_salts_uses_given_salts() { + let stale = (0x0101_0101, 0x0202_0202); + let bytes = SyntheticWal::new(PAGE, false, SALT) + .frame_with_salts(1, &page_of(0x11), stale) + .bytes(); + let mut head = [0u8; 24]; + head.copy_from_slice(&bytes[32..32 + 24]); + let frame = FrameHeader::parse(&head); + assert_eq!(frame.salt, stale, "explicit salts are written verbatim"); + } + + #[test] + fn flip_bit_and_truncate_mutate_bytes() { + let pristine = SyntheticWal::new(PAGE, false, SALT) + .commit_frame(1, &page_of(0x11), 1) + .bytes(); + let flipped = SyntheticWal::new(PAGE, false, SALT) + .commit_frame(1, &page_of(0x11), 1) + .flip_bit(40, 3) + .bytes(); + assert_eq!(pristine.len(), flipped.len()); + assert_eq!( + pristine[40] ^ (1 << 3), + flipped[40], + "exactly one bit flips" + ); + let differing = pristine + .iter() + .zip(&flipped) + .filter(|(a, b)| a != b) + .count(); + assert_eq!(differing, 1, "no other byte changes"); + + let truncated = SyntheticWal::new(PAGE, false, SALT) + .commit_frame(1, &page_of(0x11), 1) + .truncate_at(100) + .bytes(); + assert_eq!(truncated.len(), 100); + assert_eq!(&truncated[..], &pristine[..100]); + } +} diff --git a/plus/bencher_replica/src/testing/wal_fixture.rs b/plus/bencher_replica/src/testing/wal_fixture.rs new file mode 100644 index 000000000..70d38ec62 --- /dev/null +++ b/plus/bencher_replica/src/testing/wal_fixture.rs @@ -0,0 +1,355 @@ +//! Real `SQLite` WAL fixtures generated deterministically via rusqlite. +//! +//! Not byte-deterministic across runs (`SQLite` randomizes salts); assertions +//! must be structural/invariant-based, never golden-byte. + +use camino::{Utf8Path, Utf8PathBuf}; + +/// A live `SQLite` database in WAL mode (`wal_autocheckpoint = 0`, +/// `synchronous = NORMAL`) with helpers to script commits, spills, and +/// checkpoints. +pub struct WalFixture { + dir: Utf8PathBuf, + conn: rusqlite::Connection, +} + +/// Failures scripting a [`WalFixture`]. +#[derive(Debug, thiserror::Error)] +pub enum FixtureError { + #[error("SQLite error: {0}")] + Sqlite(#[from] rusqlite::Error), + #[error("Failed to read fixture WAL: {0}")] + Io(#[from] std::io::Error), + #[error( + "Big transaction did not spill: WAL unchanged at {wal_len} bytes; the test cannot prove anything" + )] + NoSpill { wal_len: usize }, +} + +impl WalFixture { + /// Create a fresh DB at `/fixture.db` with the given page size and + /// a default `t(id INTEGER PRIMARY KEY, data TEXT)` table. + pub fn new(dir: &Utf8Path, page_size: u32) -> Result { + let db_path = dir.join("fixture.db"); + let conn = rusqlite::Connection::open(&db_path)?; + // page_size must be set BEFORE the database is created by the switch + // to WAL mode; it is immutable afterwards + conn.pragma_update(None, "page_size", page_size)?; + let _mode: String = conn.query_row("PRAGMA journal_mode = WAL", [], |row| row.get(0))?; + disable_autocheckpoint(&conn)?; + conn.pragma_update(None, "synchronous", "NORMAL")?; + conn.execute_batch( + "CREATE TABLE t (id INTEGER PRIMARY KEY, data TEXT); + INSERT INTO t (data) VALUES ('init');", + )?; + Ok(Self { + dir: dir.to_owned(), + conn, + }) + } + + /// Path to the database file. + #[must_use] + pub fn db_path(&self) -> Utf8PathBuf { + self.dir.join("fixture.db") + } + + /// Path to the WAL file. + #[must_use] + pub fn wal_path(&self) -> Utf8PathBuf { + self.dir.join("fixture.db-wal") + } + + /// Execute the statements inside a single transaction (one commit frame). + pub fn txn(&self, statements: &[&str]) -> Result<(), rusqlite::Error> { + self.conn.execute_batch("BEGIN IMMEDIATE")?; + for statement in statements { + if let Err(err) = self.conn.execute_batch(statement) { + let _rollback = self.conn.execute_batch("ROLLBACK"); + return Err(err); + } + } + self.conn.execute_batch("COMMIT") + } + + /// A single transaction touching at least `pages` distinct pages + /// (multi-frame single commit). + pub fn txn_touching_pages(&self, pages: u32) -> Result<(), rusqlite::Error> { + let page_size: u32 = self + .conn + .query_row("PRAGMA page_size", [], |row| row.get(0))?; + // One row per page: a payload of page_size bytes forces at least one + // overflow page per row + let data = "x".repeat(page_size as usize); + self.conn.execute_batch("BEGIN IMMEDIATE")?; + let result = (|| { + let mut statement = self.conn.prepare("INSERT INTO t (data) VALUES (?1)")?; + for _ in 0..pages { + statement.execute([&data])?; + } + Ok(()) + })(); + match result { + Ok(()) => self.conn.execute_batch("COMMIT"), + Err(err) => { + let _rollback = self.conn.execute_batch("ROLLBACK"); + Err(err) + }, + } + } + + /// A large transaction with a tiny page cache so `SQLite` spills + /// uncommitted frames into the WAL mid-transaction, then commits. + /// With `commit` false, the transaction is rolled back after spilling, + /// leaving flushed-but-uncommitted frames in the WAL. + /// + /// Errors with [`FixtureError::NoSpill`] if the WAL did not grow before + /// the commit/rollback: a test relying on spilled frames must not pass + /// silently without them. + pub fn big_txn_spilling(&self, commit: bool) -> Result<(), FixtureError> { + // A 10-page cache guarantees a multi-hundred-KiB transaction spills. + // cache_spill must be set explicitly: some builds (notably Apple's + // system SQLite) raise the default spill threshold to 20000 pages, + // which would silently prevent any spilling here. + self.conn.execute_batch("PRAGMA cache_size = 10")?; + self.conn.execute_batch("PRAGMA cache_spill = 10")?; + let wal_before = self.wal_snapshot()?; + self.conn.execute_batch("BEGIN IMMEDIATE")?; + let spilled = self.insert_spilling_rows(&wal_before); + let end = if commit && spilled.is_ok() { + "COMMIT" + } else { + "ROLLBACK" + }; + let end_result = self.conn.execute_batch(end); + // Restore the default cache size regardless of the outcome + let _restore = self.conn.execute_batch("PRAGMA cache_size = -2000"); + spilled?; + end_result?; + Ok(()) + } + + /// Run a checkpoint in the given mode on a dedicated connection. + pub fn checkpoint(&self, mode: CheckpointMode) -> Result<(), rusqlite::Error> { + let conn = rusqlite::Connection::open(self.db_path())?; + conn.busy_timeout(std::time::Duration::from_secs(5))?; + // wal_checkpoint returns a (busy, log, checkpointed) row + let _busy: i64 = conn.query_row( + &format!("PRAGMA wal_checkpoint({})", mode.as_str()), + [], + |row| row.get(0), + )?; + Ok(()) + } + + /// Open a second, independent connection to the same DB (a "stray" + /// writer bypassing any app-level mutex). + pub fn stray_conn(&self) -> Result { + let conn = rusqlite::Connection::open(self.db_path())?; + conn.busy_timeout(std::time::Duration::from_secs(5))?; + disable_autocheckpoint(&conn)?; + Ok(conn) + } + + /// Raw bytes of the current WAL file. + pub fn wal_bytes(&self) -> Result, std::io::Error> { + std::fs::read(self.wal_path()) + } + + /// Raw bytes of the current DB file. + pub fn db_bytes(&self) -> Result, std::io::Error> { + std::fs::read(self.db_path()) + } + + /// Insert rows until well past the page cache, then verify the WAL + /// changed BEFORE the enclosing transaction ends (proof that `SQLite` + /// spilled uncommitted frames). + /// + /// The comparison is by content, not length: a second spill can + /// overwrite rolled-back frames from an earlier spill without growing + /// the file. A per-call nonce in the row data (from `total_changes()`, + /// which never resets, not even on rollback) guarantees the new frames + /// differ from any stale ones they overwrite. + fn insert_spilling_rows(&self, wal_before: &[u8]) -> Result<(), FixtureError> { + let nonce: i64 = self + .conn + .query_row("SELECT total_changes()", [], |row| row.get(0))?; + let data = format!("{nonce}-{}", "x".repeat(512)); + let mut statement = self.conn.prepare("INSERT INTO t (data) VALUES (?1)")?; + for _ in 0..400 { + statement.execute([&data])?; + } + let wal_mid = self.wal_snapshot()?; + if wal_mid == wal_before { + return Err(FixtureError::NoSpill { + wal_len: wal_mid.len(), + }); + } + Ok(()) + } + + /// Current WAL bytes (empty if the file does not exist yet). + fn wal_snapshot(&self) -> Result, std::io::Error> { + match std::fs::read(self.wal_path()) { + Ok(bytes) => Ok(bytes), + Err(err) if err.kind() == std::io::ErrorKind::NotFound => Ok(Vec::new()), + Err(err) => Err(err), + } + } +} + +/// Checkpoint mode for [`WalFixture::checkpoint`]. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum CheckpointMode { + Passive, + Full, + Restart, + Truncate, +} + +impl CheckpointMode { + fn as_str(self) -> &'static str { + match self { + Self::Passive => "PASSIVE", + Self::Full => "FULL", + Self::Restart => "RESTART", + Self::Truncate => "TRUNCATE", + } + } +} + +/// `wal_autocheckpoint = 0`: the fixture (and its stray connections) never +/// checkpoint behind a test's back (invariant I2). +fn disable_autocheckpoint(conn: &rusqlite::Connection) -> Result<(), rusqlite::Error> { + let _pages: i64 = conn.query_row("PRAGMA wal_autocheckpoint = 0", [], |row| row.get(0))?; + Ok(()) +} + +#[cfg(test)] +mod tests { + use pretty_assertions::assert_eq; + + use super::*; + + fn tempdir_path(dir: &tempfile::TempDir) -> &Utf8Path { + Utf8Path::from_path(dir.path()).unwrap() + } + + fn row_count(conn: &rusqlite::Connection) -> i64 { + conn.query_row("SELECT COUNT(*) FROM t", [], |row| row.get(0)) + .unwrap() + } + + #[test] + fn new_creates_wal_mode_db() { + let tmp = tempfile::tempdir().unwrap(); + let fixture = WalFixture::new(tempdir_path(&tmp), 4096).unwrap(); + assert!(fixture.db_path().exists(), "db file exists"); + assert!(fixture.wal_path().exists(), "wal file exists"); + assert!( + !fixture.wal_bytes().unwrap().is_empty(), + "table creation left frames in the WAL" + ); + let mode: String = fixture + .conn + .query_row("PRAGMA journal_mode", [], |row| row.get(0)) + .unwrap(); + assert_eq!(mode, "wal"); + let autocheckpoint: i64 = fixture + .conn + .query_row("PRAGMA wal_autocheckpoint", [], |row| row.get(0)) + .unwrap(); + assert_eq!(autocheckpoint, 0, "fixture is the sole checkpointer"); + let page_size: i64 = fixture + .conn + .query_row("PRAGMA page_size", [], |row| row.get(0)) + .unwrap(); + assert_eq!(page_size, 4096); + } + + #[test] + fn txn_commits_statements_atomically() { + let tmp = tempfile::tempdir().unwrap(); + let fixture = WalFixture::new(tempdir_path(&tmp), 4096).unwrap(); + let before = row_count(&fixture.conn); + fixture + .txn(&[ + "INSERT INTO t (data) VALUES ('a')", + "INSERT INTO t (data) VALUES ('b')", + ]) + .unwrap(); + assert_eq!(row_count(&fixture.conn), before + 2); + // A failing statement rolls the whole transaction back + fixture + .txn(&[ + "INSERT INTO t (data) VALUES ('c')", + "INSERT INTO no_such_table (data) VALUES ('d')", + ]) + .unwrap_err(); + assert_eq!(row_count(&fixture.conn), before + 2, "nothing committed"); + } + + #[test] + fn checkpoint_truncate_resets_wal() { + let tmp = tempfile::tempdir().unwrap(); + let fixture = WalFixture::new(tempdir_path(&tmp), 4096).unwrap(); + assert!(!fixture.wal_bytes().unwrap().is_empty()); + fixture.checkpoint(CheckpointMode::Truncate).unwrap(); + assert!( + fixture.wal_bytes().unwrap().is_empty(), + "TRUNCATE resets the WAL to zero bytes" + ); + // Data survives the checkpoint + assert!(row_count(&fixture.conn) >= 1, "rows persist"); + } + + #[test] + fn big_txn_spilling_rollback_leaves_uncommitted_frames() { + let tmp = tempfile::tempdir().unwrap(); + let fixture = WalFixture::new(tempdir_path(&tmp), 4096).unwrap(); + let before_rows = row_count(&fixture.conn); + let before_len = fixture.wal_bytes().unwrap().len(); + fixture.big_txn_spilling(false).unwrap(); + assert_eq!( + row_count(&fixture.conn), + before_rows, + "rollback leaves no rows" + ); + assert!( + fixture.wal_bytes().unwrap().len() > before_len, + "spilled frames remain in the WAL after rollback" + ); + } + + #[test] + fn big_txn_spilling_commit_persists_rows() { + let tmp = tempfile::tempdir().unwrap(); + let fixture = WalFixture::new(tempdir_path(&tmp), 4096).unwrap(); + let before_rows = row_count(&fixture.conn); + fixture.big_txn_spilling(true).unwrap(); + assert!( + row_count(&fixture.conn) > before_rows, + "committed spill persists its rows" + ); + } + + #[test] + fn stray_conn_writes_independently() { + let tmp = tempfile::tempdir().unwrap(); + let fixture = WalFixture::new(tempdir_path(&tmp), 4096).unwrap(); + let stray = fixture.stray_conn().unwrap(); + let autocheckpoint: i64 = stray + .query_row("PRAGMA wal_autocheckpoint", [], |row| row.get(0)) + .unwrap(); + assert_eq!(autocheckpoint, 0, "stray connections never autocheckpoint"); + let before = row_count(&fixture.conn); + stray + .execute("INSERT INTO t (data) VALUES ('stray')", []) + .unwrap(); + assert_eq!( + row_count(&fixture.conn), + before + 1, + "stray write is visible through the fixture connection" + ); + } +} diff --git a/plus/bencher_replica/src/testing/workload.rs b/plus/bencher_replica/src/testing/workload.rs new file mode 100644 index 000000000..942ea22de --- /dev/null +++ b/plus/bencher_replica/src/testing/workload.rs @@ -0,0 +1,837 @@ +//! Seeded workload generator and script runner: the property-testing +//! replacement. +//! +//! Workloads are explicit, replayable `Vec` scripts generated +//! from a fixed seed (`rand::rngs::StdRng::seed_from_u64`) with weighted op +//! selection; SQL operations and engine events interleave from the same +//! seed. [`WorkloadRunner`] applies the SQL ops to the source database via +//! rusqlite (creating `t0..t15` tables lazily on first touch) and maps the +//! engine events onto step-driven [`SyncEngine`] calls. Every failure is +//! wrapped in a [`WorkloadError`] carrying the seed, the failing op index, +//! and the op itself, so the calling test can print the full script +//! alongside. +//! +//! Out of scope by design: `ATTACH`, page-size changes, journal-mode flips, +//! `WITHOUT ROWID` tables, multi-process writers. + +use std::fmt; +use std::time::Duration; + +use bencher_json::Clock; +use camino::{Utf8Path, Utf8PathBuf}; +use rand::rngs::StdRng; +use rand::{RngExt as _, SeedableRng as _}; +use slog::Logger; + +use crate::checkpoint::CheckpointOutcome; +use crate::config::ReplicaConfig; +use crate::local::LocalStorage; +use crate::replicator::ReplicaDb; +use crate::snapshot::SnapshotStatus; +use crate::storage::ReplicaStorage; +use crate::sync::{EngineState, SyncEngine, SyncError}; + +use super::flaky::{FailurePlan, FlakyStorage}; + +/// One step of a generated workload: either a SQL-level operation on the +/// source database or an engine event. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum WorkloadOp { + Insert { + table: u8, + rows: u16, + }, + Update { + table: u8, + rows: u16, + }, + Delete { + table: u8, + rows: u16, + }, + CreateTable { + table: u8, + }, + CreateIndex { + table: u8, + }, + DropTable { + table: u8, + }, + /// Forces overflow pages. + BlobWrite { + table: u8, + len: u32, + }, + /// 1..=50 statements in a single commit. + BigTxn { + statements: u8, + }, + UserVersionBump, + /// Full-DB rewrite through the WAL; low weight, great frame-burst + /// stress. + Vacuum, + // Engine events, interleaved by the same seed: + Sync, + Checkpoint, + Snapshot, + RestartReplicator, + /// A write through a second connection bypassing the app writer mutex. + StrayWrite, +} + +impl WorkloadOp { + /// Number of [`WorkloadOp`] variants (kept in lockstep with the weight + /// table by a unit test). + pub const VARIANT_COUNT: usize = 15; + + /// The variant name, for coverage accounting in tests. + #[must_use] + pub const fn kind_name(&self) -> &'static str { + match self { + Self::Insert { .. } => "Insert", + Self::Update { .. } => "Update", + Self::Delete { .. } => "Delete", + Self::CreateTable { .. } => "CreateTable", + Self::CreateIndex { .. } => "CreateIndex", + Self::DropTable { .. } => "DropTable", + Self::BlobWrite { .. } => "BlobWrite", + Self::BigTxn { .. } => "BigTxn", + Self::UserVersionBump => "UserVersionBump", + Self::Vacuum => "Vacuum", + Self::Sync => "Sync", + Self::Checkpoint => "Checkpoint", + Self::Snapshot => "Snapshot", + Self::RestartReplicator => "RestartReplicator", + Self::StrayWrite => "StrayWrite", + } + } +} + +impl fmt::Display for WorkloadOp { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::Insert { table, rows } => write!(f, "Insert(t{table} x{rows})"), + Self::Update { table, rows } => write!(f, "Update(t{table} x{rows})"), + Self::Delete { table, rows } => write!(f, "Delete(t{table} x{rows})"), + Self::CreateTable { table } => write!(f, "CreateTable(t{table})"), + Self::CreateIndex { table } => write!(f, "CreateIndex(t{table})"), + Self::DropTable { table } => write!(f, "DropTable(t{table})"), + Self::BlobWrite { table, len } => write!(f, "BlobWrite(t{table}, {len} bytes)"), + Self::BigTxn { statements } => write!(f, "BigTxn({statements} statements)"), + Self::UserVersionBump + | Self::Vacuum + | Self::Sync + | Self::Checkpoint + | Self::Snapshot + | Self::RestartReplicator + | Self::StrayWrite => f.write_str(self.kind_name()), + } + } +} + +/// Generate a deterministic workload of `len` ops from `seed`. +#[must_use] +pub fn generate_workload(seed: u64, len: usize) -> Vec { + let mut rng = StdRng::seed_from_u64(seed); + std::iter::repeat_with(|| sample_op(&mut rng)) + .take(len) + .collect() +} + +/// The `t0..t15` table namespace size. +const TABLE_COUNT: u8 = 16; + +/// The weight table: mostly `Insert`/`Update`/`Sync`, every variant +/// reachable well within 200 ops across typical seeds (validated by unit +/// tests). Weights sum to [`WEIGHT_TOTAL`]. +const WEIGHTED: [(Shape, u32); WorkloadOp::VARIANT_COUNT] = [ + (Shape::Insert, 24), + (Shape::Sync, 22), + (Shape::Update, 14), + (Shape::Delete, 6), + (Shape::Checkpoint, 6), + (Shape::BlobWrite, 5), + (Shape::BigTxn, 4), + (Shape::CreateTable, 4), + (Shape::CreateIndex, 3), + (Shape::DropTable, 3), + (Shape::UserVersionBump, 2), + (Shape::Snapshot, 2), + (Shape::RestartReplicator, 2), + (Shape::StrayWrite, 2), + (Shape::Vacuum, 1), +]; + +/// Sum of the weights in [`WEIGHTED`] (checked by a unit test). +const WEIGHT_TOTAL: u32 = 100; + +/// A [`WorkloadOp`] variant without its parameters, for weighted sampling. +#[derive(Debug, Clone, Copy)] +enum Shape { + Insert, + Update, + Delete, + CreateTable, + CreateIndex, + DropTable, + BlobWrite, + BigTxn, + UserVersionBump, + Vacuum, + Sync, + Checkpoint, + Snapshot, + RestartReplicator, + StrayWrite, +} + +/// Draw one weighted op from the rng. +fn sample_op(rng: &mut StdRng) -> WorkloadOp { + let mut roll = rng.random_range(0..WEIGHT_TOTAL); + for (shape, weight) in WEIGHTED { + if roll < weight { + return instantiate(shape, rng); + } + roll -= weight; + } + // Unreachable while the weights sum to WEIGHT_TOTAL (unit-tested); fall + // back to the cheapest op rather than panicking. + WorkloadOp::Sync +} + +/// Fill in the parameters for a sampled op shape. +fn instantiate(shape: Shape, rng: &mut StdRng) -> WorkloadOp { + match shape { + Shape::Insert => WorkloadOp::Insert { + table: rng.random_range(0..TABLE_COUNT), + rows: rng.random_range(1..=20), + }, + Shape::Update => WorkloadOp::Update { + table: rng.random_range(0..TABLE_COUNT), + rows: rng.random_range(1..=10), + }, + Shape::Delete => WorkloadOp::Delete { + table: rng.random_range(0..TABLE_COUNT), + rows: rng.random_range(1..=10), + }, + Shape::CreateTable => WorkloadOp::CreateTable { + table: rng.random_range(0..TABLE_COUNT), + }, + Shape::CreateIndex => WorkloadOp::CreateIndex { + table: rng.random_range(0..TABLE_COUNT), + }, + Shape::DropTable => WorkloadOp::DropTable { + table: rng.random_range(0..TABLE_COUNT), + }, + Shape::BlobWrite => WorkloadOp::BlobWrite { + table: rng.random_range(0..TABLE_COUNT), + // Well past any page size, forcing overflow pages. + len: rng.random_range(4096..=64 * 1024), + }, + Shape::BigTxn => WorkloadOp::BigTxn { + statements: rng.random_range(1..=50), + }, + Shape::UserVersionBump => WorkloadOp::UserVersionBump, + Shape::Vacuum => WorkloadOp::Vacuum, + Shape::Sync => WorkloadOp::Sync, + Shape::Checkpoint => WorkloadOp::Checkpoint, + Shape::Snapshot => WorkloadOp::Snapshot, + Shape::RestartReplicator => WorkloadOp::RestartReplicator, + Shape::StrayWrite => WorkloadOp::StrayWrite, + } +} + +/// Everything needed to (re)build the engine over the same directories: +/// `RestartReplicator` simulates a process crash by dropping the engine and +/// resuming a fresh one against the same replica root. +pub struct WorkloadEnv { + pub log: Logger, + pub config: ReplicaConfig, + pub db: ReplicaDb, + pub clock: Clock, + /// Root of the local replica directory; rebuilt engines wrap it in a + /// fresh `Flaky(Local)` storage with an empty failure plan. + pub replica_root: Utf8PathBuf, +} + +/// Applies a workload script: SQL ops through dedicated rusqlite writer +/// connections, engine events through the step-driven [`SyncEngine`]. +pub struct WorkloadRunner { + seed: u64, + script: Vec, + next_index: usize, + env: WorkloadEnv, + engine: SyncEngine, + /// The scripted "app" writer connection. + conn: rusqlite::Connection, + /// A second connection bypassing the app writer mutex (`StrayWrite`). + stray: rusqlite::Connection, + /// Bitset of which `t0..t15` tables currently exist. + tables: u16, +} + +/// A workload failure, carrying the seed and (for op failures) the failing +/// op index and op. The calling test prints the full script alongside. +#[derive(Debug, thiserror::Error)] +pub enum WorkloadError { + #[error("workload seed {seed} setup: {source}")] + Setup { seed: u64, source: rusqlite::Error }, + #[error("workload seed {seed} op {index} ({op}): {source}")] + Op { + seed: u64, + index: usize, + op: WorkloadOp, + source: WorkloadOpError, + }, + #[error("workload seed {seed} drain: {source}")] + Drain { seed: u64, source: WorkloadOpError }, +} + +/// Why a single workload op failed. +#[derive(Debug, thiserror::Error)] +pub enum WorkloadOpError { + #[error("SQLite: {0}")] + Sqlite(#[from] rusqlite::Error), + #[error("engine step failed: {0}")] + Engine(#[from] SyncError), + #[error("sync tick reported an error: {0}")] + Tick(SyncError), + #[error("snapshot still in progress after {steps} steps")] + SnapshotStalled { steps: u32 }, + #[error("engine failed to quiesce after {ticks} ticks (state: {})", state_name(*state))] + DrainStalled { ticks: u32, state: EngineState }, +} + +/// Upper bound on [`SyncEngine::snapshot_step`] calls per `Snapshot` op. +const SNAPSHOT_MAX_STEPS: u32 = 10_000; + +/// Upper bound on [`SyncEngine::sync_once`] calls in +/// [`WorkloadRunner::drain`]. +const DRAIN_MAX_TICKS: u32 = 10_000; + +impl WorkloadRunner { + /// Open the writer connections over the (already existing, WAL-mode) + /// source database and take ownership of the engine. + pub fn new( + seed: u64, + script: Vec, + env: WorkloadEnv, + engine: SyncEngine, + ) -> Result { + let setup = |source| WorkloadError::Setup { seed, source }; + let conn = open_writer(&env.db.db_path).map_err(setup)?; + let stray = open_writer(&env.db.db_path).map_err(setup)?; + let tables = scan_tables(&conn).map_err(setup)?; + Ok(Self { + seed, + script, + next_index: 0, + env, + engine, + conn, + stray, + tables, + }) + } + + /// Apply the next script op; `Ok(None)` once the script is exhausted. + pub async fn step(&mut self) -> Result, WorkloadError> { + let Some(op) = self.script.get(self.next_index).cloned() else { + return Ok(None); + }; + let index = self.next_index; + self.next_index = index.saturating_add(1); + let checkpoint = self + .apply(index, &op) + .await + .map_err(|source| WorkloadError::Op { + seed: self.seed, + index, + op: op.clone(), + source, + })?; + Ok(Some(AppliedOp { + index, + op, + checkpoint, + })) + } + + /// Apply every remaining script op. + pub async fn run(&mut self) -> Result<(), WorkloadError> { + while self.step().await?.is_some() {} + Ok(()) + } + + /// Final drain: drive `sync_once` until any in-flight snapshot has + /// finished and a tick ships nothing, so the replica holds every + /// committed transaction. + pub async fn drain(&mut self) -> Result<(), WorkloadError> { + self.drain_inner() + .await + .map_err(|source| WorkloadError::Drain { + seed: self.seed, + source, + }) + } + + /// The engine under test. + #[must_use] + pub fn engine(&self) -> &SyncEngine { + &self.engine + } + + /// Mutable access to the engine under test. + pub fn engine_mut(&mut self) -> &mut SyncEngine { + &mut self.engine + } + + async fn drain_inner(&mut self) -> Result<(), WorkloadOpError> { + for _tick in 0..DRAIN_MAX_TICKS { + let progress = self.engine.sync_once().await?; + if let Some(error) = progress.error { + return Err(WorkloadOpError::Tick(error)); + } + if progress.backing_off { + continue; + } + // AwaitingEpoch counts as drained: it means everything shipped + // and checkpointed, with no new frames to bind yet. + let quiescent = progress.shipped_segments == 0 + && progress.snapshot.is_none() + && matches!( + self.engine.state(), + EngineState::Streaming | EngineState::AwaitingEpoch + ); + if quiescent { + return Ok(()); + } + } + Err(WorkloadOpError::DrainStalled { + ticks: DRAIN_MAX_TICKS, + state: self.engine.state(), + }) + } + + /// Dispatch one op. Returns the checkpoint outcome for `Checkpoint`. + async fn apply( + &mut self, + index: usize, + op: &WorkloadOp, + ) -> Result, WorkloadOpError> { + match op { + WorkloadOp::Insert { table, rows } => self.insert(index, *table, *rows)?, + WorkloadOp::Update { table, rows } => self.update(index, *table, *rows)?, + WorkloadOp::Delete { table, rows } => self.delete(*table, *rows)?, + WorkloadOp::CreateTable { table } => self.ensure_table(*table)?, + WorkloadOp::CreateIndex { table } => self.create_index(*table)?, + WorkloadOp::DropTable { table } => self.drop_table(*table)?, + WorkloadOp::BlobWrite { table, len } => self.blob_write(index, *table, *len)?, + WorkloadOp::BigTxn { statements } => self.big_txn(index, *statements)?, + WorkloadOp::UserVersionBump => self.user_version_bump()?, + WorkloadOp::Vacuum => self.conn.execute_batch("VACUUM")?, + WorkloadOp::Sync => self.sync().await?, + WorkloadOp::Checkpoint => return Ok(Some(self.checkpoint().await?)), + WorkloadOp::Snapshot => self.snapshot().await?, + WorkloadOp::RestartReplicator => self.restart().await?, + WorkloadOp::StrayWrite => self.stray_write(index)?, + } + Ok(None) + } + + async fn sync(&mut self) -> Result<(), WorkloadOpError> { + let progress = self.engine.sync_once().await?; + if let Some(error) = progress.error { + return Err(WorkloadOpError::Tick(error)); + } + Ok(()) + } + + /// Production `sync_once` ships before a due checkpoint; mirror that + /// order so a scripted checkpoint can actually complete (I1). + async fn checkpoint(&mut self) -> Result { + match self.engine.state() { + EngineState::Streaming | EngineState::AwaitingEpoch => { + self.engine.ship_once().await?; + }, + EngineState::PendingSnapshot | EngineState::Snapshotting => {}, + } + Ok(self.engine.checkpoint_once().await?) + } + + /// Trigger a new-generation snapshot and drive it to completion. + async fn snapshot(&mut self) -> Result<(), WorkloadOpError> { + self.engine.trigger_snapshot(); + for _step in 0..SNAPSHOT_MAX_STEPS { + if self.engine.snapshot_step().await? == SnapshotStatus::Finished { + return Ok(()); + } + } + Err(WorkloadOpError::SnapshotStalled { + steps: SNAPSHOT_MAX_STEPS, + }) + } + + /// Process-crash simulation: rebuild the engine over the same replica + /// root and let the resume logic pick the position back up. The + /// step-driven engine is inert between steps, so building the successor + /// before the old engine drops is equivalent to drop-then-rebuild (the + /// resume reads only on-disk and replica state). + async fn restart(&mut self) -> Result<(), WorkloadOpError> { + let storage = ReplicaStorage::Flaky(Box::new(FlakyStorage::new( + ReplicaStorage::Local(LocalStorage::new(self.env.replica_root.clone())), + FailurePlan::new(), + ))); + self.engine = SyncEngine::new_with_storage( + self.env.log.clone(), + self.env.config.clone(), + self.env.db.clone(), + self.env.clock.clone(), + false, + storage, + ) + .await?; + Ok(()) + } + + fn insert(&mut self, index: usize, table: u8, rows: u16) -> Result<(), WorkloadOpError> { + self.ensure_table(table)?; + let sql = format!("INSERT INTO {} (data) VALUES (?1)", table_name(table)); + let seed = self.seed; + self.in_txn(|conn| { + let mut statement = conn.prepare(&sql)?; + for row in 0..rows { + statement.execute(rusqlite::params![format!("seed{seed}-op{index}-row{row}")])?; + } + Ok(()) + })?; + Ok(()) + } + + fn update(&mut self, index: usize, table: u8, rows: u16) -> Result<(), WorkloadOpError> { + if !self.table_exists(table) { + // Deterministic fallback: updating an absent table becomes an + // insert (which creates it), keeping scripts replayable. + return self.insert(index, table, rows); + } + let sql = format!( + "UPDATE {t} SET data = ?1 WHERE id IN (SELECT id FROM {t} ORDER BY id LIMIT ?2)", + t = table_name(table) + ); + self.conn.execute( + &sql, + rusqlite::params![ + format!("seed{}-op{index}-update", self.seed), + i64::from(rows) + ], + )?; + Ok(()) + } + + fn delete(&mut self, table: u8, rows: u16) -> Result<(), WorkloadOpError> { + if !self.table_exists(table) { + // Deterministic no-op on an absent table. + return Ok(()); + } + let sql = format!( + "DELETE FROM {t} WHERE id IN (SELECT id FROM {t} ORDER BY id LIMIT ?1)", + t = table_name(table) + ); + self.conn + .execute(&sql, rusqlite::params![i64::from(rows)])?; + Ok(()) + } + + /// Create the table if it does not exist yet (lazy first touch). + fn ensure_table(&mut self, table: u8) -> Result<(), WorkloadOpError> { + if self.table_exists(table) { + return Ok(()); + } + self.conn.execute_batch(&create_table_sql(table))?; + self.set_table(table, true); + Ok(()) + } + + fn create_index(&mut self, table: u8) -> Result<(), WorkloadOpError> { + self.ensure_table(table)?; + let slot = table_slot(table); + self.conn.execute_batch(&format!( + "CREATE INDEX IF NOT EXISTS idx_t{slot} ON t{slot} (data)" + ))?; + Ok(()) + } + + fn drop_table(&mut self, table: u8) -> Result<(), WorkloadOpError> { + if !self.table_exists(table) { + // Deterministic no-op on an absent table. + return Ok(()); + } + self.conn + .execute_batch(&format!("DROP TABLE IF EXISTS {}", table_name(table)))?; + self.set_table(table, false); + Ok(()) + } + + fn blob_write(&mut self, index: usize, table: u8, len: u32) -> Result<(), WorkloadOpError> { + self.ensure_table(table)?; + // u32 always fits usize on supported (64-bit) targets. + let len = usize::try_from(len).unwrap_or_default(); + // The mask keeps the value in u8 range, so the conversion is total. + let byte = u8::try_from(index & 0xff).unwrap_or_default(); + let blob = vec![byte; len]; + self.conn.execute( + &format!("INSERT INTO {} (bin) VALUES (?1)", table_name(table)), + rusqlite::params![blob], + )?; + Ok(()) + } + + /// `statements` inserts in one commit, always into table `t0`. + fn big_txn(&mut self, index: usize, statements: u8) -> Result<(), WorkloadOpError> { + const BIG_TABLE: u8 = 0; + self.ensure_table(BIG_TABLE)?; + let sql = format!("INSERT INTO {} (data) VALUES (?1)", table_name(BIG_TABLE)); + let seed = self.seed; + self.in_txn(|conn| { + let mut statement = conn.prepare(&sql)?; + for n in 0..statements { + statement.execute(rusqlite::params![format!("seed{seed}-op{index}-big{n}")])?; + } + Ok(()) + })?; + Ok(()) + } + + fn user_version_bump(&self) -> Result<(), WorkloadOpError> { + let version: i64 = self + .conn + .query_row("PRAGMA user_version", [], |row| row.get(0))?; + self.conn + .pragma_update(None, "user_version", version.saturating_add(1))?; + Ok(()) + } + + /// One insert through the second connection, always into table `t0`. + fn stray_write(&mut self, index: usize) -> Result<(), WorkloadOpError> { + const STRAY_TABLE: u8 = 0; + if !self.table_exists(STRAY_TABLE) { + self.stray.execute_batch(&create_table_sql(STRAY_TABLE))?; + self.set_table(STRAY_TABLE, true); + } + self.stray.execute( + &format!("INSERT INTO {} (data) VALUES (?1)", table_name(STRAY_TABLE)), + rusqlite::params![format!("seed{}-op{index}-stray", self.seed)], + )?; + Ok(()) + } + + /// Run `body` inside a single `BEGIN IMMEDIATE` transaction (one commit + /// frame set), rolling back on error. + fn in_txn(&self, body: F) -> Result<(), rusqlite::Error> + where + F: FnOnce(&rusqlite::Connection) -> Result<(), rusqlite::Error>, + { + self.conn.execute_batch("BEGIN IMMEDIATE")?; + match body(&self.conn) { + Ok(()) => self.conn.execute_batch("COMMIT"), + Err(error) => { + let _rollback = self.conn.execute_batch("ROLLBACK"); + Err(error) + }, + } + } + + fn table_exists(&self, table: u8) -> bool { + self.tables & table_bit(table) != 0 + } + + fn set_table(&mut self, table: u8, exists: bool) { + if exists { + self.tables |= table_bit(table); + } else { + self.tables &= !table_bit(table); + } + } +} + +/// One applied script op, with the checkpoint outcome when the op was +/// [`WorkloadOp::Checkpoint`]. +#[derive(Debug)] +pub struct AppliedOp { + pub index: usize, + pub op: WorkloadOp, + pub checkpoint: Option, +} + +/// Build a runner and apply the whole script, returning the runner for the +/// final drain and assertions. +pub async fn run_workload( + seed: u64, + script: Vec, + env: WorkloadEnv, + engine: SyncEngine, +) -> Result, WorkloadError> { + let mut runner = WorkloadRunner::new(seed, script, env, engine)?; + runner.run().await?; + Ok(runner) +} + +/// A writer connection with the same discipline as the app: 5s busy timeout +/// and `wal_autocheckpoint = 0` (invariant I2 applies to every writer). +fn open_writer(db_path: &Utf8Path) -> Result { + let conn = rusqlite::Connection::open(db_path)?; + conn.busy_timeout(Duration::from_secs(5))?; + let _pages: i64 = conn.query_row("PRAGMA wal_autocheckpoint = 0", [], |row| row.get(0))?; + Ok(conn) +} + +/// Which `t0..t15` tables already exist (the runner may be built over a +/// database with earlier workload state). +fn scan_tables(conn: &rusqlite::Connection) -> Result { + let mut tables = 0u16; + let mut statement = conn.prepare("SELECT name FROM sqlite_master WHERE type = 'table'")?; + let names = statement.query_map([], |row| row.get::<_, String>(0))?; + for name in names { + let name = name?; + if let Some(rest) = name.strip_prefix('t') + && let Ok(slot) = rest.parse::() + && slot < TABLE_COUNT + { + tables |= table_bit(slot); + } + } + Ok(tables) +} + +/// Map any u8 into the `t0..t15` namespace (mask instead of modulo). +const fn table_slot(table: u8) -> u8 { + table & (TABLE_COUNT - 1) +} + +const fn table_bit(table: u8) -> u16 { + 1 << table_slot(table) +} + +fn table_name(table: u8) -> String { + format!("t{}", table_slot(table)) +} + +fn create_table_sql(table: u8) -> String { + format!( + "CREATE TABLE IF NOT EXISTS {} (id INTEGER PRIMARY KEY, data TEXT, bin BLOB)", + table_name(table) + ) +} + +/// [`EngineState`] does not implement `Display`; name it for error messages. +fn state_name(state: EngineState) -> &'static str { + match state { + EngineState::Streaming => "Streaming", + EngineState::AwaitingEpoch => "AwaitingEpoch", + EngineState::PendingSnapshot => "PendingSnapshot", + EngineState::Snapshotting => "Snapshotting", + } +} + +#[cfg(test)] +mod tests { + use std::collections::{BTreeMap, BTreeSet}; + + use pretty_assertions::assert_eq; + + use super::{WEIGHT_TOTAL, WEIGHTED, WorkloadOp, generate_workload}; + + const LEN: usize = 200; + + const ALL_VARIANTS: [&str; WorkloadOp::VARIANT_COUNT] = [ + "BigTxn", + "BlobWrite", + "Checkpoint", + "CreateIndex", + "CreateTable", + "Delete", + "DropTable", + "Insert", + "RestartReplicator", + "Snapshot", + "StrayWrite", + "Sync", + "Update", + "UserVersionBump", + "Vacuum", + ]; + + #[test] + fn generate_is_deterministic() { + assert_eq!( + generate_workload(42, 500), + generate_workload(42, 500), + "the same seed must generate the identical script" + ); + } + + #[test] + fn generate_has_requested_length() { + assert_eq!(generate_workload(7, 123).len(), 123, "script length"); + } + + #[test] + fn different_seeds_differ() { + assert_ne!( + generate_workload(0, LEN), + generate_workload(1, LEN), + "different seeds must generate different scripts" + ); + } + + #[test] + fn weights_sum_to_total_and_cover_every_variant() { + let sum: u32 = WEIGHTED.iter().map(|(_, weight)| *weight).sum(); + assert_eq!(sum, WEIGHT_TOTAL, "the weight table must sum to the total"); + let shapes: BTreeSet<&'static str> = WEIGHTED + .iter() + .map(|(shape, _)| { + // Round-trip through instantiate via a fixed rng so the + // weight table provably reaches every WorkloadOp variant. + use rand::SeedableRng as _; + let mut rng = rand::rngs::StdRng::seed_from_u64(0); + super::instantiate(*shape, &mut rng).kind_name() + }) + .collect(); + let all: BTreeSet<&'static str> = ALL_VARIANTS.iter().copied().collect(); + assert_eq!(shapes, all, "the weight table must list every variant"); + } + + #[test] + fn seeds_0_through_7_collectively_cover_every_variant() { + let mut seen = BTreeSet::new(); + for seed in 0..8u64 { + for op in generate_workload(seed, LEN) { + seen.insert(op.kind_name()); + } + } + let all: BTreeSet<&'static str> = ALL_VARIANTS.iter().copied().collect(); + assert_eq!( + seen, all, + "seeds 0..8 x {LEN} ops must collectively cover every variant" + ); + } + + #[test] + fn no_variant_dominates() { + for seed in 0..8u64 { + let script = generate_workload(seed, LEN); + let mut counts = BTreeMap::new(); + for op in &script { + *counts.entry(op.kind_name()).or_insert(0usize) += 1; + } + for (kind, count) in counts { + assert!( + count * 10 <= LEN * 6, + "seed {seed}: {kind} appears {count}/{LEN} times, over the 60% cap" + ); + } + } + } +} diff --git a/plus/bencher_replica/src/verify.rs b/plus/bencher_replica/src/verify.rs new file mode 100644 index 000000000..159955501 --- /dev/null +++ b/plus/bencher_replica/src/verify.rs @@ -0,0 +1,536 @@ +//! Restore-and-compare verification: prove the replica reproduces the +//! source database at a known position. +//! +//! The full choreography (ship tail, record position P, pin a read snapshot, +//! then compare off-lock) is wired into the sync loop separately; this +//! module provides the pieces: a logical fingerprint over a pinned +//! connection, restore-to-position, and the comparison. +//! +//! The row/schema serialization helpers here are `pub(crate)` and reused by +//! the test-only equivalence oracle in [`crate::testing`] (the reverse +//! dependency is impossible: production code cannot depend on the testing +//! module). + +use camino::Utf8Path; +use rusqlite::types::ValueRef; +use sha2::{Digest as _, Sha256}; + +use crate::position::Position; +use crate::storage::ReplicaStorage; + +/// The result of one verification run. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum VerifyReport { + /// Source and restored fingerprints match at the verified position. + Pass, + /// Fingerprints differ: the replica lineage is divergent. + Fail { + /// Human-readable description of the first difference. + detail: String, + }, +} + +#[derive(Debug, thiserror::Error)] +pub enum VerifyError { + #[error("Failed to fingerprint database: {0}")] + Fingerprint(rusqlite::Error), + #[error("Failed to restore replica for verification: {0}")] + Restore(#[from] crate::restore::RestoreError), + #[error("Verification scratch directory: {0}")] + TempDir(std::io::Error), + #[error("Verification task panicked: {0}")] + Join(tokio::task::JoinError), +} + +/// Compute a logical fingerprint of the database visible through `conn` +/// (which should hold a pinned read snapshot): the `user_version`, every +/// `sqlite_master` row (`ORDER BY type, name`), per user table (internal +/// `sqlite_%` tables skipped) a row count and a SHA-256 over its rows, and +/// the `sqlite_sequence` AUTOINCREMENT counters (`ORDER BY name`) when that +/// internal table exists. Rows are serialized with type-tagged, +/// delimiter-escaped values in a deterministic order (`rowid`, or every +/// column for `WITHOUT ROWID` tables that have none). +/// +/// The output is a deterministic line-oriented byte vector, so a failed +/// comparison can name the first differing line. +pub fn fingerprint_database(conn: &rusqlite::Connection) -> Result, VerifyError> { + let mut lines = Vec::new(); + let version = user_version(conn).map_err(VerifyError::Fingerprint)?; + lines.push(format!("user_version={version}")); + for row in schema_rows(conn).map_err(VerifyError::Fingerprint)? { + lines.push(format!("schema={row}")); + } + for table in user_table_names(conn).map_err(VerifyError::Fingerprint)? { + let (rows, sha256) = table_digest(conn, &table).map_err(VerifyError::Fingerprint)?; + lines.push(format!("table={table}|rows={rows}|sha256={sha256}")); + } + // `sqlite_sequence` holds AUTOINCREMENT high-water marks; it is skipped + // by `user_table_names` as an internal table, so a restore that corrupted + // it would otherwise be invisible. + for row in sqlite_sequence_rows(conn).map_err(VerifyError::Fingerprint)? { + lines.push(format!("sqlite_sequence={row}")); + } + let mut fingerprint = lines.join("\n"); + fingerprint.push('\n'); + Ok(fingerprint.into_bytes()) +} + +/// Restore the replica into `scratch_dir` up to `position` and compare its +/// fingerprint against `source_fingerprint`. The restore itself runs the +/// same `quick_check` hard gate as a startup restore. +pub async fn verify_against_replica( + log: &slog::Logger, + storage: &ReplicaStorage, + position: &Position, + source_fingerprint: &[u8], + scratch_dir: &Utf8Path, +) -> Result { + std::fs::create_dir_all(scratch_dir).map_err(VerifyError::TempDir)?; + let scratch_db = scratch_dir.join("verify.db"); + let restored = crate::restore::restore_to(log, storage, &scratch_db, Some(position)).await?; + if restored.is_none() { + // An empty (or vanished) replica cannot reproduce the source; the + // caller reacts the same way as to a content mismatch. + return Ok(VerifyReport::Fail { + detail: "no valid generation found on the replica".to_owned(), + }); + } + let fingerprint = tokio::task::spawn_blocking(move || -> Result, VerifyError> { + // No CREATE flag: a missing scratch database is a bug, not an empty + // database. + let conn = rusqlite::Connection::open_with_flags( + &scratch_db, + rusqlite::OpenFlags::SQLITE_OPEN_READ_WRITE | rusqlite::OpenFlags::SQLITE_OPEN_NO_MUTEX, + ) + .map_err(VerifyError::Fingerprint)?; + fingerprint_database(&conn) + }) + .await + .map_err(VerifyError::Join)??; + if fingerprint == source_fingerprint { + Ok(VerifyReport::Pass) + } else { + Ok(VerifyReport::Fail { + detail: first_differing_line(source_fingerprint, &fingerprint), + }) + } +} + +/// The `sqlite_master` rows serialized deterministically +/// (`ORDER BY type, name`; columns `type, name, tbl_name, sql`). +pub(crate) fn schema_rows(conn: &rusqlite::Connection) -> Result, rusqlite::Error> { + query_serialized_rows( + conn, + "SELECT type, name, tbl_name, sql FROM sqlite_master ORDER BY type, name", + ) +} + +/// User table names in deterministic order; internal `sqlite_%` tables are +/// skipped (they cannot be queried like user tables). +pub(crate) fn user_table_names( + conn: &rusqlite::Connection, +) -> Result, rusqlite::Error> { + let mut statement = conn.prepare( + "SELECT name FROM sqlite_master \ + WHERE type = 'table' AND name NOT LIKE 'sqlite_%' ORDER BY name", + )?; + let names = statement.query_map([], |row| row.get(0))?; + names.collect() +} + +/// The `sqlite_sequence` AUTOINCREMENT counters (`name, seq` `ORDER BY +/// name`), serialized like any other row. Empty when the table does not +/// exist (no AUTOINCREMENT column anywhere), so a schema without +/// AUTOINCREMENT never errors. +pub(crate) fn sqlite_sequence_rows( + conn: &rusqlite::Connection, +) -> Result, rusqlite::Error> { + if !table_exists(conn, "sqlite_sequence")? { + return Ok(Vec::new()); + } + query_serialized_rows(conn, "SELECT name, seq FROM sqlite_sequence ORDER BY name") +} + +/// Whether a table of the given name exists (`sqlite_master`). +fn table_exists(conn: &rusqlite::Connection, name: &str) -> Result { + let count: i64 = conn.query_row( + "SELECT count(*) FROM sqlite_master WHERE type = 'table' AND name = ?1", + [name], + |row| row.get(0), + )?; + Ok(count > 0) +} + +/// `PRAGMA user_version`. +pub(crate) fn user_version(conn: &rusqlite::Connection) -> Result { + conn.query_row("PRAGMA user_version", [], |row| row.get(0)) +} + +/// Run `sql` and serialize every result row with [`serialize_row`]. +pub(crate) fn query_serialized_rows( + conn: &rusqlite::Connection, + sql: &str, +) -> Result, rusqlite::Error> { + let mut statement = conn.prepare(sql)?; + let mut rows = statement.query([])?; + let mut serialized = Vec::new(); + while let Some(row) = rows.next()? { + serialized.push(serialize_row(row)?); + } + Ok(serialized) +} + +/// Serialize one result row as `|`-joined tagged values (see [`tag_value`]), +/// with each value's `|` characters escaped first so a value containing the +/// delimiter cannot merge with an adjacent column: `("a|b", "c")` and +/// `("a", "b|c")` must never serialize identically. +pub(crate) fn serialize_row(row: &rusqlite::Row) -> Result { + let column_count = row.as_ref().column_count(); + let mut values = Vec::with_capacity(column_count); + for index in 0..column_count { + values.push(escape_delimiter(&tag_value(row.get_ref(index)?))); + } + Ok(values.join("|")) +} + +/// Escape the `|` column delimiter inside an already-tagged value. `tag_value` +/// backslash-escapes TEXT via `escape_default` (which doubles `\` and never +/// emits `|`) and every other tag is delimiter-free, so an escaped `\|` is +/// unambiguous: after a join, a raw `|` is preceded by an even run of +/// backslashes exactly when it is a genuine column boundary. +fn escape_delimiter(value: &str) -> String { + value.replace('|', "\\|") +} + +/// Double-quote an identifier for interpolation into SQL. +pub(crate) fn quote_ident(name: &str) -> String { + format!("\"{}\"", name.replace('"', "\"\"")) +} + +/// Serialize one `SQLite` value with an explicit type tag so `1`, `'1'`, and +/// `x'31'` never collide. Deterministic across platforms: REALs carry their +/// exact bit pattern, TEXT is escaped to a single line, BLOBs are hex. +fn tag_value(value: ValueRef<'_>) -> String { + match value { + ValueRef::Null => "NULL".to_owned(), + ValueRef::Integer(int) => format!("INTEGER:{int}"), + ValueRef::Real(real) => format!("REAL:{real}:{:016x}", real.to_bits()), + ValueRef::Text(text) => match std::str::from_utf8(text) { + Ok(text) => format!("TEXT:{}", text.escape_default()), + Err(_) => format!("TEXT(hex):{}", hex::encode(text)), + }, + ValueRef::Blob(blob) => format!("BLOB:{}", hex::encode(blob)), + } +} + +/// Row count plus SHA-256 over the serialized rows of one table, streamed +/// row by row (a production fingerprint never materializes a whole table). +fn table_digest( + conn: &rusqlite::Connection, + table: &str, +) -> Result<(u64, String), rusqlite::Error> { + let order = row_order_clause(conn, table)?; + let mut statement = conn.prepare(&format!( + "SELECT * FROM {} ORDER BY {order}", + quote_ident(table) + ))?; + let mut rows = statement.query([])?; + let mut hasher = Sha256::new(); + let mut count = 0u64; + while let Some(row) = rows.next()? { + hasher.update(serialize_row(row)?.as_bytes()); + hasher.update(b"\n"); + count = count.saturating_add(1); + } + Ok((count, hex::encode(hasher.finalize()))) +} + +/// A deterministic `ORDER BY` clause for a table's digest. Ordinary tables +/// order by `rowid`; a `WITHOUT ROWID` table has none, so it orders by every +/// column (its PRIMARY KEY is a subset, so the order is canonical). Ordering +/// by `rowid` on a `WITHOUT ROWID` table would raise "no such column: rowid" +/// on every verification, forever. +fn row_order_clause(conn: &rusqlite::Connection, table: &str) -> Result { + if table_has_rowid(conn, table)? { + return Ok("rowid".to_owned()); + } + let clause = table_columns(conn, table)? + .iter() + .map(|column| quote_ident(column)) + .collect::>() + .join(", "); + // Every SQLite table has at least one column, and a WITHOUT ROWID table + // has a PRIMARY KEY, so the clause is never empty. + Ok(clause) +} + +/// Whether the table exposes a `rowid` (true for every ordinary table, +/// including `INTEGER PRIMARY KEY`; false for `WITHOUT ROWID`). Probing with +/// a zero-row `SELECT rowid` is the reliable signal: the column resolves +/// (prepare succeeds) for rowid tables and fails only for `WITHOUT ROWID` +/// ones. Any other preparation failure propagates. +fn table_has_rowid(conn: &rusqlite::Connection, table: &str) -> Result { + match conn.prepare(&format!("SELECT rowid FROM {} LIMIT 0", quote_ident(table))) { + Ok(_statement) => Ok(true), + Err(rusqlite::Error::SqliteFailure(_, Some(message))) if message.contains("rowid") => { + Ok(false) + }, + Err(error) => Err(error), + } +} + +/// Column names in declaration order (the order `SELECT *` yields). +fn table_columns(conn: &rusqlite::Connection, table: &str) -> Result, rusqlite::Error> { + let mut statement = conn.prepare(&format!("PRAGMA table_info({})", quote_ident(table)))?; + let names = statement.query_map([], |row| row.get::<_, String>(1))?; + names.collect() +} + +/// Name the first line where two fingerprints diverge. +fn first_differing_line(source: &[u8], restored: &[u8]) -> String { + let source_text = String::from_utf8_lossy(source); + let restored_text = String::from_utf8_lossy(restored); + let mut source_lines = source_text.lines(); + let mut restored_lines = restored_text.lines(); + loop { + match (source_lines.next(), restored_lines.next()) { + (Some(source_line), Some(restored_line)) if source_line == restored_line => {}, + (Some(source_line), Some(restored_line)) => { + return format!("source: {source_line} | restored: {restored_line}"); + }, + (Some(source_line), None) => return format!("source has extra line: {source_line}"), + (None, Some(restored_line)) => { + return format!("restored has extra line: {restored_line}"); + }, + (None, None) => return "fingerprints are equal".to_owned(), + } + } +} + +#[cfg(test)] +mod tests { + use pretty_assertions::assert_eq; + + use super::{fingerprint_database, first_differing_line}; + + fn seeded_db() -> rusqlite::Connection { + let conn = rusqlite::Connection::open_in_memory().unwrap(); + conn.execute_batch( + "CREATE TABLE t (id INTEGER PRIMARY KEY, data TEXT); + CREATE TABLE mixed (id INTEGER PRIMARY KEY, n INTEGER, r REAL, s TEXT, b BLOB); + INSERT INTO t (data) VALUES ('alpha'), ('beta'); + INSERT INTO mixed (n, r, s, b) VALUES + (42, 1.5, 'line1\nline2', x'00ff'), + (NULL, NULL, NULL, NULL);", + ) + .unwrap(); + conn + } + + #[test] + fn fingerprint_is_deterministic() { + let first = fingerprint_database(&seeded_db()).unwrap(); + let second = fingerprint_database(&seeded_db()).unwrap(); + assert_eq!( + first, second, + "identical databases produce identical fingerprints" + ); + let text = String::from_utf8(first).unwrap(); + assert!( + text.starts_with("user_version=0\n"), + "fingerprint starts with the user_version line: {text}" + ); + assert!( + text.contains("schema=TEXT:table|TEXT:t|TEXT:t|"), + "fingerprint carries the schema rows: {text}" + ); + assert!( + text.contains("table=t|rows=2|sha256="), + "fingerprint carries per-table digests: {text}" + ); + assert!( + text.contains("table=mixed|rows=2|sha256="), + "fingerprint covers every user table: {text}" + ); + } + + #[test] + fn fingerprint_reflects_row_changes() { + let conn = seeded_db(); + let before = fingerprint_database(&conn).unwrap(); + conn.execute("UPDATE t SET data = 'gamma' WHERE id = 1", []) + .unwrap(); + let after = fingerprint_database(&conn).unwrap(); + assert_ne!(before, after, "a row change must change the fingerprint"); + let detail = first_differing_line(&before, &after); + assert!( + detail.contains("table=t|rows=2|sha256="), + "the differing line names the changed table: {detail}" + ); + } + + #[test] + fn fingerprint_reflects_value_type_not_just_display() { + // The INTEGER 1 and the TEXT '1' must never fingerprint equally. + // The column is untyped (BLOB affinity) so no coercion happens. + let int_conn = rusqlite::Connection::open_in_memory().unwrap(); + int_conn + .execute_batch( + "CREATE TABLE v (id INTEGER PRIMARY KEY, x); INSERT INTO v (x) VALUES (1);", + ) + .unwrap(); + let text_conn = rusqlite::Connection::open_in_memory().unwrap(); + text_conn + .execute_batch( + "CREATE TABLE v (id INTEGER PRIMARY KEY, x); INSERT INTO v (x) VALUES ('1');", + ) + .unwrap(); + assert_ne!( + fingerprint_database(&int_conn).unwrap(), + fingerprint_database(&text_conn).unwrap(), + "tagged serialization distinguishes value types" + ); + } + + #[test] + fn fingerprint_reflects_user_version() { + let conn = seeded_db(); + let before = fingerprint_database(&conn).unwrap(); + conn.pragma_update(None, "user_version", 7).unwrap(); + let after = fingerprint_database(&conn).unwrap(); + assert_ne!(before, after, "user_version is part of the fingerprint"); + assert_eq!( + first_differing_line(&before, &after), + "source: user_version=0 | restored: user_version=7", + "the differing line is the user_version line" + ); + } + + #[test] + fn fingerprint_skips_sqlite_internal_tables() { + let conn = rusqlite::Connection::open_in_memory().unwrap(); + // AUTOINCREMENT creates the internal sqlite_sequence table. + conn.execute_batch( + "CREATE TABLE a (id INTEGER PRIMARY KEY AUTOINCREMENT, data TEXT); + INSERT INTO a (data) VALUES ('x');", + ) + .unwrap(); + let text = String::from_utf8(fingerprint_database(&conn).unwrap()).unwrap(); + assert!( + !text.contains("table=sqlite_sequence"), + "internal tables are not fingerprinted as user tables: {text}" + ); + assert!( + text.contains("TEXT:sqlite_sequence"), + "internal tables still appear in the schema rows: {text}" + ); + } + + #[test] + fn fingerprint_escapes_column_delimiter() { + // Without escaping, ("a|TEXT:b", "c") and ("a", "b|TEXT:c") both + // serialize to "TEXT:a|TEXT:b|TEXT:c", so a divergent table hashes + // equal: a false verification PASS. + let left = two_text_columns("a|TEXT:b", "c"); + let right = two_text_columns("a", "b|TEXT:c"); + assert_ne!( + fingerprint_database(&left).unwrap(), + fingerprint_database(&right).unwrap(), + "the column delimiter must be escaped so the two rows differ" + ); + } + + #[test] + fn fingerprint_reflects_sqlite_sequence() { + // Two DBs with an identical (empty) user table but different + // AUTOINCREMENT high-water marks must not fingerprint equally: a + // restore that corrupted sqlite_sequence has to be caught. + let low = autoincrement_db(1); + let high = autoincrement_db(1000); + assert_ne!( + fingerprint_database(&low).unwrap(), + fingerprint_database(&high).unwrap(), + "sqlite_sequence counters are part of the fingerprint" + ); + let text = String::from_utf8(fingerprint_database(&high).unwrap()).unwrap(); + assert!( + text.contains("sqlite_sequence=TEXT:a|INTEGER:1000"), + "the AUTOINCREMENT counter appears in the fingerprint: {text}" + ); + } + + #[test] + fn fingerprint_handles_without_rowid_tables() { + // `ORDER BY rowid` errors forever on a WITHOUT ROWID table; the + // fingerprint must still compute and still reflect row changes. + let conn = rusqlite::Connection::open_in_memory().unwrap(); + conn.execute_batch( + "CREATE TABLE wr (k TEXT PRIMARY KEY, v INTEGER) WITHOUT ROWID; + INSERT INTO wr (k, v) VALUES ('a', 1), ('b', 2);", + ) + .unwrap(); + let before = fingerprint_database(&conn).unwrap(); + let text = String::from_utf8(before.clone()).unwrap(); + assert!( + text.contains("table=wr|rows=2|sha256="), + "the WITHOUT ROWID table is fingerprinted: {text}" + ); + conn.execute("UPDATE wr SET v = 99 WHERE k = 'a'", []) + .unwrap(); + let after = fingerprint_database(&conn).unwrap(); + assert_ne!( + before, after, + "a change in a WITHOUT ROWID table changes the fingerprint" + ); + } + + fn two_text_columns(a: &str, b: &str) -> rusqlite::Connection { + let conn = rusqlite::Connection::open_in_memory().unwrap(); + conn.execute_batch("CREATE TABLE t (a TEXT, b TEXT);") + .unwrap(); + conn.execute( + "INSERT INTO t (a, b) VALUES (?1, ?2)", + rusqlite::params![a, b], + ) + .unwrap(); + conn + } + + fn autoincrement_db(seq: i64) -> rusqlite::Connection { + let conn = rusqlite::Connection::open_in_memory().unwrap(); + // Insert then delete leaves the table empty but seeds sqlite_sequence; + // the explicit UPDATE then sets the counter deterministically. + conn.execute_batch( + "CREATE TABLE a (id INTEGER PRIMARY KEY AUTOINCREMENT, data TEXT); + INSERT INTO a (data) VALUES ('x'); + DELETE FROM a;", + ) + .unwrap(); + conn.execute( + "UPDATE sqlite_sequence SET seq = ?1 WHERE name = 'a'", + [seq], + ) + .unwrap(); + conn + } + + #[test] + fn first_differing_line_reports_each_shape() { + assert_eq!( + first_differing_line(b"a\nb\n", b"a\nc\n"), + "source: b | restored: c" + ); + assert_eq!( + first_differing_line(b"a\nb\n", b"a\n"), + "source has extra line: b" + ); + assert_eq!( + first_differing_line(b"a\n", b"a\nb\n"), + "restored has extra line: b" + ); + assert_eq!( + first_differing_line(b"a\n", b"a\n"), + "fingerprints are equal" + ); + } +} diff --git a/plus/bencher_replica/src/wal.rs b/plus/bencher_replica/src/wal.rs new file mode 100644 index 000000000..8adfbaa85 --- /dev/null +++ b/plus/bencher_replica/src/wal.rs @@ -0,0 +1,1541 @@ +//! `SQLite` WAL file parsing: header and frame validation with salt and +//! cumulative checksum verification. +//! +//! Format reference: +//! +//! Layout: a 32-byte header followed by zero or more frames. Each frame is a +//! 24-byte header plus one page of payload (`page_size` bytes). +//! +//! Only bytes up through the last *commit* frame may ever be shipped to the +//! replica: frames after the last commit belong to an open (or torn) +//! transaction and are not durable. + +use std::io::{Read, Seek, SeekFrom}; + +/// Size of the WAL file header in bytes. +pub const WAL_HEADER_SIZE: u64 = 32; +/// Size of each frame header in bytes. +pub const FRAME_HEADER_SIZE: u64 = 24; + +/// [`WAL_HEADER_SIZE`] as a `usize` for buffer sizing. +const WAL_HEADER_LEN: usize = 32; +/// [`FRAME_HEADER_SIZE`] as a `usize` for buffer sizing. +const FRAME_HEADER_LEN: usize = 24; +/// Minimum valid `SQLite` page size. +const MIN_PAGE_SIZE: u32 = 512; +/// Maximum valid `SQLite` page size (65536). +const MAX_PAGE_SIZE: u32 = 0x0001_0000; + +/// WAL magic when checksums are computed over little-endian words. +pub const WAL_MAGIC_LE: u32 = 0x377f_0682; +/// WAL magic when checksums are computed over big-endian words. +pub const WAL_MAGIC_BE: u32 = 0x377f_0683; +/// The only supported WAL format version. +pub const WAL_FORMAT: u32 = 3_007_000; + +/// Parsed WAL file header (all fields stored big-endian on disk). +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct WalHeader { + pub magic: u32, + pub format: u32, + pub page_size: u32, + pub checkpoint_seq: u32, + pub salt: (u32, u32), + /// Header self-checksum over the first 24 bytes, seeded with (0, 0). + pub checksum: (u32, u32), +} + +impl WalHeader { + /// Whether frame checksums are computed over big-endian 32-bit words + /// (magic `0x377f0683`) or little-endian words (magic `0x377f0682`). + #[must_use] + pub fn big_endian_checksum(&self) -> bool { + self.magic == WAL_MAGIC_BE + } + + /// Total on-disk size of one frame: header plus one page. + #[must_use] + pub fn frame_size(&self) -> u64 { + FRAME_HEADER_SIZE + u64::from(self.page_size) + } +} + +#[derive(Debug, thiserror::Error)] +pub enum WalError { + #[error("Failed to read WAL: {0}")] + Read(#[source] std::io::Error), + #[error("Failed to seek WAL: {0}")] + Seek(#[source] std::io::Error), + #[error("WAL header is truncated: {0} bytes")] + TruncatedHeader(usize), + #[error("Invalid WAL magic: {0:#010x}")] + BadMagic(u32), + #[error("Unsupported WAL format version: {0}")] + UnsupportedFormat(u32), + #[error("Invalid WAL page size: {0}")] + InvalidPageSize(u32), + #[error("WAL header checksum mismatch: stored {stored:?}, computed {computed:?}")] + HeaderChecksum { + stored: (u32, u32), + computed: (u32, u32), + }, + #[error( + "WAL resume position is invalid: offset {offset} is not frame-aligned for page size {page_size}" + )] + MisalignedOffset { offset: u64, page_size: u32 }, + #[error( + "A WAL run of {bytes} bytes since the last commit exceeds the scan bound of {max_bytes}; refusing to buffer or ship it" + )] + TransactionTooLarge { bytes: u64, max_bytes: u64 }, +} + +/// Parse and validate a 32-byte WAL header. +/// +/// Errors on truncation, bad magic, unsupported format, invalid page size +/// (must be a power of two between 512 and 65536), or a header self-checksum +/// mismatch. +pub fn parse_wal_header(bytes: &[u8]) -> Result { + let header: &[u8; WAL_HEADER_LEN] = bytes + .get(..WAL_HEADER_LEN) + .and_then(|prefix| prefix.try_into().ok()) + .ok_or(WalError::TruncatedHeader(bytes.len()))?; + let magic = be_u32(header[0], header[1], header[2], header[3]); + if magic != WAL_MAGIC_LE && magic != WAL_MAGIC_BE { + return Err(WalError::BadMagic(magic)); + } + let format = be_u32(header[4], header[5], header[6], header[7]); + if format != WAL_FORMAT { + return Err(WalError::UnsupportedFormat(format)); + } + let page_size = be_u32(header[8], header[9], header[10], header[11]); + if !page_size.is_power_of_two() || !(MIN_PAGE_SIZE..=MAX_PAGE_SIZE).contains(&page_size) { + return Err(WalError::InvalidPageSize(page_size)); + } + let checkpoint_seq = be_u32(header[12], header[13], header[14], header[15]); + let salt = ( + be_u32(header[16], header[17], header[18], header[19]), + be_u32(header[20], header[21], header[22], header[23]), + ); + let stored = ( + be_u32(header[24], header[25], header[26], header[27]), + be_u32(header[28], header[29], header[30], header[31]), + ); + // The self-checksum covers header bytes 0..24 (everything before it) + let (covered, _stored_bytes) = header.split_at(24); + let computed = wal_checksum(magic == WAL_MAGIC_BE, (0, 0), covered); + if stored != computed { + return Err(WalError::HeaderChecksum { stored, computed }); + } + Ok(WalHeader { + magic, + format, + page_size, + checkpoint_seq, + salt, + checksum: stored, + }) +} + +/// `SQLite`'s custom cumulative WAL checksum. +/// +/// `data.len()` must be a multiple of 8. Words are read as big-endian when +/// `big_endian` is true, little-endian otherwise. The running pair `seed` +/// (`s1`, `s2`) is folded as: `s1 += x[i] + s2; s2 += x[i+1] + s1;` with +/// wrapping 32-bit arithmetic. +/// +/// - Header self-checksum: input is header bytes `0..24`, seed `(0, 0)`. +/// - Frame checksum: input is frame-header bytes `0..8` followed by the full +/// page payload, seeded with the previous frame's checksum (or the header +/// checksum for the first frame). Stored in frame-header bytes `16..24`. +#[must_use] +pub fn wal_checksum(big_endian: bool, seed: (u32, u32), data: &[u8]) -> (u32, u32) { + assert!( + data.len().is_multiple_of(8), + "WAL checksum input must be a multiple of 8 bytes" + ); + let (mut s1, mut s2) = seed; + for pair in data.chunks_exact(8) { + let mut word = [0u8; 8]; + word.copy_from_slice(pair); + let (x0, x1) = if big_endian { + ( + be_u32(word[0], word[1], word[2], word[3]), + be_u32(word[4], word[5], word[6], word[7]), + ) + } else { + ( + le_u32(word[0], word[1], word[2], word[3]), + le_u32(word[4], word[5], word[6], word[7]), + ) + }; + s1 = s1.wrapping_add(x0).wrapping_add(s2); + s2 = s2.wrapping_add(x1).wrapping_add(s1); + } + (s1, s2) +} + +/// Parsed 24-byte frame header (all fields big-endian on disk). +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct FrameHeader { + pub page_no: u32, + /// For commit frames, the size of the database in pages after the commit; + /// zero for all non-commit frames. + pub db_size: u32, + pub salt: (u32, u32), + pub checksum: (u32, u32), +} + +impl FrameHeader { + /// Extract the fields from a raw 24-byte frame header. + #[must_use] + pub fn parse(bytes: &[u8; 24]) -> Self { + Self { + page_no: be_u32(bytes[0], bytes[1], bytes[2], bytes[3]), + db_size: be_u32(bytes[4], bytes[5], bytes[6], bytes[7]), + salt: ( + be_u32(bytes[8], bytes[9], bytes[10], bytes[11]), + be_u32(bytes[12], bytes[13], bytes[14], bytes[15]), + ), + checksum: ( + be_u32(bytes[16], bytes[17], bytes[18], bytes[19]), + be_u32(bytes[20], bytes[21], bytes[22], bytes[23]), + ), + } + } + + /// Whether this frame commits a transaction. + #[must_use] + pub fn is_commit(&self) -> bool { + self.db_size != 0 + } +} + +/// A run of verified, committed WAL bytes ending exactly on a commit frame. +#[derive(Debug, Clone)] +pub struct CommittedChunk { + /// Byte offset in the WAL file where this chunk starts (inclusive). + pub start_offset: u64, + /// Byte offset just past the last commit frame in this chunk (exclusive). + pub end_offset: u64, + /// Running checksum at `end_offset`, for resuming the chain. + pub checksum_at_end: (u32, u32), + /// Number of commit frames contained in this chunk. + pub commit_count: u64, + /// Database size in pages recorded by the last commit frame. + pub db_size_pages: u32, + /// Raw WAL bytes `[start_offset, end_offset)`. + pub bytes: Vec, +} + +/// Incremental, checksum-verified scanner over a WAL file. +/// +/// The scanner walks frames from a known position, verifying that each +/// frame's salts match the header salts and that the cumulative checksum +/// chain is intact. It stops (returning `None`) at the first frame that is +/// torn, checksum-broken, salt-stale, pgno-0, or past EOF; frames after the +/// last commit are never surfaced. A stop is not an error: the valid prefix +/// up to the stop point remains shippable. +pub struct WalScanner { + reader: R, + header: WalHeader, + /// Next unread byte offset; always frame-aligned: + /// `(offset - 32) % frame_size == 0`. + offset: u64, + /// Running checksum at `offset`. + checksum: (u32, u32), +} + +impl WalScanner { + /// Open a WAL from the start: read and validate the header, position the + /// scanner at the first frame. + /// + /// Returns `Ok(None)` when the file is shorter than a full header (an + /// empty or freshly truncated WAL): there is nothing to ship and nothing + /// wrong. + pub fn open(mut reader: R) -> Result, WalError> { + reader.seek(SeekFrom::Start(0)).map_err(WalError::Seek)?; + let mut header_buf = [0u8; WAL_HEADER_LEN]; + if !read_full(&mut reader, &mut header_buf)? { + return Ok(None); + } + let header = parse_wal_header(&header_buf)?; + Ok(Some(Self { + reader, + header, + offset: WAL_HEADER_SIZE, + checksum: header.checksum, + })) + } + + /// Resume scanning from a previously verified position. + /// + /// `offset` must be frame-aligned and `checksum` must be the running + /// checksum at that offset (as previously returned in + /// [`CommittedChunk::checksum_at_end`]). + pub fn resume( + reader: R, + header: WalHeader, + offset: u64, + checksum: (u32, u32), + ) -> Result { + let aligned = offset + .checked_sub(WAL_HEADER_SIZE) + .is_some_and(|frame_bytes| frame_bytes.is_multiple_of(header.frame_size())); + if !aligned { + return Err(WalError::MisalignedOffset { + offset, + page_size: header.page_size, + }); + } + Ok(Self { + reader, + header, + offset, + checksum, + }) + } + + /// The validated WAL header. + #[must_use] + pub fn header(&self) -> &WalHeader { + &self.header + } + + /// Current frame-aligned offset (next unread byte). + #[must_use] + pub fn offset(&self) -> u64 { + self.offset + } + + /// Running checksum at [`Self::offset`]. + #[must_use] + pub fn checksum(&self) -> (u32, u32) { + self.checksum + } + + /// Read the next run of committed frames, up to `max_bytes` of raw WAL + /// bytes (always ending on a commit-frame boundary; a single transaction + /// larger than `max_bytes` is returned whole, so chunks may exceed + /// `max_bytes`). + /// + /// Returns `Ok(None)` when no complete committed transaction lies beyond + /// the current offset (clean EOF, torn tail, checksum break, or stale + /// salts). The scanner does not advance past the last commit it returns. + /// + /// MEMORY: this retaining scan applies no transaction bound (returning an + /// oversized transaction whole is the contract), so an arbitrarily large + /// open or rolled-back tail is heap-buffered before being discarded. + /// Callers facing an untrusted tail must front-run with the bounded + /// [`Self::scan_committed_extent`] (as the ship path does) and cap this + /// read to the discovered extent. + pub fn next_committed(&mut self, max_bytes: u64) -> Result, WalError> { + // Retain page bytes; no transaction bound (the public contract is to + // return an oversized transaction whole). + self.scan(max_bytes, u64::MAX, Retention::Retain) + } + + /// Scan-only variant of [`Self::next_committed`]: validate the committed + /// prefix and advance to the last commit within `max_bytes` while + /// DISCARDING page bytes (checksums require reading every byte, not + /// retaining it). `max_txn_bytes` bounds the bytes read since the last + /// commit, so an oversized open transaction errors with + /// [`WalError::TransactionTooLarge`] instead of scanning (and, for the + /// retaining scan, allocating) without bound. Returns the absolute WAL + /// offset just past the last commit (equal to [`Self::offset`] on entry + /// when nothing new is committed). + pub fn scan_committed_extent( + &mut self, + max_bytes: u64, + max_txn_bytes: u64, + ) -> Result { + self.scan(max_bytes, max_txn_bytes, Retention::Discard)?; + Ok(self.offset) + } + + /// The shared scan loop behind [`Self::next_committed`] and + /// [`Self::scan_committed_extent`]. `retention` selects whether page + /// bytes are accumulated (a heap buffer for the returned chunk) or + /// discarded (offsets and checksums only). `max_txn_bytes` is the bytes + /// permitted since the last commit before the scan aborts, so a huge + /// open transaction can never buffer or scan unboundedly. + fn scan( + &mut self, + max_bytes: u64, + max_txn_bytes: u64, + retention: Retention, + ) -> Result, WalError> { + self.reader + .seek(SeekFrom::Start(self.offset)) + .map_err(WalError::Seek)?; + let frame_size = self.header.frame_size(); + let big_endian = self.header.big_endian_checksum(); + let mut frame_header_buf = [0u8; FRAME_HEADER_LEN]; + let mut page_buf = vec![0u8; self.header.page_size as usize]; + let retain = matches!(retention, Retention::Retain); + + // Frames read so far, possibly extending past the last commit (empty + // in discard mode). + let mut pending: Vec = Vec::new(); + let mut cursor = self.offset; + let mut running = self.checksum; + + // High-water mark of the last commit frame seen + let mut committed_len = 0usize; + let mut committed_end = self.offset; + let mut committed_checksum = self.checksum; + let mut commit_count = 0u64; + let mut db_size_pages = 0u32; + + loop { + // A torn read anywhere in the frame ends the valid prefix + if !read_full(&mut self.reader, &mut frame_header_buf)? + || !read_full(&mut self.reader, &mut page_buf)? + { + break; + } + let frame = FrameHeader::parse(&frame_header_buf); + // Stale salts (leftovers from before a WAL restart) end the prefix + if frame.salt != self.header.salt { + break; + } + // `SQLite`'s walDecodeFrame rejects page number 0; a checksum-valid + // pgno-0 frame is not a real frame and ends the valid prefix + // exactly like a checksum mismatch. + if frame.page_no == 0 { + break; + } + // The frame checksum covers frame-header bytes 0..8 plus the page, + // seeded from the previous frame (or the WAL header) + let (covered_head, _rest) = frame_header_buf.split_at(8); + let after_head = wal_checksum(big_endian, running, covered_head); + let computed = wal_checksum(big_endian, after_head, &page_buf); + if computed != frame.checksum { + break; + } + running = computed; + if retain { + pending.extend_from_slice(&frame_header_buf); + pending.extend_from_slice(&page_buf); + } + cursor += frame_size; + // Bound the run since the last commit BEFORE it grows further, so + // an open (or committed) transaction beyond the cap aborts here + // instead of buffering or scanning the whole thing. + let since_commit = cursor.saturating_sub(committed_end); + if since_commit > max_txn_bytes { + return Err(WalError::TransactionTooLarge { + bytes: since_commit, + max_bytes: max_txn_bytes, + }); + } + if frame.is_commit() { + committed_len = pending.len(); + committed_end = cursor; + committed_checksum = running; + commit_count += 1; + db_size_pages = frame.db_size; + if committed_end - self.offset >= max_bytes { + break; + } + } + } + + if commit_count == 0 { + return Ok(None); + } + // Drop any valid-but-uncommitted frames read past the last commit + // (a no-op in discard mode, where `pending` stayed empty). + pending.truncate(committed_len); + let chunk = CommittedChunk { + start_offset: self.offset, + end_offset: committed_end, + checksum_at_end: committed_checksum, + commit_count, + db_size_pages, + bytes: pending, + }; + self.offset = committed_end; + self.checksum = committed_checksum; + Ok(Some(chunk)) + } +} + +/// Whether [`WalScanner::scan`] retains page bytes for a returned chunk or +/// discards them (offsets and checksums only). +#[derive(Clone, Copy)] +enum Retention { + Retain, + Discard, +} + +/// Fill `buf` from `reader`. Returns `Ok(true)` when completely filled, +/// `Ok(false)` on a clean or torn EOF before the buffer is full. +fn read_full(reader: &mut R, buf: &mut [u8]) -> Result { + let mut filled = 0usize; + while filled < buf.len() { + let Some(rest) = buf.get_mut(filled..) else { + break; + }; + match reader.read(rest) { + Ok(0) => return Ok(false), + Ok(n) => filled += n, + Err(err) if err.kind() == std::io::ErrorKind::Interrupted => {}, + Err(err) => return Err(WalError::Read(err)), + } + } + Ok(true) +} + +/// Assemble a big-endian `u32` (`SQLite` WAL integer fields are big-endian). +fn be_u32(b0: u8, b1: u8, b2: u8, b3: u8) -> u32 { + (u32::from(b0) << 24) | (u32::from(b1) << 16) | (u32::from(b2) << 8) | u32::from(b3) +} + +/// Assemble a little-endian `u32` (checksum words under magic `0x377f0682`). +fn le_u32(b0: u8, b1: u8, b2: u8, b3: u8) -> u32 { + u32::from(b0) | (u32::from(b1) << 8) | (u32::from(b2) << 16) | (u32::from(b3) << 24) +} + +#[cfg(test)] +mod tests { + use std::fmt::Write as _; + use std::io::Cursor; + + use camino::Utf8Path; + use pretty_assertions::assert_eq; + + use super::*; + use crate::testing::{CheckpointMode, SyntheticWal, WalFixture}; + + /// Small page size keeps synthetic WALs compact. + const PAGE: u32 = 512; + const SALT: (u32, u32) = (0x1122_3344, 0x5566_7788); + + fn page_of(fill: u8) -> Vec { + vec![fill; usize::try_from(PAGE).unwrap()] + } + + fn push_be(bytes: &mut Vec, value: u32) { + bytes.push(u8::try_from(value >> 24).unwrap()); + bytes.push(u8::try_from((value >> 16) & 0xff).unwrap()); + bytes.push(u8::try_from((value >> 8) & 0xff).unwrap()); + bytes.push(u8::try_from(value & 0xff).unwrap()); + } + + /// Hand-build a 32-byte WAL header with a correct self-checksum. + fn build_header( + magic: u32, + format: u32, + page_size: u32, + seq: u32, + salt: (u32, u32), + ) -> Vec { + let mut header = Vec::with_capacity(32); + push_be(&mut header, magic); + push_be(&mut header, format); + push_be(&mut header, page_size); + push_be(&mut header, seq); + push_be(&mut header, salt.0); + push_be(&mut header, salt.1); + let checksum = wal_checksum(magic == WAL_MAGIC_BE, (0, 0), &header); + push_be(&mut header, checksum.0); + push_be(&mut header, checksum.1); + header + } + + fn scan_all(bytes: &[u8]) -> (WalHeader, Vec) { + let mut scanner = WalScanner::open(Cursor::new(bytes.to_vec())) + .unwrap() + .unwrap(); + let header = *scanner.header(); + let mut chunks = Vec::new(); + while let Some(chunk) = scanner.next_committed(u64::MAX).unwrap() { + chunks.push(chunk); + } + (header, chunks) + } + + /// Total shippable end offset over a whole WAL byte string. + fn shippable_end(bytes: &[u8]) -> u64 { + let (_, chunks) = scan_all(bytes); + chunks + .last() + .map_or(WAL_HEADER_SIZE, |chunk| chunk.end_offset) + } + + fn chunk_frames(header: &WalHeader, chunk: &CommittedChunk) -> Vec { + let frame_size = usize::try_from(header.frame_size()).unwrap(); + chunk + .bytes + .chunks_exact(frame_size) + .map(|frame| { + let mut head = [0u8; 24]; + head.copy_from_slice(&frame[..24]); + FrameHeader::parse(&head) + }) + .collect() + } + + fn tempdir_path(dir: &tempfile::TempDir) -> &Utf8Path { + Utf8Path::from_path(dir.path()).unwrap() + } + + // 1. Header parsing and the checksum fold + + #[test] + fn checksum_fold_known_vectors() { + // One pair, little-endian words: x = (1, 2) + let data = [1, 0, 0, 0, 2, 0, 0, 0]; + assert_eq!(wal_checksum(false, (0, 0), &data), (1, 3)); + // Same bytes as big-endian words + assert_eq!( + wal_checksum(true, (0, 0), &data), + (0x0100_0000, 0x0300_0000) + ); + // Two pairs, big-endian words: x = (1, 2, 3, 4) + // s1 = 1, s2 = 3; then s1 = 1 + 3 + 3 = 7, s2 = 3 + 4 + 7 = 14 + let data = [0, 0, 0, 1, 0, 0, 0, 2, 0, 0, 0, 3, 0, 0, 0, 4]; + assert_eq!(wal_checksum(true, (0, 0), &data), (7, 14)); + // Seeding continues the chain + assert_eq!(wal_checksum(true, (1, 3), &data[8..]), (7, 14)); + // Empty input returns the seed + assert_eq!(wal_checksum(false, (9, 9), &[]), (9, 9)); + } + + #[test] + fn reject_truncated_header() { + for len in [0usize, 1, 31] { + let bytes = vec![0u8; len]; + let err = parse_wal_header(&bytes).unwrap_err(); + assert!( + matches!(err, WalError::TruncatedHeader(n) if n == len), + "expected TruncatedHeader({len}), got {err:?}" + ); + } + } + + #[test] + fn reject_bad_magic() { + let header = build_header(0xdead_beef, WAL_FORMAT, PAGE, 0, SALT); + let err = parse_wal_header(&header).unwrap_err(); + assert!( + matches!(err, WalError::BadMagic(0xdead_beef)), + "expected BadMagic, got {err:?}" + ); + // WalScanner::open propagates header errors (32 bytes is not "short") + let Err(err) = WalScanner::open(Cursor::new(header)) else { + panic!("expected BadMagic from open") + }; + assert!( + matches!(err, WalError::BadMagic(0xdead_beef)), + "expected BadMagic from open, got {err:?}" + ); + } + + #[test] + fn reject_unsupported_format() { + let header = build_header(WAL_MAGIC_LE, 3_007_001, PAGE, 0, SALT); + let err = parse_wal_header(&header).unwrap_err(); + assert!( + matches!(err, WalError::UnsupportedFormat(3_007_001)), + "expected UnsupportedFormat, got {err:?}" + ); + } + + #[test] + fn reject_invalid_page_size() { + for page_size in [0u32, 1, 100, 256, 1024 + 1, 2 * MAX_PAGE_SIZE] { + let header = build_header(WAL_MAGIC_LE, WAL_FORMAT, page_size, 0, SALT); + let err = parse_wal_header(&header).unwrap_err(); + assert!( + matches!(err, WalError::InvalidPageSize(n) if n == page_size), + "expected InvalidPageSize({page_size}), got {err:?}" + ); + } + } + + #[test] + fn reject_header_checksum_mismatch() { + // Corrupt a covered byte (a salt byte, so magic/format/page_size + // still parse) after the checksum was computed + let mut header = build_header(WAL_MAGIC_LE, WAL_FORMAT, PAGE, 0, SALT); + header[17] ^= 0x01; + let err = parse_wal_header(&header).unwrap_err(); + assert!( + matches!(err, WalError::HeaderChecksum { stored, computed } if stored != computed), + "expected HeaderChecksum, got {err:?}" + ); + // Corrupt a stored checksum byte (24..32) + let mut header = build_header(WAL_MAGIC_LE, WAL_FORMAT, PAGE, 0, SALT); + header[25] ^= 0x80; + let err = parse_wal_header(&header).unwrap_err(); + assert!( + matches!(err, WalError::HeaderChecksum { .. }), + "expected HeaderChecksum, got {err:?}" + ); + } + + // 2. Empty and header-only WALs + + #[test] + fn parse_empty_file() { + for len in [0usize, 1, 31] { + let scanner = WalScanner::open(Cursor::new(vec![0u8; len])).unwrap(); + assert!( + scanner.is_none(), + "expected Ok(None) for a {len}-byte WAL file" + ); + } + } + + #[test] + fn parse_header_only_wal() { + let bytes = SyntheticWal::new(PAGE, false, SALT).bytes(); + assert_eq!(bytes.len(), 32); + let mut scanner = WalScanner::open(Cursor::new(bytes)).unwrap().unwrap(); + assert_eq!(scanner.header().page_size, PAGE); + assert_eq!(scanner.header().salt, SALT); + assert_eq!(scanner.offset(), WAL_HEADER_SIZE); + assert!( + scanner.next_committed(u64::MAX).unwrap().is_none(), + "a header-only WAL has nothing to ship" + ); + assert_eq!(scanner.offset(), WAL_HEADER_SIZE); + } + + // 3. Real-WAL fixtures: header and frame iteration + + #[test] + fn parse_valid_header_all_page_sizes() { + for page_size in [512u32, 4096, 8192, MAX_PAGE_SIZE] { + let tmp = tempfile::tempdir().unwrap(); + let fixture = WalFixture::new(tempdir_path(&tmp), page_size).unwrap(); + let wal = fixture.wal_bytes().unwrap(); + assert!( + wal.len() > 32, + "fixture WAL for page size {page_size} should have frames" + ); + let header = parse_wal_header(&wal).unwrap(); + assert_eq!(header.page_size, page_size, "page size {page_size}"); + assert_eq!(header.format, WAL_FORMAT, "page size {page_size}"); + assert!( + header.magic == WAL_MAGIC_LE || header.magic == WAL_MAGIC_BE, + "unexpected magic {:#010x}", + header.magic + ); + // Everything the fixture wrote is committed: the shippable + // prefix covers the whole file. + let (_, chunks) = scan_all(&wal); + let end = chunks.last().map(|chunk| chunk.end_offset); + assert_eq!( + end, + Some(u64::try_from(wal.len()).unwrap()), + "page size {page_size}" + ); + } + } + + #[test] + fn parse_single_commit_frame() { + let tmp = tempfile::tempdir().unwrap(); + let fixture = WalFixture::new(tempdir_path(&tmp), 4096).unwrap(); + // Reset the WAL so it contains exactly one transaction + fixture.checkpoint(CheckpointMode::Truncate).unwrap(); + fixture + .txn(&["INSERT INTO t (data) VALUES ('one')"]) + .unwrap(); + let wal = fixture.wal_bytes().unwrap(); + let (header, chunks) = scan_all(&wal); + assert_eq!(chunks.len(), 1, "one committed chunk"); + let chunk = &chunks[0]; + assert_eq!(chunk.commit_count, 1, "one commit frame"); + assert_eq!(chunk.start_offset, WAL_HEADER_SIZE); + assert_eq!(chunk.end_offset, u64::try_from(wal.len()).unwrap()); + let frames = chunk_frames(&header, chunk); + assert!(!frames.is_empty(), "at least one frame"); + let last = frames.last().unwrap(); + assert!(last.is_commit(), "last frame carries the commit flag"); + assert_eq!(chunk.db_size_pages, last.db_size); + for frame in &frames[..frames.len() - 1] { + assert!(!frame.is_commit(), "only the last frame commits"); + } + } + + #[test] + fn parse_multi_frame_single_commit() { + let tmp = tempfile::tempdir().unwrap(); + let fixture = WalFixture::new(tempdir_path(&tmp), 4096).unwrap(); + fixture.checkpoint(CheckpointMode::Truncate).unwrap(); + fixture.txn_touching_pages(8).unwrap(); + let wal = fixture.wal_bytes().unwrap(); + let (header, chunks) = scan_all(&wal); + assert_eq!(chunks.len(), 1, "one committed chunk"); + let chunk = &chunks[0]; + assert_eq!(chunk.commit_count, 1, "a single commit"); + let frames = chunk_frames(&header, chunk); + assert!( + frames.len() > 1, + "expected a multi-frame transaction, got {} frame(s)", + frames.len() + ); + assert_eq!(chunk.end_offset, u64::try_from(wal.len()).unwrap()); + } + + // 4. The crate's most important test: uncommitted frames never ship + + #[test] + fn uncommitted_tail_frames_not_shipped() { + let tmp = tempfile::tempdir().unwrap(); + let fixture = WalFixture::new(tempdir_path(&tmp), 4096).unwrap(); + fixture.checkpoint(CheckpointMode::Truncate).unwrap(); + fixture + .txn(&["INSERT INTO t (data) VALUES ('committed')"]) + .unwrap(); + let committed_len = u64::try_from(fixture.wal_bytes().unwrap().len()).unwrap(); + + // Spill a big transaction into the WAL, then roll it back: the WAL + // now ends with flushed-but-uncommitted frames. + fixture.big_txn_spilling(false).unwrap(); + let wal = fixture.wal_bytes().unwrap(); + assert!( + u64::try_from(wal.len()).unwrap() > committed_len, + "rolled-back transaction must have spilled frames into the WAL" + ); + + // The shippable prefix ends exactly at the last commit; not one byte + // of the rolled-back tail is surfaced. + assert_eq!(shippable_end(&wal), committed_len); + + // Small max_bytes chunking never crosses the boundary either + let mut scanner = WalScanner::open(Cursor::new(wal.clone())).unwrap().unwrap(); + while let Some(chunk) = scanner.next_committed(1).unwrap() { + assert!( + chunk.end_offset <= committed_len, + "chunk end {} crosses the last commit at {committed_len}", + chunk.end_offset + ); + } + assert_eq!(scanner.offset(), committed_len); + + // Once a spilling transaction commits, its frames become shippable + fixture.big_txn_spilling(true).unwrap(); + let wal = fixture.wal_bytes().unwrap(); + let end = shippable_end(&wal); + assert!( + end > committed_len, + "committed spill must ship: end {end} <= {committed_len}" + ); + assert!(end <= u64::try_from(wal.len()).unwrap(), "end within file"); + } + + // 5. Chain breaks: corruption ends the trusted prefix + + #[test] + fn chain_break_bit_flip_in_frame_header() { + let frame_size = usize::try_from(FRAME_HEADER_SIZE + u64::from(PAGE)).unwrap(); + // Flip a page_no bit in the second frame's header + let bytes = SyntheticWal::new(PAGE, false, SALT) + .commit_frame(1, &page_of(0xaa), 1) + .commit_frame(2, &page_of(0xbb), 2) + .commit_frame(3, &page_of(0xcc), 3) + .flip_bit(32 + frame_size + 3, 0) + .bytes(); + let (_, chunks) = scan_all(&bytes); + assert_eq!(chunks.len(), 1, "only the frame before the break ships"); + assert_eq!( + chunks[0].end_offset, + 32 + u64::try_from(frame_size).unwrap() + ); + assert_eq!(chunks[0].commit_count, 1, "one commit before the break"); + } + + #[test] + fn chain_break_bit_flip_in_page_data() { + let frame_size = usize::try_from(FRAME_HEADER_SIZE + u64::from(PAGE)).unwrap(); + // Flip a payload bit in the second frame + let bytes = SyntheticWal::new(PAGE, false, SALT) + .commit_frame(1, &page_of(0xaa), 1) + .commit_frame(2, &page_of(0xbb), 2) + .commit_frame(3, &page_of(0xcc), 3) + .flip_bit(32 + frame_size + 24 + 100, 7) + .bytes(); + let (_, chunks) = scan_all(&bytes); + assert_eq!(chunks.len(), 1, "only the frame before the break ships"); + assert_eq!( + chunks[0].end_offset, + 32 + u64::try_from(frame_size).unwrap() + ); + } + + // 6. Checksum byte orders + + #[test] + fn checksum_verified_both_byte_orders() { + let mut ends = Vec::new(); + for big_endian in [false, true] { + let bytes = SyntheticWal::new(PAGE, big_endian, SALT) + .frame(1, &page_of(0x11)) + .commit_frame(2, &page_of(0x22), 2) + .commit_frame(1, &page_of(0x33), 2) + .bytes(); + let (header, chunks) = scan_all(&bytes); + assert_eq!( + header.big_endian_checksum(), + big_endian, + "magic selects the checksum word order" + ); + let end = chunks.last().unwrap().end_offset; + assert_eq!( + end, + u64::try_from(bytes.len()).unwrap(), + "all frames verify in {} order", + if big_endian { + "big-endian" + } else { + "little-endian" + } + ); + let commits: u64 = chunks.iter().map(|chunk| chunk.commit_count).sum(); + assert_eq!(commits, 2, "both commits found"); + ends.push(end); + } + assert_eq!(ends[0], ends[1], "same logical content, same end offset"); + } + + // 7. Stale salts (leftovers from before a WAL restart) + + #[test] + fn reject_stale_salt_frames() { + let frame_size = FRAME_HEADER_SIZE + u64::from(PAGE); + let stale_salt = (SALT.0 ^ 1, SALT.1); + let bytes = SyntheticWal::new(PAGE, false, SALT) + .commit_frame(1, &page_of(0xaa), 1) + .frame_with_salts(2, &page_of(0xbb), stale_salt) + .commit_frame(3, &page_of(0xcc), 3) + .bytes(); + let (_, chunks) = scan_all(&bytes); + assert_eq!(chunks.len(), 1, "the valid prefix ends at the stale salt"); + assert_eq!(chunks[0].end_offset, 32 + frame_size); + assert_eq!(chunks[0].commit_count, 1, "one commit before the break"); + } + + // 8. Torn tails + + #[test] + fn frame_exactly_at_eof() { + let frame_size = FRAME_HEADER_SIZE + u64::from(PAGE); + // A complete but uncommitted frame at EOF is not shipped + let bytes = SyntheticWal::new(PAGE, false, SALT) + .commit_frame(1, &page_of(0xaa), 1) + .frame(2, &page_of(0xbb)) + .bytes(); + assert_eq!(shippable_end(&bytes), 32 + frame_size); + // A commit frame ending exactly at EOF is shipped in full + let bytes = SyntheticWal::new(PAGE, false, SALT) + .commit_frame(1, &page_of(0xaa), 1) + .commit_frame(2, &page_of(0xbb), 2) + .bytes(); + assert_eq!(shippable_end(&bytes), 32 + 2 * frame_size); + } + + #[test] + fn mid_frame_eof() { + let frame_size = FRAME_HEADER_SIZE + u64::from(PAGE); + let torn_at = usize::try_from(32 + frame_size + 300).unwrap(); + let bytes = SyntheticWal::new(PAGE, false, SALT) + .commit_frame(1, &page_of(0xaa), 1) + .commit_frame(2, &page_of(0xbb), 2) + .truncate_at(torn_at) + .bytes(); + assert_eq!( + shippable_end(&bytes), + 32 + frame_size, + "the torn frame is not shipped" + ); + // Torn inside the 24-byte frame header + let bytes = SyntheticWal::new(PAGE, false, SALT) + .commit_frame(1, &page_of(0xaa), 1) + .commit_frame(2, &page_of(0xbb), 2) + .truncate_at(usize::try_from(32 + frame_size + 10).unwrap()) + .bytes(); + assert_eq!(shippable_end(&bytes), 32 + frame_size); + } + + // 9. Commit metadata and chunking + + #[test] + fn commit_frame_db_size_tracked() { + let bytes = SyntheticWal::new(PAGE, false, SALT) + .commit_frame(1, &page_of(0xaa), 3) + .commit_frame(2, &page_of(0xbb), 7) + .bytes(); + // One big chunk reports the last commit's db_size + let (_, chunks) = scan_all(&bytes); + assert_eq!(chunks.len(), 1, "single chunk"); + assert_eq!(chunks[0].commit_count, 2, "two commits"); + assert_eq!(chunks[0].db_size_pages, 7, "last commit wins"); + // Per-commit chunks report each commit's db_size + let mut scanner = WalScanner::open(Cursor::new(bytes)).unwrap().unwrap(); + let first = scanner.next_committed(1).unwrap().unwrap(); + assert_eq!(first.db_size_pages, 3); + let second = scanner.next_committed(1).unwrap().unwrap(); + assert_eq!(second.db_size_pages, 7); + } + + #[test] + fn max_bytes_chunking_ends_on_commit_boundaries() { + let frame_size = FRAME_HEADER_SIZE + u64::from(PAGE); + let bytes = SyntheticWal::new(PAGE, false, SALT) + .commit_frame(1, &page_of(0xaa), 1) + .commit_frame(2, &page_of(0xbb), 2) + .commit_frame(3, &page_of(0xcc), 3) + .bytes(); + let mut scanner = WalScanner::open(Cursor::new(bytes)).unwrap().unwrap(); + let mut previous_end = WAL_HEADER_SIZE; + let mut previous_checksum = scanner.checksum(); + for index in 0..3u64 { + let chunk = scanner.next_committed(1).unwrap().unwrap(); + assert_eq!(chunk.start_offset, previous_end, "chunks are contiguous"); + assert_eq!( + chunk.end_offset, + 32 + (index + 1) * frame_size, + "each chunk ends on the next commit boundary" + ); + assert_eq!(chunk.commit_count, 1, "one commit per chunk"); + assert_eq!( + chunk.bytes.len(), + usize::try_from(frame_size).unwrap(), + "chunk bytes cover exactly the returned frames" + ); + assert_ne!( + chunk.checksum_at_end, previous_checksum, + "the running checksum advances" + ); + previous_end = chunk.end_offset; + previous_checksum = chunk.checksum_at_end; + } + assert!(scanner.next_committed(1).unwrap().is_none(), "clean EOF"); + assert!( + scanner.next_committed(1).unwrap().is_none(), + "EOF is idempotent" + ); + } + + #[test] + fn oversized_transaction_returned_whole() { + let frame_size = FRAME_HEADER_SIZE + u64::from(PAGE); + // Four non-commit frames then the commit: one transaction + let bytes = SyntheticWal::new(PAGE, false, SALT) + .frame(1, &page_of(0x01)) + .frame(2, &page_of(0x02)) + .frame(3, &page_of(0x03)) + .frame(4, &page_of(0x04)) + .commit_frame(5, &page_of(0x05), 5) + .bytes(); + let mut scanner = WalScanner::open(Cursor::new(bytes)).unwrap().unwrap(); + // max_bytes far below the transaction size + let chunk = scanner.next_committed(100).unwrap().unwrap(); + assert_eq!(chunk.commit_count, 1, "the whole transaction is returned"); + assert_eq!(chunk.end_offset, 32 + 5 * frame_size); + assert!( + chunk.bytes.len() > 100, + "a single transaction may exceed max_bytes" + ); + } + + // 10. Resume + + #[test] + fn resume_continues_chain() { + // Synthetic: chunk, then resume from the recorded position + let bytes = SyntheticWal::new(PAGE, false, SALT) + .commit_frame(1, &page_of(0xaa), 1) + .commit_frame(2, &page_of(0xbb), 2) + .commit_frame(3, &page_of(0xcc), 3) + .bytes(); + let (header, continuous) = scan_all(&bytes); + let continuous_bytes: Vec = continuous + .iter() + .flat_map(|chunk| chunk.bytes.clone()) + .collect(); + + let mut scanner = WalScanner::open(Cursor::new(bytes.clone())) + .unwrap() + .unwrap(); + let first = scanner.next_committed(1).unwrap().unwrap(); + drop(scanner); + let mut resumed = WalScanner::resume( + Cursor::new(bytes), + header, + first.end_offset, + first.checksum_at_end, + ) + .unwrap(); + let mut resumed_bytes = first.bytes.clone(); + while let Some(chunk) = resumed.next_committed(u64::MAX).unwrap() { + resumed_bytes.extend_from_slice(&chunk.bytes); + } + assert_eq!( + resumed_bytes, continuous_bytes, + "resume yields the same bytes as one continuous scan" + ); + + // Real WAL: scan, append more commits, resume from the old position + let tmp = tempfile::tempdir().unwrap(); + let fixture = WalFixture::new(tempdir_path(&tmp), 4096).unwrap(); + let wal = fixture.wal_bytes().unwrap(); + let (header, chunks) = scan_all(&wal); + let position = chunks.last().unwrap(); + let (offset, checksum) = (position.end_offset, position.checksum_at_end); + fixture + .txn(&["INSERT INTO t (data) VALUES ('after-resume')"]) + .unwrap(); + let wal = fixture.wal_bytes().unwrap(); + let mut resumed = + WalScanner::resume(Cursor::new(wal.clone()), header, offset, checksum).unwrap(); + let tail = resumed.next_committed(u64::MAX).unwrap().unwrap(); + assert_eq!(tail.start_offset, offset); + assert_eq!(tail.end_offset, u64::try_from(wal.len()).unwrap()); + // Identical to the tail of one continuous scan + let (_, full) = scan_all(&wal); + let full_bytes: Vec = full.iter().flat_map(|chunk| chunk.bytes.clone()).collect(); + let split = usize::try_from(offset - WAL_HEADER_SIZE).unwrap(); + assert_eq!(&full_bytes[split..], &tail.bytes[..]); + } + + #[test] + fn resume_rejects_misaligned_offset() { + let bytes = SyntheticWal::new(PAGE, false, SALT) + .commit_frame(1, &page_of(0xaa), 1) + .bytes(); + let header = parse_wal_header(&bytes).unwrap(); + let frame_size = header.frame_size(); + for offset in [0, 31, 33, 32 + frame_size - 1, 32 + frame_size + 1] { + let Err(err) = WalScanner::resume(Cursor::new(bytes.clone()), header, offset, (0, 0)) + else { + panic!("expected MisalignedOffset for offset {offset}") + }; + assert!( + matches!(err, WalError::MisalignedOffset { offset: o, page_size } if o == offset && page_size == PAGE), + "expected MisalignedOffset for {offset}, got {err:?}" + ); + } + // Frame-aligned offsets are accepted + for offset in [32, 32 + frame_size] { + WalScanner::resume(Cursor::new(bytes.clone()), header, offset, (0, 0)).unwrap(); + } + } + + // 10b. pgno-0 frames end the prefix (walDecodeFrame rejects pgno 0) + + #[test] + fn pgno_zero_frame_stops_prefix() { + let frame_size = FRAME_HEADER_SIZE + u64::from(PAGE); + // A checksum-valid frame carrying page number 0 sits between two + // otherwise-valid commits. The scan must stop before it. + let bytes = SyntheticWal::new(PAGE, false, SALT) + .commit_frame(1, &page_of(0xaa), 1) + .commit_frame(0, &page_of(0xbb), 2) + .commit_frame(3, &page_of(0xcc), 3) + .bytes(); + let (_, chunks) = scan_all(&bytes); + assert_eq!(chunks.len(), 1, "the valid prefix ends at the pgno-0 frame"); + assert_eq!(chunks[0].end_offset, 32 + frame_size); + assert_eq!( + chunks[0].commit_count, 1, + "one commit before the pgno-0 frame" + ); + } + + // 10c. Scan-only (discard) mode + + #[test] + fn scan_committed_extent_agrees_with_next_committed() { + // A multi-transaction WAL, some multi-frame. The discard scan must + // land on the exact offset and checksum the retaining scan reaches. + let bytes = SyntheticWal::new(PAGE, false, SALT) + .commit_frame(1, &page_of(0xaa), 1) + .frame(2, &page_of(0xbb)) + .commit_frame(3, &page_of(0xcc), 3) + .commit_frame(4, &page_of(0xdd), 4) + .bytes(); + // Retaining scan: walk to the end, recording the final offset/checksum. + let mut retaining = WalScanner::open(Cursor::new(bytes.clone())) + .unwrap() + .unwrap(); + let mut retained_end = WAL_HEADER_SIZE; + let mut retained_checksum = retaining.checksum(); + while let Some(chunk) = retaining.next_committed(u64::MAX).unwrap() { + retained_end = chunk.end_offset; + retained_checksum = chunk.checksum_at_end; + } + // Discard scan: one pass with no size cap covers the whole WAL. + let mut discarding = WalScanner::open(Cursor::new(bytes)).unwrap().unwrap(); + let extent = discarding + .scan_committed_extent(u64::MAX, u64::MAX) + .unwrap(); + assert_eq!(extent, retained_end, "discard scan reaches the same offset"); + assert_eq!(discarding.offset(), retained_end); + assert_eq!( + discarding.checksum(), + retained_checksum, + "discard scan recovers the same running checksum" + ); + } + + #[test] + fn scan_committed_extent_bounds_oversized_open_transaction() { + let frame_size = FRAME_HEADER_SIZE + u64::from(PAGE); + // One commit, then a long OPEN (never committed) run of valid frames. + let bytes = SyntheticWal::new(PAGE, false, SALT) + .commit_frame(1, &page_of(0x01), 1) + .frame(2, &page_of(0x02)) + .frame(3, &page_of(0x03)) + .frame(4, &page_of(0x04)) + .frame(5, &page_of(0x05)) + .bytes(); + // A cap just above a single frame aborts once the open run since the + // last commit exceeds it, instead of scanning the whole tail. + let cap = frame_size + 1; + let mut scanner = WalScanner::open(Cursor::new(bytes)).unwrap().unwrap(); + let err = scanner.scan_committed_extent(u64::MAX, cap).unwrap_err(); + assert!( + matches!(err, WalError::TransactionTooLarge { max_bytes, .. } if max_bytes == cap), + "expected TransactionTooLarge, got {err:?}" + ); + } + + // 10d. Read errors mid-scan propagate (never a silent short extent) + + /// A `Read + Seek` that serves `fail_after` bytes total, then errors on + /// every subsequent read (a disk fault partway through the WAL). + struct ErroringReader { + inner: Cursor>, + fail_after: u64, + read_total: u64, + } + + impl Read for ErroringReader { + fn read(&mut self, buf: &mut [u8]) -> std::io::Result { + if self.read_total >= self.fail_after { + return Err(std::io::Error::other("injected mid-WAL read error")); + } + let budget = usize::try_from(self.fail_after - self.read_total).unwrap_or(usize::MAX); + let cap = buf.len().min(budget); + let read = self.inner.read(&mut buf[..cap])?; + self.read_total += u64::try_from(read).unwrap_or(0); + Ok(read) + } + } + + impl Seek for ErroringReader { + fn seek(&mut self, pos: SeekFrom) -> std::io::Result { + self.inner.seek(pos) + } + } + + #[test] + fn scan_propagates_mid_wal_read_error() { + let frame_size = FRAME_HEADER_SIZE + u64::from(PAGE); + let bytes = SyntheticWal::new(PAGE, false, SALT) + .commit_frame(1, &page_of(0xaa), 1) + .commit_frame(2, &page_of(0xbb), 2) + .commit_frame(3, &page_of(0xcc), 3) + .bytes(); + // Serve the header plus the first frame, then error partway through + // the second frame: the first commit is valid but the scan must NOT + // return that short extent, it must surface the read error. + let fail_after = WAL_HEADER_SIZE + frame_size + 5; + let make_scanner = || { + WalScanner::open(ErroringReader { + inner: Cursor::new(bytes.clone()), + fail_after, + read_total: 0, + }) + .unwrap() + .unwrap() + }; + let err = make_scanner().next_committed(u64::MAX).unwrap_err(); + assert!( + matches!(err, WalError::Read(_)), + "next_committed must propagate the read error, got {err:?}" + ); + let err = make_scanner() + .scan_committed_extent(u64::MAX, u64::MAX) + .unwrap_err(); + assert!( + matches!(err, WalError::Read(_)), + "scan_committed_extent must propagate the read error, got {err:?}" + ); + } + + // 11. Cross-validation against SQLite itself + + #[test] + fn parser_agrees_with_sqlite_recovery() { + let tmp = tempfile::tempdir().unwrap(); + let fixture = WalFixture::new(tempdir_path(&tmp), 4096).unwrap(); + fixture + .txn(&[ + "INSERT INTO t (data) VALUES ('alpha')", + "INSERT INTO t (data) VALUES ('beta')", + ]) + .unwrap(); + fixture.txn_touching_pages(8).unwrap(); + fixture.txn(&["DELETE FROM t WHERE data = 'beta'"]).unwrap(); + // Leave a rolled-back tail in the WAL + fixture.big_txn_spilling(false).unwrap(); + let db = fixture.db_bytes().unwrap(); + let wal = fixture.wal_bytes().unwrap(); + + // SQLite's own recovery: copy db + wal, checkpoint, read + let sqlite_dir = tempfile::tempdir().unwrap(); + let sqlite_db = sqlite_dir.path().join("recovered.db"); + std::fs::write(&sqlite_db, &db).unwrap(); + std::fs::write(sqlite_dir.path().join("recovered.db-wal"), &wal).unwrap(); + let conn = rusqlite::Connection::open(&sqlite_db).unwrap(); + let _row: i64 = conn + .query_row("PRAGMA wal_checkpoint(TRUNCATE)", [], |row| row.get(0)) + .unwrap(); + let sqlite_rows = read_rows(&conn); + drop(conn); + + // Our parser: apply each committed frame's page to the db image + let mut scanner = WalScanner::open(Cursor::new(wal)).unwrap().unwrap(); + let header = *scanner.header(); + let page_size = usize::try_from(header.page_size).unwrap(); + let frame_size = usize::try_from(header.frame_size()).unwrap(); + let mut image = db; + let mut db_pages = 0u32; + while let Some(chunk) = scanner.next_committed(u64::MAX).unwrap() { + for frame in chunk.bytes.chunks_exact(frame_size) { + let mut head = [0u8; 24]; + head.copy_from_slice(&frame[..24]); + let frame_header = FrameHeader::parse(&head); + // pgno 0 is rejected by the scanner, so page_no >= 1 here; + // checked_sub keeps the helper safe if that ever regresses. + let page_index = frame_header.page_no.checked_sub(1).expect("page_no >= 1"); + let at = usize::try_from(page_index).unwrap() * page_size; + if image.len() < at + page_size { + image.resize(at + page_size, 0); + } + image[at..at + page_size].copy_from_slice(&frame[24..]); + } + db_pages = chunk.db_size_pages; + } + assert!(db_pages > 0, "at least one commit applied"); + image.truncate(usize::try_from(db_pages).unwrap() * page_size); + + let parsed_dir = tempfile::tempdir().unwrap(); + let parsed_db = parsed_dir.path().join("parsed.db"); + std::fs::write(&parsed_db, &image).unwrap(); + let conn = rusqlite::Connection::open(&parsed_db).unwrap(); + let check: String = conn + .query_row("PRAGMA integrity_check", [], |row| row.get(0)) + .unwrap(); + assert_eq!(check, "ok", "recovered image passes integrity_check"); + let parsed_rows = read_rows(&conn); + assert_eq!( + parsed_rows, sqlite_rows, + "applying parsed committed frames reproduces SQLite's own recovery" + ); + } + + #[test] + fn synthetic_wal_accepted_by_sqlite() { + // Build a real WAL, then rebuild it byte-for-byte with SyntheticWal + let tmp = tempfile::tempdir().unwrap(); + let fixture = WalFixture::new(tempdir_path(&tmp), 4096).unwrap(); + fixture + .txn(&["INSERT INTO t (data) VALUES ('alpha')"]) + .unwrap(); + fixture + .txn(&["INSERT INTO t (data) VALUES ('beta')"]) + .unwrap(); + let db = fixture.db_bytes().unwrap(); + let wal = fixture.wal_bytes().unwrap(); + + let header = parse_wal_header(&wal).unwrap(); + assert_eq!( + header.checkpoint_seq, 0, + "a fresh database's first WAL starts at checkpoint_seq 0" + ); + let page_size = usize::try_from(header.page_size).unwrap(); + let frame_size = usize::try_from(header.frame_size()).unwrap(); + assert!( + (wal.len() - 32).is_multiple_of(frame_size), + "fixture WAL is frame-aligned" + ); + + let mut synthetic = + SyntheticWal::new(header.page_size, header.big_endian_checksum(), header.salt); + for frame in wal[32..].chunks_exact(frame_size) { + let mut head = [0u8; 24]; + head.copy_from_slice(&frame[..24]); + let frame_header = FrameHeader::parse(&head); + let page = &frame[24..24 + page_size]; + synthetic = if frame_header.is_commit() { + synthetic.commit_frame(frame_header.page_no, page, frame_header.db_size) + } else { + synthetic.frame(frame_header.page_no, page) + }; + } + let rebuilt = synthetic.bytes(); + assert_eq!(rebuilt.len(), wal.len(), "rebuilt WAL has the same length"); + assert!( + rebuilt == wal, + "SyntheticWal reproduces the real WAL byte-for-byte (checksum chain matches SQLite)" + ); + + // SQLite recovers the synthetic WAL next to the db file image + let recover_dir = tempfile::tempdir().unwrap(); + let recover_db = recover_dir.path().join("synthetic.db"); + std::fs::write(&recover_db, &db).unwrap(); + std::fs::write(recover_dir.path().join("synthetic.db-wal"), &rebuilt).unwrap(); + let conn = rusqlite::Connection::open(&recover_db).unwrap(); + let rows = read_rows(&conn); + let data: Vec<&str> = rows.iter().map(|(_, data)| data.as_str()).collect(); + assert_eq!( + data, + vec!["init", "alpha", "beta"], + "SQLite recovers the expected rows from the synthetic WAL" + ); + } + + fn read_rows(conn: &rusqlite::Connection) -> Vec<(i64, String)> { + let mut stmt = conn.prepare("SELECT id, data FROM t ORDER BY id").unwrap(); + stmt.query_map([], |row| Ok((row.get(0)?, row.get(1)?))) + .unwrap() + .collect::, _>>() + .unwrap() + } + + // 12. Golden files: a committed real {db, wal, dump} triple + + /// Rows the golden fixture contains ('init' comes from `WalFixture::new`). + const GOLDEN_ROWS: &[(i64, &str)] = &[(1, "init"), (2, "alpha"), (3, "beta"), (4, "gamma")]; + + /// Render a human-readable, stable dump of every committed frame: + /// offset, `page_no`, `db_size`, and commit flag, plus header fields and + /// end-of-scan totals. + fn render_dump(wal: &[u8]) -> String { + let mut scanner = WalScanner::open(Cursor::new(wal.to_vec())) + .unwrap() + .unwrap(); + let header = *scanner.header(); + let mut out = String::new(); + writeln!( + out, + "magic {}", + if header.big_endian_checksum() { + "big_endian" + } else { + "little_endian" + } + ) + .unwrap(); + writeln!(out, "page_size {}", header.page_size).unwrap(); + writeln!(out, "checkpoint_seq {}", header.checkpoint_seq).unwrap(); + writeln!(out, "salt {:#010x} {:#010x}", header.salt.0, header.salt.1).unwrap(); + let frame_size = usize::try_from(header.frame_size()).unwrap(); + let mut offset = WAL_HEADER_SIZE; + let mut end_offset = WAL_HEADER_SIZE; + let mut commit_count = 0u64; + let mut db_size_pages = 0u32; + while let Some(chunk) = scanner.next_committed(u64::MAX).unwrap() { + for frame in chunk.bytes.chunks_exact(frame_size) { + let mut head = [0u8; 24]; + head.copy_from_slice(&frame[..24]); + let frame_header = FrameHeader::parse(&head); + writeln!( + out, + "frame offset={offset} page_no={} db_size={} commit={}", + frame_header.page_no, + frame_header.db_size, + frame_header.is_commit() + ) + .unwrap(); + offset += header.frame_size(); + } + end_offset = chunk.end_offset; + commit_count += chunk.commit_count; + db_size_pages = chunk.db_size_pages; + } + writeln!(out, "end_offset {end_offset}").unwrap(); + writeln!(out, "commit_count {commit_count}").unwrap(); + writeln!(out, "db_size_pages {db_size_pages}").unwrap(); + out + } + + #[test] + fn golden_wal_parses_to_expected_dump() { + let wal: &[u8] = include_bytes!(concat!( + env!("CARGO_MANIFEST_DIR"), + "/golden/wal_le_4096.wal" + )); + let db: &[u8] = include_bytes!(concat!( + env!("CARGO_MANIFEST_DIR"), + "/golden/wal_le_4096.db" + )); + let dump = include_str!(concat!( + env!("CARGO_MANIFEST_DIR"), + "/golden/wal_le_4096.dump.txt" + )); + + // The parser reproduces the committed dump exactly + assert_eq!(render_dump(wal), dump, "golden dump drifted"); + let header = parse_wal_header(wal).unwrap(); + assert_eq!(header.magic, WAL_MAGIC_LE, "golden WAL is little-endian"); + assert_eq!(header.page_size, 4096); + + // Cross-SQLite-version canary: whatever SQLite is linked today must + // recover the same rows from the committed byte images + let tmp = tempfile::tempdir().unwrap(); + let db_path = tmp.path().join("golden.db"); + std::fs::write(&db_path, db).unwrap(); + std::fs::write(tmp.path().join("golden.db-wal"), wal).unwrap(); + let conn = rusqlite::Connection::open(&db_path).unwrap(); + let rows = read_rows(&conn); + let expected: Vec<(i64, String)> = GOLDEN_ROWS + .iter() + .map(|&(id, data)| (id, data.to_owned())) + .collect(); + assert_eq!(rows, expected, "SQLite recovers the golden rows"); + } + + /// Regenerate the committed golden files. Run manually during + /// development when the golden triple must change: + /// + /// ```text + /// cargo nextest run -p bencher_replica --features plus,testing \ + /// --run-ignored ignored-only -E 'test(generate_golden_files)' + /// ``` + #[test] + #[ignore = "regenerates the committed golden files under golden/"] + fn generate_golden_files() { + let golden_dir = Utf8Path::new(env!("CARGO_MANIFEST_DIR")).join("golden"); + std::fs::create_dir_all(&golden_dir).unwrap(); + let tmp = tempfile::tempdir().unwrap(); + // WalFixture::new commits the schema and the 'init' row + let fixture = WalFixture::new(tempdir_path(&tmp), 4096).unwrap(); + fixture + .txn(&[ + "INSERT INTO t (data) VALUES ('alpha')", + "INSERT INTO t (data) VALUES ('beta')", + ]) + .unwrap(); + fixture + .txn(&["INSERT INTO t (data) VALUES ('gamma')"]) + .unwrap(); + let db = fixture.db_bytes().unwrap(); + let wal = fixture.wal_bytes().unwrap(); + let header = parse_wal_header(&wal).unwrap(); + assert_eq!( + header.magic, WAL_MAGIC_LE, + "golden files must be generated on a little-endian host" + ); + std::fs::write(golden_dir.join("wal_le_4096.db"), &db).unwrap(); + std::fs::write(golden_dir.join("wal_le_4096.wal"), &wal).unwrap(); + std::fs::write(golden_dir.join("wal_le_4096.dump.txt"), render_dump(&wal)).unwrap(); + } +} diff --git a/plus/bencher_replica/tests/checkpoint_section.rs b/plus/bencher_replica/tests/checkpoint_section.rs new file mode 100644 index 000000000..d531c8c44 --- /dev/null +++ b/plus/bencher_replica/tests/checkpoint_section.rs @@ -0,0 +1,448 @@ +#![cfg(all(feature = "plus", feature = "testing"))] +//! Checkpoint critical section suite: ship-before-checkpoint (I1), the +//! PASSIVE-under-BEGIN-IMMEDIATE trick that closes the sneak-in race, and +//! the prime directive that app writers are never blocked for O(database) +//! work (I5). +//! +//! All tests use the multi-thread runtime: the engine's `spawn_blocking` +//! rusqlite calls run concurrently with real second-connection writers. +//! +//! NOTE: `unused_crate_dependencies` cannot be handled with a crate-level +//! `#![expect]` here (see `tests/storage_contract.rs`); unused package +//! dependencies are referenced explicitly instead, as rustc recommends. + +use async_compression as _; +use aws_credential_types as _; +use aws_sdk_s3 as _; +use bytes as _; +use futures as _; +use hex as _; +use rand as _; +use serde as _; +use serde_json as _; +use sha2 as _; +use thiserror as _; +use uuid as _; +use zstd as _; +// Optional dependency enabled by the otel feature; unused by tests. +#[cfg(feature = "otel")] +use bencher_otel as _; + +/// A slimmer sibling of the `tests/sync_engine.rs` harness: real fixture, +/// plain local replica (no fault injection), deterministic clock. +#[cfg(test)] +pub(crate) mod harness { + use std::sync::Arc; + use std::sync::atomic::{AtomicI64, Ordering}; + + use bencher_json::system::config::{JsonReplication, ReplicationTarget}; + use bencher_json::{Clock, DateTime}; + use bencher_replica::testing::{WalFixture, assert_replica_equivalent}; + use bencher_replica::{ + EngineState, ReplicaConfig, ReplicaDb, RestoreOutcome, SyncEngine, restore_if_missing, + }; + use camino::Utf8Path; + + /// Page size for every fixture database in this suite. + pub(crate) const PAGE_SIZE: u32 = 4096; + /// 2026-07-10T14:59:00Z, the deterministic clock start. + const BASE_SECS: i64 = 1_783_695_540; + + pub(crate) fn dir_path(tmp: &tempfile::TempDir) -> &Utf8Path { + Utf8Path::from_path(tmp.path()).expect("tempdir path is UTF-8") + } + + pub(crate) fn logger() -> slog::Logger { + slog::Logger::root(slog::Discard, slog::o!()) + } + + pub(crate) struct Harness { + pub fixture: WalFixture, + pub engine: SyncEngine<()>, + pub config: ReplicaConfig, + pub _clock_secs: Arc, + pub _fixture_tmp: tempfile::TempDir, + pub _replica_tmp: tempfile::TempDir, + } + + impl Harness { + pub(crate) async fn new() -> Self { + let fixture_tmp = tempfile::tempdir().expect("fixture tempdir"); + let replica_tmp = tempfile::tempdir().expect("replica tempdir"); + let fixture = WalFixture::new(dir_path(&fixture_tmp), PAGE_SIZE).expect("fixture"); + let replica_root = dir_path(&replica_tmp).to_path_buf(); + let config = ReplicaConfig::try_from(JsonReplication { + target: ReplicationTarget::File { + path: replica_root.into_std_path_buf(), + }, + sync_interval_secs: None, + checkpoint_interval_secs: None, + min_checkpoint_pages: None, + snapshot_interval_secs: None, + snapshot_throttle_mib: None, + retention_generations: None, + verification_interval_secs: None, + shutdown_sync_timeout_secs: None, + }) + .expect("config"); + let clock_secs = Arc::new(AtomicI64::new(BASE_SECS)); + let secs = Arc::clone(&clock_secs); + let clock = Clock::Custom(Arc::new(move || { + DateTime::try_from(secs.load(Ordering::SeqCst)).expect("valid clock seconds") + })); + let db = ReplicaDb { + db_path: fixture.db_path(), + writer: Arc::new(tokio::sync::Mutex::new(())), + busy_timeout_ms: 5000, + }; + let engine = SyncEngine::new(logger(), config.clone(), db, clock, false) + .await + .expect("engine"); + let mut harness = Self { + fixture, + engine, + config, + _clock_secs: clock_secs, + _fixture_tmp: fixture_tmp, + _replica_tmp: replica_tmp, + }; + harness.ready().await; + harness + } + + /// Drive the bootstrap snapshot to Streaming and ship the backlog. + async fn ready(&mut self) { + for _ in 0..64 { + if self.engine.state() == EngineState::Streaming { + self.engine.sync_once().await.expect("backlog sync"); + return; + } + self.engine + .sync_once() + .await + .expect("sync_once during startup"); + } + panic!( + "engine never reached Streaming; state: {:?}", + self.engine.state() + ); + } + + /// Restore the replica into a scratch directory and assert logical + /// equivalence with the live source database. + pub(crate) async fn assert_restore_equivalent(&self) { + let target_tmp = tempfile::tempdir().expect("restore target tempdir"); + let target_db = dir_path(&target_tmp).join("restored.db"); + let outcome = restore_if_missing(&logger(), &self.config, &target_db) + .await + .expect("restore"); + assert!( + matches!(outcome, RestoreOutcome::Restored { .. }), + "expected Restored, got {outcome:?}" + ); + assert_replica_equivalent(&self.fixture.db_path(), &target_db); + } + } +} + +#[cfg(test)] +mod cases { + use std::io::Cursor; + use std::sync::Arc; + use std::sync::atomic::{AtomicBool, Ordering}; + use std::time::{Duration, Instant}; + + use bencher_replica::testing::{ProbeResult, WriteProbe}; + use bencher_replica::{CheckpointOutcome, ReplicaMeta, WalScanner}; + use pretty_assertions::assert_eq; + use tokio::sync::oneshot; + + use super::harness::Harness; + + /// Read the live WAL header salts of the fixture. + fn wal_salt(harness: &Harness) -> (u32, u32) { + let wal = harness.fixture.wal_bytes().expect("wal bytes"); + WalScanner::open(Cursor::new(wal)) + .expect("wal header") + .expect("wal is not empty") + .header() + .salt + } + + /// Every probe write must have succeeded; returns the longest observed + /// block time. + fn assert_all_writes_succeeded(results: Vec) -> Duration { + assert!(!results.is_empty(), "the probe recorded at least one write"); + let mut max_blocked = Duration::ZERO; + for probe in results { + probe + .result + .expect("no probe write may fail (zero busy errors)"); + max_blocked = max_blocked.max(probe.blocked); + } + max_blocked + } + + #[tokio::test(flavor = "multi_thread")] + async fn checkpoint_only_after_ship() { + let mut harness = Harness::new().await; + harness + .fixture + .txn(&["INSERT INTO t (data) VALUES ('unshipped')"]) + .expect("txn"); + let outcome = harness.engine.checkpoint_once().await.expect("checkpoint"); + assert_eq!( + outcome, + CheckpointOutcome::SkippedUnshipped, + "unshipped committed frames block the checkpoint (I1)" + ); + let shipped = harness.engine.ship_once().await.expect("ship"); + assert!(shipped >= 1, "the tail ships"); + let outcome = harness.engine.checkpoint_once().await.expect("checkpoint"); + assert_eq!( + outcome, + CheckpointOutcome::Completed, + "once shipped, the checkpoint completes" + ); + } + + #[tokio::test(flavor = "multi_thread")] + async fn checkpoint_completes_and_sets_meta_flag() { + let mut harness = Harness::new().await; + harness + .fixture + .txn(&["INSERT INTO t (data) VALUES ('to-checkpoint')"]) + .expect("txn"); + harness.engine.ship_once().await.expect("ship"); + let meta = ReplicaMeta::load(&harness.fixture.db_path()) + .expect("load meta") + .expect("meta written by ship"); + assert!( + !meta.epoch_shipped_through_checkpoint, + "the flag is clear before any checkpoint" + ); + + let outcome = harness.engine.checkpoint_once().await.expect("checkpoint"); + assert_eq!(outcome, CheckpointOutcome::Completed, "full backfill"); + let meta = ReplicaMeta::load(&harness.fixture.db_path()) + .expect("load meta") + .expect("meta written by checkpoint"); + assert!( + meta.epoch_shipped_through_checkpoint, + "Completed sets the epoch-shipped-through-checkpoint flag" + ); + + // Full backfill means the NEXT writer restarts the WAL: new salts. + let salt_before = wal_salt(&harness); + harness + .fixture + .txn(&["INSERT INTO t (data) VALUES ('restart')"]) + .expect("txn after checkpoint"); + let salt_after = wal_salt(&harness); + assert!( + salt_after != salt_before, + "the WAL restarted after a completed checkpoint ({salt_before:?} vs {salt_after:?})" + ); + } + + #[tokio::test(flavor = "multi_thread")] + async fn passive_checkpoint_never_blocks_writers() { + let mut harness = Harness::new().await; + let probe = WriteProbe::new(&harness.fixture.db_path(), Duration::from_secs(5)); + let (stop_tx, stop_rx) = oneshot::channel::<()>(); + let hammer = tokio::spawn(async move { + probe + .hammer(async move { + drop(stop_rx.await); + }) + .await + }); + + // 5+ full ship+checkpoint cycles while the hammer runs. Outcomes + // vary (the hammer commits between ship and checkpoint), but no + // probe write may ever fail. + for _cycle in 0..6 { + harness.engine.ship_once().await.expect("ship"); + let _outcome = harness.engine.checkpoint_once().await.expect("checkpoint"); + } + stop_tx.send(()).expect("stop the hammer"); + let results = hammer.await.expect("hammer task"); + let _max_blocked = assert_all_writes_succeeded(results); + } + + #[tokio::test(flavor = "multi_thread")] + async fn critical_section_bounded() { + let mut harness = Harness::new().await; + harness.fixture.txn_touching_pages(64).expect("seed pages"); + + // A tokio-side probe AND a raw std::thread stray writer hammer the + // database through the checkpoint cycles. + let probe = WriteProbe::new(&harness.fixture.db_path(), Duration::from_secs(5)); + let (stop_tx, stop_rx) = oneshot::channel::<()>(); + let hammer = tokio::spawn(async move { + probe + .hammer(async move { + drop(stop_rx.await); + }) + .await + }); + let stray_stop = Arc::new(AtomicBool::new(false)); + let stray_flag = Arc::clone(&stray_stop); + let stray_db_path = harness.fixture.db_path().to_string(); + let stray = std::thread::spawn(move || { + let conn = rusqlite::Connection::open(&stray_db_path).expect("stray conn"); + conn.busy_timeout(Duration::from_secs(5)).expect("busy"); + let _pages: i64 = conn + .query_row("PRAGMA wal_autocheckpoint = 0", [], |row| row.get(0)) + .expect("autocheckpoint off"); + conn.execute_batch( + "CREATE TABLE IF NOT EXISTS stray_writes(id INTEGER PRIMARY KEY, at TEXT)", + ) + .expect("stray table"); + while !stray_flag.load(Ordering::SeqCst) { + conn.execute("INSERT INTO stray_writes (at) VALUES ('stray')", []) + .expect("stray insert never times out"); + // A zero-gap loop would monopolize the write lock and + // starve the other writers (a plain two-writer SQLite + // contention problem, nothing to do with the checkpoint); + // real stray writers (stats/credit sweeps) are paced. + std::thread::sleep(Duration::from_millis(2)); + } + }); + + for _ in 0..6 { + harness.engine.ship_once().await.expect("ship"); + harness.engine.checkpoint_once().await.expect("checkpoint"); + } + stop_tx.send(()).expect("stop the hammer"); + stray_stop.store(true, Ordering::SeqCst); + stray.join().expect("stray writer thread"); + let results = hammer.await.expect("hammer task"); + let max_blocked = assert_all_writes_succeeded(results); + // Smoke bound with a generous margin: the write lock is only ever + // held for O(WAL-tail) work (I5), so no writer waits anywhere near + // its 5s busy budget. + assert!( + max_blocked < Duration::from_secs(2), + "longest probe block {max_blocked:?} must stay far below the busy timeout" + ); + } + + #[tokio::test(flavor = "multi_thread")] + async fn sneak_in_stray_commit_detected() { + let mut harness = Harness::new().await; + harness + .fixture + .txn(&["INSERT INTO t (data) VALUES ('shipped')"]) + .expect("txn"); + harness.engine.ship_once().await.expect("ship"); + + // The litestream race: a stray commit lands AFTER the ship pass and + // BEFORE the checkpoint. The locked section re-verifies the tail + // under BEGIN IMMEDIATE, so the frame cannot be lost. + let stray = harness.fixture.stray_conn().expect("stray conn"); + stray + .execute("INSERT INTO t (data) VALUES ('sneak-in')", []) + .expect("sneak-in insert"); + let outcome = harness.engine.checkpoint_once().await.expect("checkpoint"); + assert_eq!( + outcome, + CheckpointOutcome::SkippedUnshipped, + "the sneak-in commit defers the checkpoint, it is never backfilled unshipped" + ); + + let shipped = harness.engine.ship_once().await.expect("ship sneak-in"); + assert!(shipped >= 1, "the sneak-in commit ships"); + let outcome = harness.engine.checkpoint_once().await.expect("checkpoint"); + assert_eq!(outcome, CheckpointOutcome::Completed, "then checkpoints"); + harness.assert_restore_equivalent().await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn long_reader_degrades_to_partial() { + let mut harness = Harness::new().await; + // Pin a read snapshot (like the online backup's long read). + let reader = harness.fixture.stray_conn().expect("reader conn"); + reader.execute_batch("BEGIN").expect("begin read txn"); + let _count: i64 = reader + .query_row("SELECT COUNT(*) FROM t", [], |row| row.get(0)) + .expect("pin the read snapshot"); + + // Commits past the pinned mark, fully shipped. + harness + .fixture + .txn(&["INSERT INTO t (data) VALUES ('past-the-mark')"]) + .expect("txn"); + harness + .fixture + .txn(&["INSERT INTO t (data) VALUES ('and-another')"]) + .expect("txn"); + harness.engine.ship_once().await.expect("ship"); + + let probe = WriteProbe::new(&harness.fixture.db_path(), Duration::from_secs(5)); + let outcome = harness.engine.checkpoint_once().await.expect("checkpoint"); + assert_eq!( + outcome, + CheckpointOutcome::Partial, + "a long reader pins the backfill mark: PASSIVE degrades to Partial" + ); + let meta = ReplicaMeta::load(&harness.fixture.db_path()) + .expect("load meta") + .expect("meta present"); + assert!( + !meta.epoch_shipped_through_checkpoint, + "Partial must NOT set the checkpoint flag" + ); + // Writers are unaffected by the partial backfill. + let result = probe.write_once().await; + result.result.expect("probe write during the long read"); + + // Release the reader: the next checkpoint completes. The probe + // write above must ship first (I1). + reader.execute_batch("COMMIT").expect("release the reader"); + harness.engine.ship_once().await.expect("ship probe write"); + let outcome = harness.engine.checkpoint_once().await.expect("checkpoint"); + assert_eq!( + outcome, + CheckpointOutcome::Completed, + "full backfill once the reader releases its mark" + ); + } + + #[tokio::test(flavor = "multi_thread")] + async fn busy_stray_writer_yields_skipped_busy() { + let mut harness = Harness::new().await; + harness + .fixture + .txn(&["INSERT INTO t (data) VALUES ('busy')"]) + .expect("txn"); + harness.engine.ship_once().await.expect("ship"); + + // A stray writer holds the SQLite write lock across the checkpoint. + let stray = harness.fixture.stray_conn().expect("stray conn"); + stray + .execute_batch("BEGIN IMMEDIATE") + .expect("hold the write lock"); + let started = Instant::now(); + let outcome = harness.engine.checkpoint_once().await.expect("checkpoint"); + let elapsed = started.elapsed(); + assert_eq!( + outcome, + CheckpointOutcome::SkippedBusy, + "a held write lock skips the checkpoint, it does not escalate" + ); + // The lock connection's busy budget is 1s; generous smoke margin. + assert!( + elapsed < Duration::from_secs(4), + "SkippedBusy within the busy budget, took {elapsed:?}" + ); + + stray.execute_batch("ROLLBACK").expect("release the lock"); + let outcome = harness.engine.checkpoint_once().await.expect("checkpoint"); + assert_eq!( + outcome, + CheckpointOutcome::Completed, + "the next attempt succeeds once the stray writer releases" + ); + } +} diff --git a/plus/bencher_replica/tests/crash_recovery.rs b/plus/bencher_replica/tests/crash_recovery.rs new file mode 100644 index 000000000..7168fac91 --- /dev/null +++ b/plus/bencher_replica/tests/crash_recovery.rs @@ -0,0 +1,838 @@ +#![cfg(all(feature = "plus", feature = "testing"))] +//! Crash-recovery matrix: six kill points, each simulated deterministically. +//! +//! A "process crash" is the DROP of the step-driven engine at a precise step +//! boundary (optionally after an injected storage failure); recovery is a +//! rebuild via `SyncEngine::new_with_storage` over the same fixture and +//! replica directories with a fresh `FlakyStorage` (journal reset, no plan). +//! Every kill point shares one postcondition: the rebuilt engine quiesces, +//! ships one more transaction, a restore reproduces the source exactly, and +//! no epoch on the replica holds duplicate or overlapping segment ranges. +//! +//! NOTE: `unused_crate_dependencies` cannot be handled with a crate-level +//! `#![expect]` here (see `tests/storage_contract.rs`); unused package +//! dependencies are referenced explicitly instead, as rustc recommends. + +use async_compression as _; +use aws_credential_types as _; +use aws_sdk_s3 as _; +use bytes as _; +use futures as _; +use hex as _; +use rand as _; +use rusqlite as _; +use serde as _; +use serde_json as _; +use sha2 as _; +use thiserror as _; +use uuid as _; +use zstd as _; +// Optional dependency enabled by the otel feature; unused by tests. +#[cfg(feature = "otel")] +use bencher_otel as _; + +/// Shared fixtures: a scripted source database, a fault-injectable replica, +/// and an engine (behind an `Option`, so a test can crash it at any step +/// boundary) built over both with a deterministic clock. +#[cfg(test)] +pub(crate) mod harness { + use std::collections::BTreeMap; + use std::sync::Arc; + use std::sync::atomic::{AtomicI64, Ordering}; + + use bencher_json::system::config::{JsonReplication, ReplicationTarget}; + use bencher_json::{Clock, DateTime}; + use bencher_replica::testing::{ + FailurePlan, FlakyStorage, OpKind, OpOutcome, WalFixture, assert_replica_equivalent, + }; + use bencher_replica::{ + EngineState, GenerationId, LocalStorage, ReplicaConfig, ReplicaDb, ReplicaStorage, + RestoreOutcome, SyncEngine, restore_if_missing, + }; + use camino::{Utf8Path, Utf8PathBuf}; + + /// Page size for every fixture database in this suite. + const PAGE_SIZE: u32 = 4096; + /// 2026-07-10T14:59:00Z, the deterministic clock start. + const BASE_SECS: i64 = 1_783_695_540; + + pub(crate) fn dir_path(tmp: &tempfile::TempDir) -> &Utf8Path { + Utf8Path::from_path(tmp.path()).expect("tempdir path is UTF-8") + } + + pub(crate) fn logger() -> slog::Logger { + slog::Logger::root(slog::Discard, slog::o!()) + } + + fn clock_for(secs: &Arc) -> Clock { + let secs = Arc::clone(secs); + Clock::Custom(Arc::new(move || { + DateTime::try_from(secs.load(Ordering::SeqCst)).expect("valid clock seconds") + })) + } + + /// Flaky(Local) storage rooted at `root`, with an empty failure plan. + fn flaky_over(root: &Utf8Path) -> ReplicaStorage { + ReplicaStorage::Flaky(Box::new(FlakyStorage::new( + ReplicaStorage::Local(LocalStorage::new(root.to_path_buf())), + FailurePlan::new(), + ))) + } + + /// Parse the `[start, end)` byte range out of a segment object key. + fn segment_range(key: &str) -> (u64, u64) { + let (_, file) = key.rsplit_once('/').expect("segment key has a directory"); + let range = file.strip_suffix(".wal.zst").expect("segment key suffix"); + let (start, end) = range.split_once('-').expect("segment range separator"); + ( + start.parse().expect("segment start offset"), + end.parse().expect("segment end offset"), + ) + } + + pub(crate) struct Harness { + fixture: WalFixture, + /// `None` between [`Harness::crash`] and [`Harness::rebuild_engine`]. + engine: Option>, + config: ReplicaConfig, + db: ReplicaDb<()>, + clock_secs: Arc, + replica_root: Utf8PathBuf, + _fixture_tmp: tempfile::TempDir, + _replica_tmp: tempfile::TempDir, + } + + impl Harness { + pub(crate) async fn new() -> Self { + Self::with_config(|_| {}).await + } + + pub(crate) async fn with_config(overrides: F) -> Self + where + F: FnOnce(&mut JsonReplication), + { + let fixture_tmp = tempfile::tempdir().expect("fixture tempdir"); + let replica_tmp = tempfile::tempdir().expect("replica tempdir"); + let fixture = WalFixture::new(dir_path(&fixture_tmp), PAGE_SIZE).expect("fixture"); + let replica_root = dir_path(&replica_tmp).to_path_buf(); + let mut json = JsonReplication { + target: ReplicationTarget::File { + path: replica_root.clone().into_std_path_buf(), + }, + sync_interval_secs: None, + checkpoint_interval_secs: None, + min_checkpoint_pages: None, + snapshot_interval_secs: None, + snapshot_throttle_mib: None, + retention_generations: None, + verification_interval_secs: None, + shutdown_sync_timeout_secs: None, + }; + overrides(&mut json); + let config = ReplicaConfig::try_from(json).expect("config"); + let clock_secs = Arc::new(AtomicI64::new(BASE_SECS)); + let db = ReplicaDb { + db_path: fixture.db_path(), + writer: Arc::new(tokio::sync::Mutex::new(())), + busy_timeout_ms: 5000, + }; + let engine = SyncEngine::new_with_storage( + logger(), + config.clone(), + db.clone(), + clock_for(&clock_secs), + false, + flaky_over(&replica_root), + ) + .await + .expect("engine"); + Self { + fixture, + engine: Some(engine), + config, + db, + clock_secs, + replica_root, + _fixture_tmp: fixture_tmp, + _replica_tmp: replica_tmp, + } + } + + /// The scripted source database. + pub(crate) fn fixture(&self) -> &WalFixture { + &self.fixture + } + + /// The live engine; panics after a crash and before a rebuild. + pub(crate) fn engine(&self) -> &SyncEngine<()> { + self.engine.as_ref().expect("the engine is running") + } + + /// The live engine, mutably. + pub(crate) fn engine_mut(&mut self) -> &mut SyncEngine<()> { + self.engine.as_mut().expect("the engine is running") + } + + /// The kill point: drop the engine mid-flight. Its open checkpoint + /// connections, in-flight multipart upload, backoff state, and + /// in-memory position all vanish, exactly as in a killed process. + pub(crate) fn crash(&mut self) { + let engine = self.engine.take(); + assert!(engine.is_some(), "a crash requires a running engine"); + drop(engine); + } + + /// The recovery: rebuild the engine over the same fixture and + /// replica directories. The storage wrapper is brand new (empty + /// journal, no failure plan), so resume sees a healed replica. + pub(crate) async fn rebuild_engine(&mut self) { + assert!( + self.engine.is_none(), + "rebuild models a process restart: crash first" + ); + self.engine = Some( + SyncEngine::new_with_storage( + logger(), + self.config.clone(), + self.db.clone(), + clock_for(&self.clock_secs), + false, + flaky_over(&self.replica_root), + ) + .await + .expect("engine rebuild"), + ); + } + + /// The fault-injection wrapper the engine was built with. + pub(crate) fn flaky(&self) -> &FlakyStorage { + if let ReplicaStorage::Flaky(flaky) = self.engine().storage() { + flaky + } else { + panic!("harness engine always wraps Flaky storage") + } + } + + /// Advance the injected clock by `secs`. + pub(crate) fn advance(&self, secs: i64) { + self.clock_secs.fetch_add(secs, Ordering::SeqCst); + } + + /// Number of put operations attempted so far (any outcome). + pub(crate) fn put_attempts(&self) -> usize { + self.flaky() + .journal() + .iter() + .filter(|(op, _)| op.kind == OpKind::Put) + .count() + } + + /// Segment object keys successfully shipped so far, in journal order. + pub(crate) fn shipped_segment_keys(&self) -> Vec { + self.flaky() + .journal() + .iter() + .filter(|(op, outcome)| { + op.kind == OpKind::Put + && *outcome == OpOutcome::Ok + && op.key.ends_with(".wal.zst") + }) + .map(|(op, _)| op.key.clone()) + .collect() + } + + /// Segment object keys whose put was injected to fail, in journal + /// order (these never reached the backend). + pub(crate) fn injected_segment_keys(&self) -> Vec { + self.flaky() + .journal() + .iter() + .filter(|(op, outcome)| { + op.kind == OpKind::Put + && *outcome == OpOutcome::Injected + && op.key.ends_with(".wal.zst") + }) + .map(|(op, _)| op.key.clone()) + .collect() + } + + /// Drive `sync_once` until the engine is streaming (the fresh-replica + /// bootstrap snapshot has completed). + pub(crate) async fn until_streaming(&mut self) { + for _ in 0..64 { + if self.engine().state() == EngineState::Streaming { + return; + } + self.engine_mut() + .sync_once() + .await + .expect("sync_once during startup"); + } + panic!( + "engine never reached Streaming; state: {:?}", + self.engine().state() + ); + } + + /// Bootstrap plus one sync tick, so the initial WAL backlog is + /// shipped and the engine is quiescent. + pub(crate) async fn ready(&mut self) { + self.until_streaming().await; + self.engine_mut().sync_once().await.expect("backlog sync"); + } + + /// Drive `sync_once` until a fully quiet streaming tick: no shipped + /// segments, no snapshot work, no backoff, no error. This absorbs + /// whatever recovery work a rebuild scheduled (backlog re-ship or a + /// forced new-generation snapshot). + pub(crate) async fn drive_to_quiescence(&mut self) { + for _ in 0..128 { + let progress = self + .engine_mut() + .sync_once() + .await + .expect("quiescence sync"); + assert!( + progress.error.is_none(), + "recovery ticks over healed storage are clean: {progress:?}" + ); + if self.engine().state() == EngineState::Streaming + && progress.shipped_segments == 0 + && progress.snapshot.is_none() + && !progress.backing_off + { + return; + } + } + panic!("engine never quiesced; state: {:?}", self.engine().state()); + } + + /// Restore the replica into a scratch directory, assert logical + /// equivalence with the live source database, and return the + /// generation the restore was served from. + pub(crate) async fn assert_restore_equivalent(&self) -> GenerationId { + let target_tmp = tempfile::tempdir().expect("restore target tempdir"); + let target_db = dir_path(&target_tmp).join("restored.db"); + let outcome = restore_if_missing(&logger(), &self.config, &target_db) + .await + .expect("restore"); + let RestoreOutcome::Restored { generation, .. } = outcome else { + panic!("expected Restored, got {outcome:?}"); + }; + assert_replica_equivalent(&self.fixture.db_path(), &target_db); + generation + } + + /// List every shipped segment on the replica and assert that within + /// each (generation, epoch) directory the byte ranges are contiguous + /// from offset 0 with strictly increasing starts: no duplicate, + /// overlapping, or missing segment ranges survive a crash. + pub(crate) async fn assert_contiguous_segments(&self) { + let keys = self + .engine() + .storage() + .list("generations") + .await + .expect("list replica"); + let mut epochs: BTreeMap<&str, Vec<(u64, u64)>> = BTreeMap::new(); + for key in &keys { + if !key.ends_with(".wal.zst") { + continue; + } + let (dir, _file) = key.rsplit_once('/').expect("segment key has a directory"); + epochs.entry(dir).or_default().push(segment_range(key)); + } + assert!(!epochs.is_empty(), "at least one epoch shipped segments"); + for (dir, mut ranges) in epochs { + ranges.sort_unstable(); + let mut expected_start = 0u64; + for (start, end) in ranges { + assert_eq!( + start, expected_start, + "{dir}: strictly increasing contiguous segment starts \ + (no duplicate, overlapping, or missing ranges)" + ); + assert!(end > start, "{dir}: segment [{start}, {end}) is non-empty"); + expected_start = end; + } + } + } + + /// The common crash postcondition: quiesce, prove the pipeline still + /// ships (one more transaction plus a sync), and prove the replica + /// restores to an exact copy with a duplicate-free segment layout. + /// Returns the generation the restore was served from. + pub(crate) async fn assert_crash_postcondition(&mut self) -> GenerationId { + self.drive_to_quiescence().await; + self.fixture + .txn(&["INSERT INTO t (data) VALUES ('post-crash probe')"]) + .expect("post-crash txn"); + let progress = self + .engine_mut() + .sync_once() + .await + .expect("post-crash sync"); + assert!( + progress.error.is_none(), + "the post-crash sync is clean: {progress:?}" + ); + let generation = self.assert_restore_equivalent().await; + self.assert_contiguous_segments().await; + generation + } + } +} + +#[cfg(test)] +mod cases { + use bencher_replica::testing::{CheckpointMode, FailurePlan, OpKind}; + use bencher_replica::{CheckpointOutcome, EngineState, SnapshotStatus, StorageError}; + use pretty_assertions::{assert_eq, assert_ne}; + + use super::harness::Harness; + + /// Drive `sync_once` until the scheduled snapshot cycle reports + /// `Finished`. + async fn drive_snapshot_cycle(harness: &mut Harness) { + for _ in 0..64 { + let progress = harness + .engine_mut() + .sync_once() + .await + .expect("snapshot cycle sync"); + assert!( + progress.error.is_none(), + "snapshot cycle ticks are clean: {progress:?}" + ); + if progress.snapshot == Some(SnapshotStatus::Finished) { + return; + } + } + panic!("the snapshot cycle never finished"); + } + + /// Count replica objects whose key ends with `suffix`. + async fn object_count(harness: &Harness, suffix: &str) -> usize { + harness + .engine() + .storage() + .list("generations") + .await + .expect("list replica") + .iter() + .filter(|key| key.ends_with(suffix)) + .count() + } + + /// `K_a`: crash with committed frames that never shipped. The frames + /// survive in the local WAL; the rebuilt engine resumes at the replica + /// tip and ships them on its very first sync. + #[tokio::test] + async fn crash_before_any_ship() { + let mut harness = Harness::new().await; + harness.ready().await; + let shipped_offset = harness.engine().position().expect("position").offset; + harness + .fixture() + .txn(&["INSERT INTO t (data) VALUES ('committed, never shipped')"]) + .expect("txn"); + harness.crash(); + harness.rebuild_engine().await; + + let position = harness.engine().position().expect("resumed position"); + assert_eq!( + position.offset, shipped_offset, + "resume from the replica LIST lands at the last shipped offset" + ); + let progress = harness.engine_mut().sync_once().await.expect("first sync"); + assert!( + progress.shipped_segments >= 1, + "the crash-surviving frames ship on the FIRST sync after rebuild: {progress:?}" + ); + harness.assert_crash_postcondition().await; + } + + /// `K_b`: crash mid-segment-upload. The put is injected to fail (backoff + /// armed), the process dies mid-backoff, and the rebuilt engine re-ships + /// the segment under the exact same key: idempotent (epoch, start, end) + /// naming means a retried upload can never duplicate a range. + #[tokio::test] + async fn crash_mid_segment_upload() { + let mut harness = Harness::new().await; + harness.ready().await; + harness + .fixture() + .txn(&["INSERT INTO t (data) VALUES ('upload will fail')"]) + .expect("txn"); + harness + .flaky() + .set_plan(FailurePlan::new().fail_nth(OpKind::Put, 1)); + + let progress = harness + .engine_mut() + .sync_once() + .await + .expect("failing sync"); + assert!( + progress.error.is_some(), + "the injected put surfaces as a tick error: {progress:?}" + ); + let progress = harness.engine_mut().sync_once().await.expect("gated sync"); + assert!( + progress.backing_off, + "the backoff is armed when the process dies: {progress:?}" + ); + let injected = harness.injected_segment_keys(); + assert_eq!( + injected.len(), + 1, + "exactly one injected segment put: {injected:?}" + ); + let failed_key = injected.first().expect("the failed segment key").clone(); + let error = harness + .engine() + .storage() + .get(&failed_key) + .await + .expect_err("the failed segment never reached the replica"); + assert!( + matches!(error, StorageError::NotFound { .. }), + "expected NotFound for the failed key, got {error}" + ); + + harness.crash(); + harness.rebuild_engine().await; + let progress = harness + .engine_mut() + .sync_once() + .await + .expect("re-ship sync"); + assert!( + progress.shipped_segments >= 1, + "the lost segment re-ships after rebuild: {progress:?}" + ); + assert_eq!( + harness.shipped_segment_keys(), + vec![failed_key], + "the re-shipped segment reuses the exact key that failed mid-upload \ + (same epoch, start, end)" + ); + harness.assert_crash_postcondition().await; + } + + /// `K_c`: crash after a successful ship but before any checkpoint. The + /// rebuilt engine re-derives its position from the replica LIST alone: + /// a no-write sync issues zero puts (nothing re-ships, nothing is lost). + #[tokio::test] + async fn crash_after_ship_before_checkpoint() { + let mut harness = Harness::new().await; + harness.ready().await; + harness + .fixture() + .txn(&["INSERT INTO t (data) VALUES ('shipped, not checkpointed')"]) + .expect("txn"); + harness.engine_mut().sync_once().await.expect("sync"); + let position = harness.engine().position().cloned().expect("position"); + + harness.crash(); + harness.rebuild_engine().await; + assert_eq!( + harness.engine().position(), + Some(&position), + "the position is re-derived from the replica LIST alone \ + (exact offset and checksum)" + ); + let progress = harness + .engine_mut() + .sync_once() + .await + .expect("no-write sync"); + assert_eq!(progress.shipped_segments, 0, "nothing new to ship"); + assert_eq!( + harness.put_attempts(), + 0, + "a no-write sync after rebuild issues zero storage puts: \ + no segment is ever shipped twice" + ); + harness.assert_crash_postcondition().await; + } + + /// `K_d`: crash mid-snapshot-upload. The half-built generation has no + /// snapshot.json, so it is invisible to resume and restore; 25 hours + /// later the scheduled snapshot cycle prunes the stale incomplete + /// generation prefix. + #[tokio::test] + async fn crash_mid_snapshot_upload() { + let mut harness = Harness::with_config(|json| { + json.snapshot_throttle_mib = Some(1); + }) + .await; + harness.ready().await; + // Grow the database FILE (not just the WAL) so the snapshot copy + // spans multiple steps: big transaction, ship, full backfill. + harness + .fixture() + .txn_touching_pages(800) + .expect("big transaction"); + while harness + .engine_mut() + .ship_once() + .await + .expect("ship big txn") + > 0 + {} + let outcome = harness + .engine_mut() + .checkpoint_once() + .await + .expect("checkpoint"); + assert_eq!( + outcome, + CheckpointOutcome::Completed, + "backfill the db file" + ); + let bootstrap = harness.engine().generation().cloned().expect("generation"); + + // Start a new-generation snapshot and kill it mid-upload: after two + // steps the multipart upload into the new generation is open but + // nothing is finalized. + harness.engine_mut().trigger_snapshot(); + for step in 0..2 { + let status = harness + .engine_mut() + .snapshot_step() + .await + .expect("snapshot step"); + assert_eq!( + status, + SnapshotStatus::InProgress, + "step {step} leaves the snapshot in flight" + ); + } + harness.crash(); + harness.rebuild_engine().await; + + // The half generation is a directory without snapshot.json: + // invisible to both resume and restore. + let dirs = harness + .engine() + .storage() + .list_dirs("generations") + .await + .expect("list generations"); + assert_eq!( + dirs.len(), + 2, + "the bootstrap generation plus the crashed half generation: {dirs:?}" + ); + let half = dirs + .iter() + .find(|dir| dir.as_str() != bootstrap.as_str()) + .expect("the crashed generation directory") + .clone(); + let error = harness + .engine() + .storage() + .get(&format!("generations/{half}/snapshot.json")) + .await + .expect_err("the crashed generation must have no snapshot.json"); + assert!( + matches!(error, StorageError::NotFound { .. }), + "snapshot.json is absent, not another error: {error}" + ); + + let restored = harness.assert_crash_postcondition().await; + assert_ne!( + restored.as_str(), + half.as_str(), + "the crashed generation is never the restore source" + ); + harness + .engine() + .storage() + .get(&format!("generations/{}/snapshot.json", restored.as_str())) + .await + .expect("the restore source generation has snapshot.json"); + + // 25 hours later the scheduled snapshot cycle runs; its prune reaps + // the stale (> 24h) incomplete generation prefix. + harness.advance(25 * 60 * 60); + drive_snapshot_cycle(&mut harness).await; + let dirs = harness + .engine() + .storage() + .list_dirs("generations") + .await + .expect("list after prune"); + assert!( + !dirs.contains(&half), + "the stale incomplete generation was pruned: {dirs:?}" + ); + // The new lineage still restores to an exact copy. + harness.drive_to_quiescence().await; + harness.assert_restore_equivalent().await; + harness.assert_contiguous_segments().await; + } + + /// `K_e`: crash after a completed checkpoint, after the WAL restarted + /// (new salts), but before any new-epoch segment shipped. The advisory + /// meta proves the old epoch was fully shipped through the checkpoint, + /// so resume continues as epoch+1 in the SAME generation: no + /// re-snapshot, no re-shipped old segments. + #[tokio::test] + async fn crash_after_checkpoint_before_new_epoch_ship() { + let mut harness = Harness::new().await; + harness.ready().await; + harness + .fixture() + .txn(&["INSERT INTO t (data) VALUES ('seals the epoch')"]) + .expect("txn"); + harness.engine_mut().sync_once().await.expect("sync"); + let sealed = harness.engine().position().cloned().expect("position"); + let outcome = harness + .engine_mut() + .checkpoint_once() + .await + .expect("checkpoint"); + assert_eq!( + outcome, + CheckpointOutcome::Completed, + "the full backfill seals the epoch" + ); + let generation = harness.engine().generation().cloned().expect("generation"); + + // The WAL is fully backfilled, so this write restarts it (new + // salts); the process dies before the engine ever sees the new + // cycle. + harness + .fixture() + .txn(&["INSERT INTO t (data) VALUES ('first write of the new cycle')"]) + .expect("txn after checkpoint"); + harness.crash(); + harness.rebuild_engine().await; + + assert_eq!( + harness.engine().state(), + EngineState::Streaming, + "the meta-verified resume streams immediately: no re-snapshot" + ); + assert_eq!( + harness.engine().generation(), + Some(&generation), + "the generation is unchanged across the crash" + ); + let position = harness.engine().position().expect("resumed position"); + assert_eq!( + position.epoch, + sealed.epoch + 1, + "resume continues as the next epoch" + ); + assert_eq!(position.offset, 0, "the new epoch starts at offset 0"); + let salt = position.salt; + + let progress = harness + .engine_mut() + .sync_once() + .await + .expect("ship the new cycle"); + assert!( + progress.shipped_segments >= 1, + "the new cycle ships after rebuild: {progress:?}" + ); + let epoch_dir = format!("{:010}-{:08x}{:08x}", sealed.epoch + 1, salt.0, salt.1); + let keys = harness.shipped_segment_keys(); + assert!( + !keys.is_empty() && keys.iter().all(|key| key.contains(&epoch_dir)), + "every post-rebuild put lands in the NEW epoch directory {epoch_dir}; \ + no old segment re-ships: {keys:?}" + ); + assert_eq!( + object_count(&harness, "snapshot.json").await, + 1, + "no new generation was created" + ); + harness.assert_crash_postcondition().await; + } + + /// `K_f`: crash, then (with the engine down) a stray writer restarts the + /// WAL, a stray TRUNCATE checkpoint backfills those never-shipped + /// frames, and another write restarts the WAL again, wiping them. The + /// rebuilt engine cannot prove continuity across the buried cycle and + /// must force a NEW generation; the fresh snapshot recaptures the + /// buried transaction so nothing is lost from the replica. + /// + /// This test originally exposed a real engine gap: the meta-verified + /// epoch+1 resume said nothing about how many WAL restarts happened + /// while the process was down, so the buried cycle's commits (living + /// only in the db file, backfilled by the stray TRUNCATE) silently never + /// shipped. Resume now also proves salt CONTINUITY: `SQLite` increments + /// WAL `salt1` by exactly one per restart, so the epoch+1 path requires + /// `header.salt1 == meta.salt1 + 1` and diverges to a new generation on + /// any larger jump. + #[tokio::test] + async fn crash_with_unshipped_frames_lost_to_wal_reset() { + let mut harness = Harness::new().await; + harness.ready().await; + harness + .fixture() + .txn(&["INSERT INTO t (data) VALUES ('shipped and checkpointed')"]) + .expect("txn"); + harness.engine_mut().sync_once().await.expect("sync"); + let outcome = harness + .engine_mut() + .checkpoint_once() + .await + .expect("checkpoint"); + assert_eq!( + outcome, + CheckpointOutcome::Completed, + "the epoch is sealed by a full backfill" + ); + let old_generation = harness.engine().generation().cloned().expect("generation"); + harness.crash(); + + // With the engine down: a stray transaction restarts the fully + // backfilled WAL (legitimately: everything so far was shipped). Its + // frames are never shipped. CREATE TABLE makes the eventual loss + // visible: no later transaction rewrites the schema this touches, + // so nothing downstream can mask a silently dropped cycle. + let stray = harness.fixture().stray_conn().expect("stray conn"); + stray + .execute_batch( + "BEGIN IMMEDIATE; + CREATE TABLE buried (id INTEGER PRIMARY KEY, data TEXT); + INSERT INTO buried (data) VALUES ('never shipped'); + COMMIT;", + ) + .expect("stray txn"); + // A stray TRUNCATE checkpoint backfills the unshipped frames into + // the db file, and the next write restarts the WAL a second time, + // wiping every trace of them from the local WAL. + harness + .fixture() + .checkpoint(CheckpointMode::Truncate) + .expect("stray TRUNCATE"); + harness + .fixture() + .txn(&["INSERT INTO t (data) VALUES ('second restart cycle')"]) + .expect("txn after truncate"); + + harness.rebuild_engine().await; + assert_eq!( + harness.engine().state(), + EngineState::PendingSnapshot, + "resume cannot prove continuity across the buried WAL cycle; \ + a new generation must be forced, never a silent epoch+1 resume" + ); + let restored = harness.assert_crash_postcondition().await; + assert_ne!( + restored, old_generation, + "a second generation appears and becomes the restore source" + ); + assert_eq!( + object_count(&harness, "snapshot.json").await, + 2, + "both the old and the replacement generation are complete" + ); + } +} diff --git a/plus/bencher_replica/tests/equivalence.rs b/plus/bencher_replica/tests/equivalence.rs new file mode 100644 index 000000000..b0b1451b0 --- /dev/null +++ b/plus/bencher_replica/tests/equivalence.rs @@ -0,0 +1,347 @@ +#![cfg(all(feature = "plus", feature = "testing"))] +//! End-to-end equivalence suite: seeded workloads drive the source database +//! and the step-driven engine together, then prove the governing invariant +//! `restore(replica) == source` at every commit boundary that matters (after +//! every completed checkpoint and after the final drain). +//! +//! Scripts come from `bencher_replica::testing::generate_workload`; the +//! smoke test's script is hand-written so a broken op type fails readably +//! before the seeded scripts do. On any failure the harness prints the seed, +//! the failing op index, and the full script. +//! +//! NOTE: `unused_crate_dependencies` cannot be handled with a crate-level +//! `#![expect]` here (see `tests/storage_contract.rs`); unused package +//! dependencies are referenced explicitly instead, as rustc recommends. + +use async_compression as _; +use aws_credential_types as _; +use aws_sdk_s3 as _; +use bytes as _; +use futures as _; +use hex as _; +use rand as _; +use rusqlite as _; +use serde as _; +use serde_json as _; +use sha2 as _; +use thiserror as _; +use uuid as _; +use zstd as _; +// Optional dependency enabled by the otel feature; unused by tests. +#[cfg(feature = "otel")] +use bencher_otel as _; + +/// Shared fixtures: a scripted source database, a workload environment over +/// a local replica, and restore-equivalence assertions. +#[cfg(test)] +pub(crate) mod harness { + use std::sync::Arc; + use std::sync::atomic::{AtomicI64, Ordering}; + + use bencher_json::system::config::{JsonReplication, ReplicationTarget}; + use bencher_json::{Clock, DateTime}; + use bencher_replica::testing::{ + FailurePlan, FlakyStorage, WalFixture, WorkloadEnv, WorkloadError, WorkloadOp, + assert_replica_equivalent, + }; + use bencher_replica::{ + EngineState, LocalStorage, ReplicaConfig, ReplicaDb, ReplicaStorage, RestoreOutcome, + SyncEngine, restore_if_missing, + }; + use camino::{Utf8Path, Utf8PathBuf}; + + /// Page size for every fixture database in this suite. + pub(crate) const PAGE_SIZE: u32 = 4096; + /// 2026-07-10T14:59:00Z, the deterministic clock start. + pub(crate) const BASE_SECS: i64 = 1_783_695_540; + + pub(crate) fn dir_path(tmp: &tempfile::TempDir) -> &Utf8Path { + Utf8Path::from_path(tmp.path()).expect("tempdir path is UTF-8") + } + + pub(crate) fn logger() -> slog::Logger { + slog::Logger::root(slog::Discard, slog::o!()) + } + + pub(crate) fn clock_for(secs: &Arc) -> Clock { + let secs = Arc::clone(secs); + Clock::Custom(Arc::new(move || { + DateTime::try_from(secs.load(Ordering::SeqCst)).expect("valid clock seconds") + })) + } + + /// Flaky(Local) storage rooted at `root`, with an empty failure plan. + pub(crate) fn flaky_over(root: &Utf8Path) -> ReplicaStorage { + ReplicaStorage::Flaky(Box::new(FlakyStorage::new( + ReplicaStorage::Local(LocalStorage::new(root.to_path_buf())), + FailurePlan::new(), + ))) + } + + /// The equivalence rig: the source database and everything a + /// `WorkloadRunner` needs, minus the engine itself (the runner owns it, + /// so `RestartReplicator` can rebuild it in place). + pub(crate) struct Rig { + pub fixture: WalFixture, + pub config: ReplicaConfig, + pub db: ReplicaDb<()>, + pub clock_secs: Arc, + pub replica_root: Utf8PathBuf, + pub _fixture_tmp: tempfile::TempDir, + pub _replica_tmp: tempfile::TempDir, + } + + impl Rig { + pub(crate) fn new() -> Self { + let fixture_tmp = tempfile::tempdir().expect("fixture tempdir"); + let replica_tmp = tempfile::tempdir().expect("replica tempdir"); + let fixture = WalFixture::new(dir_path(&fixture_tmp), PAGE_SIZE).expect("fixture"); + let replica_root = dir_path(&replica_tmp).to_path_buf(); + let json = JsonReplication { + target: ReplicationTarget::File { + path: replica_root.clone().into_std_path_buf(), + }, + sync_interval_secs: None, + checkpoint_interval_secs: None, + min_checkpoint_pages: None, + snapshot_interval_secs: None, + snapshot_throttle_mib: None, + retention_generations: None, + verification_interval_secs: None, + shutdown_sync_timeout_secs: None, + }; + let config = ReplicaConfig::try_from(json).expect("config"); + let clock_secs = Arc::new(AtomicI64::new(BASE_SECS)); + let db = ReplicaDb { + db_path: fixture.db_path(), + writer: Arc::new(tokio::sync::Mutex::new(())), + busy_timeout_ms: 5000, + }; + Self { + fixture, + config, + db, + clock_secs, + replica_root, + _fixture_tmp: fixture_tmp, + _replica_tmp: replica_tmp, + } + } + + /// The pieces a `WorkloadRunner` needs to drive (and rebuild) the + /// engine over this rig's directories. + pub(crate) fn env(&self) -> WorkloadEnv<()> { + WorkloadEnv { + log: logger(), + config: self.config.clone(), + db: self.db.clone(), + clock: clock_for(&self.clock_secs), + replica_root: self.replica_root.clone(), + } + } + + /// Build an engine over the replica and drive the fresh-replica + /// bootstrap snapshot to completion, ending quiescent in Streaming. + pub(crate) async fn ready_engine(&self) -> SyncEngine<()> { + let mut engine = SyncEngine::new_with_storage( + logger(), + self.config.clone(), + self.db.clone(), + clock_for(&self.clock_secs), + false, + flaky_over(&self.replica_root), + ) + .await + .expect("engine"); + for _ in 0..256 { + if engine.state() == EngineState::Streaming { + break; + } + engine.sync_once().await.expect("bootstrap sync"); + } + assert!( + engine.state() == EngineState::Streaming, + "engine never reached Streaming during bootstrap; state: {:?}", + engine.state() + ); + engine.sync_once().await.expect("backlog sync"); + engine + } + + /// Restore the replica into a scratch directory and assert logical + /// equivalence with the live source database. + pub(crate) async fn assert_restore_equivalent(&self) { + let target_tmp = tempfile::tempdir().expect("restore target tempdir"); + let target_db = dir_path(&target_tmp).join("restored.db"); + let outcome = restore_if_missing(&logger(), &self.config, &target_db) + .await + .expect("restore"); + assert!( + matches!(outcome, RestoreOutcome::Restored { .. }), + "expected Restored, got {outcome:?}" + ); + assert_replica_equivalent(&self.fixture.db_path(), &target_db); + } + } + + /// Unwrap a workload result; on failure print the error (which carries + /// the seed, the failing op index, and the op) plus the full script. + pub(crate) fn expect_workload(result: Result, script: &[WorkloadOp]) -> T { + match result { + Ok(value) => value, + Err(error) => panic!( + "{error}\nfull script ({} ops):\n{}", + script.len(), + format_script(script) + ), + } + } + + fn format_script(script: &[WorkloadOp]) -> String { + use std::fmt::Write as _; + script + .iter() + .enumerate() + .fold(String::new(), |mut out, (index, op)| { + writeln!(out, " {index:>4}: {op}").expect("write to string"); + out + }) + } +} + +#[cfg(test)] +mod cases { + use bencher_replica::CheckpointOutcome; + use bencher_replica::testing::{WorkloadOp, WorkloadRunner, generate_workload, run_workload}; + use pretty_assertions::assert_eq; + + use super::harness::{Rig, expect_workload}; + + /// Ops per generated script (the soak uses more). + const SCRIPT_LEN: usize = 200; + + /// A hand-written script exercising every `WorkloadOp` variant once (in + /// a sensible order), so a broken op type fails here, readably, before + /// the seeded scripts bury it. + #[tokio::test] + async fn equivalence_all_op_types_smoke() { + let script = vec![ + WorkloadOp::CreateTable { table: 0 }, + WorkloadOp::Insert { table: 0, rows: 8 }, + WorkloadOp::CreateIndex { table: 0 }, + WorkloadOp::BlobWrite { + table: 0, + len: 64 * 1024, + }, + WorkloadOp::Sync, + WorkloadOp::Update { table: 0, rows: 4 }, + WorkloadOp::BigTxn { statements: 12 }, + WorkloadOp::UserVersionBump, + WorkloadOp::Sync, + WorkloadOp::Checkpoint, + WorkloadOp::StrayWrite, + WorkloadOp::Insert { table: 1, rows: 5 }, + WorkloadOp::Delete { table: 0, rows: 2 }, + WorkloadOp::Vacuum, + WorkloadOp::Sync, + WorkloadOp::Snapshot, + WorkloadOp::Insert { table: 2, rows: 3 }, + WorkloadOp::RestartReplicator, + WorkloadOp::Insert { table: 0, rows: 2 }, + WorkloadOp::DropTable { table: 1 }, + WorkloadOp::Sync, + WorkloadOp::Checkpoint, + ]; + // The hand-written script must exercise every variant. + let mut kinds: Vec<&'static str> = script.iter().map(WorkloadOp::kind_name).collect(); + kinds.sort_unstable(); + kinds.dedup(); + assert_eq!( + kinds.len(), + WorkloadOp::VARIANT_COUNT, + "smoke script covers every WorkloadOp variant" + ); + + let rig = Rig::new(); + let engine = rig.ready_engine().await; + let seed = 424_242; + let mut runner = expect_workload( + WorkloadRunner::new(seed, script.clone(), rig.env(), engine), + &script, + ); + expect_workload(runner.run().await, &script); + expect_workload(runner.drain().await, &script); + rig.assert_restore_equivalent().await; + } + + /// Generate, run, drain, restore, compare: one full equivalence pass. + async fn run_seed(seed: u64, len: usize) { + let rig = Rig::new(); + let engine = rig.ready_engine().await; + let script = generate_workload(seed, len); + let mut runner = expect_workload( + run_workload(seed, script.clone(), rig.env(), engine).await, + &script, + ); + expect_workload(runner.drain().await, &script); + rig.assert_restore_equivalent().await; + } + + macro_rules! equivalence_seed { + ($($name:ident => $seed:literal,)+) => {$( + #[tokio::test] + async fn $name() { + // Boxed: the workload future is large (clippy::large_futures). + Box::pin(run_seed($seed, SCRIPT_LEN)).await; + } + )+}; + } + + equivalence_seed! { + equivalence_seed_0 => 0, + equivalence_seed_1 => 1, + equivalence_seed_2 => 2, + equivalence_seed_3 => 3, + equivalence_seed_4 => 4, + equivalence_seed_5 => 5, + equivalence_seed_6 => 6, + equivalence_seed_7 => 7, + } + + /// After every checkpoint that reports `Completed`, the replica must + /// already restore to the exact committed source state: `Completed` + /// means every WAL frame was shipped (I1) and backfilled, and the runner + /// is between ops, so the live committed state IS the shipped state. + #[tokio::test] + async fn equivalence_restore_at_every_checkpoint() { + let seed = 42; + let rig = Rig::new(); + let engine = rig.ready_engine().await; + let script = generate_workload(seed, 120); + let mut runner = expect_workload( + WorkloadRunner::new(seed, script.clone(), rig.env(), engine), + &script, + ); + let mut completed = 0u32; + while let Some(applied) = expect_workload(runner.step().await, &script) { + if applied.checkpoint == Some(CheckpointOutcome::Completed) { + completed += 1; + rig.assert_restore_equivalent().await; + } + } + assert!( + completed >= 1, + "seed {seed} must complete at least one checkpoint, got {completed}" + ); + expect_workload(runner.drain().await, &script); + rig.assert_restore_equivalent().await; + } + + /// Long soak: thousands of ops across multiple generations and + /// replicator restarts. + #[tokio::test] + #[ignore = "soak: run manually or nightly"] + async fn equivalence_soak() { + Box::pin(run_seed(1337, 3000)).await; + } +} diff --git a/plus/bencher_replica/tests/fault_matrix.rs b/plus/bencher_replica/tests/fault_matrix.rs new file mode 100644 index 000000000..7b560acfc --- /dev/null +++ b/plus/bencher_replica/tests/fault_matrix.rs @@ -0,0 +1,1049 @@ +#![cfg(all(feature = "plus", feature = "testing"))] +//! Fault-injection matrix: eight storage-failure scenarios (F1-F8). +//! +//! Every scenario follows arrange -> act -> assert: arrange a scripted +//! [`bencher_replica::testing::FailurePlan`] on the Flaky(Local) replica +//! storage, act by driving the step-driven engine (`sync_once`, +//! `checkpoint_once`, `snapshot_step`, `prune_once`) with all scheduling +//! under an injected `Clock` (never a sleep), and assert on the operation +//! journal, the on-disk replica state, and end-to-end restore equivalence +//! where meaningful. +//! +//! NOTE: `unused_crate_dependencies` cannot be handled with a crate-level +//! `#![expect]` here (see `tests/storage_contract.rs`); unused package +//! dependencies are referenced explicitly instead, as rustc recommends. + +use async_compression as _; +use aws_credential_types as _; +use aws_sdk_s3 as _; +use bytes as _; +use futures as _; +use hex as _; +use rand as _; +use rusqlite as _; +use serde as _; +use serde_json as _; +use sha2 as _; +use thiserror as _; +use uuid as _; +use zstd as _; +// Optional dependency enabled by the otel feature; unused by tests. +#[cfg(feature = "otel")] +use bencher_otel as _; + +/// Shared fixtures (private copy of the `tests/sync_engine.rs` harness): a +/// scripted source database, a fault-injectable replica, and an engine built +/// over both with a deterministic clock. +#[cfg(test)] +pub(crate) mod harness { + use std::sync::Arc; + use std::sync::atomic::{AtomicI64, Ordering}; + + use bencher_json::system::config::{JsonReplication, ReplicationTarget}; + use bencher_json::{Clock, DateTime}; + use bencher_replica::testing::{ + FailurePlan, FlakyStorage, OpKind, OpOutcome, WalFixture, assert_replica_equivalent, + }; + use bencher_replica::{ + EngineState, GenerationId, LocalStorage, ReplicaConfig, ReplicaDb, ReplicaStorage, + RestoreOutcome, SnapshotStatus, SyncEngine, SyncError, restore_if_missing, + }; + use camino::{Utf8Path, Utf8PathBuf}; + + /// Page size for every fixture database in this suite. + pub(crate) const PAGE_SIZE: u32 = 4096; + /// 2026-07-10T14:59:00Z, the deterministic clock start. + pub(crate) const BASE_SECS: i64 = 1_783_695_540; + + pub(crate) fn dir_path(tmp: &tempfile::TempDir) -> &Utf8Path { + Utf8Path::from_path(tmp.path()).expect("tempdir path is UTF-8") + } + + pub(crate) fn logger() -> slog::Logger { + slog::Logger::root(slog::Discard, slog::o!()) + } + + pub(crate) fn clock_for(secs: &Arc) -> Clock { + let secs = Arc::clone(secs); + Clock::Custom(Arc::new(move || { + DateTime::try_from(secs.load(Ordering::SeqCst)).expect("valid clock seconds") + })) + } + + /// Flaky(Local) storage rooted at `root` with the given failure plan. + pub(crate) fn flaky_with(root: &Utf8Path, plan: FailurePlan) -> ReplicaStorage { + ReplicaStorage::Flaky(Box::new(FlakyStorage::new( + ReplicaStorage::Local(LocalStorage::new(root.to_path_buf())), + plan, + ))) + } + + /// Parse the `[start, end)` byte range out of a segment object key. + pub(crate) fn segment_range(key: &str) -> (u64, u64) { + let (_, file) = key.rsplit_once('/').expect("segment key has a directory"); + let range = file.strip_suffix(".wal.zst").expect("segment key suffix"); + let (start, end) = range.split_once('-').expect("segment range separator"); + ( + start.parse().expect("segment start offset"), + end.parse().expect("segment end offset"), + ) + } + + /// Every file under `root`, as sorted root-relative path strings. + /// Used to prove a failed resume uploaded nothing at all. + pub(crate) fn replica_file_set(root: &Utf8Path) -> Vec { + fn walk(dir: &std::path::Path, out: &mut Vec) { + let Ok(entries) = std::fs::read_dir(dir) else { + return; + }; + for entry in entries.flatten() { + let path = entry.path(); + if path.is_dir() { + walk(&path, out); + } else { + out.push(path); + } + } + } + let mut paths = Vec::new(); + walk(root.as_std_path(), &mut paths); + let mut files: Vec = paths + .iter() + .map(|path| { + path.strip_prefix(root.as_std_path()) + .unwrap_or(path) + .to_string_lossy() + .into_owned() + }) + .collect(); + files.sort(); + files + } + + pub(crate) struct Harness { + pub fixture: WalFixture, + pub engine: SyncEngine<()>, + pub config: ReplicaConfig, + pub db: ReplicaDb<()>, + pub clock_secs: Arc, + pub replica_root: Utf8PathBuf, + pub _fixture_tmp: tempfile::TempDir, + pub _replica_tmp: tempfile::TempDir, + } + + impl Harness { + pub(crate) async fn new() -> Self { + Self::with_config(|_| {}).await + } + + pub(crate) async fn with_config(overrides: F) -> Self + where + F: FnOnce(&mut JsonReplication), + { + let fixture_tmp = tempfile::tempdir().expect("fixture tempdir"); + let replica_tmp = tempfile::tempdir().expect("replica tempdir"); + let fixture = WalFixture::new(dir_path(&fixture_tmp), PAGE_SIZE).expect("fixture"); + let replica_root = dir_path(&replica_tmp).to_path_buf(); + let mut json = JsonReplication { + target: ReplicationTarget::File { + path: replica_root.clone().into_std_path_buf(), + }, + sync_interval_secs: None, + checkpoint_interval_secs: None, + min_checkpoint_pages: None, + snapshot_interval_secs: None, + snapshot_throttle_mib: None, + retention_generations: None, + verification_interval_secs: None, + shutdown_sync_timeout_secs: None, + }; + overrides(&mut json); + let config = ReplicaConfig::try_from(json).expect("config"); + let clock_secs = Arc::new(AtomicI64::new(BASE_SECS)); + let db = ReplicaDb { + db_path: fixture.db_path(), + writer: Arc::new(tokio::sync::Mutex::new(())), + busy_timeout_ms: 5000, + }; + let engine = SyncEngine::new_with_storage( + logger(), + config.clone(), + db.clone(), + clock_for(&clock_secs), + false, + flaky_with(&replica_root, FailurePlan::new()), + ) + .await + .expect("engine"); + Self { + fixture, + engine, + config, + db, + clock_secs, + replica_root, + _fixture_tmp: fixture_tmp, + _replica_tmp: replica_tmp, + } + } + + /// The fault-injection wrapper the engine was built with. + pub(crate) fn flaky(&self) -> &FlakyStorage { + if let ReplicaStorage::Flaky(flaky) = self.engine.storage() { + flaky + } else { + panic!("harness engine always wraps Flaky storage") + } + } + + /// Advance the injected clock by `secs`. + pub(crate) fn advance(&self, secs: i64) { + self.clock_secs.fetch_add(secs, Ordering::SeqCst); + } + + /// Current injected clock second. + pub(crate) fn now_secs(&self) -> i64 { + self.clock_secs.load(Ordering::SeqCst) + } + + /// Number of put operations attempted so far (any outcome). + pub(crate) fn put_attempts(&self) -> usize { + self.flaky() + .journal() + .iter() + .filter(|(op, _)| op.kind == OpKind::Put) + .count() + } + + /// Segment object keys successfully shipped so far, in journal order. + pub(crate) fn shipped_segment_keys(&self) -> Vec { + self.flaky() + .journal() + .iter() + .filter(|(op, outcome)| { + op.kind == OpKind::Put + && *outcome == OpOutcome::Ok + && op.key.ends_with(".wal.zst") + }) + .map(|(op, _)| op.key.clone()) + .collect() + } + + /// Drive `sync_once` until the engine is streaming (the fresh-replica + /// bootstrap snapshot has completed). + pub(crate) async fn until_streaming(&mut self) { + for _ in 0..64 { + if self.engine.state() == EngineState::Streaming { + return; + } + self.engine + .sync_once() + .await + .expect("sync_once during startup"); + } + panic!( + "engine never reached Streaming; state: {:?}", + self.engine.state() + ); + } + + /// Bootstrap plus one sync tick, so the initial WAL backlog is + /// shipped and the engine is quiescent. + pub(crate) async fn ready(&mut self) { + self.until_streaming().await; + self.engine.sync_once().await.expect("backlog sync"); + } + + /// Replace the engine (dropping the old one on success) with a + /// freshly resumed engine over the same replica, built on storage + /// with the given failure plan (the plan is live BEFORE + /// construction). On `Err` the previous engine is untouched; the + /// flaky journal starts empty either way. + pub(crate) async fn rebuild_engine_with_plan( + &mut self, + plan: FailurePlan, + ) -> Result<(), SyncError> { + self.engine = SyncEngine::new_with_storage( + logger(), + self.config.clone(), + self.db.clone(), + clock_for(&self.clock_secs), + false, + flaky_with(&self.replica_root, plan), + ) + .await?; + Ok(()) + } + + /// Trigger a new-generation snapshot and drive it to completion. + pub(crate) async fn drive_snapshot(&mut self) { + self.engine.trigger_snapshot(); + for _ in 0..1000 { + let status = self.engine.snapshot_step().await.expect("snapshot step"); + if status == SnapshotStatus::Finished { + return; + } + } + panic!("snapshot never finished"); + } + + /// On-disk directory of a generation under the local replica root. + pub(crate) fn generation_dir(&self, generation: &GenerationId) -> Utf8PathBuf { + self.replica_root + .join("generations") + .join(generation.as_str()) + } + + /// Restore the replica into a scratch directory, assert logical + /// equivalence with the live source database, and return the + /// generation the restore picked. + pub(crate) async fn assert_restore_equivalent(&self) -> GenerationId { + let target_tmp = tempfile::tempdir().expect("restore target tempdir"); + let target_db = dir_path(&target_tmp).join("restored.db"); + let outcome = restore_if_missing(&logger(), &self.config, &target_db) + .await + .expect("restore"); + let RestoreOutcome::Restored { generation, .. } = outcome else { + panic!("expected Restored, got {outcome:?}"); + }; + assert_replica_equivalent(&self.fixture.db_path(), &target_db); + generation + } + } +} + +#[cfg(test)] +mod cases { + use bencher_replica::testing::{FailurePlan, OpKind, OpOutcome}; + use bencher_replica::{ + CheckpointOutcome, EngineState, RestoreError, SnapshotStatus, StorageError, SyncError, + restore_if_missing, + }; + use pretty_assertions::assert_eq; + + use super::harness::{Harness, dir_path, logger, replica_file_set, segment_range}; + + /// Count objects on the replica whose key ends with `suffix`. + async fn object_count(harness: &Harness, suffix: &str) -> usize { + harness + .engine + .storage() + .list("generations") + .await + .expect("list replica") + .iter() + .filter(|key| key.ends_with(suffix)) + .count() + } + + /// F1: a single injected put arms the backoff; once the clock passes the + /// delay the SAME segment ships exactly once (one injected attempt, one + /// successful put for the same key) and the replica restores equivalent. + #[tokio::test] + async fn put_fails_once_then_heals() { + let mut harness = Harness::new().await; + harness.ready().await; + harness + .flaky() + .set_plan(FailurePlan::new().fail_nth(OpKind::Put, 1)); + harness + .fixture + .txn(&["INSERT INTO t (data) VALUES ('flaky-once')"]) + .expect("txn"); + + let progress = harness.engine.sync_once().await.expect("failing sync"); + assert!( + progress.error.is_some(), + "the injected put arms a backoff: {progress:?}" + ); + assert_eq!( + progress.shipped_segments, 0, + "nothing ships on the failed tick" + ); + + let progress = harness.engine.sync_once().await.expect("gated sync"); + assert!( + progress.backing_off, + "the next tick is gated until the clock advances: {progress:?}" + ); + + harness.advance(1); + let progress = harness.engine.sync_once().await.expect("retry sync"); + assert!(progress.error.is_none(), "the retry succeeds: {progress:?}"); + assert_eq!( + progress.shipped_segments, 1, + "the pending segment ships exactly once" + ); + + let key = harness + .shipped_segment_keys() + .pop() + .expect("a shipped segment"); + let outcomes: Vec = harness + .flaky() + .journal() + .into_iter() + .filter(|(op, _)| op.kind == OpKind::Put && op.key == key) + .map(|(_, outcome)| outcome) + .collect(); + assert_eq!( + outcomes, + vec![OpOutcome::Injected, OpOutcome::Ok], + "the same key sees exactly one injected attempt then one successful ship" + ); + harness.assert_restore_equivalent().await; + } + + /// F2: under a persistent put outage the spacing between real put + /// attempts follows the exact capped doubling sequence 1, 2, 4, ..., + /// 256, 300, 300; healing recovers, and a success resets the backoff so + /// the next failure starts back at 1s. + #[tokio::test] + async fn backoff_caps_at_max() { + let mut harness = Harness::new().await; + harness.ready().await; + harness + .flaky() + .set_plan(FailurePlan::new().fail_all(OpKind::Put)); + harness + .fixture + .txn(&["INSERT INTO t (data) VALUES ('cap')"]) + .expect("txn"); + + // Step the clock one second at a time; a tick either attempts a put + // (journal grows) or is backoff-gated. Record the clock second of + // every real attempt. + let mut attempt_secs = Vec::new(); + let mut seen = harness.put_attempts(); + for _ in 0..2000 { + if attempt_secs.len() == 12 { + break; + } + let progress = harness.engine.sync_once().await.expect("outage sync"); + let attempts = harness.put_attempts(); + if attempts > seen { + seen = attempts; + assert!( + !progress.backing_off, + "an attempt tick is never backoff-gated: {progress:?}" + ); + assert!( + progress.error.is_some(), + "every attempt during the outage fails: {progress:?}" + ); + attempt_secs.push(harness.now_secs()); + } else { + assert!( + progress.backing_off, + "a no-attempt tick during the outage is backoff-gated: {progress:?}" + ); + harness.advance(1); + } + } + assert_eq!(attempt_secs.len(), 12, "twelve attempts were observed"); + let deltas: Vec = attempt_secs + .iter() + .zip(attempt_secs.iter().skip(1)) + .map(|(previous, next)| next - previous) + .collect(); + assert_eq!( + deltas, + vec![1, 2, 4, 8, 16, 32, 64, 128, 256, 300, 300], + "attempt spacing doubles from 1s and caps at 300s" + ); + + // Heal and wait out the final (capped) delay: the backlog ships. + harness.flaky().heal(); + harness.advance(301); + let progress = harness.engine.sync_once().await.expect("healed sync"); + assert!( + progress.error.is_none(), + "the healed tick succeeds: {progress:?}" + ); + assert!( + progress.shipped_segments >= 1, + "the backlog ships after healing: {progress:?}" + ); + + // The success reset the backoff: a fresh failure starts at 1s again, + // not at the 300s cap. + harness + .flaky() + .set_plan(FailurePlan::new().fail_nth(OpKind::Put, 1)); + harness + .fixture + .txn(&["INSERT INTO t (data) VALUES ('reset')"]) + .expect("txn"); + let progress = harness.engine.sync_once().await.expect("failing sync"); + assert!( + progress.error.is_some(), + "the induced failure arms a fresh backoff: {progress:?}" + ); + let progress = harness.engine.sync_once().await.expect("gated sync"); + assert!( + progress.backing_off, + "gated within the fresh 1s delay: {progress:?}" + ); + harness.advance(1); + let progress = harness.engine.sync_once().await.expect("retry sync"); + assert!( + progress.error.is_none() && progress.shipped_segments >= 1, + "one second later the retry ships: the backoff restarted at 1s, not 300s: {progress:?}" + ); + harness.assert_restore_equivalent().await; + } + + /// F3: during an outage the WAL is the buffer (it grows every round and + /// no checkpoint may complete while unshipped frames exist); after + /// healing the backlog ships strictly in order, the checkpoint + /// completes, and the replica restores equivalent. + #[tokio::test] + async fn outage_wal_grows_then_catchup() { + let mut harness = Harness::new().await; + harness.ready().await; + harness + .flaky() + .set_plan(FailurePlan::new().fail_all(OpKind::Put)); + + let mut wal_len = harness.fixture.wal_bytes().expect("wal").len(); + for round in 0..5u32 { + // Past the backoff cap: every round makes a real attempt. + harness.advance(400); + let statement = format!("INSERT INTO t (data) VALUES ('outage-{round}')"); + harness.fixture.txn(&[statement.as_str()]).expect("txn"); + let grown = harness.fixture.wal_bytes().expect("wal").len(); + assert!( + grown > wal_len, + "round {round}: the WAL buffers the outage, growing monotonically ({wal_len} -> {grown})" + ); + wal_len = grown; + let progress = harness.engine.sync_once().await.expect("outage sync"); + assert!( + progress.error.is_some(), + "round {round}: the ship attempt fails: {progress:?}" + ); + let outcome = harness.engine.checkpoint_once().await.expect("checkpoint"); + assert_eq!( + outcome, + CheckpointOutcome::SkippedUnshipped, + "round {round}: no checkpoint completes while unshipped frames exist (I1)" + ); + } + + // Heal, clear the backoff, and drive to quiescence. + harness.flaky().heal(); + harness.advance(400); + let mut quiescent = false; + for _ in 0..100 { + let progress = harness.engine.sync_once().await.expect("catch-up sync"); + assert!( + progress.error.is_none(), + "healed ticks succeed: {progress:?}" + ); + if progress.shipped_segments == 0 { + quiescent = true; + break; + } + } + assert!(quiescent, "the backlog fully ships after healing"); + + // Every successful ship, bootstrap included, landed strictly in + // (epoch, start) order and contiguously. + let keys = harness.shipped_segment_keys(); + assert!(!keys.is_empty(), "segments were shipped: {keys:?}"); + for (previous, next) in keys.iter().zip(keys.iter().skip(1)) { + assert!( + previous < next, + "segments ship in strictly increasing key order: {keys:?}" + ); + } + let mut expected_start = None; + for key in &keys { + let (start, end) = segment_range(key); + if let Some(expected) = expected_start { + assert_eq!(start, expected, "segments are contiguous: {keys:?}"); + } + expected_start = Some(end); + } + + let outcome = harness.engine.checkpoint_once().await.expect("checkpoint"); + assert_eq!( + outcome, + CheckpointOutcome::Completed, + "with everything shipped the checkpoint completes" + ); + harness.assert_restore_equivalent().await; + } + + /// F4: a directory-listing failure at resume (the first replica read the + /// engine makes) surfaces as a retryable error, creates no generation, and + /// uploads nothing (an unreachable replica must never read as an empty + /// one); a healed retry resumes the EXISTING generation at the exact + /// position. + #[tokio::test] + async fn list_error_at_resume_no_new_generation() { + let mut harness = Harness::new().await; + harness.ready().await; + harness + .fixture + .txn(&["INSERT INTO t (data) VALUES ('resume me')"]) + .expect("txn"); + harness.engine.sync_once().await.expect("sync"); + let generation = harness.engine.generation().cloned().expect("generation"); + let position = harness.engine.position().cloned().expect("position"); + + // "Process crash": the old engine is dropped on rebuild; the failing + // construction below never touches it. + let files_before = replica_file_set(&harness.replica_root); + let error = harness + .rebuild_engine_with_plan(FailurePlan::new().fail_all(OpKind::ListDirs)) + .await + .expect_err("resume against an unreachable replica must fail"); + assert!( + error.is_retryable(), + "an unreachable replica at boot is retryable, never fatal: {error}" + ); + assert!( + matches!( + error, + SyncError::Storage(StorageError::Injected { + op: "list_dirs", + .. + }) + ), + "the injected directory-listing error surfaces unchanged: {error}" + ); + assert_eq!( + replica_file_set(&harness.replica_root), + files_before, + "a failed resume creates no generation and uploads nothing" + ); + + // Heal (same replica root) and retry: the existing generation is + // resumed, nothing re-ships, and no new generation dir appears. + harness + .rebuild_engine_with_plan(FailurePlan::new()) + .await + .expect("healed resume"); + assert_eq!( + harness.engine.state(), + EngineState::Streaming, + "the healed resume streams immediately (no snapshot)" + ); + assert_eq!( + harness.engine.generation(), + Some(&generation), + "the EXISTING generation is resumed" + ); + assert_eq!( + harness.engine.position(), + Some(&position), + "the exact position is recovered from the replica LIST" + ); + let progress = harness.engine.sync_once().await.expect("quiet sync"); + assert_eq!(progress.shipped_segments, 0, "nothing re-ships"); + assert_eq!(harness.put_attempts(), 0, "no puts after a clean resume"); + let generation_dirs = + std::fs::read_dir(harness.replica_root.join("generations").as_std_path()) + .expect("generations dir") + .count(); + assert_eq!(generation_dirs, 1, "exactly one generation dir on disk"); + harness.assert_restore_equivalent().await; + } + + /// F5: orphaned junk on the replica (a crashed `.partial-` upload and a + /// bogus half-object under the WAL prefix) is invisible to shipping and + /// restore, is never deleted by shipping, and is reaped only when its + /// whole generation is pruned. + #[tokio::test] + async fn orphan_partial_upload_ignored() { + let mut harness = Harness::with_config(|json| { + json.retention_generations = Some(1); + }) + .await; + harness.ready().await; + let old_generation = harness.engine.generation().cloned().expect("generation"); + let old_dir = harness.generation_dir(&old_generation); + + // Plant the orphans by hand, as a crashed uploader would leave them. + let partial_orphan = old_dir.join("snapshot.db.zst.partial-garbage"); + std::fs::write(partial_orphan.as_std_path(), b"crashed multipart part") + .expect("plant partial orphan"); + let wal_dir = old_dir.join("wal"); + std::fs::create_dir_all(wal_dir.as_std_path()).expect("wal dir"); + let half_object = wal_dir.join("0000000000-halfobject"); + std::fs::write(half_object.as_std_path(), b"half an object").expect("plant half object"); + + // Shipping proceeds unbothered. + harness + .fixture + .txn(&["INSERT INTO t (data) VALUES ('after orphans')"]) + .expect("txn"); + let progress = harness.engine.sync_once().await.expect("sync"); + assert!( + progress.error.is_none() && progress.shipped_segments >= 1, + "shipping ignores the orphans: {progress:?}" + ); + + // Restore ignores them too (list filters the partial infix; the + // unparseable key is skipped with a warning). + let restored = harness.assert_restore_equivalent().await; + assert_eq!( + restored, old_generation, + "restore uses the real generation despite the orphans" + ); + assert!( + partial_orphan.as_std_path().exists(), + "shipping never deletes the partial orphan" + ); + assert!( + half_object.as_std_path().exists(), + "shipping never deletes the half object" + ); + + // A new generation under retention 1 prunes the old generation + // wholesale, orphans included. + harness.drive_snapshot().await; + let new_generation = harness.engine.generation().cloned().expect("generation"); + assert!( + new_generation != old_generation, + "the snapshot created a new generation" + ); + assert!( + !old_dir.as_std_path().exists(), + "pruning removes the whole old generation directory" + ); + assert!( + !partial_orphan.as_std_path().exists(), + "pruning reaps the partial orphan with its generation" + ); + assert!( + !half_object.as_std_path().exists(), + "pruning reaps the half object with its generation" + ); + + harness.engine.sync_once().await.expect("backlog sync"); + let restored = harness.assert_restore_equivalent().await; + assert_eq!(restored, new_generation, "restore picks the new generation"); + } + + /// F6: a corrupt (truncated) snapshot body makes the startup restore + /// fail HARD, leaving no file at the target path and no `.restore` / + /// `.restore-wal` / `.restore-shm` scratch leftovers behind. + #[tokio::test] + async fn get_fails_mid_restore_no_half_db() { + let mut harness = Harness::new().await; + harness.ready().await; + harness + .fixture + .txn(&["INSERT INTO t (data) VALUES ('tamper target')"]) + .expect("txn"); + harness.engine.sync_once().await.expect("sync"); + let generation = harness.engine.generation().cloned().expect("generation"); + + // Corrupt the snapshot body in place: truncate it to half. The + // restore fails on the broken zstd stream (or, failing that, on the + // sha256 verification). + let snapshot_path = harness.generation_dir(&generation).join("snapshot.db.zst"); + let body = std::fs::read(snapshot_path.as_std_path()).expect("snapshot body"); + assert!(body.len() >= 2, "the snapshot body is non-trivial"); + std::fs::write( + snapshot_path.as_std_path(), + body.get(..body.len() >> 1).expect("half the body"), + ) + .expect("truncate snapshot"); + + let target_tmp = tempfile::tempdir().expect("target tempdir"); + let target_db = dir_path(&target_tmp).join("restored.db"); + let error = restore_if_missing(&logger(), &harness.config, &target_db) + .await + .expect_err("a truncated snapshot must fail the restore"); + assert!( + matches!( + error, + RestoreError::SnapshotDownload { .. } | RestoreError::SnapshotChecksum { .. } + ), + "the failure names the snapshot, not a generic IO error: {error}" + ); + + assert!( + !target_db.as_std_path().exists(), + "no half-restored database at the target path" + ); + let leftovers: Vec = std::fs::read_dir(dir_path(&target_tmp).as_std_path()) + .expect("read target dir") + .flatten() + .map(|entry| entry.file_name().to_string_lossy().into_owned()) + .collect(); + assert_eq!( + leftovers, + Vec::::new(), + "a failed restore leaves no scratch files behind" + ); + } + + /// F7: delete failures during pruning are logged and swallowed (the + /// snapshot completes and sync keeps working); after healing, the next + /// prune deletes the old generation and restore picks the newest. + #[tokio::test] + async fn delete_fails_during_prune() { + let mut harness = Harness::with_config(|json| { + json.retention_generations = Some(1); + }) + .await; + harness.ready().await; + let old_generation = harness.engine.generation().cloned().expect("generation"); + harness + .fixture + .txn(&["INSERT INTO t (data) VALUES ('pre-prune')"]) + .expect("txn"); + harness.engine.sync_once().await.expect("sync"); + + harness + .flaky() + .set_plan(FailurePlan::new().fail_all(OpKind::DeletePrefix)); + // The snapshot completes even though its finalize-time prune fails. + harness.drive_snapshot().await; + let new_generation = harness.engine.generation().cloned().expect("generation"); + assert!( + new_generation != old_generation, + "the snapshot created a new generation" + ); + let injected_deletes = harness + .flaky() + .journal() + .iter() + .filter(|(op, outcome)| { + op.kind == OpKind::DeletePrefix && *outcome == OpOutcome::Injected + }) + .count(); + assert!( + injected_deletes >= 1, + "the prune attempted a delete_prefix and it was injected" + ); + assert!( + harness + .generation_dir(&old_generation) + .as_std_path() + .exists(), + "the old generation survives the failed prune" + ); + // Marker-first prune: the snapshot.json delete precedes the injected + // delete_prefix, so the surviving old generation is markerless + // (invisible to resume/restore), never a marker without a body. + assert!( + !harness + .generation_dir(&old_generation) + .join("snapshot.json") + .as_std_path() + .exists(), + "the failed prune already removed the old generation's marker" + ); + + // Sync keeps working through the prune failure. + assert_eq!( + harness.engine.state(), + EngineState::Streaming, + "the engine keeps streaming" + ); + harness + .fixture + .txn(&["INSERT INTO t (data) VALUES ('post-prune-failure')"]) + .expect("txn"); + let progress = harness + .engine + .sync_once() + .await + .expect("sync keeps working"); + assert!( + progress.error.is_none() && progress.shipped_segments >= 1, + "prune failures never poison the sync loop: {progress:?}" + ); + + // Heal and prune again. The marker-first failed prune left the old + // generation markerless, so it is classified incomplete + // (indistinguishable from a still-uploading snapshot) and is reaped + // as stale-incomplete rather than immediately: advance past the 24h + // stale cutoff so the healed prune reaps it. + harness.flaky().heal(); + harness.advance(25 * 60 * 60); + harness.engine.prune_once().await.expect("healed prune"); + assert!( + !harness + .generation_dir(&old_generation) + .as_std_path() + .exists(), + "the old generation is reaped once deletes heal and it goes stale" + ); + assert!( + harness + .generation_dir(&new_generation) + .as_std_path() + .exists(), + "the current generation is never pruned" + ); + let restored = harness.assert_restore_equivalent().await; + assert_eq!( + restored, new_generation, + "restore picks the newest generation" + ); + } + + /// F8: an injected multipart write aborts the snapshot; the aborted + /// generation stays invisible (no snapshot.json, no visible body), the + /// engine retriggers on its own, and after healing a fresh snapshot + /// completes and restore picks it. + #[tokio::test] + async fn multipart_write_fails_mid_snapshot() { + // 1 MiB copy budget per step forces multiple multipart writes. + let mut harness = Harness::with_config(|json| { + json.snapshot_throttle_mib = Some(1); + }) + .await; + harness.ready().await; + // Grow the database FILE (not just the WAL) so the copy spans + // several parts: big transaction, ship, full checkpoint. + harness + .fixture + .txn_touching_pages(800) + .expect("big transaction"); + harness.engine.ship_once().await.expect("ship big txn"); + let outcome = harness.engine.checkpoint_once().await.expect("checkpoint"); + assert_eq!( + outcome, + CheckpointOutcome::Completed, + "backfill the db file" + ); + assert_eq!( + object_count(&harness, "snapshot.json").await, + 1, + "one bootstrap generation before the fault" + ); + + harness + .flaky() + .set_plan(FailurePlan::new().fail_nth(OpKind::MultipartWrite, 2)); + harness.engine.trigger_snapshot(); + let mut surfaced = None; + for _ in 0..64 { + match harness.engine.snapshot_step().await { + Ok(SnapshotStatus::InProgress) => {}, + Ok(SnapshotStatus::Finished) => { + panic!("the snapshot must not finish past an injected multipart write") + }, + Err(error) => { + surfaced = Some(error); + break; + }, + } + } + let error = surfaced.expect("the injected multipart write must surface"); + assert!( + matches!( + error, + SyncError::Storage(StorageError::Injected { + op: "multipart_write", + .. + }) + ), + "the snapshot aborts on the injected multipart write: {error}" + ); + assert_eq!( + object_count(&harness, "snapshot.json").await, + 1, + "no snapshot.json for the aborted generation: it stays invisible" + ); + assert_eq!( + object_count(&harness, "snapshot.db.zst").await, + 1, + "the aborted body upload never became visible" + ); + + // The abort scheduled a retrigger: the next tick starts over. + harness.flaky().heal(); + harness.advance(5); + harness.engine.sync_once().await.expect("retrigger tick"); + assert_eq!( + harness.engine.state(), + EngineState::Snapshotting, + "the aborted snapshot retriggers after healing" + ); + for _ in 0..64 { + if object_count(&harness, "snapshot.json").await == 2 { + break; + } + harness.engine.sync_once().await.expect("recovery sync"); + } + assert_eq!( + object_count(&harness, "snapshot.json").await, + 2, + "the retried snapshot completes" + ); + let generation = harness.engine.generation().cloned().expect("generation"); + harness.engine.sync_once().await.expect("backlog sync"); + let restored = harness.assert_restore_equivalent().await; + assert_eq!( + restored, generation, + "restore picks the completed generation" + ); + } + + /// F9: a segment put lands server-side but reports failure (a lost 200). + /// Seeing failure, the engine retries and re-ships the SAME segment; the + /// key is put twice (applied-but-reported-failed, then a clean overwrite), + /// and the idempotent re-ship leaves the replica correct and restorable. + /// Exercises retry idempotency under ambiguous success. + #[tokio::test] + async fn put_reports_failure_after_applying_reships() { + let mut harness = Harness::new().await; + harness.ready().await; + // The next segment put reaches the backend but returns an error. + harness.flaky().set_plan( + FailurePlan::new() + .fail_matching(Some(OpKind::Put), ".wal.zst", 1) + .after(), + ); + harness + .fixture + .txn(&["INSERT INTO t (data) VALUES ('lost-200')"]) + .expect("txn"); + + let progress = harness.engine.sync_once().await.expect("ambiguous sync"); + assert!( + progress.error.is_some(), + "the lost-200 arms a backoff even though the segment reached the backend: {progress:?}" + ); + assert_eq!( + progress.shipped_segments, 0, + "nothing counts as shipped on the failed tick" + ); + + harness.advance(1); + let progress = harness.engine.sync_once().await.expect("retry sync"); + assert!(progress.error.is_none(), "the retry succeeds: {progress:?}"); + assert_eq!( + progress.shipped_segments, 1, + "the pending segment re-ships exactly once" + ); + + // The same key was put twice: applied-but-reported-failed, then a + // clean re-ship. Put is idempotent, so the replica ends correct. + let key = harness + .shipped_segment_keys() + .pop() + .expect("a shipped segment"); + let outcomes: Vec = harness + .flaky() + .journal() + .into_iter() + .filter(|(op, _)| op.kind == OpKind::Put && op.key == key) + .map(|(_, outcome)| outcome) + .collect(); + assert_eq!( + outcomes, + vec![OpOutcome::InjectedAfter, OpOutcome::Ok], + "the key is put twice: applied-but-reported-failed, then a clean re-ship" + ); + harness.assert_restore_equivalent().await; + } +} diff --git a/plus/bencher_replica/tests/live_s3.rs b/plus/bencher_replica/tests/live_s3.rs new file mode 100644 index 000000000..8ee222bba --- /dev/null +++ b/plus/bencher_replica/tests/live_s3.rs @@ -0,0 +1,321 @@ +#![cfg(all(feature = "plus", feature = "testing"))] +//! Live-S3 leg of the storage contract suite (`#[ignore]`-gated). +//! +//! Runs the SAME case table as `tests/storage_contract.rs` (included below +//! as a module) against a real S3-compatible bucket, following the +//! connectivity-probe pattern of `plus/bencher_oci_storage/tests/conformance.rs`: +//! when the environment is not configured, every test skips cleanly. +//! +//! # Running +//! +//! ```sh +//! export BENCHER_REPLICA_TEST_S3_BUCKET=my-test-bucket +//! export BENCHER_REPLICA_TEST_S3_ACCESS_KEY_ID=... +//! export BENCHER_REPLICA_TEST_S3_SECRET_ACCESS_KEY=... +//! # Optional (MinIO/R2 need an endpoint; region defaults to us-east-1): +//! export BENCHER_REPLICA_TEST_S3_ENDPOINT=http://localhost:9000 +//! export BENCHER_REPLICA_TEST_S3_REGION=us-east-1 +//! export BENCHER_REPLICA_TEST_S3_PREFIX=ci-scratch +//! cargo nextest run -p bencher_replica --features plus,testing \ +//! --run-ignored ignored-only -E 'binary(live_s3)' +//! ``` +//! +//! CI runs this tier against dockerized `MinIO` and `RustFS` on every rust +//! change (the `replica_live_s3` job in `.github/workflows/test.yml`). To +//! run locally without AWS, start either server the same way, e.g.: +//! +//! ```sh +//! docker run -d -p 9000:9000 \ +//! -e MINIO_ROOT_USER=bencher-test \ +//! -e MINIO_ROOT_PASSWORD=bencher-test-secret \ +//! minio/minio server /data +//! aws --endpoint-url http://127.0.0.1:9000 s3api create-bucket \ +//! --bucket bencher-replica-test +//! ``` +//! +//! Known server differences the suite tolerates: `RustFS` (1.0.0-beta.10) +//! returns an empty `ListMultipartUploads`, so the orphan-sweep case skips +//! there (`MinIO` and AWS exercise it). +//! +//! Every test isolates itself under a randomized per-run key prefix and +//! deletes that prefix in teardown (even when the case fails). A case that +//! drops an unfinished multipart upload leaves an invisible uncompleted +//! upload behind; use a bucket lifecycle rule to expire those. + +// The contract suite is shared by including the file as a module; its test +// modules also compile (and run) in this binary, which keeps the shared +// cases exercised even when the live environment is absent. +// Optional dependency enabled by the otel feature; unused by tests. +#[cfg(feature = "otel")] +use bencher_otel as _; +#[macro_use] +#[path = "storage_contract.rs"] +mod storage_contract; + +#[cfg(test)] +mod live_s3_backend { + use std::pin::Pin; + + use bencher_replica::{ReplicaStorage, S3Storage}; + use uuid::Uuid; + + /// Environment configuration for the live-S3 leg. + struct LiveS3Config { + bucket: String, + endpoint: Option, + region: Option, + access_key_id: String, + secret_access_key: String, + prefix: Option, + } + + impl LiveS3Config { + /// Read `BENCHER_REPLICA_TEST_S3_*`; `None` when the required + /// variables are unset (the test then skips cleanly). + fn from_env() -> Option { + let var = |name: &str| std::env::var(name).ok().filter(|value| !value.is_empty()); + Some(Self { + bucket: var("BENCHER_REPLICA_TEST_S3_BUCKET")?, + access_key_id: var("BENCHER_REPLICA_TEST_S3_ACCESS_KEY_ID")?, + secret_access_key: var("BENCHER_REPLICA_TEST_S3_SECRET_ACCESS_KEY")?, + endpoint: var("BENCHER_REPLICA_TEST_S3_ENDPOINT"), + region: var("BENCHER_REPLICA_TEST_S3_REGION"), + prefix: var("BENCHER_REPLICA_TEST_S3_PREFIX"), + }) + } + + /// A storage handle scoped to `run_prefix` under the optional + /// configured base prefix. + fn storage(&self, run_prefix: &str) -> ReplicaStorage { + self.s3_storage(run_prefix, None) + } + + /// A storage handle scoped to `run_prefix`, optionally capping the + /// `ListObjectsV2` page size so pagination is exercised over a small + /// object count. + fn s3_storage(&self, run_prefix: &str, max_keys: Option) -> ReplicaStorage { + let prefix = match &self.prefix { + Some(base) => format!("{base}/{run_prefix}"), + None => run_prefix.to_owned(), + }; + let mut s3 = S3Storage::new( + self.bucket.clone(), + Some(prefix), + self.endpoint.clone(), + self.region.clone(), + self.access_key_id.clone(), + &self.secret_access_key, + ); + if let Some(max_keys) = max_keys { + s3.set_max_keys(max_keys); + } + ReplicaStorage::S3(Box::new(s3)) + } + } + + #[expect( + clippy::print_stderr, + reason = "report the clean skip when the live environment is absent" + )] + fn report_skip() { + eprintln!("BENCHER_REPLICA_TEST_S3_* not set, skipping live S3 test"); + } + + #[expect( + clippy::print_stderr, + reason = "report the clean skip when the server cannot express the case" + )] + fn report_incomplete_uploads_unsupported() { + eprintln!( + "server does not report incomplete multipart uploads \ + (ListMultipartUploads returned none); skipping the sweep case" + ); + } + + /// Run one contract case against a live bucket under a fresh randomized + /// prefix, then delete the prefix in teardown, failure or not. The case + /// runs in a spawned task so a panic still reaches teardown; the + /// original panic message is printed by the panic hook and the test is + /// failed afterwards. + async fn run_live_case(case: &str, case_fn: CaseFn) + where + CaseFn: for<'s> FnOnce(&'s ReplicaStorage) -> Pin + Send + 's>> + + Send + + 'static, + { + let Some(config) = LiveS3Config::from_env() else { + report_skip(); + return; + }; + let run_prefix = format!( + "bencher-replica-contract/{case}-{}", + Uuid::new_v4().simple() + ); + let storage = config.storage(&run_prefix); + let outcome = tokio::spawn(async move { case_fn(&storage).await }).await; + // Teardown: remove everything this case wrote, even on failure. + config + .storage(&run_prefix) + .delete_prefix("") + .await + .unwrap_or_else(|error| { + panic!("failed to clean up live S3 prefix {run_prefix}: {error}") + }); + if let Err(error) = outcome { + panic!("live S3 contract case {case} failed: {error}"); + } + } + + macro_rules! live_case { + ($case:ident) => { + #[tokio::test] + #[ignore = "requires BENCHER_REPLICA_TEST_S3_* and a live bucket"] + async fn $case() { + run_live_case(stringify!($case), |storage| { + Box::pin(crate::storage_contract::cases::$case(storage)) + }) + .await; + } + }; + } + for_each_contract_case!(live_case); + + /// Continuation-token pagination: with a tiny page size, both `list` and + /// `list_dirs` must reassemble the complete, sorted result across several + /// `ListObjectsV2` round-trips. Objects are spread over several epoch + /// directories so `list_dirs` (delimiter-scoped) paginates too. + #[tokio::test] + #[ignore = "requires BENCHER_REPLICA_TEST_S3_* and a live bucket"] + async fn pagination_spans_continuation_tokens() { + use bytes::Bytes; + use pretty_assertions::assert_eq; + + const EPOCHS: u32 = 6; + const PER_EPOCH: u32 = 4; + + let Some(config) = LiveS3Config::from_env() else { + report_skip(); + return; + }; + let run_prefix = format!( + "bencher-replica-contract/pagination-{}", + Uuid::new_v4().simple() + ); + // Page size 5 over 6 epoch dirs x 4 segments = 24 objects: `list` + // spans ~5 pages and `list_dirs` (6 common prefixes) spans 2. + let storage = config.s3_storage(&run_prefix, Some(5)); + let outcome = tokio::spawn(async move { + let mut expected_keys = Vec::new(); + for epoch in 0..EPOCHS { + for seq in 0..PER_EPOCH { + let key = format!("generations/g/wal/epoch{epoch:02}/seg{seq:02}.wal.zst"); + storage + .put(&key, Bytes::from(format!("{epoch}-{seq}"))) + .await + .expect("paginated put"); + expected_keys.push(key); + } + } + expected_keys.sort(); + + let listed = storage + .list("generations/g/wal/") + .await + .expect("paginated list"); + assert_eq!( + listed, expected_keys, + "paginated list must return every key, sorted, across pages" + ); + + let dirs = storage + .list_dirs("generations/g/wal/") + .await + .expect("paginated list_dirs"); + let expected_dirs: Vec = (0..EPOCHS) + .map(|epoch| format!("epoch{epoch:02}")) + .collect(); + assert_eq!( + dirs, expected_dirs, + "paginated list_dirs must return every directory, sorted, across pages" + ); + }) + .await; + // Teardown: remove everything this case wrote, even on failure. + config + .s3_storage(&run_prefix, None) + .delete_prefix("") + .await + .unwrap_or_else(|error| { + panic!("failed to clean up live S3 prefix {run_prefix}: {error}") + }); + if let Err(error) = outcome { + panic!("live S3 pagination case failed: {error}"); + } + } + + /// Crash-orphan reconciliation: an incomplete multipart upload left behind + /// by a killed process is reclaimed by the best-effort sweep. Start a + /// multipart upload, drop it unfinished, sweep, and assert no incomplete + /// upload remains under the prefix. + #[tokio::test] + #[ignore = "requires BENCHER_REPLICA_TEST_S3_* and a live bucket"] + async fn abort_incomplete_uploads_sweeps_orphans() { + use bencher_replica::ReplicaStorage; + + let Some(config) = LiveS3Config::from_env() else { + report_skip(); + return; + }; + let run_prefix = format!( + "bencher-replica-contract/orphan-sweep-{}", + Uuid::new_v4().simple() + ); + let log = slog::Logger::root(slog::Discard, slog::o!()); + let storage = config.s3_storage(&run_prefix, None); + let outcome = tokio::spawn(async move { + // Start a multipart upload and drop it unfinished: a crash-orphaned + // incomplete upload now exists server-side under the prefix. + let upload = storage + .start_multipart("snap/db.zst") + .await + .expect("start multipart"); + drop(upload); + let ReplicaStorage::S3(s3) = &storage else { + panic!("expected the S3 backend"); + }; + // Some S3-compatible servers (rustfs 1.0.0-beta.10) return an + // empty ListMultipartUploads even when an incomplete upload + // exists, so the sweep has nothing observable to reclaim there: + // skip rather than fail, and rely on the AWS S3/MinIO legs to + // exercise the sweep. + if s3.incomplete_upload_count().await == 0 { + report_incomplete_uploads_unsupported(); + return; + } + // Sweep through the enum method: callers need no backend match. + storage.abort_incomplete_uploads(&log).await; + assert_eq!( + s3.incomplete_upload_count().await, + 0, + "the sweep aborts every incomplete upload under the prefix" + ); + }) + .await; + // Teardown: abort any straggler upload and remove written objects. + let cleanup_log = slog::Logger::root(slog::Discard, slog::o!()); + config + .s3_storage(&run_prefix, None) + .abort_incomplete_uploads(&cleanup_log) + .await; + config + .s3_storage(&run_prefix, None) + .delete_prefix("") + .await + .unwrap_or_else(|error| { + panic!("failed to clean up live S3 prefix {run_prefix}: {error}") + }); + if let Err(error) = outcome { + panic!("live S3 orphan-sweep case failed: {error}"); + } + } +} diff --git a/plus/bencher_replica/tests/restore_matrix.rs b/plus/bencher_replica/tests/restore_matrix.rs new file mode 100644 index 000000000..79cce4490 --- /dev/null +++ b/plus/bencher_replica/tests/restore_matrix.rs @@ -0,0 +1,1048 @@ +#![cfg(all(feature = "plus", feature = "testing"))] +//! Restore matrix: startup restore and restore-to-position verification. +//! +//! Every replica in this suite is fabricated BY HAND from the foundation +//! APIs (real WALs from `WalFixture`, chunks from `WalScanner`, zstd via +//! `compress_segment`, objects via `ReplicaStorage`): restore must be +//! correct independently of the sync engine that normally produces +//! replicas. +//! +//! NOTE: `unused_crate_dependencies` cannot be handled with a crate-level +//! `#![expect]` here (see `tests/storage_contract.rs`); unused package +//! dependencies are referenced explicitly instead, as rustc recommends. + +use async_compression as _; +use aws_credential_types as _; +use aws_sdk_s3 as _; +use futures as _; +use rand as _; +use serde as _; +use serde_json as _; +use thiserror as _; +use uuid as _; +use zstd as _; +// Optional dependency enabled by the otel feature; unused by tests. +#[cfg(feature = "otel")] +use bencher_otel as _; + +/// Shared fixtures: a scripted source database, a local replica, and a +/// restore target, plus hand-rolled snapshot/segment shippers. +#[cfg(test)] +pub(crate) mod harness { + use std::io::Cursor; + + use bencher_json::DateTime; + use bencher_json::system::config::{JsonReplication, ReplicationTarget}; + use bencher_replica::testing::{CheckpointMode, WalFixture}; + use bencher_replica::{ + GenerationId, LocalStorage, ReplicaConfig, ReplicaStorage, RestoreError, RestoreOutcome, + SnapshotMeta, WAL_HEADER_SIZE, WalBoundary, WalScanner, compress_segment, + decompress_segment, restore_if_missing, + }; + use bytes::Bytes; + use camino::{Utf8Path, Utf8PathBuf}; + use sha2::{Digest as _, Sha256}; + + /// Page size for every fixture database in this suite. + pub(crate) const PAGE_SIZE: u32 = 4096; + /// 2026-07-10T14:59:00Z, the base second for deterministic generations. + const BASE_SECS: i64 = 1_783_695_540; + + pub(crate) fn dir_path(tmp: &tempfile::TempDir) -> &Utf8Path { + Utf8Path::from_path(tmp.path()).expect("tempdir path is UTF-8") + } + + pub(crate) fn logger() -> slog::Logger { + slog::Logger::root(slog::Discard, slog::o!()) + } + + /// Deterministic generation ids: `generation(n)` sorts before + /// `generation(n + 1)`. + pub(crate) fn generation(index: u32) -> GenerationId { + let created = DateTime::try_from(BASE_SECS + i64::from(index)).expect("valid timestamp"); + GenerationId::new(created, index) + } + + /// A scripted source database, a local replica root, and a separate + /// restore target directory. + pub(crate) struct Harness { + pub fixture: WalFixture, + pub storage: ReplicaStorage, + pub config: ReplicaConfig, + /// Restore target: `/bencher.db` (does not exist yet). + pub db_path: Utf8PathBuf, + /// Directory holding the fixture (for boundary-state copies). + pub fixture_dir: Utf8PathBuf, + pub _fixture_tmp: tempfile::TempDir, + pub _replica_tmp: tempfile::TempDir, + pub _target_tmp: tempfile::TempDir, + } + + impl Harness { + pub(crate) fn new() -> Self { + let fixture_tmp = tempfile::tempdir().expect("fixture tempdir"); + let replica_tmp = tempfile::tempdir().expect("replica tempdir"); + let target_tmp = tempfile::tempdir().expect("target tempdir"); + let fixture_dir = dir_path(&fixture_tmp).to_path_buf(); + let fixture = WalFixture::new(&fixture_dir, PAGE_SIZE).expect("fixture"); + // Start from a clean, fully checkpointed state: the WAL is empty + // and the database file alone is the snapshot boundary. + fixture + .checkpoint(CheckpointMode::Truncate) + .expect("initial checkpoint"); + let replica_root = dir_path(&replica_tmp).to_path_buf(); + let storage = ReplicaStorage::Local(LocalStorage::new(replica_root.clone())); + let config = ReplicaConfig::try_from(JsonReplication { + target: ReplicationTarget::File { + path: replica_root.into_std_path_buf(), + }, + sync_interval_secs: None, + checkpoint_interval_secs: None, + min_checkpoint_pages: None, + snapshot_interval_secs: None, + snapshot_throttle_mib: None, + retention_generations: None, + verification_interval_secs: None, + shutdown_sync_timeout_secs: None, + }) + .expect("config"); + let db_path = dir_path(&target_tmp).join("bencher.db"); + Self { + fixture, + storage, + config, + db_path, + fixture_dir, + _fixture_tmp: fixture_tmp, + _replica_tmp: replica_tmp, + _target_tmp: target_tmp, + } + } + + pub(crate) async fn restore(&self) -> Result { + restore_if_missing(&logger(), &self.config, &self.db_path).await + } + + /// Checkpoint (TRUNCATE) the fixture and copy its database file as a + /// named boundary state for later equivalence assertions. The + /// truncation also ends the current WAL salt cycle: the next write + /// starts a new epoch. + pub(crate) fn boundary_copy(&self, name: &str) -> Utf8PathBuf { + self.fixture + .checkpoint(CheckpointMode::Truncate) + .expect("boundary checkpoint"); + let path = self.fixture_dir.join(name); + std::fs::copy(self.fixture.db_path(), &path).expect("boundary copy"); + path + } + + /// Write raw database bytes as a named boundary state. + pub(crate) fn write_state(&self, name: &str, db: &[u8]) -> Utf8PathBuf { + let path = self.fixture_dir.join(name); + std::fs::write(&path, db).expect("write state"); + path + } + + /// Run `count` single-insert transactions (one commit frame each). + pub(crate) fn txns(&self, count: usize, tag: &str) { + for index in 0..count { + self.fixture + .txn(&[&format!("INSERT INTO t (data) VALUES ('{tag}-{index}')")]) + .expect("txn"); + } + } + } + + pub(crate) fn snapshot_key(generation: &GenerationId) -> String { + format!("generations/{}/snapshot.db.zst", generation.as_str()) + } + + pub(crate) fn snapshot_meta_key(generation: &GenerationId) -> String { + format!("generations/{}/snapshot.json", generation.as_str()) + } + + /// Segment keys constructed independently of `bencher_replica`'s own key + /// builder: a deliberate cross-check of the pinned naming scheme. + pub(crate) fn segment_key( + generation: &GenerationId, + epoch: u64, + salt: (u32, u32), + start: u64, + end: u64, + ) -> String { + format!( + "generations/{}/wal/{epoch:010}-{:08x}{:08x}/{start:020}-{end:020}.wal.zst", + generation.as_str(), + salt.0, + salt.1 + ) + } + + /// Upload a snapshot body plus its `snapshot.json` commit marker. + pub(crate) async fn put_snapshot( + storage: &ReplicaStorage, + generation: &GenerationId, + db: &[u8], + boundary_salt: (u32, u32), + ) { + let compressed = compress_segment(db).expect("compress snapshot"); + let sha256 = hex::encode(Sha256::digest(&compressed)); + storage + .put(&snapshot_key(generation), Bytes::from(compressed)) + .await + .expect("put snapshot body"); + let meta = SnapshotMeta { + version: 1, + created: "2026-07-10T14:59:00Z".to_owned(), + db_bytes: u64::try_from(db.len()).expect("db size"), + page_size: PAGE_SIZE, + sha256, + wal_boundary: WalBoundary { + epoch: 0, + salt1: boundary_salt.0, + salt2: boundary_salt.1, + offset: 0, + }, + }; + storage + .put( + &snapshot_meta_key(generation), + Bytes::from(meta.to_bytes().expect("snapshot meta bytes")), + ) + .await + .expect("put snapshot.json"); + } + + /// Everything shipped for one epoch: per-segment keys, end offsets, and + /// running checksums (for building `Position`s). + pub(crate) struct ShippedEpoch { + pub salt: (u32, u32), + pub keys: Vec, + pub ends: Vec, + pub checksums: Vec<(u32, u32)>, + } + + /// Ship every committed transaction of `wal` as one segment per commit + /// under `epoch`. The first segment starts at offset 0 and carries the + /// 32-byte WAL header, so restore can rebuild the `-wal` file by + /// decompress-and-concatenate. + pub(crate) async fn ship_epoch( + storage: &ReplicaStorage, + generation: &GenerationId, + epoch: u64, + wal: &[u8], + ) -> ShippedEpoch { + let mut scanner = WalScanner::open(Cursor::new(wal.to_vec())) + .expect("WAL header") + .expect("WAL is not empty"); + let salt = scanner.header().salt; + let mut shipped = ShippedEpoch { + salt, + keys: Vec::new(), + ends: Vec::new(), + checksums: Vec::new(), + }; + while let Some(chunk) = scanner.next_committed(1).expect("scan WAL") { + let (start, raw) = if chunk.start_offset == WAL_HEADER_SIZE { + let mut with_header = wal[..32].to_vec(); + with_header.extend_from_slice(&chunk.bytes); + (0u64, with_header) + } else { + (chunk.start_offset, chunk.bytes.clone()) + }; + let key = segment_key(generation, epoch, salt, start, chunk.end_offset); + let compressed = compress_segment(&raw).expect("compress segment"); + storage + .put(&key, Bytes::from(compressed)) + .await + .expect("put segment"); + shipped.keys.push(key); + shipped.ends.push(chunk.end_offset); + shipped.checksums.push(chunk.checksum_at_end); + } + shipped + } + + /// Flip one byte in the middle of a stored object. + pub(crate) async fn corrupt_object(storage: &ReplicaStorage, key: &str) { + let bytes = storage.get(key).await.expect("get object to corrupt"); + let mut bytes = bytes.to_vec(); + // A shift instead of `/ 2` keeps the integer_division lint quiet. + let middle = bytes.len() >> 1; + bytes[middle] ^= 0xff; + storage + .put(key, Bytes::from(bytes)) + .await + .expect("put corrupted object"); + } + + /// Tamper a stored segment so the OBJECT stays fully valid (zstd content + /// checksum recomputed, byte length preserved) while the WAL checksum + /// chain inside it breaks: decompress, flip a payload byte, recompress. + /// This models a shipped-then-forked or bit-rotted lineage that + /// per-object integrity checks cannot see. + pub(crate) async fn recompress_tampered_segment(storage: &ReplicaStorage, key: &str) { + let compressed = storage.get(key).await.expect("get segment to tamper"); + let mut raw = decompress_segment(&compressed).expect("decompress segment"); + // Flip one byte in the FIRST frame's page payload (frame header is + // 24 bytes; non-first segments carry no WAL header). + raw[24 + 50] ^= 0xff; + let recompressed = compress_segment(&raw).expect("recompress tampered segment"); + storage + .put(key, Bytes::from(recompressed)) + .await + .expect("put tampered segment"); + } +} + +#[cfg(test)] +mod cases { + use bencher_replica::testing::{CheckpointMode, assert_replica_equivalent}; + use bencher_replica::{ + Position, ReplicaMeta, RestoreError, RestoreOutcome, SnapshotMeta, VerifyReport, + WalBoundary, compress_segment, fingerprint_database, verify_against_replica, + }; + use bytes::Bytes; + use camino::Utf8PathBuf; + use pretty_assertions::assert_eq; + use sha2::{Digest as _, Sha256}; + + use super::harness::{ + Harness, corrupt_object, dir_path, generation, logger, put_snapshot, + recompress_tampered_segment, ship_epoch, snapshot_key, snapshot_meta_key, + }; + + /// Arbitrary boundary salts for snapshot-only replicas (no WAL cycle ever + /// started after the snapshot). + const NO_WAL_SALT: (u32, u32) = (0x1111_2222, 0x3333_4444); + + fn assert_restored(outcome: &RestoreOutcome) -> (&bencher_replica::GenerationId, u64, u64) { + let RestoreOutcome::Restored { + generation, + db_bytes, + segments, + } = outcome + else { + panic!("expected Restored, got {outcome:?}") + }; + (generation, *db_bytes, *segments) + } + + #[tokio::test] + async fn restore_snapshot_only() { + let harness = Harness::new(); + let boundary = harness.boundary_copy("boundary.db"); + let db = harness.fixture.db_bytes().expect("db bytes"); + put_snapshot(&harness.storage, &generation(1), &db, NO_WAL_SALT).await; + + let outcome = harness.restore().await.expect("restore"); + let (restored_generation, db_bytes, segments) = assert_restored(&outcome); + assert_eq!(restored_generation, &generation(1), "generation"); + assert_eq!(segments, 0, "snapshot only: no segments replayed"); + assert_eq!( + db_bytes, + u64::try_from(db.len()).expect("db size"), + "restored size equals the uncompressed snapshot" + ); + assert_replica_equivalent(&boundary, &harness.db_path); + } + + #[tokio::test] + async fn restore_snapshot_plus_n_segments() { + for n in [1usize, 2, 17] { + let harness = Harness::new(); + let db = harness.fixture.db_bytes().expect("db bytes"); + harness.txns(n, "seg"); + let wal = harness.fixture.wal_bytes().expect("wal bytes"); + let shipped = ship_epoch(&harness.storage, &generation(1), 0, &wal).await; + assert_eq!( + shipped.ends.len(), + n, + "one segment per committed transaction (n = {n})" + ); + put_snapshot(&harness.storage, &generation(1), &db, shipped.salt).await; + let boundary = harness.boundary_copy("boundary.db"); + + let outcome = harness.restore().await.expect("restore"); + let (_, _, segments) = assert_restored(&outcome); + assert_eq!( + segments, + u64::try_from(n).expect("segment count"), + "all {n} segments replayed" + ); + assert_replica_equivalent(&boundary, &harness.db_path); + } + } + + #[tokio::test] + async fn restore_multi_epoch() { + let harness = Harness::new(); + let db = harness.fixture.db_bytes().expect("db bytes"); + // Epoch 0: three commits. + harness.txns(3, "epoch0"); + let wal = harness.fixture.wal_bytes().expect("wal bytes"); + let epoch0 = ship_epoch(&harness.storage, &generation(1), 0, &wal).await; + put_snapshot(&harness.storage, &generation(1), &db, epoch0.salt).await; + // The boundary checkpoint truncates the WAL; the next write starts + // epoch 1 with fresh salts. + harness.boundary_copy("epoch0.db"); + harness.txns(2, "epoch1"); + let wal = harness.fixture.wal_bytes().expect("wal bytes"); + let epoch1 = ship_epoch(&harness.storage, &generation(1), 1, &wal).await; + assert_ne!(epoch0.salt, epoch1.salt, "a WAL restart changes the salts"); + let boundary = harness.boundary_copy("epoch1.db"); + + let outcome = harness.restore().await.expect("restore"); + let (_, _, segments) = assert_restored(&outcome); + assert_eq!(segments, 5, "both epochs replayed in full"); + assert_replica_equivalent(&boundary, &harness.db_path); + + let meta = ReplicaMeta::load(&harness.db_path) + .expect("load meta") + .expect("meta written"); + assert_eq!(meta.epoch, 1, "meta records the last applied epoch"); + assert_eq!( + (meta.salt1, meta.salt2), + epoch1.salt, + "meta records the last epoch's salts" + ); + assert_eq!( + Some(&meta.shipped_offset), + epoch1.ends.last(), + "meta records the last applied segment end" + ); + } + + #[tokio::test] + async fn restore_picks_latest_generation() { + let harness = Harness::new(); + harness.boundary_copy("state_a.db"); + let db_a = harness.fixture.db_bytes().expect("db bytes"); + put_snapshot(&harness.storage, &generation(1), &db_a, NO_WAL_SALT).await; + + harness.txns(1, "later"); + let boundary_b = harness.boundary_copy("state_b.db"); + let db_b = harness.fixture.db_bytes().expect("db bytes"); + put_snapshot(&harness.storage, &generation(2), &db_b, NO_WAL_SALT).await; + + let outcome = harness.restore().await.expect("restore"); + let (restored_generation, ..) = assert_restored(&outcome); + assert_eq!( + restored_generation, + &generation(2), + "the newest complete generation wins" + ); + assert_replica_equivalent(&boundary_b, &harness.db_path); + } + + #[tokio::test] + async fn restore_ignores_generation_without_snapshot_json() { + let harness = Harness::new(); + let boundary_a = harness.boundary_copy("state_a.db"); + let db_a = harness.fixture.db_bytes().expect("db bytes"); + put_snapshot(&harness.storage, &generation(1), &db_a, NO_WAL_SALT).await; + + // A newer generation with a snapshot BODY but no snapshot.json: a + // crash mid-snapshot. It must be invisible to restore. + harness.txns(1, "crashed"); + harness.boundary_copy("state_b.db"); + let db_b = harness.fixture.db_bytes().expect("db bytes"); + let compressed = compress_segment(&db_b).expect("compress"); + harness + .storage + .put(&snapshot_key(&generation(2)), Bytes::from(compressed)) + .await + .expect("put orphan snapshot body"); + + let outcome = harness.restore().await.expect("restore"); + let (restored_generation, ..) = assert_restored(&outcome); + assert_eq!( + restored_generation, + &generation(1), + "a generation without snapshot.json is invisible" + ); + assert_replica_equivalent(&boundary_a, &harness.db_path); + } + + #[tokio::test] + async fn restore_empty_replica_no_replica() { + let harness = Harness::new(); + // Stale advisory meta from a previous life of this volume. + stale_meta() + .store(&harness.db_path) + .expect("store stale meta"); + + let outcome = harness.restore().await.expect("restore"); + assert!( + matches!(outcome, RestoreOutcome::NoReplica), + "expected NoReplica, got {outcome:?}" + ); + assert!( + !harness.db_path.exists(), + "no database file is fabricated for a fresh server" + ); + assert_eq!( + ReplicaMeta::load(&harness.db_path).expect("load meta"), + None, + "stale advisory meta is removed when the database is missing" + ); + } + + #[tokio::test] + async fn restore_db_exists_skipped() { + let harness = Harness::new(); + // A perfectly valid replica exists, so a skip can only come from the + // database file being present. + let db = harness.fixture.db_bytes().expect("db bytes"); + put_snapshot(&harness.storage, &generation(1), &db, NO_WAL_SALT).await; + std::fs::write(&harness.db_path, b"pre-existing database").expect("write db"); + let meta = stale_meta(); + meta.store(&harness.db_path).expect("store meta"); + + let outcome = harness.restore().await.expect("restore"); + assert!( + matches!(outcome, RestoreOutcome::Skipped), + "expected Skipped, got {outcome:?}" + ); + assert_eq!( + std::fs::read(&harness.db_path).expect("read db"), + b"pre-existing database".to_vec(), + "an existing database file is untouched" + ); + assert_eq!( + ReplicaMeta::load(&harness.db_path).expect("load meta"), + Some(meta), + "the meta file is untouched too" + ); + } + + #[tokio::test] + async fn restore_corrupted_segment_stops_at_last_good_epoch() { + let harness = Harness::new(); + let db = harness.fixture.db_bytes().expect("db bytes"); + harness.txns(2, "epoch0"); + let wal = harness.fixture.wal_bytes().expect("wal bytes"); + let epoch0 = ship_epoch(&harness.storage, &generation(1), 0, &wal).await; + put_snapshot(&harness.storage, &generation(1), &db, epoch0.salt).await; + let boundary0 = harness.boundary_copy("epoch0.db"); + harness.txns(3, "epoch1"); + let wal = harness.fixture.wal_bytes().expect("wal bytes"); + let epoch1 = ship_epoch(&harness.storage, &generation(1), 1, &wal).await; + assert_eq!(epoch1.keys.len(), 3, "epoch 1 has three segments"); + // Flip a byte inside the middle stored segment of epoch 1. + corrupt_object(&harness.storage, &epoch1.keys[1]).await; + + let outcome = harness.restore().await.expect("restore boots anyway"); + let (_, _, segments) = assert_restored(&outcome); + assert_eq!( + segments, 2, + "only epoch 0's segments count; the corrupt epoch is discarded" + ); + assert_replica_equivalent(&boundary0, &harness.db_path); + + let meta = ReplicaMeta::load(&harness.db_path) + .expect("load meta") + .expect("meta written"); + assert_eq!(meta.epoch, 0, "meta stops at the last good epoch"); + assert_eq!( + Some(&meta.shipped_offset), + epoch0.ends.last(), + "meta records epoch 0's end" + ); + } + + /// A stored segment whose object is valid but decompresses to far more + /// than its key's declared byte range: restore must reject it via the + /// exact-size cap (never inflating toward the multi-gigabyte ceiling) and + /// soft-stop, keeping the last good epoch. + #[tokio::test] + async fn restore_oversized_segment_stops_at_last_good_epoch() { + let harness = Harness::new(); + let db = harness.fixture.db_bytes().expect("db bytes"); + harness.txns(2, "epoch0"); + let wal = harness.fixture.wal_bytes().expect("wal bytes"); + let epoch0 = ship_epoch(&harness.storage, &generation(1), 0, &wal).await; + put_snapshot(&harness.storage, &generation(1), &db, epoch0.salt).await; + let boundary0 = harness.boundary_copy("epoch0.db"); + + harness.txns(2, "epoch1"); + let wal = harness.fixture.wal_bytes().expect("wal bytes"); + let epoch1 = ship_epoch(&harness.storage, &generation(1), 1, &wal).await; + assert_eq!(epoch1.keys.len(), 2, "epoch 1 has two segments"); + // Replace the first segment body with a valid object that decompresses + // to 512 KiB, orders of magnitude beyond its key's declared range. + let bloated = compress_segment(&vec![0u8; 512 * 1024]).expect("compress bloated segment"); + harness + .storage + .put(&epoch1.keys[0], Bytes::from(bloated)) + .await + .expect("put bloated segment"); + + let outcome = harness.restore().await.expect("restore boots anyway"); + let (_, _, segments) = assert_restored(&outcome); + assert_eq!( + segments, 2, + "only epoch 0 applies; the oversized segment is discarded" + ); + assert_replica_equivalent(&boundary0, &harness.db_path); + + let meta = ReplicaMeta::load(&harness.db_path) + .expect("load meta") + .expect("meta written"); + assert_eq!(meta.epoch, 0, "meta stops at the last good epoch"); + } + + /// A tampered segment whose OBJECT is fully valid (recompressed, so the + /// zstd checksum and the byte length both pass) but whose WAL checksum + /// chain is broken inside the epoch: the chain pre-validation must + /// soft-stop BEFORE application, discarding that epoch AND every later + /// one, instead of letting `SQLite` recovery silently truncate at the + /// break and replay later epochs on top of the wrong base. + #[tokio::test] + async fn restore_chain_break_inside_epoch_stops_before_application() { + let harness = Harness::new(); + let db = harness.fixture.db_bytes().expect("db bytes"); + harness.txns(2, "epoch0"); + let wal = harness.fixture.wal_bytes().expect("wal bytes"); + let epoch0 = ship_epoch(&harness.storage, &generation(1), 0, &wal).await; + put_snapshot(&harness.storage, &generation(1), &db, epoch0.salt).await; + let boundary0 = harness.boundary_copy("epoch0.db"); + + // Epoch 1: three segments; break the chain inside the middle one + // with an object-valid recompression tamper. + harness.txns(3, "epoch1"); + let wal = harness.fixture.wal_bytes().expect("wal bytes"); + let epoch1 = ship_epoch(&harness.storage, &generation(1), 1, &wal).await; + assert_eq!(epoch1.keys.len(), 3, "epoch 1 has three segments"); + recompress_tampered_segment(&harness.storage, &epoch1.keys[1]).await; + + // Epoch 2 exists and is intact: it must STILL not apply, because + // its page images build on epoch 1's lost frames. + harness.txns(1, "epoch2"); + let wal = harness.fixture.wal_bytes().expect("wal bytes"); + let _epoch2 = ship_epoch(&harness.storage, &generation(1), 2, &wal).await; + + let outcome = harness.restore().await.expect("restore boots anyway"); + let (_, _, segments) = assert_restored(&outcome); + assert_eq!( + segments, 2, + "only epoch 0 applies; the broken epoch and everything after are discarded" + ); + assert_replica_equivalent(&boundary0, &harness.db_path); + + let meta = ReplicaMeta::load(&harness.db_path) + .expect("load meta") + .expect("meta written"); + assert_eq!(meta.epoch, 0, "meta stops at the last good epoch"); + } + + /// A single generation that holds the SAME epoch number twice under + /// different salts: two epoch-1 directories, each a complete, well-formed + /// WAL lineage. This is the salt-collision state a restore must never be + /// poisoned by. `plan_epochs` sorts by `(epoch, start, end)` with salt + /// EXCLUDED, so the two lineages collapse into one epoch-1 group and the + /// salt collision discards the epoch wholesale. Restore soft-stops at + /// epoch 0: it boots on the last unambiguous state, never a torn mixture of + /// the two lineages and never silently picking one. This is safe but + /// pessimistic (a complete valid lineage is dropped), which is why the + /// engine's resume path is designed never to CREATE the collision (see + /// `resume_after_soft_stop_below_corrupt_epoch_forces_new_generation` in + /// `tests/sync_engine.rs`, and the `meta_matches` note in `src/sync.rs`). + #[tokio::test] + async fn restore_duplicate_epoch_salt_collision_soft_stops_at_last_unambiguous_epoch() { + let harness = Harness::new(); + let db = harness.fixture.db_bytes().expect("db bytes"); + // Epoch 0: two commits, shipped, plus the snapshot boundary. + harness.txns(2, "epoch0"); + let wal0 = harness.fixture.wal_bytes().expect("wal bytes"); + let epoch0 = ship_epoch(&harness.storage, &generation(1), 0, &wal0).await; + put_snapshot(&harness.storage, &generation(1), &db, epoch0.salt).await; + let boundary0 = harness.boundary_copy("epoch0.db"); + + // Epoch 1, first lineage: a real WAL cycle shipped under epoch 1. + harness.txns(2, "epoch1-first"); + let wal_first = harness.fixture.wal_bytes().expect("wal bytes"); + let first_lineage = ship_epoch(&harness.storage, &generation(1), 1, &wal_first).await; + // Restart the WAL (fresh salts) and ship a SECOND, distinct lineage + // that also claims epoch 1: the same epoch number, different salts. + harness + .fixture + .checkpoint(CheckpointMode::Truncate) + .expect("truncate to restart the WAL cycle"); + harness.txns(2, "epoch1-second"); + let wal_second = harness.fixture.wal_bytes().expect("wal bytes"); + let second_lineage = ship_epoch(&harness.storage, &generation(1), 1, &wal_second).await; + assert_ne!( + first_lineage.salt, second_lineage.salt, + "the two epoch-1 lineages carry distinct salts (a real directory collision)" + ); + + let outcome = harness.restore().await.expect("restore boots anyway"); + let (_, _, segments) = assert_restored(&outcome); + assert_eq!( + segments, 2, + "only epoch 0 applies; the duplicate-epoch collision discards both epoch-1 lineages" + ); + assert_replica_equivalent(&boundary0, &harness.db_path); + + let meta = ReplicaMeta::load(&harness.db_path) + .expect("load meta") + .expect("meta written"); + assert_eq!(meta.epoch, 0, "meta stops at the last unambiguous epoch"); + assert_eq!( + Some(&meta.shipped_offset), + epoch0.ends.last(), + "meta records epoch 0's end" + ); + } + + #[tokio::test] + async fn restore_missing_middle_segment_stops() { + let harness = Harness::new(); + let db = harness.fixture.db_bytes().expect("db bytes"); + let boundary = harness.write_state("snapshot-boundary.db", &db); + harness.txns(3, "epoch0"); + let wal = harness.fixture.wal_bytes().expect("wal bytes"); + let epoch0 = ship_epoch(&harness.storage, &generation(1), 0, &wal).await; + assert_eq!(epoch0.keys.len(), 3, "epoch 0 has three segments"); + put_snapshot(&harness.storage, &generation(1), &db, epoch0.salt).await; + // Delete segment 2 of 3: the offset chain has a hole. + harness + .storage + .delete(&epoch0.keys[1]) + .await + .expect("delete middle segment"); + + let outcome = harness.restore().await.expect("restore boots anyway"); + let (_, _, segments) = assert_restored(&outcome); + assert_eq!( + segments, 0, + "an epoch with a segment gap is discarded entirely" + ); + assert_replica_equivalent(&boundary, &harness.db_path); + + let meta = ReplicaMeta::load(&harness.db_path) + .expect("load meta") + .expect("meta written"); + assert_eq!(meta.epoch, 0, "meta falls back to the snapshot boundary"); + assert_eq!(meta.shipped_offset, 0, "nothing shipped in the epoch"); + assert_eq!( + (meta.salt1, meta.salt2), + epoch0.salt, + "boundary salts from snapshot.json" + ); + } + + #[tokio::test] + async fn restore_corrupted_snapshot_sha256_hard_error() { + // Tampered recorded hash: the download succeeds but the checksum + // comparison must hard-fail and leave no files behind. + let harness = Harness::new(); + let db = harness.fixture.db_bytes().expect("db bytes"); + put_snapshot(&harness.storage, &generation(1), &db, NO_WAL_SALT).await; + let meta_key = snapshot_meta_key(&generation(1)); + let meta_bytes = harness.storage.get(&meta_key).await.expect("get meta"); + let mut snapshot_meta = SnapshotMeta::from_bytes(&meta_bytes).expect("parse meta"); + snapshot_meta.sha256 = "00".repeat(32); + harness + .storage + .put( + &meta_key, + Bytes::from(snapshot_meta.to_bytes().expect("meta bytes")), + ) + .await + .expect("put tampered meta"); + + let error = harness.restore().await.expect_err("checksum must fail"); + assert!( + matches!(error, RestoreError::SnapshotChecksum { .. }), + "expected SnapshotChecksum, got {error}" + ); + assert_no_db_left_behind(&harness); + + // A corrupted snapshot BODY is also a hard error (zstd decode or + // checksum), never a half-restored database. + let harness = Harness::new(); + let db = harness.fixture.db_bytes().expect("db bytes"); + put_snapshot(&harness.storage, &generation(1), &db, NO_WAL_SALT).await; + corrupt_object(&harness.storage, &snapshot_key(&generation(1))).await; + harness + .restore() + .await + .expect_err("corrupt snapshot body must fail"); + assert_no_db_left_behind(&harness); + } + + #[tokio::test] + async fn restore_snapshot_body_larger_than_db_bytes_fails() { + // A body that decompresses to far more than the recorded db_bytes: + // the streamed restore must fail on the size cap, before the checksum + // verdict, without writing unboundedly. + let harness = Harness::new(); + let oversized = vec![0u8; 256 * 1024]; + let compressed = compress_segment(&oversized).expect("compress oversized body"); + let sha256 = hex::encode(Sha256::digest(&compressed)); + harness + .storage + .put(&snapshot_key(&generation(1)), Bytes::from(compressed)) + .await + .expect("put oversized snapshot body"); + let meta = SnapshotMeta { + version: 1, + created: "2026-07-10T14:59:00Z".to_owned(), + db_bytes: 4096, + page_size: 4096, + sha256, + wal_boundary: WalBoundary { + epoch: 0, + salt1: NO_WAL_SALT.0, + salt2: NO_WAL_SALT.1, + offset: 0, + }, + }; + harness + .storage + .put( + &snapshot_meta_key(&generation(1)), + Bytes::from(meta.to_bytes().expect("meta bytes")), + ) + .await + .expect("put snapshot.json"); + + let error = harness + .restore() + .await + .expect_err("oversize snapshot body must fail"); + assert!( + matches!(error, RestoreError::SnapshotTooLarge { .. }), + "expected SnapshotTooLarge, got {error}" + ); + assert_no_db_left_behind(&harness); + } + + #[tokio::test] + async fn restore_falls_through_to_older_generation_when_newest_body_missing() { + // The newest generation has a valid snapshot.json marker but no body + // (the partial-prune state): restore must fall through to the intact + // older generation instead of refusing to boot. + let harness = Harness::new(); + let boundary_old = harness.boundary_copy("old.db"); + let db_old = harness.fixture.db_bytes().expect("db bytes"); + put_snapshot(&harness.storage, &generation(1), &db_old, NO_WAL_SALT).await; + + // Newer generation: upload ONLY snapshot.json (no snapshot.db.zst). + // The recorded sha256 is never consulted because the body is missing. + let marker = SnapshotMeta { + version: 1, + created: "2026-07-10T14:59:00Z".to_owned(), + db_bytes: 4096, + page_size: 4096, + sha256: "00".repeat(32), + wal_boundary: WalBoundary { + epoch: 0, + salt1: NO_WAL_SALT.0, + salt2: NO_WAL_SALT.1, + offset: 0, + }, + }; + harness + .storage + .put( + &snapshot_meta_key(&generation(2)), + Bytes::from(marker.to_bytes().expect("marker bytes")), + ) + .await + .expect("put newer snapshot.json marker"); + + let outcome = harness + .restore() + .await + .expect("restore succeeds from the older generation"); + let (restored_generation, ..) = assert_restored(&outcome); + assert_eq!( + restored_generation, + &generation(1), + "falls through to the intact older generation" + ); + assert_replica_equivalent(&boundary_old, &harness.db_path); + } + + #[tokio::test] + async fn restore_hard_fails_on_corrupt_newest_body_without_falling_through() { + // A newest generation whose body is PRESENT but whose recorded sha256 + // is wrong: unlike a missing body, this is corruption and must HARD + // fail, never silently regress to the older generation. + let harness = Harness::new(); + harness.boundary_copy("old.db"); + let db_old = harness.fixture.db_bytes().expect("db bytes"); + put_snapshot(&harness.storage, &generation(1), &db_old, NO_WAL_SALT).await; + + put_snapshot(&harness.storage, &generation(2), &db_old, NO_WAL_SALT).await; + let meta_key = snapshot_meta_key(&generation(2)); + let meta_bytes = harness.storage.get(&meta_key).await.expect("get meta"); + let mut snapshot_meta = SnapshotMeta::from_bytes(&meta_bytes).expect("parse meta"); + snapshot_meta.sha256 = "00".repeat(32); + harness + .storage + .put( + &meta_key, + Bytes::from(snapshot_meta.to_bytes().expect("meta bytes")), + ) + .await + .expect("put tampered meta"); + + let error = harness + .restore() + .await + .expect_err("a corrupt newest body must hard-fail, not fall through"); + assert!( + matches!(error, RestoreError::SnapshotChecksum { .. }), + "expected SnapshotChecksum, got {error}" + ); + assert_no_db_left_behind(&harness); + } + + #[tokio::test] + async fn restore_writes_fresh_meta() { + // Snapshot only: the meta sits at the wal_boundary with offset 0. + let harness = Harness::new(); + let db = harness.fixture.db_bytes().expect("db bytes"); + put_snapshot(&harness.storage, &generation(1), &db, NO_WAL_SALT).await; + harness.restore().await.expect("restore"); + let meta = ReplicaMeta::load(&harness.db_path) + .expect("load meta") + .expect("meta written"); + assert_eq!( + meta, + ReplicaMeta { + version: 1, + generation: generation(1).as_str().to_owned(), + epoch: 0, + salt1: NO_WAL_SALT.0, + salt2: NO_WAL_SALT.1, + shipped_offset: 0, + epoch_shipped_through_checkpoint: true, + shadow: false, + }, + "snapshot-only restore writes a boundary meta" + ); + + // Snapshot plus segments: the meta records the applied end position. + let harness = Harness::new(); + let db = harness.fixture.db_bytes().expect("db bytes"); + harness.txns(2, "meta"); + let wal = harness.fixture.wal_bytes().expect("wal bytes"); + let shipped = ship_epoch(&harness.storage, &generation(1), 0, &wal).await; + put_snapshot(&harness.storage, &generation(1), &db, shipped.salt).await; + harness.restore().await.expect("restore"); + let meta = ReplicaMeta::load(&harness.db_path) + .expect("load meta") + .expect("meta written"); + assert_eq!( + meta, + ReplicaMeta { + version: 1, + generation: generation(1).as_str().to_owned(), + epoch: 0, + salt1: shipped.salt.0, + salt2: shipped.salt.1, + shipped_offset: *shipped.ends.last().expect("segments shipped"), + epoch_shipped_through_checkpoint: true, + shadow: false, + }, + "restore writes the applied position with checkpoint-verified epoch" + ); + } + + #[tokio::test] + async fn restore_to_position_replays_prefix_only() { + let harness = Harness::new(); + let db = harness.fixture.db_bytes().expect("db bytes"); + harness.txns(2, "at-p"); + // Fingerprint the source at position P (after two commits). + let source_conn = + rusqlite::Connection::open(harness.fixture.db_path()).expect("open source"); + let fingerprint_p = fingerprint_database(&source_conn).expect("fingerprint at P"); + // One more commit past P. + harness.txns(1, "after-p"); + let fingerprint_after = fingerprint_database(&source_conn).expect("fingerprint after P"); + assert_ne!( + fingerprint_p, fingerprint_after, + "the fingerprints must differ across the extra commit" + ); + let wal = harness.fixture.wal_bytes().expect("wal bytes"); + let shipped = ship_epoch(&harness.storage, &generation(1), 0, &wal).await; + assert_eq!(shipped.ends.len(), 3, "three commits shipped"); + put_snapshot(&harness.storage, &generation(1), &db, shipped.salt).await; + + let position = Position { + generation: generation(1), + epoch: 0, + salt: shipped.salt, + offset: shipped.ends[1], + checksum: shipped.checksums[1], + }; + let scratch_pass = tempfile::tempdir().expect("scratch tempdir"); + let report = verify_against_replica( + &logger(), + &harness.storage, + &position, + &fingerprint_p, + dir_path(&scratch_pass), + ) + .await + .expect("verify at P"); + assert_eq!( + report, + VerifyReport::Pass, + "the replica restored to P matches the source fingerprint at P" + ); + + let scratch_fail = tempfile::tempdir().expect("scratch tempdir"); + let report = verify_against_replica( + &logger(), + &harness.storage, + &position, + &fingerprint_after, + dir_path(&scratch_fail), + ) + .await + .expect("verify after P"); + let VerifyReport::Fail { detail } = report else { + panic!("a fingerprint taken past P must not verify at P, got {report:?}") + }; + assert!( + detail.contains("table=t"), + "the failure names the differing table line: {detail}" + ); + } + + fn stale_meta() -> ReplicaMeta { + ReplicaMeta { + version: 1, + generation: generation(9).as_str().to_owned(), + epoch: 7, + salt1: 0xdead_beef, + salt2: 0xcafe_f00d, + shipped_offset: 4128, + epoch_shipped_through_checkpoint: false, + shadow: false, + } + } + + fn assert_no_db_left_behind(harness: &Harness) { + assert!( + !harness.db_path.exists(), + "no database file after a hard restore error" + ); + let scratch = Utf8PathBuf::from(format!("{}.restore", harness.db_path)); + assert!( + !scratch.exists(), + "the partial .restore scratch file is deleted" + ); + assert!( + !Utf8PathBuf::from(format!("{scratch}-wal")).exists(), + "no scratch -wal left behind" + ); + } +} diff --git a/plus/bencher_replica/tests/storage_contract.rs b/plus/bencher_replica/tests/storage_contract.rs new file mode 100644 index 000000000..fc551a079 --- /dev/null +++ b/plus/bencher_replica/tests/storage_contract.rs @@ -0,0 +1,588 @@ +#![cfg(all(feature = "plus", feature = "testing"))] +//! Storage backend contract suite. +//! +//! Every [`ReplicaStorage`] backend must satisfy the same observable +//! semantics (documented in `src/storage.rs`). The suite is table-driven: +//! the [`for_each_contract_case`] macro is the single case table, and each +//! backend module instantiates one test per case against a fresh, isolated +//! backend. The suite runs here against (a) the local-filesystem backend and +//! (b) the fault-injection wrapper around it with an empty failure plan +//! (proving the wrapper is transparent). The live-S3 leg in +//! `tests/live_s3.rs` includes this file as a module and runs the same table +//! against a real bucket. +//! +//! NOTE: `unused_crate_dependencies` cannot be handled with a crate-level +//! `#![expect]` here because this file is also compiled as a module of +//! `tests/live_s3.rs`, where such an expectation is ignored and reported as +//! unfulfilled. Unused package dependencies are referenced explicitly +//! instead, as rustc recommends. + +use async_compression as _; +use aws_credential_types as _; +use aws_sdk_s3 as _; +use bencher_json as _; +use hex as _; +use rand as _; +use rusqlite as _; +use serde as _; +use serde_json as _; +use sha2 as _; +use slog as _; +use thiserror as _; +use uuid as _; +use zstd as _; +// Optional dependency enabled by the otel feature; unused by tests. +#[cfg(feature = "otel")] +use bencher_otel as _; + +/// The contract case table: expands `$per_case!(case_name);` once per case. +/// +/// Each case is an `async fn(&ReplicaStorage)` in [`cases`]. Backend modules +/// pass a macro that wraps one case into one `#[tokio::test]`, so the case +/// name is part of the failing test's name and panic message. +macro_rules! for_each_contract_case { + ($per_case:ident) => { + $per_case!(put_then_get_roundtrip); + $per_case!(get_missing_is_not_found_error); + $per_case!(put_overwrite_replaces); + $per_case!(list_prefix_returns_only_prefix); + $per_case!(list_ordering_lexicographic); + $per_case!(list_empty_prefix_ok_empty_vec); + $per_case!(list_dirs_immediate_components_sorted); + $per_case!(delete_then_get_not_found); + $per_case!(delete_missing_idempotent_ok); + $per_case!(delete_prefix_removes_all); + $per_case!(multipart_roundtrip_large_object); + $per_case!(multipart_abort_leaves_nothing); + $per_case!(multipart_unfinished_invisible); + $per_case!(key_charset_roundtrip); + $per_case!(concurrent_puts_distinct_keys); + $per_case!(get_stream_roundtrip); + }; +} + +/// The generic contract cases. Each takes only a `&ReplicaStorage` rooted at +/// a fresh, empty, isolated namespace. +#[cfg(test)] +pub(crate) mod cases { + use bencher_replica::{ReplicaStorage, StorageError}; + use bytes::Bytes; + use pretty_assertions::assert_eq; + use tokio::io::AsyncReadExt as _; + + /// One mebibyte, for sizing multipart writes. + const MIB: usize = 1024 * 1024; + + /// Deterministic payload of `len` bytes (no RNG, no wall clock). + fn pattern(len: usize) -> Vec { + let unit: Vec = (0..=u8::MAX).collect(); + let mut data = unit.repeat(len.div_ceil(unit.len())); + data.truncate(len); + data + } + + async fn put_str(storage: &ReplicaStorage, key: &str, value: &str) { + storage + .put(key, Bytes::copy_from_slice(value.as_bytes())) + .await + .unwrap_or_else(|error| panic!("put {key} failed: {error}")); + } + + fn assert_not_found(result: Result, key: &str, context: &str) { + match result { + Ok(bytes) => panic!( + "{context}: get {key} returned Ok with {} bytes, expected NotFound", + bytes.len() + ), + Err(StorageError::NotFound { key: found }) => { + assert_eq!(found, key, "{context}: NotFound names the wrong key"); + }, + Err(error) => panic!("{context}: get {key} returned unexpected error: {error}"), + } + } + + pub(crate) async fn put_then_get_roundtrip(storage: &ReplicaStorage) { + let key = "roundtrip/hello.txt"; + let value = Bytes::from_static(b"hello replica"); + storage.put(key, value.clone()).await.expect("put failed"); + let got = storage.get(key).await.expect("get failed"); + assert_eq!(got, value, "roundtrip value mismatch"); + } + + pub(crate) async fn get_missing_is_not_found_error(storage: &ReplicaStorage) { + let result = storage.get("missing/never-written.txt").await; + assert_not_found( + result, + "missing/never-written.txt", + "get_missing_is_not_found_error", + ); + } + + pub(crate) async fn put_overwrite_replaces(storage: &ReplicaStorage) { + let key = "overwrite/value.txt"; + put_str(storage, key, "first").await; + put_str(storage, key, "second, longer than the first").await; + let got = storage.get(key).await.expect("get failed"); + assert_eq!( + got.as_ref(), + b"second, longer than the first".as_slice(), + "overwrite did not replace the value" + ); + } + + pub(crate) async fn list_prefix_returns_only_prefix(storage: &ReplicaStorage) { + put_str(storage, "gen/a/1", "1").await; + put_str(storage, "gen/a/2", "2").await; + put_str(storage, "gen/b/1", "3").await; + put_str(storage, "other/x", "4").await; + let keys = storage.list("gen/a/").await.expect("list failed"); + assert_eq!( + keys, + vec!["gen/a/1".to_owned(), "gen/a/2".to_owned()], + "list gen/a/ returned keys outside the prefix" + ); + let keys = storage.list("gen/").await.expect("list failed"); + assert_eq!( + keys, + vec![ + "gen/a/1".to_owned(), + "gen/a/2".to_owned(), + "gen/b/1".to_owned() + ], + "list gen/ returned keys outside the prefix" + ); + } + + pub(crate) async fn list_ordering_lexicographic(storage: &ReplicaStorage) { + // Insert in scrambled order; list must sort lexicographically, + // including "a" < "aa" < "b" and descent into nested keys. + put_str(storage, "seg/b", "1").await; + put_str(storage, "seg/aa", "2").await; + put_str(storage, "seg/c/d", "3").await; + put_str(storage, "seg/a", "4").await; + let keys = storage.list("seg/").await.expect("list failed"); + assert_eq!( + keys, + vec![ + "seg/a".to_owned(), + "seg/aa".to_owned(), + "seg/b".to_owned(), + "seg/c/d".to_owned() + ], + "list is not sorted lexicographically" + ); + } + + pub(crate) async fn list_empty_prefix_ok_empty_vec(storage: &ReplicaStorage) { + let keys = storage + .list("") + .await + .expect("list of empty storage failed"); + assert_eq!(keys, Vec::::new(), "empty storage must list empty"); + let keys = storage + .list("no/such/prefix/") + .await + .expect("list of missing prefix failed"); + assert_eq!(keys, Vec::::new(), "missing prefix must list empty"); + } + + pub(crate) async fn list_dirs_immediate_components_sorted(storage: &ReplicaStorage) { + put_str( + storage, + "generations/20260710T145900Z-3f8a2c1d/snapshot.json", + "a", + ) + .await; + put_str( + storage, + "generations/20260101T000000Z-0b1c2d3e/snapshot.json", + "b", + ) + .await; + put_str( + storage, + "generations/20260315T120000Z-99aabbcc/wal/0000000000-9d2f1c4a8b3e6f70/x", + "c", + ) + .await; + // A plain object at the listed level must not appear as a directory. + put_str(storage, "rootfile", "d").await; + + let expected = vec![ + "20260101T000000Z-0b1c2d3e".to_owned(), + "20260315T120000Z-99aabbcc".to_owned(), + "20260710T145900Z-3f8a2c1d".to_owned(), + ]; + let dirs = storage + .list_dirs("generations/") + .await + .expect("list_dirs failed"); + assert_eq!(dirs, expected, "list_dirs returned wrong components"); + // A prefix without a trailing slash behaves the same. + let dirs = storage + .list_dirs("generations") + .await + .expect("list_dirs without trailing slash failed"); + assert_eq!( + dirs, expected, + "list_dirs must normalize the trailing slash" + ); + + let root_dirs = storage.list_dirs("").await.expect("list_dirs root failed"); + assert_eq!( + root_dirs, + vec!["generations".to_owned()], + "root list_dirs must contain only directories, not objects" + ); + } + + pub(crate) async fn delete_then_get_not_found(storage: &ReplicaStorage) { + let key = "delete/me.txt"; + put_str(storage, key, "ephemeral").await; + storage.delete(key).await.expect("delete failed"); + assert_not_found(storage.get(key).await, key, "delete_then_get_not_found"); + } + + pub(crate) async fn delete_missing_idempotent_ok(storage: &ReplicaStorage) { + storage + .delete("never/existed.txt") + .await + .expect("delete of missing key must be Ok"); + storage + .delete("never/existed.txt") + .await + .expect("repeated delete of missing key must be Ok"); + } + + pub(crate) async fn delete_prefix_removes_all(storage: &ReplicaStorage) { + put_str(storage, "gens/old/snapshot.db.zst", "s").await; + put_str(storage, "gens/old/wal/0/seg1", "w1").await; + put_str(storage, "gens/old/wal/1/seg2", "w2").await; + put_str(storage, "gens/new/snapshot.db.zst", "keep").await; + storage + .delete_prefix("gens/old/") + .await + .expect("delete_prefix failed"); + let keys = storage + .list("") + .await + .expect("list after delete_prefix failed"); + assert_eq!( + keys, + vec!["gens/new/snapshot.db.zst".to_owned()], + "delete_prefix must remove exactly the prefixed keys" + ); + storage + .delete_prefix("gens/missing/") + .await + .expect("delete_prefix of missing prefix must be Ok"); + } + + pub(crate) async fn multipart_roundtrip_large_object(storage: &ReplicaStorage) { + // ~12 MiB plus an unaligned tail: exercises the S3 backend's 5 MiB + // minimum-part buffering across several parts. + let data = pattern(12 * MIB + 513); + let key = "snapshots/large/snapshot.db.zst"; + let mut upload = storage.start_multipart(key).await.expect("start failed"); + let part_sizes = [MIB, 6 * MIB, 4 * MIB, MIB, 513]; + assert_eq!( + part_sizes.iter().sum::(), + data.len(), + "part sizes must cover the payload" + ); + let mut remaining = Bytes::from(data.clone()); + for size in part_sizes { + let part = remaining.split_to(size); + upload.write_part(part).await.expect("write_part failed"); + } + upload.finish().await.expect("finish failed"); + + let got = storage.get(key).await.expect("get of large object failed"); + assert_eq!(got.len(), data.len(), "large object length mismatch"); + assert!( + got.as_ref() == data.as_slice(), + "large object content mismatch" + ); + let keys = storage.list("snapshots/").await.expect("list failed"); + assert_eq!(keys, vec![key.to_owned()], "large object missing from list"); + } + + pub(crate) async fn multipart_abort_leaves_nothing(storage: &ReplicaStorage) { + let key = "snapshots/aborted/snapshot.db.zst"; + let mut upload = storage.start_multipart(key).await.expect("start failed"); + upload + .write_part(Bytes::from(pattern(MIB))) + .await + .expect("write_part failed"); + upload.abort().await.expect("abort failed"); + assert_not_found( + storage.get(key).await, + key, + "multipart_abort_leaves_nothing", + ); + let keys = storage.list("").await.expect("list after abort failed"); + assert_eq!(keys, Vec::::new(), "abort must leave no objects"); + } + + pub(crate) async fn multipart_unfinished_invisible(storage: &ReplicaStorage) { + let key = "snapshots/unfinished/snapshot.db.zst"; + let mut upload = storage.start_multipart(key).await.expect("start failed"); + upload + .write_part(Bytes::from(pattern(MIB))) + .await + .expect("write_part failed"); + // Drop without finish: the object must never appear under its final + // key (local: only a partial temp file; S3: uncompleted upload). + drop(upload); + assert_not_found( + storage.get(key).await, + key, + "multipart_unfinished_invisible", + ); + let keys = storage.list("").await.expect("list after drop failed"); + assert_eq!( + keys, + Vec::::new(), + "an unfinished multipart upload must not be listable" + ); + } + + pub(crate) async fn key_charset_roundtrip(storage: &ReplicaStorage) { + // A real replica key: generation id, epoch dir with salts, and a + // zero-padded offset range segment name. + let key = "generations/20260710T145900Z-3f8a2c1d/wal/0000000000-9d2f1c4a8b3e6f70/00000000000000000000-00000000000000524320.wal.zst"; + let value = Bytes::from(pattern(1024)); + storage.put(key, value.clone()).await.expect("put failed"); + let got = storage.get(key).await.expect("get failed"); + assert_eq!(got, value, "key charset roundtrip value mismatch"); + let keys = storage.list("generations/").await.expect("list failed"); + assert_eq!( + keys, + vec![key.to_owned()], + "key charset key missing from list" + ); + } + + pub(crate) async fn concurrent_puts_distinct_keys(storage: &ReplicaStorage) { + let puts = (0..8) + .map(|index| { + let key = format!("concurrent/{index:02}.txt"); + let value = Bytes::from(format!("value {index}")); + async move { + storage + .put(&key, value) + .await + .expect("concurrent put failed"); + key + } + }) + .collect::>(); + let mut expected = futures::future::join_all(puts).await; + expected.sort(); + let keys = storage.list("concurrent/").await.expect("list failed"); + assert_eq!(keys, expected, "concurrent puts must all be visible"); + for (index, key) in keys.iter().enumerate() { + let got = storage.get(key).await.expect("get failed"); + assert_eq!( + got, + Bytes::from(format!("value {index}")), + "concurrent put content mismatch" + ); + } + } + + pub(crate) async fn get_stream_roundtrip(storage: &ReplicaStorage) { + let data = pattern(3 * MIB + 7); + let key = "stream/blob.bin"; + storage + .put(key, Bytes::from(data.clone())) + .await + .expect("put failed"); + let mut stream = storage.get_stream(key).await.expect("get_stream failed"); + let mut got = Vec::new(); + stream + .read_to_end(&mut got) + .await + .expect("reading stream failed"); + assert_eq!(got.len(), data.len(), "streamed length mismatch"); + assert!(got == data, "streamed content mismatch"); + + match storage.get_stream("stream/missing.bin").await { + Ok(_) => panic!("get_stream of a missing key must be NotFound"), + Err(StorageError::NotFound { key: found }) => { + assert_eq!(found, "stream/missing.bin", "NotFound names the wrong key"); + }, + Err(error) => panic!("get_stream returned unexpected error: {error}"), + } + } +} + +/// Shared backend constructors for the suites in this file. +#[cfg(test)] +pub(crate) mod harness { + use bencher_replica::testing::{FailurePlan, FlakyStorage}; + use bencher_replica::{LocalStorage, ReplicaStorage}; + use camino::Utf8Path; + + pub(crate) fn local_storage(tmp: &tempfile::TempDir) -> ReplicaStorage { + let root = Utf8Path::from_path(tmp.path()) + .expect("tempdir path is UTF-8") + .to_path_buf(); + ReplicaStorage::Local(LocalStorage::new(root)) + } + + pub(crate) fn flaky_storage(tmp: &tempfile::TempDir, plan: FailurePlan) -> ReplicaStorage { + ReplicaStorage::Flaky(Box::new(FlakyStorage::new(local_storage(tmp), plan))) + } +} + +/// The contract suite against the local-filesystem backend in a tempdir. +#[cfg(test)] +mod local_backend { + macro_rules! local_case { + ($case:ident) => { + #[tokio::test] + async fn $case() { + let tmp = tempfile::tempdir().expect("tempdir failed"); + let storage = super::harness::local_storage(&tmp); + super::cases::$case(&storage).await; + } + }; + } + for_each_contract_case!(local_case); +} + +/// The contract suite against `Flaky(Local)` with an empty failure plan: +/// proves the fault-injection wrapper is transparent when no rule fires. +#[cfg(test)] +mod flaky_passthrough { + macro_rules! flaky_case { + ($case:ident) => { + #[tokio::test] + async fn $case() { + let tmp = tempfile::tempdir().expect("tempdir failed"); + let storage = super::harness::flaky_storage( + &tmp, + bencher_replica::testing::FailurePlan::new(), + ); + super::cases::$case(&storage).await; + } + }; + } + for_each_contract_case!(flaky_case); +} + +/// Flaky-only contract cases: these require a scripted failure plan, so they +/// cannot run against a plain backend. +#[cfg(test)] +mod flaky_only { + use bencher_replica::testing::{FailurePlan, OpKind}; + use bencher_replica::{ReplicaStorage, StorageError}; + use bytes::Bytes; + use pretty_assertions::assert_eq; + + /// The load-bearing case: a backend failure during `list` MUST surface + /// as an `Err`, never as `Ok(vec![])`. Conflating the two would make an + /// unreachable replica look empty and trigger a spurious new generation. + #[tokio::test] + async fn list_error_is_error_not_empty() { + let tmp = tempfile::tempdir().expect("tempdir failed"); + let storage = + super::harness::flaky_storage(&tmp, FailurePlan::new().fail_all(OpKind::List)); + storage + .put("generations/g1/snapshot.json", Bytes::from_static(b"{}")) + .await + .expect("put must pass through: only List is failed"); + + match storage.list("").await { + Ok(keys) => panic!( + "list during an outage returned Ok({keys:?}): an unreachable \ + replica must never look empty" + ), + Err(StorageError::Injected { op, .. }) => { + assert_eq!(op, "list", "injected error names the wrong op"); + }, + Err(error) => panic!("list returned unexpected error: {error}"), + } + + // After healing, the object put before the outage is visible again. + let ReplicaStorage::Flaky(flaky) = &storage else { + panic!("flaky_storage must build the Flaky variant"); + }; + flaky.heal(); + let keys = storage.list("").await.expect("list after heal failed"); + assert_eq!( + keys, + vec!["generations/g1/snapshot.json".to_owned()], + "healed list must show the pre-outage object" + ); + } +} + +/// Local-only behaviors that diverge from S3 by design (documented in +/// `src/storage.rs`). Pinned here for the local backend so they cannot change +/// silently; the engine only ever uses directory-aligned prefixes and never +/// relies on the divergent cases. +#[cfg(test)] +mod local_divergences { + use bencher_replica::StorageError; + use bytes::Bytes; + use pretty_assertions::assert_eq; + + /// Local `get` of a key that is actually a directory errors (`EISDIR`), + /// where S3 (no directories) would return `NotFound`. + #[tokio::test] + async fn get_of_directory_key_errors() { + let tmp = tempfile::tempdir().expect("tempdir failed"); + let storage = super::harness::local_storage(&tmp); + storage + .put("dir/child", Bytes::from_static(b"x")) + .await + .expect("put failed"); + match storage.get("dir").await { + Err(StorageError::Local(_)) => {}, + other => panic!("local get of a directory key must be a local error, got: {other:?}"), + } + } + + /// Local `list` errors when a prefix names a path whose component is a + /// regular file, where S3 returns an empty listing. + #[tokio::test] + async fn list_prefix_through_a_file_errors() { + let tmp = tempfile::tempdir().expect("tempdir failed"); + let storage = super::harness::local_storage(&tmp); + storage + .put("a", Bytes::from_static(b"x")) + .await + .expect("put failed"); + match storage.list("a/").await { + Err(StorageError::Local(_)) => {}, + other => { + panic!("local list through a regular file must be a local error, got: {other:?}") + }, + } + } + + /// Local `delete_prefix` treats its argument as a directory path, so a + /// prefix that ends mid-component matches no directory and no-ops, where + /// S3 would delete every key with that raw string prefix. + #[tokio::test] + async fn delete_prefix_mid_component_noops() { + let tmp = tempfile::tempdir().expect("tempdir failed"); + let storage = super::harness::local_storage(&tmp); + storage + .put("generations/g/x", Bytes::from_static(b"keep")) + .await + .expect("put failed"); + // "gen" is not a directory-aligned prefix of "generations/...". + storage + .delete_prefix("gen") + .await + .expect("mid-component delete_prefix must be Ok"); + assert_eq!( + storage.list("").await.expect("list failed"), + vec!["generations/g/x".to_owned()], + "a mid-component prefix must delete nothing on the local backend" + ); + } +} diff --git a/plus/bencher_replica/tests/sync_engine.rs b/plus/bencher_replica/tests/sync_engine.rs new file mode 100644 index 000000000..c559d1f2c --- /dev/null +++ b/plus/bencher_replica/tests/sync_engine.rs @@ -0,0 +1,1751 @@ +#![cfg(all(feature = "plus", feature = "testing"))] +//! Sync engine suite: shipping, resume, backoff, snapshots, and pruning. +//! +//! Every test drives the step-driven engine core directly (`sync_once`, +//! `ship_once`, `checkpoint_once`, `snapshot_step`, `prune_once`) against a +//! real `SQLite` database (`WalFixture`) and a fault-injectable replica +//! (`FlakyStorage` over `LocalStorage`), with all scheduling under an +//! injected `Clock`. +//! +//! NOTE: `unused_crate_dependencies` cannot be handled with a crate-level +//! `#![expect]` here (see `tests/storage_contract.rs`); unused package +//! dependencies are referenced explicitly instead, as rustc recommends. + +use async_compression as _; +use aws_credential_types as _; +use aws_sdk_s3 as _; +use futures as _; +use hex as _; +use rand as _; +use rusqlite as _; +use serde as _; +use serde_json as _; +use sha2 as _; +use thiserror as _; +use uuid as _; +use zstd as _; +// Optional dependency enabled by the otel feature; unused by tests. +#[cfg(feature = "otel")] +use bencher_otel as _; + +/// Shared fixtures: a scripted source database, a fault-injectable replica, +/// and an engine built over both with a deterministic clock. +#[cfg(test)] +pub(crate) mod harness { + use std::sync::Arc; + use std::sync::atomic::{AtomicI64, Ordering}; + + use bencher_json::system::config::{JsonReplication, ReplicationTarget}; + use bencher_json::{Clock, DateTime}; + use bencher_replica::testing::{ + FailurePlan, FlakyStorage, OpKind, OpOutcome, WalFixture, assert_replica_equivalent, + }; + use bencher_replica::{ + EngineState, LocalStorage, ReplicaConfig, ReplicaDb, ReplicaStorage, RestoreOutcome, + SyncEngine, restore_if_missing, + }; + use camino::{Utf8Path, Utf8PathBuf}; + + /// Page size for every fixture database in this suite. + pub(crate) const PAGE_SIZE: u32 = 4096; + /// 2026-07-10T14:59:00Z, the deterministic clock start. + pub(crate) const BASE_SECS: i64 = 1_783_695_540; + + pub(crate) fn dir_path(tmp: &tempfile::TempDir) -> &Utf8Path { + Utf8Path::from_path(tmp.path()).expect("tempdir path is UTF-8") + } + + pub(crate) fn logger() -> slog::Logger { + slog::Logger::root(slog::Discard, slog::o!()) + } + + pub(crate) fn clock_for(secs: &Arc) -> Clock { + let secs = Arc::clone(secs); + Clock::Custom(Arc::new(move || { + DateTime::try_from(secs.load(Ordering::SeqCst)).expect("valid clock seconds") + })) + } + + /// Flaky(Local) storage rooted at `root`, with an empty failure plan. + pub(crate) fn flaky_over(root: &Utf8Path) -> ReplicaStorage { + ReplicaStorage::Flaky(Box::new(FlakyStorage::new( + ReplicaStorage::Local(LocalStorage::new(root.to_path_buf())), + FailurePlan::new(), + ))) + } + + /// Parse the `[start, end)` byte range out of a segment object key. + pub(crate) fn segment_range(key: &str) -> (u64, u64) { + let (_, file) = key.rsplit_once('/').expect("segment key has a directory"); + let range = file.strip_suffix(".wal.zst").expect("segment key suffix"); + let (start, end) = range.split_once('-').expect("segment range separator"); + ( + start.parse().expect("segment start offset"), + end.parse().expect("segment end offset"), + ) + } + + pub(crate) struct Harness { + pub fixture: WalFixture, + pub engine: SyncEngine<()>, + pub config: ReplicaConfig, + pub db: ReplicaDb<()>, + pub clock_secs: Arc, + pub replica_root: Utf8PathBuf, + pub shadow: bool, + pub _fixture_tmp: tempfile::TempDir, + pub _replica_tmp: tempfile::TempDir, + } + + impl Harness { + pub(crate) async fn new() -> Self { + Self::with_config(|_| {}).await + } + + pub(crate) async fn with_config(overrides: F) -> Self + where + F: FnOnce(&mut JsonReplication), + { + Self::with_shadow_config(false, overrides).await + } + + pub(crate) async fn with_shadow_config(shadow: bool, overrides: F) -> Self + where + F: FnOnce(&mut JsonReplication), + { + let fixture_tmp = tempfile::tempdir().expect("fixture tempdir"); + let replica_tmp = tempfile::tempdir().expect("replica tempdir"); + let fixture = WalFixture::new(dir_path(&fixture_tmp), PAGE_SIZE).expect("fixture"); + let replica_root = dir_path(&replica_tmp).to_path_buf(); + let mut json = JsonReplication { + target: ReplicationTarget::File { + path: replica_root.clone().into_std_path_buf(), + }, + sync_interval_secs: None, + checkpoint_interval_secs: None, + min_checkpoint_pages: None, + snapshot_interval_secs: None, + snapshot_throttle_mib: None, + retention_generations: None, + verification_interval_secs: None, + shutdown_sync_timeout_secs: None, + }; + overrides(&mut json); + let config = ReplicaConfig::try_from(json).expect("config"); + let clock_secs = Arc::new(AtomicI64::new(BASE_SECS)); + let db = ReplicaDb { + db_path: fixture.db_path(), + writer: Arc::new(tokio::sync::Mutex::new(())), + busy_timeout_ms: 5000, + }; + let engine = SyncEngine::new_with_storage( + logger(), + config.clone(), + db.clone(), + clock_for(&clock_secs), + shadow, + flaky_over(&replica_root), + ) + .await + .expect("engine"); + Self { + fixture, + engine, + config, + db, + clock_secs, + replica_root, + shadow, + _fixture_tmp: fixture_tmp, + _replica_tmp: replica_tmp, + } + } + + /// The fault-injection wrapper the engine was built with. + pub(crate) fn flaky(&self) -> &FlakyStorage { + if let ReplicaStorage::Flaky(flaky) = self.engine.storage() { + flaky + } else { + panic!("harness engine always wraps Flaky storage") + } + } + + /// Advance the injected clock by `secs`. + pub(crate) fn advance(&self, secs: i64) { + self.clock_secs.fetch_add(secs, Ordering::SeqCst); + } + + /// Number of put operations attempted so far (any outcome). + pub(crate) fn put_attempts(&self) -> usize { + self.flaky() + .journal() + .iter() + .filter(|(op, _)| op.kind == OpKind::Put) + .count() + } + + /// Segment object keys successfully shipped so far, in journal order. + pub(crate) fn shipped_segment_keys(&self) -> Vec { + self.flaky() + .journal() + .iter() + .filter(|(op, outcome)| { + op.kind == OpKind::Put + && *outcome == OpOutcome::Ok + && op.key.ends_with(".wal.zst") + }) + .map(|(op, _)| op.key.clone()) + .collect() + } + + /// Drive `sync_once` until the engine is streaming (the fresh-replica + /// bootstrap snapshot has completed). + pub(crate) async fn until_streaming(&mut self) { + for _ in 0..64 { + if self.engine.state() == EngineState::Streaming { + return; + } + self.engine + .sync_once() + .await + .expect("sync_once during startup"); + } + panic!( + "engine never reached Streaming; state: {:?}", + self.engine.state() + ); + } + + /// Bootstrap plus one sync tick, so the initial WAL backlog is + /// shipped and the engine is quiescent. + pub(crate) async fn ready(&mut self) { + self.until_streaming().await; + self.engine.sync_once().await.expect("backlog sync"); + } + + /// Replace the engine (dropping the old one) with a freshly resumed + /// engine over the same replica; the flaky journal starts empty. + pub(crate) async fn rebuild_engine(&mut self) { + self.engine = SyncEngine::new_with_storage( + logger(), + self.config.clone(), + self.db.clone(), + clock_for(&self.clock_secs), + self.shadow, + flaky_over(&self.replica_root), + ) + .await + .expect("engine rebuild"); + } + + /// Restore the replica into a scratch directory and assert logical + /// equivalence with the live source database. + pub(crate) async fn assert_restore_equivalent(&self) { + let target_tmp = tempfile::tempdir().expect("restore target tempdir"); + let target_db = dir_path(&target_tmp).join("restored.db"); + let outcome = restore_if_missing(&logger(), &self.config, &target_db) + .await + .expect("restore"); + assert!( + matches!(outcome, RestoreOutcome::Restored { .. }), + "expected Restored, got {outcome:?}" + ); + assert_replica_equivalent(&self.fixture.db_path(), &target_db); + } + } +} + +#[cfg(test)] +mod cases { + use std::io::Cursor; + use std::time::Duration; + + use bencher_json::DateTime; + use bencher_replica::testing::{CheckpointMode, FailurePlan, OpKind, WriteProbe}; + use bencher_replica::{ + CheckpointOutcome, EngineState, GenerationId, Replicator, SnapshotStatus, VerifyReport, + WalScanner, decompress_segment, + }; + use bytes::Bytes; + use pretty_assertions::assert_eq; + + use super::harness::{BASE_SECS, Harness, PAGE_SIZE, clock_for, logger, segment_range}; + + /// Read the live WAL header salts of the fixture. + fn wal_salt(harness: &Harness) -> (u32, u32) { + let wal = harness.fixture.wal_bytes().expect("wal bytes"); + WalScanner::open(Cursor::new(wal)) + .expect("wal header") + .expect("wal is not empty") + .header() + .salt + } + + /// Count objects on the replica whose key ends with `suffix`. + async fn object_count(harness: &Harness, suffix: &str) -> usize { + harness + .engine + .storage() + .list("generations") + .await + .expect("list replica") + .iter() + .filter(|key| key.ends_with(suffix)) + .count() + } + + #[tokio::test] + async fn sync_ships_nothing_when_no_new_commits() { + let mut harness = Harness::new().await; + harness.ready().await; + harness + .fixture + .txn(&["INSERT INTO t (data) VALUES ('one')"]) + .expect("txn"); + let progress = harness.engine.sync_once().await.expect("sync"); + assert_eq!(progress.shipped_segments, 1, "the new commit ships"); + + let puts_before = harness.put_attempts(); + let progress = harness.engine.sync_once().await.expect("quiet sync"); + assert_eq!(progress.shipped_segments, 0, "nothing new to ship"); + assert_eq!( + harness.put_attempts(), + puts_before, + "a quiet tick issues no storage puts" + ); + } + + #[tokio::test] + async fn sync_ships_only_delta_frames() { + let mut harness = Harness::new().await; + harness.ready().await; + harness + .fixture + .txn(&["INSERT INTO t (data) VALUES ('a')"]) + .expect("txn a"); + harness.engine.sync_once().await.expect("sync a"); + let end_a = harness.engine.position().expect("position").offset; + harness + .fixture + .txn(&["INSERT INTO t (data) VALUES ('b')"]) + .expect("txn b"); + harness.engine.sync_once().await.expect("sync b"); + let end_b = harness.engine.position().expect("position").offset; + assert!(end_b > end_a, "position advances: {end_a} -> {end_b}"); + + let keys = harness.shipped_segment_keys(); + assert!(keys.len() >= 2, "at least two segments shipped: {keys:?}"); + let (start_b, key_end_b) = segment_range(keys.last().expect("segment b")); + let (_, key_end_a) = segment_range(keys.get(keys.len() - 2).expect("segment a")); + assert_eq!( + start_b, key_end_a, + "the second segment starts exactly at the first segment's end" + ); + assert_eq!(key_end_a, end_a, "segment a covers through position a"); + assert_eq!(key_end_b, end_b, "segment b covers through position b"); + } + + #[tokio::test] + async fn first_segment_of_epoch_contains_wal_header() { + let mut harness = Harness::new().await; + harness.ready().await; + let keys = harness.shipped_segment_keys(); + let first = keys.first().expect("an epoch-0 segment shipped"); + let (start, end) = segment_range(first); + assert_eq!(start, 0, "the first segment of an epoch starts at 0"); + + let compressed = harness + .engine + .storage() + .get(first) + .await + .expect("get first segment"); + let raw = decompress_segment(&compressed).expect("decompress"); + assert_eq!( + raw.len(), + usize::try_from(end).expect("segment end"), + "the stored segment covers [0, end)" + ); + // The stored bytes parse as a self-contained, checksum-valid WAL. + let mut scanner = WalScanner::open(Cursor::new(raw.clone())) + .expect("stored segment parses as a WAL") + .expect("stored segment has a header"); + let mut last_end = 0; + while let Some(chunk) = scanner.next_committed(u64::MAX).expect("scan stored WAL") { + last_end = chunk.end_offset; + } + assert_eq!( + last_end, + u64::try_from(raw.len()).expect("raw len"), + "every stored frame verifies against the embedded header" + ); + } + + #[tokio::test] + async fn restart_resumes_from_replica_list() { + let mut harness = Harness::new().await; + harness.ready().await; + harness + .fixture + .txn(&["INSERT INTO t (data) VALUES ('pre-restart')"]) + .expect("txn"); + harness.engine.sync_once().await.expect("sync"); + let position = harness.engine.position().cloned().expect("position"); + + harness.rebuild_engine().await; + assert_eq!( + harness.engine.position(), + Some(&position), + "resume from the replica LIST reproduces the exact position (offset and checksum)" + ); + // A no-write sync re-ships nothing. + let progress = harness.engine.sync_once().await.expect("quiet sync"); + assert_eq!(progress.shipped_segments, 0, "no re-ship after restart"); + assert_eq!(harness.put_attempts(), 0, "no puts after restart"); + + // A new write ships exactly the gap. + harness + .fixture + .txn(&["INSERT INTO t (data) VALUES ('post-restart')"]) + .expect("txn"); + harness.engine.sync_once().await.expect("sync delta"); + let keys = harness.shipped_segment_keys(); + assert_eq!(keys.len(), 1, "exactly one delta segment: {keys:?}"); + let (start, _) = segment_range(keys.first().expect("delta segment")); + assert_eq!(start, position.offset, "the delta starts at the old end"); + harness.assert_restore_equivalent().await; + } + + #[tokio::test] + async fn restart_after_checkpoint_resumes_as_next_epoch() { + let mut harness = Harness::new().await; + harness.ready().await; + harness + .fixture + .txn(&["INSERT INTO t (data) VALUES ('epoch-0')"]) + .expect("txn"); + harness.engine.sync_once().await.expect("sync"); + let epoch_before = harness.engine.position().expect("position").epoch; + let outcome = harness.engine.checkpoint_once().await.expect("checkpoint"); + assert_eq!(outcome, CheckpointOutcome::Completed, "full backfill"); + + // The WAL is fully backfilled, so this write restarts it (new salts). + harness + .fixture + .txn(&["INSERT INTO t (data) VALUES ('epoch-1')"]) + .expect("txn after checkpoint"); + + harness.rebuild_engine().await; + let position = harness.engine.position().expect("resumed position"); + assert_eq!( + position.epoch, + epoch_before + 1, + "meta-verified resume continues as the next epoch, no re-snapshot" + ); + assert_eq!(position.offset, 0, "the new epoch starts at offset 0"); + assert_eq!( + position.salt, + wal_salt(&harness), + "the new epoch binds the local WAL header salts" + ); + harness.engine.sync_once().await.expect("ship new epoch"); + harness.assert_restore_equivalent().await; + } + + #[tokio::test] + async fn restart_with_unshipped_frames_lost_to_reset_forces_new_generation() { + let mut harness = Harness::new().await; + harness.ready().await; + harness + .fixture + .txn(&["INSERT INTO t (data) VALUES ('shipped')"]) + .expect("txn a"); + harness.engine.sync_once().await.expect("sync"); + let old_generation = harness.engine.generation().cloned().expect("generation"); + + // Unshipped commit, then a stray TRUNCATE checkpoint resets the WAL: + // the unshipped frames are gone from the WAL (they live only in the + // db file now) and the next write starts a new salt cycle. + harness + .fixture + .txn(&["INSERT INTO t (data) VALUES ('unshipped')"]) + .expect("txn b"); + harness + .fixture + .checkpoint(CheckpointMode::Truncate) + .expect("stray TRUNCATE"); + harness + .fixture + .txn(&["INSERT INTO t (data) VALUES ('new-cycle')"]) + .expect("txn c"); + + harness.rebuild_engine().await; + assert_eq!( + harness.engine.state(), + EngineState::PendingSnapshot, + "salt change with unshipped loss resolves to a new generation, never silently" + ); + harness.ready().await; + assert!( + harness + .engine + .generation() + .is_some_and(|generation| { generation != &old_generation }), + "a new generation was created" + ); + harness.assert_restore_equivalent().await; + } + + #[tokio::test] + async fn stray_writer_frames_picked_up() { + let mut harness = Harness::new().await; + harness.ready().await; + let stray = harness.fixture.stray_conn().expect("stray conn"); + stray + .execute("INSERT INTO t (data) VALUES ('stray')", []) + .expect("stray insert"); + let progress = harness.engine.sync_once().await.expect("sync"); + assert!( + progress.shipped_segments >= 1, + "stray-connection commits replicate: {progress:?}" + ); + harness.assert_restore_equivalent().await; + } + + #[tokio::test] + async fn storage_outage_defers_checkpoint_and_catches_up() { + let mut harness = Harness::new().await; + harness.ready().await; + harness + .flaky() + .set_plan(FailurePlan::new().fail_all(OpKind::Put)); + + harness + .fixture + .txn(&["INSERT INTO t (data) VALUES ('outage-1')"]) + .expect("txn 1"); + let wal_len_1 = harness.fixture.wal_bytes().expect("wal").len(); + let progress = harness.engine.sync_once().await.expect("sync in outage"); + assert!( + progress.error.is_some(), + "the tick reports the storage error class: {progress:?}" + ); + let progress = harness.engine.sync_once().await.expect("gated sync"); + assert!( + progress.backing_off, + "the next tick is gated by backoff: {progress:?}" + ); + + harness + .fixture + .txn(&["INSERT INTO t (data) VALUES ('outage-2')"]) + .expect("txn 2"); + let wal_len_2 = harness.fixture.wal_bytes().expect("wal").len(); + assert!( + wal_len_2 > wal_len_1, + "the WAL grows monotonically during the outage ({wal_len_1} -> {wal_len_2})" + ); + let outcome = harness.engine.checkpoint_once().await.expect("checkpoint"); + assert_eq!( + outcome, + CheckpointOutcome::SkippedUnshipped, + "no checkpoint while unshipped frames exist (I1): the WAL is the buffer" + ); + harness + .fixture + .txn(&["INSERT INTO t (data) VALUES ('outage-3')"]) + .expect("txn 3"); + let wal_len_3 = harness.fixture.wal_bytes().expect("wal").len(); + assert!(wal_len_3 > wal_len_2, "the WAL keeps growing"); + + // Heal, wait out the backoff, and catch up in order. + harness.flaky().heal(); + harness.advance(400); + let progress = harness.engine.sync_once().await.expect("catch-up sync"); + assert!(progress.error.is_none(), "healed tick: {progress:?}"); + assert!( + progress.shipped_segments >= 1, + "the backlog ships after healing: {progress:?}" + ); + let keys = harness.shipped_segment_keys(); + let mut expected_start = None; + for key in &keys { + let (start, end) = segment_range(key); + if let Some(expected) = expected_start { + assert_eq!(start, expected, "segments ship in order: {keys:?}"); + } + expected_start = Some(end); + } + harness.assert_restore_equivalent().await; + } + + #[tokio::test] + async fn backoff_delays_follow_sequence() { + let mut harness = Harness::new().await; + harness.ready().await; + harness + .flaky() + .set_plan(FailurePlan::new().fail_all(OpKind::Put)); + harness + .fixture + .txn(&["INSERT INTO t (data) VALUES ('backoff')"]) + .expect("txn"); + + // The bootstrap already issued puts; count attempts from here. + let base = harness.put_attempts(); + + // Attempt at t0 fails and arms a 1s delay. + harness.engine.sync_once().await.expect("sync t0"); + assert_eq!(harness.put_attempts(), base + 1, "first attempt at t0"); + harness.engine.sync_once().await.expect("sync t0 again"); + assert_eq!(harness.put_attempts(), base + 1, "gated before t0+1"); + + harness.advance(1); // t0+1 + harness.engine.sync_once().await.expect("sync t0+1"); + assert_eq!(harness.put_attempts(), base + 2, "second attempt at t0+1"); + + harness.advance(1); // t0+2 + harness.engine.sync_once().await.expect("sync t0+2"); + assert_eq!( + harness.put_attempts(), + base + 2, + "gated before t0+3 (2s delay)" + ); + + harness.advance(1); // t0+3 + harness.engine.sync_once().await.expect("sync t0+3"); + assert_eq!(harness.put_attempts(), base + 3, "third attempt at t0+3"); + + harness.advance(3); // t0+6 + harness.engine.sync_once().await.expect("sync t0+6"); + assert_eq!( + harness.put_attempts(), + base + 3, + "gated before t0+7 (4s delay)" + ); + + harness.advance(1); // t0+7 + harness.engine.sync_once().await.expect("sync t0+7"); + assert_eq!(harness.put_attempts(), base + 4, "fourth attempt at t0+7"); + } + + #[tokio::test] + async fn snapshot_end_to_end_restore_equivalent() { + let mut harness = Harness::new().await; + harness.ready().await; + harness + .fixture + .txn(&["INSERT INTO t (data) VALUES ('before-snapshot')"]) + .expect("txn"); + harness.engine.sync_once().await.expect("sync"); + let old_generation = harness.engine.generation().cloned().expect("generation"); + + harness.engine.trigger_snapshot(); + for step in 0..1000 { + let status = harness.engine.snapshot_step().await.expect("snapshot step"); + if status == SnapshotStatus::Finished { + break; + } + assert!(step < 999, "snapshot never finished"); + } + let new_generation = harness.engine.generation().cloned().expect("generation"); + assert!( + new_generation != old_generation, + "a snapshot always creates a new generation" + ); + assert_eq!( + harness.engine.state(), + EngineState::Streaming, + "back to streaming after the snapshot" + ); + // The accumulated WAL backlog re-ships into the new generation. + harness + .engine + .sync_once() + .await + .expect("post-snapshot sync"); + harness.assert_restore_equivalent().await; + } + + #[tokio::test] + async fn snapshot_writes_proceed_between_steps() { + // 1 MiB copy budget per step forces a multi-step copy. + let mut harness = Harness::with_config(|json| { + json.snapshot_throttle_mib = Some(1); + }) + .await; + harness.ready().await; + // Grow the database FILE (not just the WAL): big transaction, ship, + // full checkpoint. + harness + .fixture + .txn_touching_pages(800) + .expect("big transaction"); + harness.engine.ship_once().await.expect("ship big txn"); + let outcome = harness.engine.checkpoint_once().await.expect("checkpoint"); + assert_eq!( + outcome, + CheckpointOutcome::Completed, + "backfill the db file" + ); + + let probe = WriteProbe::new(&harness.fixture.db_path(), Duration::from_secs(5)); + harness.engine.trigger_snapshot(); + let mut steps = 0u32; + loop { + let status = harness.engine.snapshot_step().await.expect("snapshot step"); + steps += 1; + // The anti-Litestream regression: a real second connection + // writes IMMEDIATELY between snapshot steps; the snapshot holds + // no locks of any kind. + let result = probe.write_once().await; + result + .result + .expect("write between snapshot steps succeeds"); + assert!( + result.blocked < Duration::from_secs(2), + "write between snapshot steps must not block: {:?}", + result.blocked + ); + if status == SnapshotStatus::Finished { + break; + } + assert!(steps < 1000, "snapshot never finished"); + } + assert!( + steps > 4, + "the throttled copy must span multiple steps, got {steps}" + ); + // The probe writes accumulated in the WAL ship into the new + // generation and survive a restore. + harness + .engine + .sync_once() + .await + .expect("post-snapshot sync"); + harness.assert_restore_equivalent().await; + } + + /// An external checkpointer (even a TRUNCATE checkpoint after a stray + /// write) mid-snapshot must NOT tear or abort the snapshot: the body + /// comes from a single-step online backup taken through the pager, so + /// it is transactionally consistent no matter who checkpoints. This is + /// what makes shadow mode (Litestream checkpointing freely) viable. + #[tokio::test] + async fn snapshot_survives_external_checkpoint_mid_copy() { + let mut harness = Harness::with_config(|json| { + json.snapshot_throttle_mib = Some(1); + }) + .await; + harness.ready().await; + harness + .fixture + .txn_touching_pages(800) + .expect("big transaction"); + harness.engine.ship_once().await.expect("ship big txn"); + let outcome = harness.engine.checkpoint_once().await.expect("checkpoint"); + assert_eq!( + outcome, + CheckpointOutcome::Completed, + "backfill the db file" + ); + assert_eq!( + object_count(&harness, "snapshot.json").await, + 1, + "bootstrap" + ); + + // Enter the multi-step Copying phase (the backup already happened + // at CreateGeneration; the steps upload the scratch file). + harness.engine.trigger_snapshot(); + for _ in 0..3 { + let status = harness.engine.snapshot_step().await.expect("snapshot step"); + assert_eq!(status, SnapshotStatus::InProgress, "copy still running"); + } + // An external checkpointer mutates the LIVE db file mid-upload: a + // stray write that grows the database, then a TRUNCATE checkpoint. + // The scratch upload is unaffected. + let stray = harness.fixture.stray_conn().expect("stray conn"); + let big = "x".repeat(64 * 1024); + stray + .execute("INSERT INTO t (data) VALUES (?1)", [&big]) + .expect("stray insert"); + harness + .fixture + .checkpoint(CheckpointMode::Truncate) + .expect("external checkpoint"); + + // The snapshot completes despite the external churn. + for _ in 0..256 { + if harness.engine.state() != EngineState::Snapshotting { + break; + } + harness.engine.snapshot_step().await.expect("snapshot step"); + } + assert_eq!( + object_count(&harness, "snapshot.json").await, + 2, + "the snapshot committed" + ); + // The stray write (which the WAL restart could have buried) is + // recaptured: divergence handling plus the fresh snapshot keep the + // replica equivalent end to end. + for _ in 0..64 { + let progress = harness.engine.sync_once().await.expect("recovery sync"); + if progress.shipped_segments == 0 + && harness.engine.state() == EngineState::Streaming + && !progress.backing_off + { + break; + } + harness.advance(1); + } + if harness.engine.state() == EngineState::Snapshotting + || object_count(&harness, "snapshot.json").await > 2 + { + // A divergence-triggered replacement snapshot may be running; + // drive it home. + for _ in 0..256 { + if harness.engine.state() == EngineState::Streaming { + break; + } + harness.engine.sync_once().await.expect("replacement sync"); + } + } + harness.engine.sync_once().await.expect("final drain"); + harness.assert_restore_equivalent().await; + } + + #[tokio::test] + async fn prune_keeps_newest_n_and_reaps_stale_incomplete() { + let mut harness = Harness::with_config(|json| { + json.retention_generations = Some(2); + }) + .await; + harness.until_streaming().await; + let current = harness.engine.generation().cloned().expect("generation"); + + let old = |age_secs: i64, suffix: u32| { + GenerationId::new( + DateTime::try_from(BASE_SECS - age_secs).expect("timestamp"), + suffix, + ) + }; + let complete_old = old(200_000, 1); + let complete_mid = old(150_000, 2); + let stale_incomplete = old(100_000, 3); // > 24h without snapshot.json + let fresh_incomplete = old(3_600, 4); // 1h old, still in flight + let storage = harness.engine.storage(); + for generation in [&complete_old, &complete_mid] { + storage + .put( + &format!("generations/{}/snapshot.db.zst", generation.as_str()), + Bytes::from_static(b"body"), + ) + .await + .expect("put body"); + storage + .put( + &format!("generations/{}/snapshot.json", generation.as_str()), + Bytes::from_static(b"{}"), + ) + .await + .expect("put marker"); + } + for generation in [&stale_incomplete, &fresh_incomplete] { + storage + .put( + &format!("generations/{}/snapshot.db.zst", generation.as_str()), + Bytes::from_static(b"body"), + ) + .await + .expect("put orphan body"); + } + + harness.engine.prune_once().await.expect("prune"); + let dirs = harness + .engine + .storage() + .list_dirs("generations") + .await + .expect("list generations"); + assert_eq!( + dirs, + vec![ + complete_mid.as_str().to_owned(), + fresh_incomplete.as_str().to_owned(), + current.as_str().to_owned(), + ], + "keep the newest 2 complete generations plus fresh incomplete ones" + ); + } + + #[tokio::test] + async fn new_generation_epoch_zero_rebinding() { + let mut harness = Harness::new().await; + harness.ready().await; + harness + .fixture + .txn(&["INSERT INTO t (data) VALUES ('pre')"]) + .expect("txn"); + harness.engine.sync_once().await.expect("sync"); + let outcome = harness.engine.checkpoint_once().await.expect("checkpoint"); + assert_eq!(outcome, CheckpointOutcome::Completed, "full backfill"); + + // Snapshot with a fully backfilled WAL and no writes in between: the + // boundary binds to the old salt cycle. + harness.engine.trigger_snapshot(); + for step in 0..1000 { + let status = harness.engine.snapshot_step().await.expect("snapshot step"); + if status == SnapshotStatus::Finished { + break; + } + assert!(step < 999, "snapshot never finished"); + } + let new_generation = harness.engine.generation().cloned().expect("generation"); + let boundary_salt = harness.engine.position().expect("position").salt; + + // A post-snapshot write restarts the WAL BEFORE any epoch-0 segment + // ships: epoch 0 must rebind to the new cycle, not leave an empty + // epoch behind. + harness + .fixture + .txn(&["INSERT INTO t (data) VALUES ('post-snapshot')"]) + .expect("txn after snapshot"); + let new_salt = wal_salt(&harness); + assert!( + new_salt != boundary_salt, + "the write after a full backfill restarts the WAL" + ); + let progress = harness.engine.sync_once().await.expect("sync"); + assert!( + progress.shipped_segments >= 1, + "the restarted cycle ships: {progress:?}" + ); + let position = harness.engine.position().expect("position"); + assert_eq!(position.epoch, 0, "epoch numbering stays contiguous from 0"); + assert_eq!(position.salt, new_salt, "epoch 0 rebound to the new salts"); + + // Exactly one epoch directory exists, named for epoch 0 with the + // NEW salts. + let prefix = format!("generations/{}/wal/", new_generation.as_str()); + let keys = harness + .engine + .storage() + .list(&prefix) + .await + .expect("list epoch dirs"); + let mut epoch_dirs: Vec = keys + .iter() + .filter_map(|key| { + key.strip_prefix(&prefix) + .and_then(|rest| rest.split_once('/')) + .map(|(dir, _)| dir.to_owned()) + }) + .collect(); + epoch_dirs.dedup(); + assert_eq!( + epoch_dirs, + vec![format!("{:010}-{:08x}{:08x}", 0, new_salt.0, new_salt.1)], + "one epoch dir: epoch 0 under the new salts" + ); + harness.assert_restore_equivalent().await; + } + + #[tokio::test] + async fn replicator_shutdown_ships_tail() { + // Pre-populate a valid generation with a manually driven engine. + let mut harness = Harness::new().await; + harness.ready().await; + harness + .fixture + .txn(&["INSERT INTO t (data) VALUES ('shipped')"]) + .expect("txn"); + harness.engine.sync_once().await.expect("sync"); + harness.rebuild_engine().await; // drop the active engine + + // An unshipped tail, then the production shell: start + shutdown. + // final_sync ships the tail without any tick having elapsed. + harness + .fixture + .txn(&["INSERT INTO t (data) VALUES ('tail')"]) + .expect("tail txn"); + let handle = Replicator::start( + logger(), + harness.config.clone(), + harness.db.clone(), + clock_for(&harness.clock_secs), + false, + ); + handle + .shutdown(Duration::from_secs(10)) + .await + .expect("shutdown"); + harness.assert_restore_equivalent().await; + } + + /// A post-power-loss rewind fork: the local WAL still reaches the + /// shipped offset with a VALID checksum chain, but the content differs + /// from what the replica stored (with `synchronous = NORMAL`, committed + /// frames can be lost from the page cache and re-written differently). + /// Resume must compare content against the replica tip, not just chain + /// length, and force a new generation. + #[tokio::test] + async fn resume_rewind_fork_forces_new_generation() { + let mut harness = Harness::new().await; + harness.ready().await; + harness + .fixture + .txn(&["INSERT INTO t (data) VALUES ('pre fork')"]) + .expect("txn"); + harness.engine.sync_once().await.expect("sync"); + let before = harness.engine.generation().cloned(); + + // Simulate the fork by doctoring the LAST frame in place: same + // length, same commit flag, different payload, RECOMPUTED cumulative + // checksum, so the local chain is fully valid but the bytes differ + // from the shipped segment. + let wal_path = harness.fixture.wal_path(); + let mut wal = std::fs::read(&wal_path).expect("read wal"); + let frame_size = 24 + PAGE_SIZE as usize; + assert!(wal.len() >= 32 + 2 * frame_size, "need at least two frames"); + let last = wal.len() - frame_size; + // Seed: the cumulative checksum stored in the PREVIOUS frame header. + let prev = last - frame_size; + let prev_header: [u8; 24] = wal[prev..prev + 24].try_into().expect("frame header"); + let seed = bencher_replica::FrameHeader::parse(&prev_header).checksum; + let header: [u8; 24] = wal[..24].try_into().expect("wal header prefix"); + // Manual byte math: the workspace warns on from_be_bytes/to_be_bytes. + let big_endian = header[3] & 1 == 1; + // Flip a payload byte, then recompute the frame checksum over the + // first 8 header bytes plus the full page payload. + wal[last + 24 + 100] ^= 0xFF; + let mut input = Vec::with_capacity(8 + PAGE_SIZE as usize); + input.extend_from_slice(&wal[last..last + 8]); + input.extend_from_slice(&wal[last + 24..last + frame_size]); + let (c1, c2) = bencher_replica::wal_checksum(big_endian, seed, &input); + for (index, checksum) in [c1, c2].into_iter().enumerate() { + let at = last + 16 + index * 4; + wal[at] = u8::try_from((checksum >> 24) & 0xff).expect("byte"); + wal[at + 1] = u8::try_from((checksum >> 16) & 0xff).expect("byte"); + wal[at + 2] = u8::try_from((checksum >> 8) & 0xff).expect("byte"); + wal[at + 3] = u8::try_from(checksum & 0xff).expect("byte"); + } + std::fs::write(&wal_path, wal).expect("write doctored wal"); + + // Resume must detect the fork and start a new generation. + harness.rebuild_engine().await; + harness.until_streaming().await; + harness.engine.sync_once().await.expect("backlog"); + let after = harness.engine.generation().cloned(); + assert_ne!(before, after, "a rewind fork must force a new generation"); + harness.assert_restore_equivalent().await; + } + + /// A WAL cycle buried behind the engine's back BETWEEN ticks (stray + /// restart, stray TRUNCATE checkpoint, stray restart) jumps salt1 by + /// more than one: the live epoch transition must refuse the shortcut + /// and force a new generation, and the fresh snapshot recaptures the + /// buried commits. + #[tokio::test] + async fn live_buried_wal_cycle_forces_new_generation() { + let mut harness = Harness::new().await; + harness.ready().await; + harness + .fixture + .txn(&["INSERT INTO t (data) VALUES ('shipped')"]) + .expect("txn"); + harness.engine.sync_once().await.expect("sync"); + let outcome = harness.engine.checkpoint_once().await.expect("checkpoint"); + assert_eq!(outcome, CheckpointOutcome::Completed); + + // Bury a whole cycle: restart (fully backfilled WAL), TRUNCATE + // checkpoint backfills the never-shipped frames, restart again. + harness + .fixture + .txn(&["CREATE TABLE buried (id INTEGER PRIMARY KEY, data TEXT)"]) + .expect("buried txn"); + harness + .fixture + .checkpoint(CheckpointMode::Truncate) + .expect("stray truncate"); + harness + .fixture + .txn(&["INSERT INTO buried (data) VALUES ('after burial')"]) + .expect("post burial txn"); + + // The live transition detects the salt discontinuity and diverges. + let before = harness.engine.generation().cloned(); + harness.engine.sync_once().await.expect("sync"); + harness.until_streaming().await; + harness.engine.sync_once().await.expect("backlog"); + let after = harness.engine.generation().cloned(); + assert_ne!(before, after, "a buried cycle must force a new generation"); + // The buried table survives on the replica via the fresh snapshot. + harness.assert_restore_equivalent().await; + } + + /// The salt-collision poison scenario, end to end. A generation's epoch 1 + /// is present but corrupt (parseable keys, broken checksum chain). After a + /// volume loss the restore soft-stops at epoch 0 and pins an advisory meta + /// to epoch 0, while the replica tip is still the corrupt epoch 1. On the + /// following resume the engine must NOT bind fresh salts and extend the + /// poisoned generation (that is what would leave it with two epoch-1 + /// directories, which `plan_epochs` can only ever soft-stop on): the + /// `meta.epoch == tip.epoch` proof fails, so resume diverges to a brand-new + /// generation. This is the guard that keeps the pessimistic collision + /// handling in `plan_epochs` unreachable in practice. + #[tokio::test] + async fn resume_after_soft_stop_below_corrupt_epoch_forces_new_generation() { + use std::sync::Arc; + + use bencher_replica::{ + LocalStorage, ReplicaDb, ReplicaMeta, ReplicaStorage, RestoreOutcome, SyncEngine, + compress_segment, restore_if_missing, + }; + + use super::harness::dir_path; + + let mut harness = Harness::new().await; + harness.ready().await; + let poisoned = harness.engine.generation().cloned().expect("generation"); + + // Epoch 0: ship one commit and checkpoint it, then restart the WAL + // (new salts) into epoch 1 and ship that too. + harness + .fixture + .txn(&["INSERT INTO t (data) VALUES ('epoch-0')"]) + .expect("epoch 0 txn"); + harness.engine.sync_once().await.expect("ship epoch 0"); + let outcome = harness.engine.checkpoint_once().await.expect("checkpoint"); + assert_eq!(outcome, CheckpointOutcome::Completed, "full backfill"); + harness + .fixture + .txn(&["INSERT INTO t (data) VALUES ('epoch-1')"]) + .expect("epoch 1 txn"); + harness.engine.sync_once().await.expect("ship epoch 1"); + + // Corrupt epoch 1 IN PLACE so the object stays valid (recompressed, + // same byte length) while the WAL checksum chain breaks: parseable + // keys, invalid content. A later restore plans the epoch, then + // soft-stops when the assembled WAL fails chain validation. + let epoch1_key = harness + .shipped_segment_keys() + .into_iter() + .find(|key| key.contains("/0000000001-")) + .expect("an epoch-1 segment shipped"); + let compressed = harness + .engine + .storage() + .get(&epoch1_key) + .await + .expect("get epoch-1 segment"); + let mut raw = decompress_segment(&compressed).expect("decompress epoch-1 segment"); + raw[24 + 50] ^= 0xff; + let recompressed = compress_segment(&raw).expect("recompress tampered segment"); + harness + .engine + .storage() + .put(&epoch1_key, Bytes::from(recompressed)) + .await + .expect("put tampered epoch-1 segment"); + + // Volume loss: restore into a fresh target. Replay soft-stops at + // epoch 0 (the corrupt epoch 1 is discarded) and writes an advisory + // meta pinned to epoch 0. + let target_tmp = tempfile::tempdir().expect("restore target tempdir"); + let target_db = dir_path(&target_tmp).join("bencher.db"); + let restored = restore_if_missing(&logger(), &harness.config, &target_db) + .await + .expect("restore boots on epoch 0"); + assert!( + matches!(restored, RestoreOutcome::Restored { .. }), + "restore boots on the last good epoch, got {restored:?}" + ); + let restored_meta = ReplicaMeta::load(&target_db) + .expect("load restored meta") + .expect("restore writes an advisory meta"); + assert_eq!( + restored_meta.epoch, 0, + "restore soft-stopped below the corrupt epoch 1, pinning the meta to epoch 0" + ); + + // Resume over the restored volume against the same replica. The tip is + // the still-present epoch 1, but the meta records epoch 0: the + // meta-verified epoch+1 path is refused and the engine diverges to a + // fresh generation instead of colliding a second epoch-1 directory + // into the poisoned generation. + let db = ReplicaDb { + db_path: target_db.clone(), + writer: Arc::new(tokio::sync::Mutex::new(())), + busy_timeout_ms: 5000, + }; + let resumed = SyncEngine::new_with_storage( + logger(), + harness.config.clone(), + db, + clock_for(&harness.clock_secs), + false, + ReplicaStorage::Local(LocalStorage::new(harness.replica_root.clone())), + ) + .await + .expect("resume engine"); + assert_eq!( + resumed.state(), + EngineState::PendingSnapshot, + "resume refuses to extend the poisoned generation and forces a fresh one" + ); + assert_eq!( + resumed.generation(), + None, + "no lineage is bound; a brand-new generation will be snapshotted next, \ + so the poisoned generation {} keeps its single epoch-1 directory", + poisoned.as_str() + ); + } + + /// Verification pins the shipped position and passes on a faithful + /// replica; app writes proceed while the fingerprint and restore run. + #[tokio::test] + async fn verify_passes_after_writes() { + let mut harness = Harness::new().await; + harness.ready().await; + harness + .fixture + .txn(&["INSERT INTO t (data) VALUES ('verified')"]) + .expect("txn"); + let report = harness.engine.verify_once().await.expect("verify"); + assert_eq!(report, Some(VerifyReport::Pass)); + } + + /// A fresh engine with no bound position cannot verify yet. + #[tokio::test] + async fn verify_unavailable_without_position() { + let mut harness = Harness::new().await; + // No bootstrap: the replica is empty and no position is bound. + let report = harness.engine.verify_once().await.expect("verify"); + assert_eq!(report, None); + } + + /// A tampered replica object fails verification, and (outside the + /// rate-limit window) forces a new generation on the next tick. + #[tokio::test] + async fn verify_fail_on_tampered_replica_triggers_new_generation() { + let mut harness = Harness::new().await; + harness.ready().await; + harness + .fixture + .txn(&["INSERT INTO t (data) VALUES ('will be tampered')"]) + .expect("txn"); + harness.engine.sync_once().await.expect("sync"); + + // Corrupt the newest shipped segment directly on the local replica. + let key = harness + .shipped_segment_keys() + .pop() + .expect("a shipped segment"); + let object = harness.replica_root.join(&key); + std::fs::write(&object, b"garbage, not zstd").expect("tamper"); + + // Move past the retrigger rate limit so the failure forces a new + // generation. + harness.advance(7 * 60 * 60); + let report = harness.engine.verify_once().await.expect("verify"); + assert!( + matches!(report, Some(VerifyReport::Fail { .. })), + "expected Fail, got {report:?}" + ); + // The failure schedules a new generation: the next tick snapshots, + // and the replica heals to full equivalence again. + let progress = harness.engine.sync_once().await.expect("sync"); + assert!( + progress.snapshot.is_some() || harness.engine.state() == EngineState::Snapshotting, + "verification failure must trigger a new generation" + ); + harness.until_streaming().await; + harness.engine.sync_once().await.expect("backlog"); + harness.assert_restore_equivalent().await; + } + + /// Scheduled verification runs through `sync_once` once the interval + /// elapses on the injected clock. + #[tokio::test] + async fn verify_scheduled_via_sync_once() { + let mut harness = Harness::with_config(|json| { + json.verification_interval_secs = Some(60); + }) + .await; + harness.ready().await; + harness + .fixture + .txn(&["INSERT INTO t (data) VALUES ('scheduled')"]) + .expect("txn"); + + // Not due yet: no verification report. + let progress = harness.engine.sync_once().await.expect("sync"); + assert_eq!(progress.verify, None); + + harness.advance(61); + let progress = harness.engine.sync_once().await.expect("sync"); + assert_eq!(progress.verify, Some(VerifyReport::Pass)); + } + + /// A verification EXECUTION error (here: the replica cannot be read for + /// the restore) must not arm the global ship backoff. Shipping proceeds + /// on the failing tick and, crucially, on the very next tick. + #[tokio::test] + async fn verify_error_does_not_block_shipping() { + let mut harness = Harness::with_config(|json| { + json.verification_interval_secs = Some(60); + }) + .await; + harness.ready().await; + // Verification restores the replica by reading (Get) its objects; + // fail all Gets so verification cannot COMPLETE (an execution error, + // distinct from a content Fail). Shipping only Puts, so it is + // unaffected. + harness + .flaky() + .set_plan(FailurePlan::new().fail_all(OpKind::Get)); + harness + .fixture + .txn(&["INSERT INTO t (data) VALUES ('a')"]) + .expect("txn a"); + harness.advance(61); + let progress = harness + .engine + .sync_once() + .await + .expect("sync with verify error"); + assert!( + progress.shipped_segments >= 1, + "the write ships despite the verify error: {progress:?}" + ); + assert!( + progress.error.is_some(), + "the verify execution error is reported: {progress:?}" + ); + assert!( + !progress.backing_off, + "the failing tick is not itself gated" + ); + + // The decoupling: the verify error must NOT arm the global ship + // backoff, so the very next tick ships a new write immediately. + harness + .fixture + .txn(&["INSERT INTO t (data) VALUES ('b')"]) + .expect("txn b"); + let progress = harness.engine.sync_once().await.expect("next tick"); + assert!( + !progress.backing_off, + "a verify error must not throttle WAL shipping: {progress:?}" + ); + assert!( + progress.shipped_segments >= 1, + "the next write ships immediately: {progress:?}" + ); + } + + /// Sole mode: an EXTERNAL checkpointer restarting the WAL between the + /// snapshot backup and finalize must ABORT the generation (no + /// snapshot.json), never rebind epoch 0 onto a body that may be missing + /// buried frames. A fresh, consistent generation then replaces it. + #[tokio::test] + async fn sole_finalize_aborts_on_external_wal_restart() { + let mut harness = Harness::with_config(|json| { + json.snapshot_throttle_mib = Some(1); + }) + .await; + harness.ready().await; + // Grow the database FILE so the copy spans multiple steps and there is + // room to interfere after the backup. + harness + .fixture + .txn_touching_pages(400) + .expect("big transaction"); + while harness.engine.ship_once().await.expect("ship big txn") > 0 {} + let outcome = harness.engine.checkpoint_once().await.expect("checkpoint"); + assert_eq!( + outcome, + CheckpointOutcome::Completed, + "backfill the db file" + ); + let bootstrap_markers = object_count(&harness, "snapshot.json").await; + + harness.engine.trigger_snapshot(); + // Drive past CreateGeneration (records the boundary salt) into Copying. + harness + .engine + .snapshot_step() + .await + .expect("ship tail step"); + harness + .engine + .snapshot_step() + .await + .expect("create generation step"); + assert_eq!(harness.engine.state(), EngineState::Snapshotting); + + // The engine's own checkpoints are suppressed while a snapshot runs, + // so this WAL restart is provably external: TRUNCATE resets the WAL, + // then a write starts a NEW salt cycle. + harness + .fixture + .checkpoint(CheckpointMode::Truncate) + .expect("external truncate"); + harness + .fixture + .txn(&["INSERT INTO t (data) VALUES ('buried new cycle')"]) + .expect("new cycle txn"); + + // The Finalize step must abort (Err) rather than commit a torn marker. + let mut aborted = false; + for _ in 0..256 { + if harness.engine.state() != EngineState::Snapshotting { + break; + } + if harness.engine.snapshot_step().await.is_err() { + aborted = true; + break; + } + } + assert!( + aborted, + "sole-mode finalize aborts on an external WAL restart" + ); + assert_eq!( + object_count(&harness, "snapshot.json").await, + bootstrap_markers, + "the aborted generation committed no snapshot.json" + ); + + // The engine scheduled a fresh generation; drive it home and confirm + // the replica restores to an exact copy of the (now-larger) source. + for _ in 0..256 { + let progress = harness.engine.sync_once().await.expect("recovery sync"); + if harness.engine.state() == EngineState::Streaming + && progress.shipped_segments == 0 + && progress.snapshot.is_none() + && !progress.backing_off + { + break; + } + harness.advance(1); + } + assert_eq!( + object_count(&harness, "snapshot.json").await, + bootstrap_markers + 1, + "a fresh consistent generation replaces the aborted one" + ); + harness.assert_restore_equivalent().await; + } + + /// Shadow mode: the same external WAL restart between backup and finalize + /// is legitimate (Litestream owns checkpoints), so the marker IS still + /// committed (the boundary rebinds; the shadow replica is disposable). + #[tokio::test] + async fn shadow_finalize_commits_despite_external_wal_restart() { + let mut harness = Harness::with_shadow_config(true, |json| { + json.snapshot_throttle_mib = Some(1); + }) + .await; + harness.until_streaming().await; + harness + .fixture + .txn_touching_pages(400) + .expect("big transaction"); + while harness.engine.ship_once().await.expect("ship big txn") > 0 {} + let before = object_count(&harness, "snapshot.json").await; + + harness.engine.trigger_snapshot(); + harness + .engine + .snapshot_step() + .await + .expect("ship tail step"); + harness + .engine + .snapshot_step() + .await + .expect("create generation step"); + assert_eq!(harness.engine.state(), EngineState::Snapshotting); + + // An external checkpointer (Litestream) legitimately restarts the WAL. + harness + .fixture + .checkpoint(CheckpointMode::Truncate) + .expect("external truncate"); + harness + .fixture + .txn(&["INSERT INTO t (data) VALUES ('litestream cycle')"]) + .expect("new cycle txn"); + + // Shadow mode keeps committing: the snapshot finishes and the marker + // lands (unlike sole mode, which would abort). + for step in 0..256 { + let status = harness + .engine + .snapshot_step() + .await + .expect("shadow snapshot commits despite the restart"); + if status == SnapshotStatus::Finished { + break; + } + assert!(step < 255, "shadow snapshot never finished"); + } + assert_eq!( + object_count(&harness, "snapshot.json").await, + before + 1, + "shadow mode commits the marker despite the external WAL restart" + ); + assert_eq!( + harness.engine.state(), + EngineState::Streaming, + "back to streaming after the shadow snapshot" + ); + } + + /// Shadow-mode burn-in: a real external checkpointer commits and TRUNCATE- + /// checkpoints the WAL repeatedly (several restarts) between engine ticks. + /// The engine ships across every restart, never checkpoints itself, and + /// the final replica restores to an exact copy of the source. + #[tokio::test] + async fn shadow_burn_in_ships_across_external_checkpoints() { + let mut harness = Harness::with_shadow_config(true, |_| {}).await; + harness.until_streaming().await; + harness.engine.sync_once().await.expect("initial backlog"); + + for round in 0..5 { + // A stray writer commits a frame in the current cycle. + let stray = harness.fixture.stray_conn().expect("stray conn"); + let data = format!("round-{round}"); + stray + .execute("INSERT INTO t (data) VALUES (?1)", [&data]) + .expect("stray insert"); + // The engine ships that frame BEFORE the checkpoint (I1: nothing + // unshipped is ever backfilled). + harness.engine.sync_once().await.expect("ship the cycle"); + // Then an external TRUNCATE checkpoint restarts the WAL. + harness + .fixture + .checkpoint(CheckpointMode::Truncate) + .expect("external truncate"); + harness + .engine + .sync_once() + .await + .expect("sync after the restart"); + } + + // The engine's own checkpoints are always skipped in shadow mode. + let outcome = harness.engine.checkpoint_once().await.expect("checkpoint"); + assert_eq!( + outcome, + CheckpointOutcome::SkippedShadow, + "shadow mode never checkpoints" + ); + + // Drain and confirm the final replica is logically equivalent. + for _ in 0..64 { + let progress = harness.engine.sync_once().await.expect("drain"); + if harness.engine.state() == EngineState::Streaming + && progress.shipped_segments == 0 + && progress.snapshot.is_none() + && !progress.backing_off + { + break; + } + harness.advance(1); + } + harness.assert_restore_equivalent().await; + } + + /// Prune deletes the `snapshot.json` marker FIRST, so a crash (here: an + /// injected failure of the unordered `delete_prefix` batch) mid-prune + /// leaves an invisible markerless generation, never a marker without its + /// body that resume/restore would trust. + #[tokio::test] + async fn prune_deletes_marker_before_body() { + let mut harness = Harness::with_config(|json| { + json.retention_generations = Some(1); + }) + .await; + harness.until_streaming().await; + let current = harness.engine.generation().cloned().expect("generation"); + + // An older COMPLETE generation that retention (1) must prune. + let old = GenerationId::new( + DateTime::try_from(BASE_SECS - 200_000).expect("timestamp"), + 1, + ); + let body_key = format!("generations/{}/snapshot.db.zst", old.as_str()); + let marker_key = format!("generations/{}/snapshot.json", old.as_str()); + let storage = harness.engine.storage(); + storage + .put(&body_key, Bytes::from_static(b"body")) + .await + .expect("put body"); + storage + .put(&marker_key, Bytes::from_static(b"{}")) + .await + .expect("put marker"); + + // Fail the unordered prefix-delete batch (a separate op from the + // single marker delete). The marker delete runs first and succeeds; + // the prefix batch then fails, and the prune surfaces the error. + harness + .flaky() + .set_plan(FailurePlan::new().fail_all(OpKind::DeletePrefix)); + let result = harness.engine.prune_once().await; + assert!(result.is_err(), "the failed prefix delete surfaces"); + + let storage = harness.engine.storage(); + let marker = storage.get(&marker_key).await; + assert!( + matches!(marker, Err(bencher_replica::StorageError::NotFound { .. })), + "the marker was deleted first (generation now invisible): {marker:?}" + ); + assert!( + storage.get(&body_key).await.is_ok(), + "the body remains after the failed prefix delete" + ); + assert_eq!( + harness.engine.generation(), + Some(¤t), + "the current generation is untouched" + ); + } + + /// The poison classification (the last-resort fatal safety net): an + /// oversized transaction qualifies, transient/control errors do not. In + /// normal operation the ship path intercepts an oversized transaction and + /// re-snapshots instead (see `oversized_committed_txn_re_snapshots_and_drains`), + /// so this classification only fires for an unconverted escape. + #[test] + fn poison_classification() { + use bencher_replica::{SyncError, WalError}; + assert!( + SyncError::TransactionTooLarge { + bytes: 5, + max_bytes: 4 + } + .is_poison() + ); + assert!( + SyncError::Wal(WalError::TransactionTooLarge { + bytes: 5, + max_bytes: 4 + }) + .is_poison() + ); + assert!(!SyncError::TaskExited.is_poison()); + assert!(!SyncError::SnapshotBoundaryDiverged.is_poison()); + } + + /// An oversized COMMITTED transaction cannot ship (restore rejects the + /// segment) and, being durable, cannot be split retroactively. It must NOT + /// wedge or crash the engine: the ship path forces a re-snapshot that + /// captures its state, pins epoch 0 at the boundary, and drains the WAL via + /// a checkpoint so shipping resumes on a fresh cycle. The replica stays + /// restorable throughout and ends logically equivalent to the source. + #[tokio::test] + async fn oversized_committed_txn_re_snapshots_and_drains() { + let mut harness = Harness::new().await; + harness.ready().await; + // Impose a small per-transaction cap and rebuild the engine to resume + // with it (there is no JSON knob, so mutate the resolved config). The + // cap is above a normal single-page write (~4 KiB) but well below the + // deliberately oversized transaction below. + harness.config.max_transaction_bytes = 16 * 1024; + harness.rebuild_engine().await; + let bootstrap = harness.engine.generation().cloned().expect("generation"); + + // A committed transaction far over the cap (dozens of pages). + harness + .fixture + .txn_touching_pages(16) + .expect("oversized committed txn"); + + // Ship detects it and schedules a re-snapshot; it is NOT a fatal error. + let progress = harness.engine.sync_once().await.expect("sync (no fatal)"); + assert!( + progress.error.is_none(), + "an oversized committed transaction is not a fatal error: {progress:?}" + ); + + // Drive the re-snapshot to completion (position pins at the boundary). + for _ in 0..256 { + if harness.engine.state() == EngineState::Streaming { + break; + } + harness.engine.sync_once().await.expect("re-snapshot tick"); + } + let drained_gen = harness.engine.generation().cloned().expect("generation"); + assert_ne!( + drained_gen, bootstrap, + "a fresh generation captured the oversized transaction" + ); + + // Pinned at the boundary: the ship path ships nothing until the drain. + let progress = harness.engine.sync_once().await.expect("pinned ship"); + assert_eq!( + progress.shipped_segments, 0, + "epoch 0 is pinned at the boundary, awaiting the drain" + ); + + // Drain: a checkpoint backfills the WAL below the boundary (driven + // directly to bypass the min-pages threshold for this small fixture). + let outcome = harness + .engine + .checkpoint_once() + .await + .expect("drain checkpoint"); + assert_eq!( + outcome, + CheckpointOutcome::Completed, + "the WAL below the boundary backfills, allowing the restart" + ); + + // A write restarts the WAL; the ship path rebinds epoch 0 to the fresh + // cycle and resumes shipping. + harness + .fixture + .txn(&["INSERT INTO t (data) VALUES ('after drain')"]) + .expect("post-drain write"); + let progress = harness.engine.sync_once().await.expect("resume shipping"); + assert!( + progress.error.is_none() && progress.shipped_segments >= 1, + "shipping resumes on a fresh cycle after the drain: {progress:?}" + ); + + // The oversized transaction lives in the snapshot body, the post-drain + // write in an epoch-0 segment: restore reproduces the source exactly. + harness.assert_restore_equivalent().await; + } + + /// The checkpoint gate bounds its unshipped-commit scan with + /// `max_transaction_bytes`: an oversized run aborts the scan (never read in + /// full inside the critical section) and is treated CONSERVATIVELY as + /// unshipped, deferring the checkpoint (invariant I1) rather than + /// backfilling an ambiguous tail. + #[tokio::test] + async fn checkpoint_gate_defers_on_oversized_unshipped_commit() { + let mut harness = Harness::new().await; + harness.ready().await; + harness.config.max_transaction_bytes = 16 * 1024; + harness.rebuild_engine().await; + // A committed transaction over the cap, sitting unshipped above the + // resumed position. + harness + .fixture + .txn_touching_pages(16) + .expect("oversized committed txn"); + let outcome = harness.engine.checkpoint_once().await.expect("checkpoint"); + assert_eq!( + outcome, + CheckpointOutcome::SkippedUnshipped, + "an oversized unshipped commit defers the checkpoint conservatively" + ); + } +} diff --git a/services/api/CLAUDE.md b/services/api/CLAUDE.md index b3b6de438..807be12bb 100644 --- a/services/api/CLAUDE.md +++ b/services/api/CLAUDE.md @@ -14,6 +14,18 @@ please alert the developer working with you and indicate that this is the case b - SQLite - Litestream +## Replication (`plus/bencher_replica` and Litestream) + +`plus/bencher_replica` is the in-process replacement for Litestream: it ships WAL segments to a file or S3 replica, snapshots via a single-step `SQLite` online backup, and restores at startup. Configure via `plus.replica` in the server config; when BOTH `plus.litestream` and `plus.replica` are set, the replica runs in shadow mode (Litestream keeps checkpoint ownership and remains the restore source). The crate's `src/lib.rs` documents the six invariants (I1 to I6) that govern every design decision; read them before touching the sync engine. Its tests require `--features plus,testing`. + +## Litestream + +Litestream (v0.5.x, LTX-based) runs as a child process of the API server and is the sole WAL checkpointer (`wal_autocheckpoint = 0`). + +- When Litestream decides it needs a full re-snapshot (WAL continuity break, salt reset, or process restart), it copies the entire database into a local LTX file while holding the SQLite write lock via its `_litestream_lock` table. While that copy runs, API writes cannot take the write lock: each one burns its `busy_timeout` and fails with "database is locked" (surfaced as retryable 503s plus Sentry tripwires). The copy scales with database size, so the larger the database the longer the stall; removing this blocking failure mode is the reason `plus/bencher_replica` exists. A burst of "database is locked" Conflicts arriving together with a large read-plus-write burst on the data volume points at a Litestream re-snapshot, not the `/v0/server/backup` endpoint. +- `truncate-page-n: 0` genuinely disables blocking TRUNCATE checkpoints in v0.5.x (`TruncatePageN` is a `*int` in Litestream's config, so explicit 0 is honored, verified against the v0.5.13 source). +- `JsonLitestream.metrics_port` enables Litestream's Prometheus endpoint (serialized as the top-level `addr` setting in `litestream.yml`). The Fly configs (`services/api/fly/*.toml`) scrape port 9090 via their `[[metrics]]` section. + ## Building & Running ```bash diff --git a/services/api/Cargo.toml b/services/api/Cargo.toml index 5d1b2c060..aeb765ed1 100644 --- a/services/api/Cargo.toml +++ b/services/api/Cargo.toml @@ -13,6 +13,8 @@ plus = [ "dep:api_oci", "dep:api_runners", "dep:api_specs", + "dep:bencher_replica", + "dep:camino", "api_auth/plus", "api_checkout/plus", "api_mcp/plus", @@ -29,6 +31,7 @@ plus = [ "bencher_json/plus", "bencher_otel?/plus", "bencher_otel_provider/plus", + "bencher_replica/plus", "bencher_schema/plus", ] sentry = [ @@ -61,6 +64,7 @@ otel = [ "api_users/otel", "bencher_config/otel", "bencher_endpoint/otel", + "bencher_replica?/otel", "bencher_schema/otel", ] @@ -83,6 +87,8 @@ bencher_logger.workspace = true bencher_schema.workspace = true bencher_otel = { workspace = true, optional = true } bencher_otel_provider = { workspace = true, optional = true } +bencher_replica = { workspace = true, optional = true } +camino = { workspace = true, optional = true } dropshot.workspace = true futures-concurrency.workspace = true futures-util.workspace = true diff --git a/services/api/fly/fly.dev.toml b/services/api/fly/fly.dev.toml index 66fadc7ba..f0e7998f9 100644 --- a/services/api/fly/fly.dev.toml +++ b/services/api/fly/fly.dev.toml @@ -49,3 +49,7 @@ memory_mb = 256 [[restart]] policy = "on-failure" retries = 6 + +[[metrics]] +port = 9090 +path = "/metrics" diff --git a/services/api/fly/fly.test.toml b/services/api/fly/fly.test.toml index 5e2653e6b..e58cfee4b 100644 --- a/services/api/fly/fly.test.toml +++ b/services/api/fly/fly.test.toml @@ -49,3 +49,7 @@ memory_mb = 256 [[restart]] policy = "on-failure" retries = 6 + +[[metrics]] +port = 9090 +path = "/metrics" diff --git a/services/api/fly/fly.toml b/services/api/fly/fly.toml index cfe9f049d..1c08dd63e 100644 --- a/services/api/fly/fly.toml +++ b/services/api/fly/fly.toml @@ -54,3 +54,7 @@ memory_mb = 2048 [[restart]] policy = "on-failure" retries = 6 + +[[metrics]] +port = 9090 +path = "/metrics" diff --git a/services/api/openapi.json b/services/api/openapi.json index fe2cf1b90..b9f9a35eb 100644 --- a/services/api/openapi.json +++ b/services/api/openapi.json @@ -8933,7 +8933,7 @@ "server" ], "summary": "View server configuration", - "description": "View the API server configuration. The user must be an admin on the server to use this route.", + "description": "View the API server configuration. Secrets in the configuration are masked in the response. The user must be an admin on the server to use this route.", "operationId": "server_config_get", "responses": { "200": { @@ -13124,6 +13124,13 @@ } ] }, + "metrics_port": { + "nullable": true, + "description": "Port for the Prometheus metrics endpoint (serialized as the Litestream `addr` setting)", + "type": "integer", + "format": "uint16", + "minimum": 0 + }, "replica": { "description": "Disaster recovery replica", "allOf": [ @@ -15334,6 +15341,14 @@ } ] }, + "replica": { + "nullable": true, + "allOf": [ + { + "$ref": "#/components/schemas/JsonReplication" + } + ] + }, "runners": { "nullable": true, "allOf": [ @@ -16005,6 +16020,80 @@ } ] }, + "JsonReplication": { + "description": "In-process disaster recovery replication (`bencher_replica`).\n\nStreams `SQLite` WAL frames to a replica target, takes periodic snapshots, and restores the latest state at API server startup. Replaces Litestream. When both `litestream` and `replica` are configured, the replica runs in shadow mode: Litestream keeps checkpoint ownership and remains the restore source.\n\nExactly one server may replicate to a given target. There is no fencing between concurrent writers, so pointing two servers at the same target interleaves their lineages and corrupts the replica.\n\n`deny_unknown_fields` is deliberate: this is disaster-recovery config, so a misspelled key (e.g. `sync_interval` for `sync_interval_secs`) must fail loudly at load time instead of silently applying a default.", + "type": "object", + "properties": { + "checkpoint_interval_secs": { + "nullable": true, + "description": "Minimum interval between checkpoints (seconds; default 60)", + "type": "integer", + "format": "uint32", + "minimum": 0 + }, + "min_checkpoint_pages": { + "nullable": true, + "description": "Minimum WAL pages before a checkpoint triggers (default 1000)", + "type": "integer", + "format": "uint32", + "minimum": 0 + }, + "retention_generations": { + "nullable": true, + "description": "Number of generations to retain (default 3)", + "type": "integer", + "format": "uint32", + "minimum": 0 + }, + "shutdown_sync_timeout_secs": { + "nullable": true, + "description": "Deadline for the final WAL ship at shutdown (seconds; default 4, which fits inside Fly's 6 second kill timeout)", + "type": "integer", + "format": "uint32", + "minimum": 0 + }, + "snapshot_interval_secs": { + "nullable": true, + "description": "How often to start a new generation with a fresh snapshot (seconds; default 86400)", + "type": "integer", + "format": "uint32", + "minimum": 0 + }, + "snapshot_throttle_mib": { + "nullable": true, + "description": "Snapshot copy throttle (MiB per second; default 32)", + "type": "integer", + "format": "uint32", + "minimum": 0 + }, + "sync_interval_secs": { + "nullable": true, + "description": "How often to ship new WAL frames (seconds; default 1)", + "type": "integer", + "format": "uint32", + "minimum": 0 + }, + "target": { + "description": "Replica target", + "allOf": [ + { + "$ref": "#/components/schemas/ReplicationTarget" + } + ] + }, + "verification_interval_secs": { + "nullable": true, + "description": "How often to run restore-and-compare verification (seconds; default 86400, 0 disables). Each run restores a full copy of the database next to the live one, so expect a transient doubling of data-volume usage while it runs.", + "type": "integer", + "format": "uint32", + "minimum": 0 + } + }, + "required": [ + "target" + ], + "additionalProperties": false + }, "JsonReport": { "type": "object", "properties": { @@ -18576,6 +18665,71 @@ } ] }, + "ReplicationTarget": { + "description": "Where the replica lives: a local directory XOR S3-compatible storage.", + "oneOf": [ + { + "type": "object", + "properties": { + "path": { + "type": "string" + }, + "scheme": { + "type": "string", + "enum": [ + "file" + ] + } + }, + "required": [ + "path", + "scheme" + ], + "additionalProperties": false + }, + { + "description": "S3-compatible object storage. Configure a lifecycle rule on the bucket to abort incomplete multipart uploads: a crashed snapshot upload leaves orphaned parts that otherwise accrue storage cost indefinitely.", + "type": "object", + "properties": { + "access_key_id": { + "type": "string" + }, + "bucket": { + "type": "string" + }, + "endpoint": { + "nullable": true, + "description": "Endpoint override for S3-compatible services (R2, `MinIO`)", + "type": "string" + }, + "path": { + "nullable": true, + "type": "string" + }, + "region": { + "nullable": true, + "type": "string" + }, + "scheme": { + "type": "string", + "enum": [ + "s3" + ] + }, + "secret_access_key": { + "$ref": "#/components/schemas/Secret" + } + }, + "required": [ + "access_key_id", + "bucket", + "scheme", + "secret_access_key" + ], + "additionalProperties": false + } + ] + }, "ReportIdempotencyKey": { "type": "string", "format": "uuid" diff --git a/services/api/src/lib.rs b/services/api/src/lib.rs index 58b6f4410..18c8159c2 100644 --- a/services/api/src/lib.rs +++ b/services/api/src/lib.rs @@ -8,6 +8,10 @@ use bencher_logger as _; use bencher_otel as _; #[cfg(any(feature = "plus", feature = "otel"))] use bencher_otel_provider as _; +#[cfg(feature = "plus")] +use bencher_replica as _; +#[cfg(feature = "plus")] +use camino as _; use futures_concurrency as _; use futures_util as _; #[cfg(feature = "sentry")] diff --git a/services/api/src/main.rs b/services/api/src/main.rs index 485194dcf..7ccc29772 100644 --- a/services/api/src/main.rs +++ b/services/api/src/main.rs @@ -3,7 +3,7 @@ reason = "dependencies used by lib but not binary" )] -#[cfg(feature = "sentry")] +#[cfg(any(feature = "plus", feature = "sentry"))] use std::path::PathBuf; use std::sync::Arc; @@ -39,6 +39,21 @@ pub enum ApiError { #[cfg(feature = "plus")] #[error("{0}")] Litestream(#[from] LitestreamError), + #[cfg(feature = "plus")] + #[error("Replica config: {0}")] + ReplicaConfig(bencher_replica::ReplicaConfigError), + #[cfg(feature = "plus")] + #[error("Replica restore: {0}")] + ReplicaRestore(bencher_replica::RestoreError), + #[cfg(feature = "plus")] + #[error("Replica: {0}")] + ReplicaSync(bencher_replica::SyncError), + #[cfg(feature = "plus")] + #[error("Failed to resolve the database path for replication: {0}")] + ReplicaDbPath(std::io::Error), + #[cfg(feature = "plus")] + #[error("Database path is not valid UTF-8: {0}")] + ReplicaDbPathUtf8(PathBuf), #[error("{0}")] ConfigTxError(bencher_config::ConfigTxError), #[error("Failed to listen for ctrl-c signal: {0}")] @@ -98,7 +113,22 @@ async fn run( .map_err(ApiError::OpenTelemetry)?; #[cfg(feature = "plus")] - if let Some(litestream) = config + let replica_config = config + .plus + .as_ref() + .and_then(|plus| plus.replica.clone()) + .map(bencher_replica::ReplicaConfig::try_from) + .transpose() + .map_err(ApiError::ReplicaConfig)?; + #[cfg(feature = "plus")] + let replica_shutdown_timeout = replica_config + .as_ref() + .map(|replica| replica.shutdown_sync_timeout); + + // Restore precedence: Litestream owns restore whenever it is configured + // (during the shadow burn-in the replica is NOT the restore source). + #[cfg(feature = "plus")] + let litestream_handle = if let Some(litestream) = config .plus .as_ref() .and_then(|plus| plus.litestream.clone()) @@ -107,29 +137,68 @@ async fn run( let litestream_handle = run_litestream(log, &config, litestream, restore_tx)?; // Wait for Litestream restore to complete (replicate starts in background) restore_rx.await.map_err(LitestreamError::RestoreRecv)?; + Some(litestream_handle) + } else { + if let Some(replica) = &replica_config { + // Same handshake as `litestream restore`: the restore fully + // completes before any SQLite connection is opened. + let db_path = replica_db_path(&config.database.file)?; + let outcome = bencher_replica::restore_if_missing(log, replica, &db_path) + .await + .map_err(ApiError::ReplicaRestore)?; + info!(log, "Replica restore: {outcome:?}"); + } + None + }; + // Both configured: the replica runs in shadow mode (no checkpoints; + // Litestream keeps checkpoint ownership until cutover). + #[cfg(feature = "plus")] + let shadow = litestream_handle.is_some(); - let server = create_api_server(config).await?; - let shutdown_wait = server.wait_for_shutdown(); - let result = ( - tokio::signal::ctrl_c().map(|r| r.map_err(ApiError::CtrlC)), - async { - litestream_handle + let server = create_api_server(config).await?; + + #[cfg(feature = "plus")] + let mut replica_handle = start_replica(log, &server, replica_config, shadow)?; + + let shutdown_wait = server.wait_for_shutdown(); + #[cfg(feature = "plus")] + let result = { + let litestream_wait = async { + match litestream_handle { + Some(litestream_handle) => litestream_handle .await .map_err(LitestreamError::JoinHandle)? - .map_err(Into::into) - }, + .map_err(Into::into), + None => std::future::pending::>().await, + } + }; + let replica_fatal = async { + match replica_handle.as_mut() { + // Resolves only if the replication task dies; storage + // outages back off internally and never resolve this. + // + // SOLE mode only: in shadow mode the replica is the unproven + // system burning in while Litestream stays authoritative, so + // a fatal replica error must never take down a healthy + // production server. The task still logs and meters the + // fatal; it just does not race the server here. + Some(replica_handle) if !shadow => replica_handle + .wait_fatal() + .await + .map_err(ApiError::ReplicaSync), + Some(_) | None => std::future::pending::>().await, + } + }; + ( + tokio::signal::ctrl_c().map(|r| r.map_err(ApiError::CtrlC)), + litestream_wait, + replica_fatal, shutdown_wait.map(|r| r.map_err(ApiError::RunServer)), ) .race() - .await; - - shutdown(log, server).await; - - return result; - } - - let server = create_api_server(config).await?; - let shutdown_wait = server.wait_for_shutdown(); + .await + }; + #[cfg(not(feature = "plus"))] let result = ( tokio::signal::ctrl_c().map(|r| r.map_err(ApiError::CtrlC)), shutdown_wait.map(|r| r.map_err(ApiError::RunServer)), @@ -139,9 +208,82 @@ async fn run( shutdown(log, server).await; + // Final replica sync AFTER the server has drained: ship the remaining WAL + // tail within the shutdown budget. On a COMPLETE drain in sole mode a + // final checkpoint then runs (after the budget, unbounded) to seal the + // epoch so the next boot resumes in place; on the deadline the tail stays + // in the local WAL (lag, never loss). + #[cfg(feature = "plus")] + if let Some(replica_handle) = replica_handle { + // `replica_shutdown_timeout` is always `Some` here (it is derived from + // the same `replica_config` that produced this handle); the fallback + // is defensive and reuses the crate default rather than a magic value. + let deadline = + replica_shutdown_timeout.unwrap_or(bencher_replica::DEFAULT_SHUTDOWN_SYNC_TIMEOUT); + if let Err(e) = replica_handle.shutdown(deadline).await { + error!(log, "Replica final sync failed: {e}"); + } + } + result } +/// Spawn the in-process replication task over the server's database. +#[cfg(feature = "plus")] +fn start_replica( + log: &Logger, + server: &HttpServer, + replica_config: Option, + shadow: bool, +) -> Result, ApiError> { + let Some(replica) = replica_config else { + return Ok(None); + }; + let ctx = server.app_private(); + let db_path = replica_db_path(&ctx.database.path)?; + Ok(Some(bencher_replica::Replicator::start( + log.clone(), + replica, + bencher_replica::ReplicaDb { + db_path, + writer: ctx.database.connection.clone(), + busy_timeout_ms: ctx.database.busy_timeout, + }, + ctx.clock.clone(), + shadow, + ))) +} + +/// The absolute, UTF-8 database path handed to the replicator. +/// +/// Canonicalized when the file exists: `SQLite` resolves a symlinked +/// database file and creates the real `-wal` next to the TARGET, so the +/// replicator must derive its WAL path from the resolved location or it +/// would silently watch a never-existing file. +#[cfg(feature = "plus")] +fn replica_db_path(file: &std::path::Path) -> Result { + let absolute = if file.is_absolute() { + file.to_path_buf() + } else { + std::env::current_dir() + .map_err(ApiError::ReplicaDbPath)? + .join(file) + }; + let resolved = match std::fs::canonicalize(&absolute) { + Ok(resolved) => resolved, + // A missing database (fresh start before the first connection) has + // nothing to resolve yet; canonicalize the parent instead so a + // symlinked data directory still lands on the real location. + Err(_missing) => match (absolute.parent(), absolute.file_name()) { + (Some(parent), Some(file_name)) => std::fs::canonicalize(parent) + .map(|parent| parent.join(file_name)) + .unwrap_or(absolute), + _ => absolute, + }, + }; + camino::Utf8PathBuf::from_path_buf(resolved).map_err(ApiError::ReplicaDbPathUtf8) +} + async fn shutdown(log: &Logger, server: HttpServer) { #[cfg(feature = "plus")] let save_rate_limiting = { diff --git a/services/cli/Dockerfile b/services/cli/Dockerfile index daf2ef0e6..55056c086 100644 --- a/services/cli/Dockerfile +++ b/services/cli/Dockerfile @@ -61,6 +61,7 @@ RUN cargo init --lib bencher_otel RUN cargo init --lib bencher_otel_provider RUN cargo init --lib bencher_rate_limiter RUN cargo init --lib bencher_recaptcha +RUN cargo init --lib bencher_replica RUN cargo init --lib bencher_rootfs RUN cargo init --lib bencher_runner RUN cargo init --lib bencher_init diff --git a/services/console/Dockerfile b/services/console/Dockerfile index e52651775..e5d2e1504 100644 --- a/services/console/Dockerfile +++ b/services/console/Dockerfile @@ -45,6 +45,7 @@ RUN cargo init --lib bencher_otel RUN cargo init --lib bencher_otel_provider RUN cargo init --lib bencher_rate_limiter RUN cargo init --lib bencher_recaptcha +RUN cargo init --lib bencher_replica RUN cargo init --lib bencher_rootfs RUN cargo init --lib bencher_runner RUN cargo init --lib bencher_init