diff --git a/docs/RESHARDING.md b/docs/RESHARDING.md index cd8193265..a14b9e039 100644 --- a/docs/RESHARDING.md +++ b/docs/RESHARDING.md @@ -16,8 +16,9 @@ RESHARD ; ``` Issued against the admin database. Parsed in [`pgdog/src/admin/reshard.rs`](../pgdog/src/admin/reshard.rs), which calls -`Orchestrator::new(source, destination, publication, slot_name)` and then -`orchestrator.replicate_and_cutover().await`. +`Orchestrator::new(source, destination, publication, slot_name)` and then starts a `ReshardTask` +([`api/resharding.rs`](../pgdog/src/api/resharding.rs)) in the background. The command replies with +the task id. `SHOW TASKS` reports the progress. > **Multi-node deployments:** Traffic cutover via `RESHARD` is supported on single-node PgDog only. > The [Enterprise Edition control plane](https://docs.pgdog.dev/enterprise_edition/control_plane/) @@ -32,7 +33,7 @@ Issued against the admin database. Parsed in [`pgdog/src/admin/reshard.rs`](../p - `publisher: Arc>` — manages replication slots, table list, and lag tracking - `replication_slot: String` — auto-generated as `__pgdog_repl_` unless overridden -`replicate_and_cutover()` is the top-level method and calls the five steps below in sequence: +`ReshardTask::run` drives the five steps below in sequence: ```mermaid flowchart LR @@ -40,7 +41,7 @@ flowchart LR B["2. schema_sync_pre
pre-data to dest
reload schema cache"] C["3. data_sync
ParallelSyncManager
binary COPY"] D["4. schema_sync_post
secondary indexes"] - E["5. replicate().cutover()
WAL drain
traffic swap"] + E["5. ReplicationTask
WAL drain
traffic swap"] A --> B --> C --> D --> E ``` @@ -129,8 +130,9 @@ index maintenance overhead during the high-throughput copy phase. ## Step 5 — Replication and cutover -`replicate()` creates a `ReplicationWaiter` that wraps a `Waiter` from `Publisher::replicate()`. -`ReplicationWaiter::cutover()` then runs two serial wait phases followed by the atomic swap. +`ReplicationTask` builds a `Migration` and calls `Migration::run` +([`api/replication.rs`](../pgdog/src/api/replication.rs)). `run` streams until a cutover signal, +cuts over, flips direction, and then streams in reverse so a rollback stays possible. ### Publisher and StreamSubscriber @@ -158,37 +160,47 @@ Two behaviours are specific to the resharding context: ### Cutover phases -**Phase 1 — `wait_for_replication()`**: polls lag every 1 second. When -`lag ≤ cutover_traffic_stop_threshold`: -1. Calls `maintenance_mode::start()` — new queries queue behind a barrier. +**Phase 1 — `CutoverPolicy::wait_for_stop_threshold()`**: polls lag every 1 second. It returns +when `lag ≤ cutover_traffic_stop_threshold`. `Migration::prepare_cutover` then: +1. Calls `MaintenanceMode::stop_traffic()`, which calls `maintenance_mode::start(None)` — new + queries queue behind a barrier. 2. Calls `cancel_all(source_db)` — cancels any queries already in flight. -**Phase 2 — `wait_for_cutover()`**: polls at 50 ms intervals. Three independent triggers can fire -cutover (whichever comes first): +**Phase 2 — `CutoverPolicy::wait_for_catchup()`**: polls at 50 ms intervals. Three independent +triggers can fire cutover (whichever comes first): | Trigger | Config key | Action | |---|---|---| | `lag ≤ threshold` | `cutover_replication_lag_threshold` | `CutoverReason::Lag` → proceed | | elapsed ≥ timeout | `cutover_timeout` | `CutoverReason::Timeout` → proceed or abort (see `cutover_timeout_action`) | -| no transactions for N ms | `cutover_last_transaction_delay` | `CutoverReason::LastTransaction` → proceed | +| no transaction applied for N ms | `cutover_last_transaction_delay` | `CutoverReason::LastTransaction` → proceed | -**Point of no return** — the `ok_or_abort!` macro wraps every subsequent call. Any failure resumes -traffic immediately via `maintenance_mode::stop()` and returns an error. Steps in order: +The `LastTransaction` trigger needs a measured transaction. A stream that has applied nothing +reports no value, so the trigger stays silent and only the timeout can fire. -1. `publisher.request_stop()` + `waiter.wait()` — stops the replication stream; drains remaining WAL. -2. `schema_sync_cutover()` — applies `SyncState::Cutover` operations (e.g. drops sequences that - won't be used in the sharded cluster). -3. `cutover(source_db, dest_db)` in [`pgdog/src/backend/databases.rs`](../pgdog/src/backend/databases.rs) — +**Phase 3 — drain**: `replicate_until_cutover()` stops the cluster task and waits for every +stream to drain. The budget is `ReplicationClusterTask::drain_timeout()` (300 s) for the cluster +and `stream_drain_timeout()` (120 s) for the streams. A stream that does not drain in time is +aborted, and its `SlotGuard` drops the replication slot on a detached task. A failed drain +returns `Error::DrainTimeout`. + +**Point of no return** — `Migration::cutover()` runs these steps in order: + +1. `Publisher::create_slots(destination)` — creates the reverse replication slots. +2. `cutover(source_db, dest_db)` in [`pgdog/src/backend/databases.rs`](../pgdog/src/backend/databases.rs) — atomically swaps the two clusters' logical identity in the routing table (and config refs via `Config::cutover`/`Users::cutover`); no data moves. Persisted to disk when `cutover_save_config = true`. -4. `orchestrator.refresh()` — re-fetches both clusters from `databases()` so the orchestrator now - treats the new cluster as source for reverse replication. -5. `schema_sync_post_cutover()` — applies `SyncState::PostCutover` (removes blockers that would - prevent reverse replication, such as unique constraints on sequence columns). -6. `orchestrator.replicate()` — starts reverse replication (new cluster → old cluster) as a - background `crate::api` task. This enables rollback without data loss. -7. `maintenance_mode::stop()` — releases the barrier; queued and new queries flow to the new cluster. +3. `Orchestrator::refresh()` — re-fetches both clusters from `databases()`. +4. `MaintenanceMode::resume_traffic()` — releases the barrier; queued and new queries flow to the + new cluster. + +`Migration::run` then flips direction and streams in reverse, from the new cluster to the old one. +The reverse phase runs in the same task, not in a separate one. A `STOP_TASK` during the reverse +phase ends the rollback window, and the task reports the migration as finished. + +The cutover schema sync (`SyncState::Cutover`, then `SyncState::PostCutover`) runs as a +`SchemaSyncTask` subtask after each stream phase ends. --- @@ -197,7 +209,7 @@ traffic immediately via `maintenance_mode::stop()` and returns an error. Steps i ### Pre-cutover failures — plain propagation Steps 1–4 (`load_schema`, `schema_sync_pre`, `data_sync`, `schema_sync_post`) propagate errors -with `?` directly from `replicate_and_cutover()`. Maintenance mode is never entered during these +with `?` directly from `Migration::run()`. Maintenance mode is never entered during these steps. A failure here leaves traffic unaffected and the source untouched, making a full restart safe. ### Schema DDL — intentional error tolerance @@ -240,38 +252,26 @@ Per-table slots created in [`Table::data_sync()`](../pgdog/src/backend/replicati automatically when the replication connection closes, including on error or panic. A failed copy task leaves no orphaned per-table slot. -### The `ok_or_abort!` macro — guaranteed traffic resumption after cutover starts - -```rust -macro_rules! ok_or_abort { - ($expr:expr) => { - match $expr { - Ok(res) => res, - Err(err) => { - maintenance_mode::stop(); - cutover_state(CutoverState::Abort { error: err.to_string() }); - return Err(err.into()); - } - } - }; -} -``` +### `MaintenanceMode` — guaranteed traffic resumption + +`Migration` owns a `MaintenanceMode` guard ([`api/replication.rs`](../pgdog/src/api/replication.rs)). +`stop_traffic()` calls `maintenance_mode::start(None)` and records that it did. +`resume_traffic()` calls `maintenance_mode::stop(None)` only when the barrier is on, so every +caller can call it safely. Three paths release the barrier: -Once `maintenance_mode::start()` is called in `wait_for_replication()`, traffic is paused. -`ok_or_abort!` is the only place that calls `maintenance_mode::stop()` for the remaining steps. -Every call after the point of no return — `waiter.wait()`, `schema_sync_cutover()`, `cutover()`, -`orchestrator.refresh()`, `schema_sync_post_cutover()`, `orchestrator.replicate()` — is wrapped -in it. This guarantees traffic always resumes, regardless of which step fails. +1. `Migration::cutover()` releases it after the swap. +2. `prepare_cutover()` releases it when the catch-up wait fails. +3. `replicate_until_cutover()` releases it when the phase ends with an error, including a + `STOP_TASK`, so the barrier does not survive the drain. -The macro also transitions the global `CutoverState` to `Abort`, which is visible via -`SHOW REPLICATION_SLOTS` in the admin database. +`ReplicationTask::run` calls `resume_traffic()` again after `Migration::run` returns. The `Drop` +impl is the last backstop, for a panic or for an aborted task future. -### AbortTimeout — the one pre-point-of-no-return stop +### AbortTimeout -When `cutover_timeout_action = "abort"` and the timeout fires in `wait_for_cutover()`, the code -explicitly calls `maintenance_mode::stop()` before returning `Err(Error::AbortTimeout)`. This is -the only code path that stops maintenance mode without being inside `ok_or_abort!` — it is the -case where the cutover was never attempted, so no data was moved and no swap occurred. +When `cutover_timeout_action = "abort"` and the timeout fires in `wait_for_catchup()`, the policy +returns `Err(Error::AbortTimeout)`. `prepare_cutover()` then resumes traffic. The cutover was +never attempted, so no data moved and no swap occurred. ### Idempotency guarantees diff --git a/integration/rust/tests/integration/admin/resharding/mod.rs b/integration/rust/tests/integration/admin/resharding/mod.rs index ec4330349..d8f42d7ad 100644 --- a/integration/rust/tests/integration/admin/resharding/mod.rs +++ b/integration/rust/tests/integration/admin/resharding/mod.rs @@ -1,12 +1,16 @@ pub mod copy_data; pub mod replication; +pub mod replication_slots; #[allow(clippy::module_inception)] pub mod resharding; pub mod schema_sync; pub mod table_copies; +use std::panic::{AssertUnwindSafe, resume_unwind}; use std::time::Duration; +use futures_util::FutureExt; + use crate::setup::connection_sqlx_direct_db; use sqlx::{Executor, Pool, Postgres, Row}; use tokio::time::{sleep, timeout}; @@ -104,13 +108,29 @@ async fn drop_test_slots(direct: &Pool) { .await; } +async fn test_slot_names(pool: &Pool) -> Vec { + sqlx::query_scalar(&format!( + "SELECT slot_name FROM pg_replication_slots WHERE {SLOT_FILTER} ORDER BY slot_name" + )) + .fetch_all(pool) + .await + .expect("replication slots must be readable") +} + +async fn drop_all_test_slots(direct: &Pool) { + drop_test_slots(direct).await; + for db in &["shard_0", "shard_1"] { + drop_test_slots(&connection_sqlx_direct_db(db).await).await; + } +} + async fn cleanup(admin: &Pool, direct: &Pool) { drain_tasks(admin).await; let _ = admin.execute("RELOAD").await; sleep(Duration::from_millis(500)).await; - drop_test_slots(direct).await; + drop_all_test_slots(direct).await; let _ = direct .execute(format!("DROP PUBLICATION IF EXISTS {TEST_PUB}").as_str()) @@ -119,6 +139,20 @@ async fn cleanup(admin: &Pool, direct: &Pool) { drop_test_schemas(direct).await; } +async fn with_cleanup( + admin: &Pool, + direct: &Pool, + body: impl Future, +) { + let result = AssertUnwindSafe(body).catch_unwind().await; + + cleanup(admin, direct).await; + + if let Err(panic) = result { + resume_unwind(panic); + } +} + async fn wait_for_task_status(admin: &Pool, task_id: i64, status: TaskProgress) { let result = timeout(Duration::from_secs(30), async { loop { diff --git a/integration/rust/tests/integration/admin/resharding/replication.rs b/integration/rust/tests/integration/admin/resharding/replication.rs index 86b8dbd04..1e398c13b 100644 --- a/integration/rust/tests/integration/admin/resharding/replication.rs +++ b/integration/rust/tests/integration/admin/resharding/replication.rs @@ -1,51 +1,106 @@ use std::time::Duration; -use crate::setup::{admin_sqlx, connection_sqlx_direct}; +use crate::setup::{ + admin_sqlx, connection_sqlx_direct, connection_sqlx_direct_db, connections_sqlx, +}; use pgdog_stats::TaskProgress; use sqlx::{Executor, Pool, Postgres, Row}; use tokio::time::{sleep, timeout}; +use super::table_copies::poll; use super::{ - POLL, TEST_PUB, Tasks, cleanup, create_publication, create_test_table, run_task_command, - seed_rows, task_status_line, wait_for_task, wait_for_task_status, + POLL, TEST_PUB, TEST_SCHEMA, TEST_TABLE, cleanup, create_publication, create_test_table, + fail_if_task_errored, run_task_command, seed_rows, task_status_line, test_slot_names, + wait_for_task, wait_for_task_status, with_cleanup, }; -async fn start_replication(admin: &Pool, direct: &Pool) -> i64 { - admin.execute("RELOAD").await.unwrap(); - sleep(Duration::from_millis(500)).await; - - direct - .execute(format!("CREATE PUBLICATION {TEST_PUB} FOR ALL TABLES").as_str()) - .await - .unwrap(); - - let row = admin - .fetch_one(format!("REPLICATE pgdog pgdog_sharded {TEST_PUB}").as_str()) - .await - .unwrap(); - let task_id: i64 = row.get::("task_id").parse().unwrap(); +pub(super) async fn prepare_replication(admin: &Pool, direct: &Pool) { + create_test_table(direct).await; + for database in ["shard_0", "shard_1"] { + create_test_table(&connection_sqlx_direct_db(database).await).await; + } + create_publication(direct).await; + admin.execute("RELOAD").await.expect("reload must succeed"); +} - let appeared = timeout(Duration::from_secs(10), async { - loop { - if Tasks::fetch(admin) - .await - .find(task_id) - .is_some_and(|t| t.kind == "replication pgdog -> pgdog_sharded") - { - return; - } - sleep(POLL).await; - } +pub(super) async fn start_replication(admin: &Pool, slot: Option<&str>) -> i64 { + let command = match slot { + Some(slot) => format!("REPLICATE pgdog pgdog_sharded {TEST_PUB} {slot}"), + None => format!("REPLICATE pgdog pgdog_sharded {TEST_PUB}"), + }; + let task_id = run_task_command(admin, &command).await; + wait_for_task(admin, "replication ready", |task| { + task.id == Some(task_id) && task.inner_status == "replicating" }) .await; - assert!( - appeared.is_ok(), - "replication task {task_id} did not appear in SHOW TASKS in time" - ); task_id } +pub(super) async fn wait_for_values( + admin: &Pool, + task_id: i64, + expected: &[(i64, &str)], +) { + let expected: Vec<(i64, String)> = expected + .iter() + .map(|(id, value)| (*id, (*value).to_owned())) + .collect(); + for database in ["shard_0", "shard_1"] { + let shard = connection_sqlx_direct_db(database).await; + poll(&format!("replicated values on {database}"), || async { + fail_if_task_errored(admin, task_id).await; + let actual = sqlx::query_as::<_, (i64, String)>(&format!( + "SELECT id, val FROM {TEST_SCHEMA}.{TEST_TABLE} ORDER BY id" + )) + .fetch_all(&shard) + .await + .expect("destination rows must be readable"); + (actual == expected).then_some(()) + }) + .await; + } +} + +#[tokio::test] +async fn test_replicate_streams_changes_without_copying_existing_rows() { + let direct = connection_sqlx_direct().await; + let admin = admin_sqlx().await; + cleanup(&admin, &direct).await; + + with_cleanup(&admin, &direct, async { + prepare_replication(&admin, &direct).await; + seed_rows(&direct, 1).await; + + let task_id = start_replication(&admin, None).await; + direct + .execute( + format!( + "INSERT INTO {TEST_SCHEMA}.{TEST_TABLE} (id, val) \ + VALUES (2, 'inserted'), (3, 'removed')" + ) + .as_str(), + ) + .await + .expect("source inserts must succeed"); + wait_for_values(&admin, task_id, &[(2, "inserted"), (3, "removed")]).await; + + direct + .execute( + format!("UPDATE {TEST_SCHEMA}.{TEST_TABLE} SET val = 'updated' WHERE id = 2") + .as_str(), + ) + .await + .expect("source update must succeed"); + direct + .execute(format!("DELETE FROM {TEST_SCHEMA}.{TEST_TABLE} WHERE id = 3").as_str()) + .await + .expect("source delete must succeed"); + wait_for_values(&admin, task_id, &[(2, "updated")]).await; + }) + .await; +} + #[tokio::test] async fn test_cutover_without_replication_task() { let direct = connection_sqlx_direct().await; @@ -66,37 +121,128 @@ async fn test_stop_task() { let admin = admin_sqlx().await; cleanup(&admin, &direct).await; - let task_id = start_replication(&admin, &direct).await; + with_cleanup(&admin, &direct, async { + prepare_replication(&admin, &direct).await; + let task_id = start_replication(&admin, None).await; - let row = admin - .fetch_one(format!("STOP_TASK {task_id}").as_str()) - .await - .unwrap(); - assert_eq!(row.get::("stop_task"), "OK"); + let slots = test_slot_names(&direct).await; + assert_eq!( + slots.len(), + 1, + "replication must create one slot on the source: {slots:?}" + ); - wait_for_task_status(&admin, task_id, TaskProgress::Cancelled).await; - cleanup(&admin, &direct).await; + let row = admin + .fetch_one(format!("STOP_TASK {task_id}").as_str()) + .await + .unwrap(); + assert_eq!(row.get::("stop_task"), "OK"); + + wait_for_task_status(&admin, task_id, TaskProgress::Cancelled).await; + + poll( + "the replication slot to be dropped on the source", + || async { test_slot_names(&direct).await.is_empty().then_some(()) }, + ) + .await; + }) + .await; } #[tokio::test] -async fn test_cutover() { +async fn test_cutover_starts_reverse_replication() { let direct = connection_sqlx_direct().await; let admin = admin_sqlx().await; cleanup(&admin, &direct).await; - create_test_table(&direct).await; - seed_rows(&direct, 20).await; - create_publication(&direct).await; + with_cleanup(&admin, &direct, async { + create_test_table(&direct).await; + seed_rows(&direct, 20).await; + create_publication(&direct).await; + + let task_id = + run_task_command(&admin, &format!("COPY_DATA pgdog pgdog_sharded {TEST_PUB}")).await; + + wait_for_task(&admin, "copy_data replicating", |t| { + t.id == Some(task_id) && t.inner_status == "replicating" + }) + .await; + + let cutover_ok = timeout(Duration::from_secs(10), async { + loop { + if let Ok(row) = admin.fetch_one("CUTOVER").await + && row.get::("cutover") == "OK" + { + return; + } + sleep(POLL).await; + } + }) + .await; + assert!( + cutover_ok.is_ok(), + "CUTOVER never returned OK ({})", + task_status_line(&admin, task_id).await + ); + + let connections = connections_sqlx().await; + poll("traffic to switch to the destination", || async { + fail_if_task_errored(&admin, task_id).await; + let database = sqlx::query_scalar::<_, String>("SELECT current_database()") + .fetch_one(&connections[0]) + .await + .ok()?; + matches!(database.as_str(), "shard_0" | "shard_1").then_some(()) + }) + .await; + connections[0] + .execute( + format!( + "INSERT INTO {TEST_SCHEMA}.{TEST_TABLE} (id, val) \ + VALUES (1001, 'written_after_cutover')" + ) + .as_str(), + ) + .await + .expect("writes through the new source must succeed"); - let task_id = - run_task_command(&admin, &format!("COPY_DATA pgdog pgdog_sharded {TEST_PUB}")).await; + for database in ["shard_0", "shard_1"] { + let shard = connection_sqlx_direct_db(database).await; + let value: String = sqlx::query_scalar(&format!( + "SELECT val FROM {TEST_SCHEMA}.{TEST_TABLE} WHERE id = 1001" + )) + .fetch_one(&shard) + .await + .expect("the new source must contain the post-cutover row"); + assert_eq!(value, "written_after_cutover"); + } - wait_for_task(&admin, "copy_data replicating", |t| { - t.id == Some(task_id) && t.inner_status == "replicating" + poll( + "the post-cutover row to replicate back to the old source", + || async { + fail_if_task_errored(&admin, task_id).await; + let value: Option = sqlx::query_scalar(&format!( + "SELECT val FROM {TEST_SCHEMA}.{TEST_TABLE} WHERE id = 1001" + )) + .fetch_optional(&direct) + .await + .expect("the old source must remain readable"); + (value.as_deref() == Some("written_after_cutover")).then_some(()) + }, + ) + .await; + + admin + .execute(format!("STOP_TASK {task_id}").as_str()) + .await + .expect("the migration task must stop"); + wait_for_task_status(&admin, task_id, TaskProgress::Finished).await; }) .await; +} - let cutover_ok = timeout(Duration::from_secs(10), async { +async fn request_cutover(admin: &Pool, task_id: i64) { + let accepted = timeout(Duration::from_secs(10), async { loop { if let Ok(row) = admin.fetch_one("CUTOVER").await && row.get::("cutover") == "OK" @@ -107,12 +253,62 @@ async fn test_cutover() { } }) .await; + assert!( - cutover_ok.is_ok(), + accepted.is_ok(), "CUTOVER never returned OK ({})", - task_status_line(&admin, task_id).await + task_status_line(admin, task_id).await ); +} + +async fn wait_for_traffic(admin: &Pool, task_id: i64, expected: &[&str]) { + let connections = connections_sqlx().await; + poll("traffic to switch", || async { + fail_if_task_errored(admin, task_id).await; + let database = sqlx::query_scalar::<_, String>("SELECT current_database()") + .fetch_one(&connections[0]) + .await + .ok()?; + expected.contains(&database.as_str()).then_some(()) + }) + .await; +} - wait_for_task_status(&admin, task_id, TaskProgress::Finished).await; +#[tokio::test] +async fn test_three_cutovers_alternate_the_traffic_target() { + let direct = connection_sqlx_direct().await; + let admin = admin_sqlx().await; cleanup(&admin, &direct).await; + + with_cleanup(&admin, &direct, async { + create_test_table(&direct).await; + seed_rows(&direct, 20).await; + create_publication(&direct).await; + + let task_id = + run_task_command(&admin, &format!("COPY_DATA pgdog pgdog_sharded {TEST_PUB}")).await; + + wait_for_task(&admin, "copy_data replicating", |t| { + t.id == Some(task_id) && t.inner_status == "replicating" + }) + .await; + + request_cutover(&admin, task_id).await; + wait_for_traffic(&admin, task_id, &["shard_0", "shard_1"]).await; + + request_cutover(&admin, task_id).await; + wait_for_traffic(&admin, task_id, &["pgdog"]).await; + + request_cutover(&admin, task_id).await; + wait_for_traffic(&admin, task_id, &["shard_0", "shard_1"]).await; + + fail_if_task_errored(&admin, task_id).await; + + admin + .execute(format!("STOP_TASK {task_id}").as_str()) + .await + .expect("the migration task must stop"); + wait_for_task_status(&admin, task_id, TaskProgress::Finished).await; + }) + .await; } diff --git a/integration/rust/tests/integration/admin/resharding/replication_slots.rs b/integration/rust/tests/integration/admin/resharding/replication_slots.rs new file mode 100644 index 000000000..b1d4145d5 --- /dev/null +++ b/integration/rust/tests/integration/admin/resharding/replication_slots.rs @@ -0,0 +1,95 @@ +use crate::setup::{admin_sqlx, connection_sqlx_direct}; +use pgdog_stats::{Lsn, TaskProgress}; +use sqlx::postgres::PgRow; +use sqlx::{Executor, Pool, Postgres, Row}; + +use super::super::assert_layout; +use super::replication::{prepare_replication, start_replication, wait_for_values}; +use super::table_copies::poll; +use super::{cleanup, seed_rows, wait_for_task_status}; + +const SLOT_PREFIX: &str = "__pgdog_repl_admin_slots"; +const SLOT_NAME: &str = "__pgdog_repl_admin_slots_0"; + +const SHOW_REPLICATION_SLOTS_LAYOUT: &[(&str, &str)] = &[ + ("host", "TEXT"), + ("port", "INT8"), + ("database_name", "TEXT"), + ("name", "TEXT"), + ("lsn", "TEXT"), + ("lag", "TEXT"), + ("lag_bytes", "INT8"), + ("copy_data", "BOOL"), + ("last_transaction", "TEXT"), + ("last_transaction_ms", "INT8"), + ("task_id", "INT8"), +]; + +async fn slot_row(admin: &Pool) -> Option { + let rows = admin + .fetch_all("SHOW REPLICATION_SLOTS") + .await + .expect("replication slots must be readable"); + if !rows.is_empty() { + assert_layout(&rows, SHOW_REPLICATION_SLOTS_LAYOUT); + } + let mut matching = rows + .into_iter() + .filter(|row| row.get::("name") == SLOT_NAME); + let row = matching.next(); + assert!(matching.next().is_none(), "slot must appear only once"); + row +} + +#[tokio::test] +async fn test_show_replication_slots_tracks_named_stream_until_stopped() { + let direct = connection_sqlx_direct().await; + let admin = admin_sqlx().await; + cleanup(&admin, &direct).await; + prepare_replication(&admin, &direct).await; + + let task_id = start_replication(&admin, Some(SLOT_PREFIX)).await; + let row = poll("the named replication slot", || slot_row(&admin)).await; + assert_eq!(row.get::("database_name"), "pgdog"); + assert!(!row.get::("copy_data")); + + let before: String = sqlx::query_scalar("SELECT pg_current_wal_lsn()::text") + .fetch_one(&direct) + .await + .expect("source WAL position must be readable"); + let before: Lsn = before.parse().expect("source WAL position must be valid"); + seed_rows(&direct, 2).await; + wait_for_values(&admin, task_id, &[(1, "v1"), (2, "v2")]).await; + + let row = poll("the slot to acknowledge new writes", || async { + let row = slot_row(&admin).await?; + let lsn: Lsn = row + .get::("lsn") + .parse() + .expect("displayed WAL position must be valid"); + (lsn.lsn > before.lsn).then_some(row) + }) + .await; + assert!(row.get::, _>("last_transaction").is_some()); + assert!(row.get::, _>("last_transaction_ms").is_some()); + assert!(row.get::, _>("task_id").is_some()); + + admin + .execute(format!("STOP_TASK {task_id}").as_str()) + .await + .expect("replication stop must succeed"); + wait_for_task_status(&admin, task_id, TaskProgress::Cancelled).await; + poll("the stopped slot to disappear", || async { + let exists: bool = sqlx::query_scalar( + "SELECT EXISTS (SELECT 1 FROM pg_replication_slots WHERE slot_name = $1)", + ) + .bind(SLOT_NAME) + .fetch_one(&direct) + .await + .expect("source slots must be readable"); + (!exists && slot_row(&admin).await.is_none()).then_some(()) + }) + .await; + + cleanup(&admin, &direct).await; +} diff --git a/pgdog-stats/src/resharding.rs b/pgdog-stats/src/resharding.rs index aa5e0c387..da8494b1c 100644 --- a/pgdog-stats/src/resharding.rs +++ b/pgdog-stats/src/resharding.rs @@ -4,7 +4,7 @@ use derive_more::Display; use pgdog_config::ServerAuth; use serde::{Deserialize, Serialize}; -use crate::{Lsn, User}; +use crate::{Lsn, TaskId, User}; /// Replication slot. #[derive(Debug, Clone, Serialize, Deserialize)] @@ -15,6 +15,7 @@ pub struct ReplicationSlot { pub copy_data: bool, pub address: Address, pub last_transaction: Option, + pub task_id: Option, } /// Server address. @@ -79,3 +80,60 @@ pub enum SyncState { PostData, Cutover, } + +#[derive(Debug, Default, Clone, Copy, PartialEq, Serialize, Deserialize, schemars::JsonSchema)] +pub struct MissedRows { + pub inserts: usize, + pub updates: usize, + pub deletes: usize, +} + +impl MissedRows { + pub fn non_zero(&self) -> bool { + self.inserts > 0 || self.updates > 0 || self.deletes > 0 + } + + pub fn merge(&mut self, other: Self) { + self.inserts += other.inserts; + self.updates += other.updates; + self.deletes += other.deletes; + } + + pub fn record(&mut self, tag: &str) { + if tag.starts_with("INSERT") { + self.inserts += 1; + } else if tag.starts_with("UPDATE") { + self.updates += 1; + } else if tag.starts_with("DELETE") { + self.deletes += 1; + } + } +} + +impl std::fmt::Display for MissedRows { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let mut written = false; + if self.inserts > 0 { + write!(f, "insert={}", self.inserts)?; + written = true; + } + if self.updates > 0 { + write!( + f, + "{}update={}", + if written { " " } else { "" }, + self.updates + )?; + written = true; + } + if self.deletes > 0 { + write!( + f, + "{}delete={}", + if written { " " } else { "" }, + self.deletes + )?; + } + Ok(()) + } +} diff --git a/pgdog-stats/src/task.rs b/pgdog-stats/src/task.rs index 8c061d808..b9f87b442 100644 --- a/pgdog-stats/src/task.rs +++ b/pgdog-stats/src/task.rs @@ -1,20 +1,26 @@ //! Task identity, status and definition reports. use std::borrow::Cow; -use std::fmt; use std::sync::Arc; use std::time::{Duration, SystemTime}; use indexmap::IndexMap; use derive_more::{Display, Error, From, FromStr}; -use pgdog_config::CopyFormat; use pgdog_postgres_types::ToDataRowColumn; use schemars::JsonSchema; use serde::{Deserialize, Serialize}; use serde_with::{TimestampMilliSeconds, serde_as, skip_serializing_none}; -use crate::{Lsn, SyncState}; +pub mod copy_data; +pub mod replication; +pub mod reshard; +pub mod schema_sync; + +pub use copy_data::*; +pub use replication::*; +pub use reshard::*; +pub use schema_sync::*; /// Identity of a task in the registry. Ids are unique per registry. #[derive( @@ -57,11 +63,12 @@ pub enum TaskStatus { // v1, in use SchemaSync(SchemaSyncStatus), SchemaShard(SchemaShardStatus), - // in progress, not used TableCopy(TableCopyStatus), CopyData(CopyDataStatus), + // in progress, not used Replication(ReplicationStatus), - ReplicationSlot(ReplicationSlotStatus), + ReplicationCluster(ReplicationClusterStatus), + ReplicationShard(ReplicationShardStatus), Reshard(ReshardStatus), /// Any other task status that is either doesn't report any status /// or is not compatible with other versions of tasks. @@ -343,11 +350,12 @@ pub enum TaskDefinitionKind { // v1, in use SchemaSync(SchemaSyncDefinition), SchemaShard(SchemaShardDefinition), - // In progress, not used yet CopyData(CopyDataDefinition), TableCopy(TableCopyDefinition), + // In progress, not used yet Replication(ReplicationDefinition), - ReplicationSlot(ReplicationSlotDefinition), + ReplicationCluster(ReplicationClusterDefinition), + ReplicationShard(ReplicationShardDefinition), Reshard(ReshardDefinition), /// No detail beyond the name, or a `kind` this build does not know. #[default] @@ -362,11 +370,12 @@ impl TaskDefinitionKind { match self { Self::Reshard(_) => "reshard", Self::CopyData(_) => "copy_data", - Self::SchemaSync(_) => "schema_sync", - Self::Replication(_) => "replication", Self::TableCopy(_) => "table_copy", - Self::ReplicationSlot(_) => "replication_slot", + Self::SchemaSync(_) => "schema_sync", Self::SchemaShard(_) => "schema_shard", + Self::Replication(_) => "replication", + Self::ReplicationCluster(_) => "replication_cluster", + Self::ReplicationShard(_) => "replication_shard", Self::Other => "other", } } @@ -380,47 +389,6 @@ pub struct Databases { pub destination: String, } -/// The full migration one reshard task runs, and which phases it was asked -/// to skip. -#[derive(Debug, Clone, PartialEq, Display, Serialize, Deserialize, JsonSchema)] -#[display("reshard {databases}")] -pub struct ReshardDefinition { - pub databases: Databases, - pub skip_schema_sync: bool, - pub replicate_only: bool, - pub sync_only: bool, - pub auto_cutover: bool, -} - -/// The bulk data copy one copy-data task runs. -#[derive(Debug, Clone, PartialEq, Display, Serialize, Deserialize, JsonSchema)] -#[display("copy_data {databases}")] -pub struct CopyDataDefinition { - pub databases: Databases, - pub format: CopyFormat, -} - -/// The schema sync one schema-sync task runs, and at which stage. -#[derive(Debug, Clone, PartialEq, Display, Serialize, Deserialize, JsonSchema)] -#[display("schema_sync({sync_state}) {databases}")] -pub struct SchemaSyncDefinition { - pub databases: Databases, - pub sync_state: SyncState, - pub ignore_errors: bool, - pub dry_run: bool, -} - -/// The replication stream one replication task drives. -#[derive(Debug, Clone, PartialEq, Display, Serialize, Deserialize, JsonSchema)] -#[display("replication {databases}{}", if *reverse { " (reverse)" } else { "" })] -pub struct ReplicationDefinition { - pub databases: Databases, - /// The post-cutover stream that backs a rollback, rather than the - /// initial migration. - pub reverse: bool, - pub auto_cutover: bool, -} - /// Generic "`done` of `total`" counter, reusable by any task that can count its /// work. #[derive(Debug, Clone, Copy, PartialEq, Eq, Display, Serialize, Deserialize, JsonSchema)] @@ -430,269 +398,11 @@ pub struct RatioProgress { pub total: u64, } -/// Stages of the migration, reported as the task's status. The fine-grained -/// schema-sync, copy, and replication stages live on the child tasks. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Display, Serialize, Deserialize, JsonSchema)] -#[serde(tag = "status", rename_all = "snake_case")] -pub enum ReshardStatus { - /// Running the pre-data schema-sync child task. - #[display("syncing schema")] - SchemaSync, - /// Running the data-copy child task. - #[display("syncing data")] - SyncingData, - /// Running the post-data schema-sync child task (indexes, constraints). - #[display("finalizing schema")] - FinalizingSchema, - /// Running the replication child task. - #[display("replicating")] - Replication, - /// A stage this build does not know. - #[display("")] - #[serde(other)] - Other, -} - -/// Stages of a bulk data copy, reported as the task's status. Per-table -/// progress lives on the [`TableCopyStatus`] child tasks. -#[skip_serializing_none] -#[derive(Debug, Clone, PartialEq, Eq, Display, Serialize, Deserialize, JsonSchema)] -#[display("{stage}")] -pub struct CopyDataStatus { - pub stage: CopyDataStage, - pub tables_per_shard: Option>, -} - -impl From for CopyDataStatus { - fn from(stage: CopyDataStage) -> Self { - Self { - stage, - tables_per_shard: None, - } - } -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, Display, Serialize, Deserialize, JsonSchema)] -#[serde(tag = "status", rename_all = "snake_case")] -pub enum CopyDataStage { - /// Fetching table and column metadata from the source. - #[display("loading table metadata")] - LoadingTableMetadata, - /// Checking that every table has a usable replica identity. - #[display("validating tables")] - ValidatingTables, - /// Creating the replication slots the copy reads from. - #[display("creating slots")] - CreatingSlots, - /// Copying table data to the destination shards. - #[display("copying tables")] - CopyingTables, - /// A stage this build does not know. - #[display("")] - #[serde(other)] - Other, -} - -/// One statement of a schema sync phase. -#[derive(Debug, Clone, PartialEq, Eq, Display, Serialize, Deserialize, JsonSchema)] -#[display("{sql}")] -pub struct SchemaSyncStatement { - pub sql: String, - /// The statement tolerates an "already exists" error from Postgres. - pub skip_if_exists: bool, -} - -impl SchemaSyncStatement { - pub fn new(sql: impl Into) -> Self { - Self { - sql: sql.into(), - skip_if_exists: false, - } - } - - pub fn set_skip_if_exists(mut self) -> Self { - self.skip_if_exists = true; - self - } -} - -/// Status of a schema sync. The phase it applies lives on -/// [`SchemaSyncDefinition`]. The plan is reported once, by the parent task. -/// Each shard subtask reports a cursor into it as a [`SchemaShardStatus`]. -#[derive(Debug, Clone, Default, PartialEq, Eq, Display, Serialize, Deserialize, JsonSchema)] -#[serde(tag = "status", rename_all = "snake_case")] -pub enum SchemaSyncStatus { - /// Dumping the schema from the source. - #[default] - #[display("loading schema")] - LoadingSchema, - /// The dump is loaded and the phase's statements are known. - #[display("applying {} statements", statements.len())] - ApplyingStatements { - #[serde(default)] - statements: Arc>, - }, - /// A status this build does not know. - #[display("")] - #[serde(other)] - Other, -} - -/// A statement one shard could not apply. `index` points into the plan the -/// parent task reported, and `message` is the error Postgres returned. The -/// display is one-based, to match the statement counter in the logs. -#[derive(Debug, Clone, PartialEq, Eq, Display, Serialize, Deserialize, JsonSchema)] -#[display("statement {}: {message}", index + 1)] -pub struct SchemaStatementFailure { - pub index: u64, - pub message: String, -} - -/// How far one destination shard got through the phase's statements. `applied` -/// counts the statements this shard ran, and `failures` records the ones it -/// could not. -#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize, JsonSchema)] -#[serde(default)] -pub struct SchemaShardStatus { - pub shard: u64, - pub total: u64, - pub applied: u64, - pub skipped: u64, - pub failures: Vec, -} - -impl SchemaShardStatus { - pub fn new(shard: u64, total: u64) -> Self { - Self { - shard, - total, - ..Default::default() - } - } - - pub fn done(&self) -> u64 { - self.applied + self.skipped + self.failures.len() as u64 - } -} - -impl fmt::Display for SchemaShardStatus { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - write!( - f, - "shard {}: {}/{} statements", - self.shard, - self.done(), - self.total - )?; - - if self.skipped > 0 { - write!(f, ", {} skipped", self.skipped)?; - } - if !self.failures.is_empty() { - write!(f, ", {} failed", self.failures.len())?; - } - - Ok(()) - } -} - -/// Stages of logical replication, reported as the task's status. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Display, Serialize, Deserialize, JsonSchema)] -#[serde(tag = "status", rename_all = "snake_case")] -pub enum ReplicationStatus { - /// Streaming changes to catch the destination up. - #[display("replicating")] - Replicating, - /// Cutting traffic over to the destination. - #[display("cutting over")] - CuttingOver, - /// Cutting traffic back to the original after a prior cutover (rollback). - #[display("rolling back")] - RollingBack, - /// Winding down on a stop request. - #[display("stopping")] - Stopping, - /// A stage this build does not know. - #[display("")] - #[serde(other)] - Other, -} - -/// The slot one per-shard replication subtask streams from. -#[derive(Debug, Clone, PartialEq, Display, Serialize, Deserialize, JsonSchema)] -#[display("{slot} on {host}:{port}/{database_name}")] -pub struct ReplicationSlotDefinition { - pub slot: String, - pub host: String, - pub port: u16, - pub database_name: String, - /// Temporary slot taken for an initial data copy, rather than a persistent - /// streaming slot. - pub copy_data: bool, -} - -/// How far one replication slot has streamed. -#[derive(Debug, Clone, Copy, PartialEq, Display, Serialize, Deserialize, JsonSchema)] -#[display("lag {lag_bytes} bytes at {lsn}")] -pub struct ReplicationSlotStatus { - pub lsn: Lsn, - /// `pg_current_wal_lsn() - confirmed_flush_lsn`. - pub lag_bytes: i64, - /// Epoch millis of the last transaction applied through this slot. - pub last_transaction: Option, -} - -/// The table one copy subtask is copying. -#[derive(Debug, Clone, PartialEq, Display, Serialize, Deserialize, JsonSchema)] -#[display("{schema}.{table}")] -pub struct TableCopyDefinition { - pub schema: String, - pub table: String, - pub source_shard: usize, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, Display, Serialize, Deserialize, JsonSchema)] -#[serde(tag = "status", rename_all = "snake_case")] -pub enum TableCopyStage { - #[display("estimating")] - Estimation, - #[display("waiting")] - WaitingForCopyHandler, - #[display("copy in progress")] - InProgress { rows: u64, bytes: u64 }, - #[display("error backoff")] - ErrorBackoff, - #[display("")] - #[serde(other)] - Other, -} - -/// How much of one table has been copied. -#[skip_serializing_none] -#[derive(Debug, Clone, PartialEq, Eq, Display, Serialize, Deserialize, JsonSchema)] -#[display("{stage}")] -pub struct TableCopyStatus { - pub stage: TableCopyStage, - pub attempt: usize, - pub estimated_rows: Option, - pub estimated_bytes: Option, - pub rows_per_sec: Option, - pub bytes_per_sec: Option, - pub last_error: Option, -} - -/// The destination shard one schema-sync subtask restores into. -#[derive(Debug, Clone, PartialEq, Display, Serialize, Deserialize, JsonSchema)] -#[display("shard {shard} of {databases} ({sync_state})")] -pub struct SchemaShardDefinition { - pub shard: u64, - pub databases: Databases, - pub sync_state: SyncState, -} - #[cfg(test)] mod test { use super::*; + use crate::{Lsn, MissedRows, SyncState}; + use pgdog_config::CopyFormat; use std::time::{Duration, UNIX_EPOCH}; fn table_copy() -> TableCopyDefinition { @@ -712,7 +422,7 @@ mod test { /// One definition per kind. The exhaustive `match` in /// [`test_definition_round_trip`] forces a new kind into this list. - fn definitions() -> [TaskDefinition; 8] { + fn definitions() -> [TaskDefinition; 9] { [ "test task".into(), ReshardDefinition { @@ -735,19 +445,23 @@ mod test { dry_run: false, } .into(), + table_copy().into(), ReplicationDefinition { databases: databases(), - reverse: true, auto_cutover: false, } .into(), - table_copy().into(), - ReplicationSlotDefinition { + ReplicationClusterDefinition { + databases: databases(), + direction: ReplicationDirection::Reverse, + } + .into(), + ReplicationShardDefinition { slot: "pgdog_0".into(), host: "127.0.0.1".into(), port: 5432, database_name: "prod".into(), - copy_data: false, + source_shard: 0, } .into(), SchemaShardDefinition { @@ -778,11 +492,12 @@ mod test { match back.kind { TaskDefinitionKind::Reshard(_) | TaskDefinitionKind::CopyData(_) - | TaskDefinitionKind::SchemaSync(_) - | TaskDefinitionKind::Replication(_) | TaskDefinitionKind::TableCopy(_) - | TaskDefinitionKind::ReplicationSlot(_) + | TaskDefinitionKind::SchemaSync(_) | TaskDefinitionKind::SchemaShard(_) + | TaskDefinitionKind::Replication(_) + | TaskDefinitionKind::ReplicationCluster(_) + | TaskDefinitionKind::ReplicationShard(_) | TaskDefinitionKind::Other => (), } } @@ -841,16 +556,20 @@ mod test { ); } - // Only the reverse stream is marked; the forward one reads plainly. - for (reverse, expected) in [ - (false, "replication prod -> prod_sharded"), - (true, "replication prod -> prod_sharded (reverse)"), + for (direction, expected) in [ + ( + ReplicationDirection::Forward, + "replication stream prod -> prod_sharded (forward)", + ), + ( + ReplicationDirection::Reverse, + "replication stream prod -> prod_sharded (reverse)", + ), ] { assert_eq!( - TaskDefinition::from(ReplicationDefinition { + TaskDefinition::from(ReplicationClusterDefinition { databases: databases(), - reverse, - auto_cutover: false, + direction, }) .to_string(), expected @@ -944,14 +663,33 @@ mod test { last_error: Some("connection reset".into()), }), TaskStatus::Replication(ReplicationStatus::Replicating), - TaskStatus::ReplicationSlot(ReplicationSlotStatus { + TaskStatus::ReplicationCluster(ReplicationClusterStatus::Replicating { + direction: ReplicationDirection::Reverse, + progress: ReplicationProgress { + lag_bytes: Some(2048), + last_transaction_ms: Some(150), + rows: 10, + bytes: 4096, + rows_per_sec: Some(5), + bytes_per_sec: Some(2048), + }, + }), + TaskStatus::ReplicationShard(ReplicationShardStatus { lsn: Lsn { high: 0, low: 16, lsn: 16, }, - lag_bytes: 4096, - last_transaction: Some(1_700_000_000_000), + lag_bytes: Some(4096), + missed_rows: MissedRows { + inserts: 1, + updates: 2, + deletes: 3, + }, + rows: 0, + bytes: 0, + rows_per_sec: None, + bytes_per_sec: None, }), TaskStatus::Other, ]; @@ -969,7 +707,8 @@ mod test { | TaskStatus::SchemaShard(_) | TaskStatus::TableCopy(_) | TaskStatus::Replication(_) - | TaskStatus::ReplicationSlot(_) + | TaskStatus::ReplicationCluster(_) + | TaskStatus::ReplicationShard(_) | TaskStatus::Other => (), } } @@ -1122,12 +861,12 @@ mod test { "public.users" ); assert_eq!( - TaskDefinitionKind::from(ReplicationSlotDefinition { + TaskDefinitionKind::from(ReplicationShardDefinition { slot: "pgdog_0".into(), host: "127.0.0.1".into(), port: 5432, database_name: "prod".into(), - copy_data: false, + source_shard: 0, }) .to_string(), "pgdog_0 on 127.0.0.1:5432/prod" diff --git a/pgdog-stats/src/task/copy_data.rs b/pgdog-stats/src/task/copy_data.rs new file mode 100644 index 000000000..f5b738a03 --- /dev/null +++ b/pgdog-stats/src/task/copy_data.rs @@ -0,0 +1,96 @@ +//! Copy-data task definition and status, and the per-table copy subtask. + +use derive_more::Display; +use pgdog_config::CopyFormat; +use schemars::JsonSchema; +use serde::{Deserialize, Serialize}; +use serde_with::skip_serializing_none; + +use crate::Databases; + +/// The bulk data copy one copy-data task runs. +#[derive(Debug, Clone, PartialEq, Display, Serialize, Deserialize, JsonSchema)] +#[display("copy_data {databases}")] +pub struct CopyDataDefinition { + pub databases: Databases, + pub format: CopyFormat, +} + +/// Stages of a bulk data copy, reported as the task's status. Per-table +/// progress lives on the [`TableCopyStatus`] child tasks. +#[skip_serializing_none] +#[derive(Debug, Clone, PartialEq, Eq, Display, Serialize, Deserialize, JsonSchema)] +#[display("{stage}")] +pub struct CopyDataStatus { + pub stage: CopyDataStage, + pub tables_per_shard: Option>, +} + +impl From for CopyDataStatus { + fn from(stage: CopyDataStage) -> Self { + Self { + stage, + tables_per_shard: None, + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Display, Serialize, Deserialize, JsonSchema)] +#[serde(tag = "status", rename_all = "snake_case")] +pub enum CopyDataStage { + /// Fetching table and column metadata from the source. + #[display("loading table metadata")] + LoadingTableMetadata, + /// Checking that every table has a usable replica identity. + #[display("validating tables")] + ValidatingTables, + /// Creating the replication slots the copy reads from. + #[display("creating slots")] + CreatingSlots, + /// Copying table data to the destination shards. + #[display("copying tables")] + CopyingTables, + /// A stage this build does not know. + #[display("")] + #[serde(other)] + Other, +} + +/// The table one copy subtask is copying. +#[derive(Debug, Clone, PartialEq, Display, Serialize, Deserialize, JsonSchema)] +#[display("{schema}.{table}")] +pub struct TableCopyDefinition { + pub schema: String, + pub table: String, + pub source_shard: usize, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Display, Serialize, Deserialize, JsonSchema)] +#[serde(tag = "status", rename_all = "snake_case")] +pub enum TableCopyStage { + #[display("estimating")] + Estimation, + #[display("waiting")] + WaitingForCopyHandler, + #[display("copy in progress")] + InProgress { rows: u64, bytes: u64 }, + #[display("error backoff")] + ErrorBackoff, + #[display("")] + #[serde(other)] + Other, +} + +/// How much of one table has been copied. +#[skip_serializing_none] +#[derive(Debug, Clone, PartialEq, Eq, Display, Serialize, Deserialize, JsonSchema)] +#[display("{stage}")] +pub struct TableCopyStatus { + pub stage: TableCopyStage, + pub attempt: usize, + pub estimated_rows: Option, + pub estimated_bytes: Option, + pub rows_per_sec: Option, + pub bytes_per_sec: Option, + pub last_error: Option, +} diff --git a/pgdog-stats/src/task/replication.rs b/pgdog-stats/src/task/replication.rs new file mode 100644 index 000000000..3376701e0 --- /dev/null +++ b/pgdog-stats/src/task/replication.rs @@ -0,0 +1,191 @@ +//! Replication task definitions and statuses: the migration, one cluster +//! stream of it, and one shard slot of that stream. + +use std::fmt; + +use derive_more::Display; +use schemars::JsonSchema; +use serde::{Deserialize, Serialize}; + +use crate::{Databases, Lsn, MissedRows}; + +/// Direction of a replication task: the initial migration (`Forward`) or the +/// post-cutover reverse stream that backs a rollback (`Reverse`). A `CUTOVER` +/// on a `Reverse` task is therefore a rollback. Affects reported status only, +/// not control flow. +#[derive( + Debug, Clone, Copy, Default, PartialEq, Eq, Display, Serialize, Deserialize, JsonSchema, +)] +#[serde(rename_all = "snake_case")] +#[display(rename_all = "snake_case")] +pub enum ReplicationDirection { + #[default] + Forward, + Reverse, +} + +/// Why the replication task stopped waiting and cut traffic over. +#[derive( + Debug, Clone, Copy, Default, PartialEq, Eq, Display, Serialize, Deserialize, JsonSchema, +)] +#[serde(rename_all = "snake_case")] +pub enum ReplicationCutoverReason { + /// Replication lag reached the configured threshold. + #[display("lag")] + Lag, + /// No transaction was applied for the configured delay. + #[display("last transaction")] + LastTransaction, + /// The configured wait expired before the other conditions were met. + #[display("timeout")] + Timeout, + /// A reason this build does not know. + #[default] + #[display("unknown")] + #[serde(other)] + Unknown, +} + +/// The migration one replication task drives, including every cutover it +/// performs. +#[derive(Debug, Clone, PartialEq, Display, Serialize, Deserialize, JsonSchema)] +#[display("replication {databases}")] +pub struct ReplicationDefinition { + pub databases: Databases, + pub auto_cutover: bool, +} + +/// Stages of logical replication, reported as the task's status. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Display, Serialize, Deserialize, JsonSchema)] +#[serde(tag = "status", rename_all = "snake_case")] +pub enum ReplicationStatus { + /// Streaming changes to catch the destination up. + #[display("replicating")] + Replicating, + /// Streaming changes back to the original source after a cutover, so a + /// rollback stays possible. + #[display("reverse replicating")] + ReverseReplicating, + #[display("stopping traffic")] + StoppingTraffic, + #[display("waiting for catch-up")] + WaitingForCatchUp, + #[display("syncing schema")] + SyncingSchema, + #[display("preparing reverse replication")] + PreparingReverseReplication, + /// Cutting traffic over to the destination. + #[display("cutting over")] + CuttingOver, + /// Cutting traffic back to the original after a prior cutover (rollback). + #[display("rolling back")] + RollingBack, + /// A stage this build does not know. + #[display("")] + #[serde(other)] + Other, +} + +/// The cluster one replication subtask streams until the parent task cuts +/// traffic over. `databases` always names the migration's original source and +/// destination; `direction` says which of them the changes flow from. +#[derive(Debug, Clone, PartialEq, Display, Serialize, Deserialize, JsonSchema)] +#[display("replication stream {databases} ({direction})")] +pub struct ReplicationClusterDefinition { + pub databases: Databases, + pub direction: ReplicationDirection, +} + +/// Stages of one replication cluster, reported as the subtask's status. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Display, Serialize, Deserialize, JsonSchema)] +#[serde(tag = "status", rename_all = "snake_case")] +pub enum ReplicationClusterStatus { + #[display("initializing replication streams")] + InitializingReplicationStreams, + /// Streaming changes to catch the destination up. + #[display("{direction} replicating, {progress}")] + Replicating { + direction: ReplicationDirection, + progress: ReplicationProgress, + }, + /// Stopped streaming so the parent task can cut traffic over. + #[display("stopped for cutover ({reason})")] + StoppedForCutover { reason: ReplicationCutoverReason }, + /// A stage this build does not know. + #[display("")] + #[serde(other)] + Other, +} + +/// How far the whole cluster has replicated: the largest lag of its shards, +/// how long ago the newest transaction was applied, and the rows and bytes +/// applied by every shard together. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize, JsonSchema)] +pub struct ReplicationProgress { + pub lag_bytes: Option, + pub last_transaction_ms: Option, + pub rows: u64, + pub bytes: u64, + pub rows_per_sec: Option, + pub bytes_per_sec: Option, +} + +impl fmt::Display for ReplicationProgress { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self.lag_bytes { + Some(lag) => write!(f, "lag {lag} bytes")?, + None => write!(f, "lag unknown")?, + } + if let Some(age) = self.last_transaction_ms { + write!(f, ", last transaction {age}ms ago")?; + } + write!(f, ", applied {} rows {} bytes", self.rows, self.bytes)?; + if let (Some(rows), Some(bytes)) = (self.rows_per_sec, self.bytes_per_sec) { + write!(f, " [since start: {rows} rows/sec, {bytes} bytes/sec]")?; + } + Ok(()) + } +} + +/// The slot one per-shard replication subtask streams from. +#[derive(Debug, Clone, PartialEq, Display, Serialize, Deserialize, JsonSchema)] +#[display("{slot} on {host}:{port}/{database_name}")] +pub struct ReplicationShardDefinition { + pub slot: String, + pub host: String, + pub port: u16, + pub database_name: String, + pub source_shard: usize, +} + +/// How far one replication slot has streamed. +#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize, JsonSchema)] +pub struct ReplicationShardStatus { + pub lsn: Lsn, + /// `pg_current_wal_lsn() - confirmed_flush_lsn`, clamped to zero. + pub lag_bytes: Option, + pub missed_rows: MissedRows, + pub rows: u64, + pub bytes: u64, + pub rows_per_sec: Option, + pub bytes_per_sec: Option, +} + +impl fmt::Display for ReplicationShardStatus { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self.lag_bytes { + Some(b) => write!(f, "lag {} bytes at {}", b, self.lsn)?, + None => write!(f, "lag unknown at {}", self.lsn)?, + } + + if self.missed_rows.non_zero() { + write!(f, ", missed {}", self.missed_rows)?; + } + + write!(f, ", applied {} rows {} bytes", self.rows, self.bytes)?; + if let (Some(rows), Some(bytes)) = (self.rows_per_sec, self.bytes_per_sec) { + write!(f, " [since start: {rows} rows/sec, {bytes} bytes/sec]")?; + } + Ok(()) + } +} diff --git a/pgdog-stats/src/task/reshard.rs b/pgdog-stats/src/task/reshard.rs new file mode 100644 index 000000000..c83b24964 --- /dev/null +++ b/pgdog-stats/src/task/reshard.rs @@ -0,0 +1,42 @@ +//! Reshard task definition and status. + +use derive_more::Display; +use schemars::JsonSchema; +use serde::{Deserialize, Serialize}; + +use crate::Databases; + +/// The full migration one reshard task runs, and which phases it was asked +/// to skip. +#[derive(Debug, Clone, PartialEq, Display, Serialize, Deserialize, JsonSchema)] +#[display("reshard {databases}")] +pub struct ReshardDefinition { + pub databases: Databases, + pub skip_schema_sync: bool, + pub replicate_only: bool, + pub sync_only: bool, + pub auto_cutover: bool, +} + +/// Stages of the migration, reported as the task's status. The fine-grained +/// schema-sync, copy, and replication stages live on the child tasks. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Display, Serialize, Deserialize, JsonSchema)] +#[serde(tag = "status", rename_all = "snake_case")] +pub enum ReshardStatus { + /// Running the pre-data schema-sync child task. + #[display("syncing schema")] + SchemaSync, + /// Running the data-copy child task. + #[display("syncing data")] + SyncingData, + /// Running the post-data schema-sync child task (indexes, constraints). + #[display("finalizing schema")] + FinalizingSchema, + /// Running the replication child task. + #[display("replicating")] + Replication, + /// A stage this build does not know. + #[display("")] + #[serde(other)] + Other, +} diff --git a/pgdog-stats/src/task/schema_sync.rs b/pgdog-stats/src/task/schema_sync.rs new file mode 100644 index 000000000..6ca8c2927 --- /dev/null +++ b/pgdog-stats/src/task/schema_sync.rs @@ -0,0 +1,132 @@ +//! Schema-sync task definitions and statuses, for the phase and its shards. + +use std::fmt; +use std::sync::Arc; + +use derive_more::Display; +use schemars::JsonSchema; +use serde::{Deserialize, Serialize}; + +use crate::{Databases, SyncState}; + +/// The schema sync one schema-sync task runs, and at which stage. +#[derive(Debug, Clone, PartialEq, Display, Serialize, Deserialize, JsonSchema)] +#[display("schema_sync({sync_state}) {databases}")] +pub struct SchemaSyncDefinition { + pub databases: Databases, + pub sync_state: SyncState, + pub ignore_errors: bool, + pub dry_run: bool, +} + +/// One statement of a schema sync phase. +#[derive(Debug, Clone, PartialEq, Eq, Display, Serialize, Deserialize, JsonSchema)] +#[display("{sql}")] +pub struct SchemaSyncStatement { + pub sql: String, + /// The statement tolerates an "already exists" error from Postgres. + pub skip_if_exists: bool, +} + +impl SchemaSyncStatement { + pub fn new(sql: impl Into) -> Self { + Self { + sql: sql.into(), + skip_if_exists: false, + } + } + + pub fn set_skip_if_exists(mut self) -> Self { + self.skip_if_exists = true; + self + } +} + +/// Status of a schema sync. The phase it applies lives on +/// [`SchemaSyncDefinition`]. The plan is reported once, by the parent task. +/// Each shard subtask reports a cursor into it as a [`SchemaShardStatus`]. +#[derive(Debug, Clone, Default, PartialEq, Eq, Display, Serialize, Deserialize, JsonSchema)] +#[serde(tag = "status", rename_all = "snake_case")] +pub enum SchemaSyncStatus { + /// Dumping the schema from the source. + #[default] + #[display("loading schema")] + LoadingSchema, + /// The dump is loaded and the phase's statements are known. + #[display("applying {} statements", statements.len())] + ApplyingStatements { + #[serde(default)] + statements: Arc>, + }, + /// A status this build does not know. + #[display("")] + #[serde(other)] + Other, +} + +/// The destination shard one schema-sync subtask restores into. +#[derive(Debug, Clone, PartialEq, Display, Serialize, Deserialize, JsonSchema)] +#[display("shard {shard} of {databases} ({sync_state})")] +pub struct SchemaShardDefinition { + pub shard: u64, + pub databases: Databases, + pub sync_state: SyncState, +} + +/// A statement one shard could not apply. `index` points into the plan the +/// parent task reported, and `message` is the error Postgres returned. The +/// display is one-based, to match the statement counter in the logs. +#[derive(Debug, Clone, PartialEq, Eq, Display, Serialize, Deserialize, JsonSchema)] +#[display("statement {}: {message}", index + 1)] +pub struct SchemaStatementFailure { + pub index: u64, + pub message: String, +} + +/// How far one destination shard got through the phase's statements. `applied` +/// counts the statements this shard ran, and `failures` records the ones it +/// could not. +#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize, JsonSchema)] +#[serde(default)] +pub struct SchemaShardStatus { + pub shard: u64, + pub total: u64, + pub applied: u64, + pub skipped: u64, + pub failures: Vec, +} + +impl SchemaShardStatus { + pub fn new(shard: u64, total: u64) -> Self { + Self { + shard, + total, + ..Default::default() + } + } + + pub fn done(&self) -> u64 { + self.applied + self.skipped + self.failures.len() as u64 + } +} + +impl fmt::Display for SchemaShardStatus { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!( + f, + "shard {}: {}/{} statements", + self.shard, + self.done(), + self.total + )?; + + if self.skipped > 0 { + write!(f, ", {} skipped", self.skipped)?; + } + if !self.failures.is_empty() { + write!(f, ", {} failed", self.failures.len())?; + } + + Ok(()) + } +} diff --git a/pgdog/src/admin/replicate.rs b/pgdog/src/admin/replicate.rs index a2749451a..660603976 100644 --- a/pgdog/src/admin/replicate.rs +++ b/pgdog/src/admin/replicate.rs @@ -68,10 +68,9 @@ impl Command for Replicate { .ignore_errors(true) .build(); - let waiter = orchestrator.replicate().await?; let task_id = run_task( ReplicationTask::builder() - .waiter(waiter) + .orchestrator(orchestrator) .schema_sync(schema_sync) .build(), ) diff --git a/pgdog/src/admin/show_replication_slots.rs b/pgdog/src/admin/show_replication_slots.rs index 7c1df83a4..368a5489c 100644 --- a/pgdog/src/admin/show_replication_slots.rs +++ b/pgdog/src/admin/show_replication_slots.rs @@ -34,6 +34,7 @@ impl Command for ShowReplicationSlots { Field::bool("copy_data"), Field::text("last_transaction"), Field::bigint("last_transaction_ms"), + Field::bigint("task_id"), ]); let mut messages = vec![rd.message()]; let now = SystemTime::now(); @@ -59,7 +60,7 @@ impl Command for ShowReplicationSlots { .add(format_bytes(slot.lag as u64).as_str()) .add(slot.lag) .add(slot.copy_data) - .add(if let Some(ref s) = last_transaction_str { + .add(if let Some(s) = &last_transaction_str { s.as_str().to_data_row_column() } else { Data::null() @@ -68,6 +69,11 @@ impl Command for ShowReplicationSlots { ms.to_data_row_column() } else { Data::null() + }) + .add(if let Some(task_id) = slot.task_id { + task_id.to_data_row_column() + } else { + Data::null() }); messages.push(row.message()); diff --git a/pgdog/src/api/copy_data.rs b/pgdog/src/api/copy_data.rs index 7cb6eb6a6..75246c012 100644 --- a/pgdog/src/api/copy_data.rs +++ b/pgdog/src/api/copy_data.rs @@ -115,16 +115,15 @@ impl Task for CopyDataTask { let ctx = ctx.clone(); let shard_number = shard.number(); let format = self.format; - // W: should we even use tasks for it? handles.push(tasks::spawn("tables copy", async move { - let table_sync_task = TableDataSyncTask { - pool, - table, - source, - dest, - format, - source_shard: shard_number, - }; + let table_sync_task = TableDataSyncTask::builder() + .pool(pool) + .table(table) + .source(source) + .dest(dest) + .format(format) + .source_shard(shard_number) + .build(); let table = ctx.run(table_sync_task).await?; @@ -146,20 +145,18 @@ impl Task for CopyDataTask { } } -#[derive(Debug)] -struct TableDataSyncTask { - pool: Arc>, - table: Table, - source: Cluster, - dest: Cluster, - format: CopyFormat, - source_shard: usize, +#[derive(Debug, bon::Builder)] +pub(crate) struct TableDataSyncTask { + pub(crate) pool: Arc>, + pub(crate) table: Table, + pub(crate) source: Cluster, + pub(crate) dest: Cluster, + pub(crate) format: CopyFormat, + pub(crate) source_shard: usize, } impl Task for TableDataSyncTask { type Status = TableCopyStatus; - - // W: table? type Output = Table; type Error = Error; @@ -180,6 +177,7 @@ impl Task for TableDataSyncTask { source: &self.source, dest: &self.dest, format: self.format, + task_id: ctx.id(), }; let cancel = ctx.cancellation_token(); let mut last_error = None; diff --git a/pgdog/src/api/replication.rs b/pgdog/src/api/replication.rs index e8b28c6da..965c7d04b 100644 --- a/pgdog/src/api/replication.rs +++ b/pgdog/src/api/replication.rs @@ -1,168 +1,738 @@ //! Logical-replication background task. -//! -//! Drives a `ReplicationWaiter` to completion. Without `auto_cutover` -//! (standalone `REPLICATE`, `copy_data`) it stops on cancellation -//! (`STOP_TASK`), cuts over on an operator `CUTOVER` addressed to this task -//! (delivered through [`ReplicationTask::cutover`]), and otherwise finishes -//! when the source slot drains (no cutover on natural drain). With -//! `auto_cutover` set (reshard) it cuts over automatically once the -//! destination has caught up. - -use std::collections::HashMap; + +use std::pin::pin; use std::sync::LazyLock; use std::time::Duration; -use parking_lot::Mutex; +use dashmap::DashMap; +use futures::future::{FusedFuture, FutureExt}; +use futures::stream::{FuturesUnordered, StreamExt}; use tokio::select; +use tokio::task::JoinHandle; use tokio_util::sync::CancellationToken; use crate::api::Task; -use crate::api::schema_sync::SchemaSyncTask; +use crate::api::schema_sync::{SchemaSyncPhase, SchemaSyncTask}; use crate::api::task::{TaskContext, TaskId}; +use crate::backend::replication::ee::{OrchestratorState, orchestrator_state}; use crate::backend::replication::logical::Error; -use crate::backend::replication::logical::orchestrator::ReplicationWaiter; -use pgdog_stats::{ReplicationDefinition, ReplicationStatus, TaskDefinition}; - -/// Direction of a replication task: the initial migration (`Forward`) or the -/// post-cutover reverse stream that backs a rollback (`Reverse`). A `CUTOVER` -/// on a `Reverse` task is therefore a rollback. Affects reported status only, -/// not control flow. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] -pub(crate) enum Direction { - #[default] - Forward, - Reverse, -} +use crate::backend::replication::logical::orchestrator::Orchestrator; +use crate::backend::replication::logical::publisher::cutover_policy::CutoverPolicy; +use crate::backend::replication::logical::publisher::replication_progress::ReplicationProgress; +use crate::backend::replication::logical::publisher::replication_stream::ReplicationStream; +use crate::backend::replication::logical::publisher::{ReplicationSlot, Table}; +use crate::backend::{ + databases::{cancel_all, cutover}, + maintenance_mode, +}; +use crate::config::config; +use crate::tasks; +use crate::util::{safe_interval, safe_timeout}; +use pgdog_stats::{ + MissedRows, ReplicationClusterDefinition, ReplicationClusterStatus, ReplicationCutoverReason, + ReplicationDefinition, ReplicationDirection, ReplicationShardDefinition, + ReplicationShardStatus, ReplicationStatus, TaskDefinition, +}; +use tracing::{info, warn}; -/// Run the replication by driving a [`ReplicationWaiter`] to completion. #[derive(Debug, bon::Builder)] pub(crate) struct ReplicationTask { - /// The running replication waiter this task drives to completion. - pub(crate) waiter: ReplicationWaiter, + pub(crate) orchestrator: Orchestrator, /// Cut over automatically once the destination has caught up, instead /// of waiting for an operator `CUTOVER`. #[builder(default)] pub(crate) auto_cutover: bool, - /// Replication direction. `Reverse` marks the post-cutover stream that - /// backs a rollback; it only affects reported status, not control flow. - #[builder(default)] - pub(crate) direction: Direction, pub(crate) schema_sync: SchemaSyncTask, } -/// Cutover tokens of the replication tasks currently awaiting an operator -/// `CUTOVER`, keyed by the root task id they belong to. A cutover token is -/// *separate* from the task's `STOP_TASK` cancellation token — signalling it -/// means "cut over", not "abandon". -static CUTOVERS: LazyLock>> = - LazyLock::new(|| Mutex::new(HashMap::new())); +/// Executes the whole replication process. It runs a replication until a cutover, +/// cuts over, then runs the opposite replication so a rollback stays possible. Each +/// cutover flips the direction, so the task never finishes on its own. +/// +/// A `STOP_TASK` during a reverse phase returns `Ok(())`, so the task reports +/// as finished: the migration is complete and the operator ended the rollback +/// window. The same signal during a forward phase returns +/// [`Error::ReplicationAborted`], which reports as cancelled. +impl Task for ReplicationTask { + type Status = ReplicationStatus; + type Output = (); + type Error = Error; -/// Guard held by a running replication task: removes its cutover -/// registration on drop. Awaiting [CutoverWaiter::requested] -/// resolves when an operator `CUTOVER` targets the task. -struct CutoverWaiter { - root_id: TaskId, - token: CancellationToken, + fn cancel_timeout() -> Duration { + // the cancellation should be handled by the task itself, + // TODO: though, some stages are not cancellable at all, + // so maybe this should be dynamic? + Duration::from_secs(600) + } + + fn definition(&self) -> impl Into { + ReplicationDefinition { + databases: self.orchestrator.databases(), + auto_cutover: self.auto_cutover, + } + } + + async fn run(self, ctx: TaskContext) -> Result<(), Error> { + let Self { + orchestrator, + schema_sync, + auto_cutover, + } = self; + + let slots = orchestrator.publication_guard(); + let mut replication = Replication::new(&ctx, orchestrator); + let result = replication.run(schema_sync, auto_cutover).await; + replication.resume_traffic(); + if let Err(err) = slots.cleanup().await { + warn!("failed to clean up replication slots: {err}"); + } + + if replication.cancelled() + && replication.direction == ReplicationDirection::Reverse + && matches!(result, Err(Error::ReplicationAborted)) + { + info!("[replication] stopped in the rollback window, migration complete"); + return Ok(()); + } + + match &result { + Ok(()) => info!("[replication] finished"), + Err(err) if replication.cancelled() => info!("[replication] cancelled: {err}"), + Err(err) => warn!("[replication] failed: {err}"), + } + + result + } } -impl CutoverWaiter { - /// Wait until a cutover is requested for this task. The token latches, so - /// a cutover that arrived earlier is delivered immediately. - async fn requested(&self) { - self.token.cancelled().await; +impl ReplicationTask { + /// Trigger a cutover on a running replication task. + pub(crate) fn trigger_cutover(target: Option) -> bool { + let token = match target { + Some(id) => CUTOVERS.get(&id).map(|entry| entry.value().clone()), + // No id: cut over the first (lowest-id) running task. + None => CUTOVERS + .iter() + .min_by_key(|entry| *entry.key()) + .map(|entry| entry.value().clone()), + }; + + match token { + Some(token) => { + token.cancel(); + true + } + None => false, + } } } -impl Drop for CutoverWaiter { - fn drop(&mut self) { - CUTOVERS.lock().remove(&self.root_id); +/// Struct to hold the replication state from the [`ReplicationTask`] +struct Replication<'a> { + ctx: &'a TaskContext, + orchestrator: Orchestrator, + direction: ReplicationDirection, + maintenance: MaintenanceMode, +} + +impl<'a> Replication<'a> { + fn new(ctx: &'a TaskContext, orchestrator: Orchestrator) -> Self { + Self { + ctx, + orchestrator, + direction: ReplicationDirection::Forward, + maintenance: MaintenanceMode::new(), + } + } + + async fn run(&mut self, schema_sync: SchemaSyncTask, auto_cutover: bool) -> Result<(), Error> { + info!( + "[replication] starting {}, auto_cutover={auto_cutover}", + self.orchestrator.databases() + ); + self.replicate_until_cutover(auto_cutover).await?; + self.sync_schema(schema_sync).await?; + + loop { + self.cutover().await?; + self.flip_direction(); + self.replicate_until_cutover(false).await?; + self.sync_schema( + // new schema sync tasks with updated orchestrator + SchemaSyncTask::builder() + .databases(self.orchestrator.databases()) + .publication(self.orchestrator.publication.clone()) + .phase(SchemaSyncPhase::Cutover) + .ignore_errors(true) + .build(), + ) + .await?; + } + } + + fn resume_traffic(&mut self) { + self.maintenance.resume_traffic(); + } + + fn cancelled(&self) -> bool { + self.ctx.cancellation_token().is_cancelled() + } + + async fn sync_schema(&self, schema_sync: SchemaSyncTask) -> Result<(), Error> { + info!("Run schema sync in {} direction", self.direction); + self.ctx.set_status(ReplicationStatus::SyncingSchema); + self.ctx.run(schema_sync).await?; + Ok(()) + } + + fn flip_direction(&mut self) { + self.direction = match self.direction { + ReplicationDirection::Reverse => ReplicationDirection::Forward, + ReplicationDirection::Forward => ReplicationDirection::Reverse, + }; + } + + /// Run the replication until we get the cutover signal and [`CutoverPolicy`] + /// waited for the stop_traffic conditions + async fn replicate_until_cutover(&mut self, auto_cutover: bool) -> Result<(), Error> { + let ctx = self.ctx; + let task_cancel = ctx.cancellation_token(); + let cutover = (!auto_cutover).then(|| CutoverWaiter::register(ctx.root_id())); + let progress = ReplicationProgress::new(self.orchestrator.source.shards().len()); + let mut cutover_reason = None; + let (cluster, stop_cluster_replication) = ReplicationClusterTask::new( + self.orchestrator.clone(), + self.direction, + progress.clone(), + ); + + info!("[replication] {} stream starting", self.direction); + orchestrator_state(OrchestratorState::Replication); + ctx.set_status(match self.direction { + ReplicationDirection::Forward => ReplicationStatus::Replicating, + ReplicationDirection::Reverse => ReplicationStatus::ReverseReplicating, + }); + let mut cluster_run = pin!(ctx.run(cluster).fuse()); + + let result = select! { + biased; + _ = task_cancel.cancelled() => { + info!("[replication] {} stream cancelled", self.direction); + Err(Error::ReplicationAborted) + }, + result = &mut cluster_run => result.and(Err(Error::ReplicationStreamStopped)), + result = async { + if let Some(cutover) = cutover.as_ref() { + cutover.requested().await; + } + self.prepare_cutover(progress).await + } => result.map(|reason| cutover_reason = Some(reason)), + }; + + if result.is_err() { + self.resume_traffic(); + } + + // stop the cluster replication and wait until it gracefully finishes, + // we should stop it despite if we succeed or not at this moment + stop_cluster_replication.stop(cutover_reason); + let drained = if cluster_run.is_terminated() { + Ok(()) + } else { + safe_timeout(ReplicationClusterTask::drain_timeout(), &mut cluster_run) + .await + .unwrap_or(Err(Error::DrainTimeout)) + }; + let result = result.and(drained); + match &result { + Ok(()) => info!("[replication] {} stream stopped", self.direction), + Err(err) => warn!("[replication] {} stream failed: {err}", self.direction), + } + result + } + + /// Wait for cutover initial conditions, stop the traffic + /// and wait until the replication catch up with the source + async fn prepare_cutover( + &mut self, + progress: ReplicationProgress, + ) -> Result { + let cutover_policy = CutoverPolicy::new(config().as_ref().into(), progress); + cutover_policy.wait_for_stop_threshold().await; + self.ctx.set_status(ReplicationStatus::StoppingTraffic); + self.maintenance.stop_traffic(); + let result = async { + cancel_all(&self.orchestrator.source.identifier().database).await?; + self.ctx.set_status(ReplicationStatus::WaitingForCatchUp); + cutover_policy.wait_for_catchup().await + } + .await; + if result.is_err() { + // in case of errors, after stop_traffic, resume it immediately + self.maintenance.resume_traffic(); + } + result + } + + /// Execute the cutover: create reverse slots, update the config, + /// refresh the orchestrator and resume traffic. + async fn cutover(&mut self) -> Result<(), Error> { + info!("Cutting over"); + self.ctx + .set_status(ReplicationStatus::PreparingReverseReplication); + self.orchestrator + .publisher() + .await + .create_slots( + &self.orchestrator.destination, + &self.ctx.cancellation_token(), + ) + .await?; + self.ctx.set_status(match self.direction { + ReplicationDirection::Forward => ReplicationStatus::CuttingOver, + ReplicationDirection::Reverse => ReplicationStatus::RollingBack, + }); + cutover( + &self.orchestrator.source.identifier().database, + &self.orchestrator.destination.identifier().database, + ) + .await?; + self.orchestrator.refresh()?; + self.maintenance.resume_traffic(); + Ok(()) } } -impl Task for ReplicationTask { - type Status = ReplicationStatus; +/// Handle that stops one replication cluster, with the cutover reason when +/// the parent task stopped it to cut traffic over. +#[derive(Debug)] +pub(crate) struct ReplicationClusterStop { + sender: tokio::sync::oneshot::Sender>, +} + +impl ReplicationClusterStop { + pub(crate) fn stop(self, cutover_reason: Option) { + let _ = self.sender.send(cutover_reason); + } +} + +/// Task that runs the replication in one direction +/// from one source cluster to another. +#[derive(Debug)] +pub(crate) struct ReplicationClusterTask { + orchestrator: Orchestrator, + progress: ReplicationProgress, + direction: ReplicationDirection, + stop: tokio::sync::oneshot::Receiver>, +} + +impl ReplicationClusterTask { + pub(crate) fn new( + orchestrator: Orchestrator, + direction: ReplicationDirection, + progress: ReplicationProgress, + ) -> (Self, ReplicationClusterStop) { + let (sender, stop) = tokio::sync::oneshot::channel(); + ( + Self { + orchestrator, + progress, + direction, + stop, + }, + ReplicationClusterStop { sender }, + ) + } +} + +impl Task for ReplicationClusterTask { + type Status = ReplicationClusterStatus; type Output = (); type Error = Error; fn cancel_timeout() -> Duration { - Duration::from_secs(60) + Duration::from_secs(120) } fn definition(&self) -> impl Into { - ReplicationDefinition { - databases: self.waiter.databases(), - reverse: self.direction == Direction::Reverse, - auto_cutover: self.auto_cutover, + ReplicationClusterDefinition { + databases: self.orchestrator.databases(), + direction: self.direction, } } - async fn run(mut self, ctx: TaskContext) -> Result<(), Error> { - let token = ctx.cancellation_token(); + async fn run(self, ctx: TaskContext) -> Result<(), Error> { + let Self { + orchestrator, + progress, + stop, + direction, + } = self; + let task_cancel = ctx.cancellation_token(); + let mut streams = ReplicationStreams::new(); + let streams_stop = CancellationToken::new(); + let _streams_stop_on_drop = streams_stop.clone().drop_guard(); + + ctx.set_status(ReplicationClusterStatus::InitializingReplicationStreams); + let init_result = Self::create_replication_shard_tasks( + &ctx, + &orchestrator, + &progress, + &streams_stop, + &mut streams, + ) + .await; + + let mut report = safe_interval(Duration::from_secs(1)); + let mut stop = stop; + let result = async { + init_result?; + loop { + ctx.set_status(ReplicationClusterStatus::Replicating { + direction, + progress: progress.snapshot(), + }); + select! { + biased; + _ = task_cancel.cancelled() => { + info!("[replication] {direction} streams cancelled, draining"); + return Ok(()); + } + stopped = &mut stop => { + match stopped { + Ok(Some(reason)) => { + info!("[replication] {direction} streams stopped for cutover ({reason}), draining"); + ctx.set_status(ReplicationClusterStatus::StoppedForCutover { reason }); + } + _ => info!("[replication] {direction} streams stopped, draining"), + } + return Ok(()); + } + // if any of streams exit early, stop the process + result = streams.next() => { + if let Some(child) = result { + child??; + } + // return an error, since it should stop by the signal + // not by itself + return Err(Error::ReplicationStreamStopped); + } + _ = report.tick() => {} + } + } + } + .await; - ctx.set_status(ReplicationStatus::Replicating); + // stop all the stream and make sure they are drained. + // If there were error on some stream it should stop other streams. + streams_stop.cancel(); + let drained = Self::drain_streams(&mut streams).await; + result.and(drained) + } +} - if self.auto_cutover { - return self.perform_cutover(&ctx, &token).await; +impl ReplicationClusterTask { + /// Create [`ReplicationShardTask`] for every source shard in the cluster + /// and track its status. + async fn create_replication_shard_tasks( + ctx: &TaskContext, + orchestrator: &Orchestrator, + progress: &ReplicationProgress, + stop: &CancellationToken, + streams: &mut ReplicationStreams, + ) -> Result<(), Error> { + let mut publisher = orchestrator.publisher().await; + publisher + .prepare_replication(&orchestrator.source, &ctx.cancellation_token()) + .await?; + for source_shard in 0..orchestrator.source.shards().len() { + let tables = publisher.pop_tables(source_shard)?; + let slot = SlotGuard::new(publisher.pop_slot(source_shard)?); + let updater = progress.updater_for_shard(source_shard); + let replication_stream = + ReplicationStream::new(&orchestrator.source, &orchestrator.destination, updater); + let task = ReplicationShardTask::builder() + .source_shard(source_shard) + .slot(slot) + .tables(tables) + .replication_stream(replication_stream) + .stop(stop.clone()) + .build(); + + streams.push(tasks::spawn("replication stream", ctx.run(task))); } - let cutover = Self::register_cutover(ctx.root_id()); + Ok(()) + } - select! { - _ = token.cancelled() => { - ctx.set_status(ReplicationStatus::Stopping); - self.waiter.stop(); - } - _ = cutover.requested() => { - self.perform_cutover(&ctx, &token).await?; + fn drain_timeout() -> Duration { + Duration::from_secs(300) + } + + fn stream_drain_timeout() -> Duration { + Duration::from_secs(120) + } + + /// Drain all the streams - make sure they are drained and generated no errors + async fn drain_streams(streams: &mut ReplicationStreams) -> Result<(), Error> { + safe_timeout(Self::stream_drain_timeout(), async { + let mut result = Ok(()); + while let Some(child) = streams.next().await { + result = result.and(child.map_err(Error::from).and_then(|result| result)); } - res = self.waiter.wait() => { - res?; + result + }) + .await + .unwrap_or_else(|_| { + streams.iter().for_each(JoinHandle::abort); + Err(Error::DrainTimeout) + }) + } +} + +/// Task for the replication stream executing on a single +/// source shard +#[derive(Debug, bon::Builder)] +pub(crate) struct ReplicationShardTask { + pub(crate) slot: SlotGuard, + pub(crate) source_shard: usize, + pub(crate) tables: Vec, + pub(crate) replication_stream: ReplicationStream, + pub(crate) stop: CancellationToken, +} + +impl Task for ReplicationShardTask { + type Status = ReplicationShardStatus; + type Output = (); + type Error = Error; + + fn cancel_timeout() -> Duration { + Duration::from_secs(60) + } + + fn definition(&self) -> impl Into { + let slot = self.slot.get(); + ReplicationShardDefinition { + slot: slot.name().to_owned(), + host: slot.addr().host.clone(), + port: slot.addr().port, + database_name: slot.addr().database_name.clone(), + source_shard: self.source_shard, + } + } + + async fn run(self, ctx: TaskContext) -> Result<(), Error> { + let Self { + mut slot, + tables, + replication_stream, + stop, + source_shard, + } = self; + + // task got cancelled + let task_cancel = ctx.cancellation_token(); + // signal to stream to stop - due to cutover or fail in other streams + let stream_stop = stop.child_token(); + + slot.slot().set_task_id(ctx.id()); + let slot_name = slot.get().name().to_owned(); + let slot_addr = slot.get().addr().clone(); + let initial_lsn = slot.get().lsn(); + ctx.set_status(ReplicationShardStatus { + lsn: initial_lsn, + lag_bytes: None, + missed_rows: MissedRows::default(), + rows: 0, + bytes: 0, + rows_per_sec: None, + bytes_per_sec: None, + }); + + info!( + shard = source_shard, + "[replication] stream starting at {initial_lsn}" + ); + let mut replication_run = + Box::pin(replication_stream.run(slot.slot(), tables, &stream_stop)); + + let report_interval = Duration::from_secs(5); + let mut report = safe_interval(report_interval); + let mut logged_rows = 0u64; + let mut logged_bytes = 0u64; + + let result = loop { + select! { + _ = task_cancel.cancelled(), if !stream_stop.is_cancelled() => { + info!(shard = source_shard, "[replication] stream cancelled"); + stream_stop.cancel(); + } + result = &mut replication_run => { + break result; + } + _ = report.tick() => { + let progress = replication_stream.progress(); + let status = progress.snapshot(initial_lsn); + let window = report_interval.as_secs_f64(); + info!( + shard = source_shard, + addr = %slot_addr, + slot = slot_name, + "[replication] origin LSN at {}, speed over the last {}s: {:.0} rows/sec, {:.3} MB/sec", + progress.origin_lsn, + report_interval.as_secs(), + (status.rows - logged_rows) as f64 / window, + (status.bytes - logged_bytes) as f64 / window / 1024.0 / 1024.0, + ); + logged_rows = status.rows; + logged_bytes = status.bytes; + ctx.set_status(status); + } } + }; + + let status = replication_stream.progress().snapshot(initial_lsn); + match &result { + Ok(()) => info!( + shard = source_shard, + "[replication] stream stopped, {status}" + ), + Err(err) => warn!( + shard = source_shard, + "[replication] stream failed: {err}, {status}" + ), } + ctx.set_status(status); + drop(replication_run); - Ok(()) + let dropped = Box::pin(slot.drop_slot()).await; + if let Err(err) = &dropped { + warn!("failed to drop replication slot {slot_name}: {err}"); + } + + result.and(dropped) } } -impl ReplicationTask { - /// Perform the actual cutover for running replication. - async fn perform_cutover( - mut self, - ctx: &TaskContext, - token: &CancellationToken, - ) -> Result<(), Error> { - ctx.set_status(match self.direction { - Direction::Forward => ReplicationStatus::CuttingOver, - Direction::Reverse => ReplicationStatus::RollingBack, - }); - self.waiter.cutover(token, ctx, self.schema_sync).await +type ReplicationStreams = FuturesUnordered>>; + +/// Owns the replication slot of one stream. +/// +/// Dropping this guard with the slot still inside schedules the drop on a +/// detached task, so an aborted stream cannot leak a permanent slot. +#[derive(Debug)] +pub(crate) struct SlotGuard { + slot: Option, +} + +impl SlotGuard { + fn new(slot: ReplicationSlot) -> Self { + Self { slot: Some(slot) } } - /// Trigger a cutover on a running replication task. - pub(crate) fn trigger_cutover(target: Option) -> bool { - let tokens = CUTOVERS.lock(); + fn slot(&mut self) -> &mut ReplicationSlot { + self.slot.as_mut().expect("slot guard owns the slot") + } - let token = match target { - Some(id) => tokens.get(&id), - // No id: cut over the first (lowest-id) running task. - None => tokens.keys().min().and_then(|id| tokens.get(id)), + fn get(&self) -> &ReplicationSlot { + self.slot.as_ref().expect("slot guard owns the slot") + } + + fn drop_timeout() -> Duration { + Duration::from_secs(30) + } + + async fn drop_slot(mut self) -> Result<(), Error> { + let mut slot = self.slot.take().expect("slot guard owns the slot"); + let name = slot.name().to_owned(); + + safe_timeout(Self::drop_timeout(), slot.drop_slot()) + .await + .unwrap_or(Err(Error::SlotDropTimeout(name))) + } +} + +impl Drop for SlotGuard { + fn drop(&mut self) { + let Some(mut slot) = self.slot.take() else { + return; }; + let name = slot.name().to_owned(); + tasks::spawn("replication slot cleanup", async move { + let dropped = safe_timeout(Self::drop_timeout(), slot.drop_slot()) + .await + .unwrap_or(Err(Error::SlotDropTimeout(name.clone()))); - match token { - Some(token) => { - token.cancel(); - true + if let Err(err) = dropped { + warn!("failed to drop replication slot {name} of an aborted stream: {err}"); } - None => false, + }); + } +} + +struct MaintenanceMode { + stopped_traffic: bool, +} + +impl MaintenanceMode { + fn new() -> Self { + Self { + stopped_traffic: false, } } - /// Register this task (by its `root_id`) to receive operator cutovers for + fn stop_traffic(&mut self) { + maintenance_mode::start(None); + self.stopped_traffic = true; + } + + fn resume_traffic(&mut self) { + if self.stopped_traffic { + maintenance_mode::stop(None); + self.stopped_traffic = false; + } + } +} + +impl Drop for MaintenanceMode { + fn drop(&mut self) { + self.resume_traffic(); + } +} + +/// Cutover tokens of the replication tasks currently awaiting an operator +/// `CUTOVER`, keyed by the root task id they belong to. A cutover token is +/// *separate* from the task's `STOP_TASK` cancellation token — signalling it +/// means "cut over", not "abandon". +static CUTOVERS: LazyLock> = LazyLock::new(DashMap::new); + +/// Guard held by a running replication task: removes its cutover +/// registration on drop. Awaiting [CutoverWaiter::requested] +/// resolves when an operator `CUTOVER` targets the task. +struct CutoverWaiter { + root_id: TaskId, + token: CancellationToken, +} + +impl CutoverWaiter { + /// Register a task (by its `root_id`) to receive operator cutovers for /// as long as the returned guard is held. - fn register_cutover(root_id: TaskId) -> CutoverWaiter { + fn register(root_id: TaskId) -> Self { let token = CancellationToken::new(); - CUTOVERS.lock().insert(root_id, token.clone()); - CutoverWaiter { root_id, token } + CUTOVERS.insert(root_id, token.clone()); + Self { root_id, token } + } + + /// Wait until a cutover is requested for this task. The token latches, so + /// a cutover that arrived earlier is delivered immediately. + async fn requested(&self) { + self.token.cancelled().await; + } +} + +impl Drop for CutoverWaiter { + fn drop(&mut self) { + CUTOVERS.remove(&self.root_id); } } @@ -181,7 +751,7 @@ mod tests { async fn cutover_delivers_even_when_buffered() { let _guard = CUTOVER_TEST_LOCK.lock().await; // Cutover lands before the task awaits: still delivered (latches). - let waiter = ReplicationTask::register_cutover(TaskId::new(1)); + let waiter = CutoverWaiter::register(TaskId::new(1)); assert!( ReplicationTask::trigger_cutover(Some(TaskId::new(1))), "the named task must receive the cutover" @@ -197,7 +767,7 @@ mod tests { let _guard = CUTOVER_TEST_LOCK.lock().await; // A cutover for one id must never disturb a task registered under a // different id — the whole point of keying by task id. - let waiter = ReplicationTask::register_cutover(TaskId::new(7)); + let waiter = CutoverWaiter::register(TaskId::new(7)); assert!( !ReplicationTask::trigger_cutover(Some(TaskId::new(8))), @@ -221,8 +791,8 @@ mod tests { let _guard = CUTOVER_TEST_LOCK.lock().await; // No id: the lowest-id (first) registered task is cut over, and only // it. - let first = ReplicationTask::register_cutover(TaskId::new(3)); - let second = ReplicationTask::register_cutover(TaskId::new(9)); + let first = CutoverWaiter::register(TaskId::new(3)); + let second = CutoverWaiter::register(TaskId::new(9)); assert!( ReplicationTask::trigger_cutover(None), @@ -246,12 +816,12 @@ mod tests { // A cutover to a task that never consumes it must die with that task, // never reaching the next one. Regression guard for the signal leak. { - let first = ReplicationTask::register_cutover(TaskId::new(1)); + let first = CutoverWaiter::register(TaskId::new(1)); assert!(ReplicationTask::trigger_cutover(Some(TaskId::new(1)))); drop(first); // ends without ever awaiting `requested()` } - let next = ReplicationTask::register_cutover(TaskId::new(2)); + let next = CutoverWaiter::register(TaskId::new(2)); assert!( tokio::time::timeout(Duration::from_millis(200), next.requested()) .await @@ -267,4 +837,38 @@ mod tests { assert!(!ReplicationTask::trigger_cutover(None)); assert!(!ReplicationTask::trigger_cutover(Some(TaskId::new(404)))); } + + static MAINTENANCE_TEST_LOCK: std::sync::LazyLock> = + std::sync::LazyLock::new(|| tokio::sync::Mutex::new(())); + + fn traffic_stopped() -> bool { + maintenance_mode::waiter("pgdog").is_some() + } + + #[tokio::test] + async fn maintenance_guard_stops_and_resumes_traffic() { + let _guard = MAINTENANCE_TEST_LOCK.lock().await; + let mut maintenance = MaintenanceMode::new(); + assert!(!traffic_stopped()); + + maintenance.stop_traffic(); + assert!(traffic_stopped()); + + maintenance.resume_traffic(); + assert!(!traffic_stopped()); + + maintenance.resume_traffic(); + assert!(!traffic_stopped()); + } + + #[tokio::test] + async fn dropping_the_maintenance_guard_resumes_traffic() { + let _guard = MAINTENANCE_TEST_LOCK.lock().await; + let mut maintenance = MaintenanceMode::new(); + maintenance.stop_traffic(); + assert!(traffic_stopped()); + + drop(maintenance); + assert!(!traffic_stopped()); + } } diff --git a/pgdog/src/api/resharding.rs b/pgdog/src/api/resharding.rs index f30bfed00..a14ac1c3c 100644 --- a/pgdog/src/api/resharding.rs +++ b/pgdog/src/api/resharding.rs @@ -124,13 +124,13 @@ impl Task for ReshardTask { orchestrator.refresh()?; // `auto_cutover` (reshard) cuts over on its own; otherwise the - // task runs until an operator `CUTOVER`/`STOP_TASK`. Both of - // those resolve to `Ok`, so awaiting surfaces only a genuine - // replication failure. - let waiter = orchestrator.replicate().await?; + // task runs until an operator `CUTOVER`/`STOP_TASK`. A stop in + // a forward phase resolves to `Err(ReplicationAborted)` and runs + // the cleanup below; a stop in a reverse phase resolves to + // `Ok`, because the migration is already complete. ctx.run( ReplicationTask::builder() - .waiter(waiter) + .orchestrator(orchestrator.clone()) .auto_cutover(self.auto_cutover) .schema_sync(schema_sync.clone().phase(SchemaSyncPhase::Cutover).build()) .build(), diff --git a/pgdog/src/api/task.rs b/pgdog/src/api/task.rs index a53feee87..a4a8070a9 100644 --- a/pgdog/src/api/task.rs +++ b/pgdog/src/api/task.rs @@ -185,7 +185,7 @@ impl TaskEntry { /// Transition the task to the specified progress state. /// No-op if the task is already in terminal state. - fn transition(&self, mut progress: TaskProgress) { + fn transition(&self, progress: TaskProgress) { let _enter = self.tracing_span.enter(); let mut state = self.state.write(); @@ -193,12 +193,6 @@ impl TaskEntry { return; } - let panicked = matches!(progress, TaskProgress::Panic { .. }); - if progress.is_terminal() && !panicked && self.cancellation_token.is_cancelled() { - info!("task is cancelled, ignoring current progress ({progress})"); - progress = TaskProgress::Cancelled; - } - debug!("task state transition to {progress}"); if progress.is_error() { @@ -420,6 +414,12 @@ impl TaskContext { Ok(output) } + Err(err) if ctx.task.cancellation_token.is_cancelled() => { + info!("task cancelled: {err}"); + ctx.transition(TaskProgress::Cancelled); + + Err(err) + } Err(err) => { ctx.transition(TaskProgress::error(err.to_string())); @@ -429,6 +429,10 @@ impl TaskContext { } } + pub(crate) fn id(&self) -> TaskId { + self.task.id + } + pub(crate) fn root_id(&self) -> TaskId { self.task.root_id } @@ -510,6 +514,11 @@ impl TaskStorage { ctx.transition(TaskProgress::Finished); let _ = sender.send(Ok(res)); } + Ok(Err(err)) if cancellation_token.is_cancelled() => { + info!("task cancelled: {err}"); + ctx.transition(TaskProgress::Cancelled); + let _ = sender.send(Err(TaskError::Failed(err))); + } Ok(Err(err)) => { ctx.transition(TaskProgress::error(err.to_string())); let _ = sender.send(Err(TaskError::Failed(err))); @@ -1123,7 +1132,7 @@ mod tests { assert_eq!(*state.lock(), "cancelled"); let entry = storage.task(task_id).unwrap(); - assert!(matches!(entry.state().progress, TaskProgress::Cancelled)); + assert!(matches!(entry.state().progress, TaskProgress::Finished)); } #[test(start_paused = true)] @@ -1189,7 +1198,7 @@ mod tests { let res = task.await; assert!(res.unwrap()); let entry = storage.task(task_id).unwrap(); - assert!(matches!(entry.state().progress, TaskProgress::Cancelled)); + assert!(matches!(entry.state().progress, TaskProgress::Finished)); } #[test(start_paused = true)] diff --git a/pgdog/src/backend/databases.rs b/pgdog/src/backend/databases.rs index 200a9c9de..5d7800e48 100644 --- a/pgdog/src/backend/databases.rs +++ b/pgdog/src/backend/databases.rs @@ -246,6 +246,7 @@ pub(crate) async fn cutover(source: &str, destination: &str) -> Result<(), Error config.config.cutover(source, destination); config.users.cutover(source, destination); + let config = crate::config::set(config)?; let databases = from_config(&config); replace_databases(databases, true)?; @@ -2010,4 +2011,42 @@ password = "testpass" assert_eq!(resolved.schema.as_deref(), Some("Public")); assert_eq!(resolved.column, "Tenant_Id"); } + + #[tokio::test] + async fn test_cutover_swaps_back_on_the_second_call() { + let mut config = ConfigAndUsers::default(); + config.config.databases.push(Database { + name: "single".into(), + host: "127.0.0.1".into(), + port: 5432, + ..Default::default() + }); + for shard in 0..2 { + config.config.databases.push(Database { + name: "sharded".into(), + host: "127.0.0.1".into(), + port: 5432, + database_name: Some(format!("shard_{shard}")), + shard, + ..Default::default() + }); + } + for database in ["single", "sharded"] { + let mut user = ConfigUser::new("pgdog", "pgdog", database); + user.schema_admin = true; + config.users.users.push(user); + } + crate::config::set(config).unwrap(); + init().unwrap(); + + let shards = |database: &str| databases().schema_owner(database).unwrap().shards().len(); + + assert_eq!((shards("single"), shards("sharded")), (1, 2)); + + cutover("single", "sharded").await.unwrap(); + assert_eq!((shards("single"), shards("sharded")), (2, 1)); + + cutover("single", "sharded").await.unwrap(); + assert_eq!((shards("single"), shards("sharded")), (1, 2)); + } } diff --git a/pgdog/src/backend/maintenance_mode.rs b/pgdog/src/backend/maintenance_mode.rs index e88e0d4f8..639bf5764 100644 --- a/pgdog/src/backend/maintenance_mode.rs +++ b/pgdog/src/backend/maintenance_mode.rs @@ -80,11 +80,6 @@ pub(crate) fn stop(database: Option<&str>) { } } -#[cfg(test)] -pub(crate) fn is_on(database: &str) -> bool { - MAINTENANCE_MODE.paused(database) -} - #[derive(Debug)] struct MaintenanceMode { state: ArcSwap, diff --git a/pgdog/src/backend/replication/logical/data_sync.rs b/pgdog/src/backend/replication/logical/data_sync.rs index 4be1ccf3d..97b8e91f0 100644 --- a/pgdog/src/backend/replication/logical/data_sync.rs +++ b/pgdog/src/backend/replication/logical/data_sync.rs @@ -5,6 +5,7 @@ use tokio::select; use tracing::{info, warn}; use pgdog_config::CopyFormat; +use pgdog_stats::TaskId; use crate::backend::pool::{Address, Request}; use crate::backend::{Cluster, ConnectReason, Server, ServerOptions}; @@ -37,6 +38,7 @@ pub(crate) struct DataSync<'a> { pub(crate) source: &'a Cluster, pub(crate) dest: &'a Cluster, pub(crate) format: CopyFormat, + pub(crate) task_id: TaskId, } impl DataSync<'_> { @@ -67,6 +69,7 @@ impl DataSync<'_> { let mut slot = ReplicationSlot::data_sync(&table.publication, address); slot.connect().await?; table.lsn = slot.create_slot().await?; + slot.set_task_id(self.task_id); // Reload table info just to be sure it's consistent. table.reload(slot.server()?).await?; diff --git a/pgdog/src/backend/replication/logical/ee/mod.rs b/pgdog/src/backend/replication/logical/ee/mod.rs index 0efadd04b..322accafc 100644 --- a/pgdog/src/backend/replication/logical/ee/mod.rs +++ b/pgdog/src/backend/replication/logical/ee/mod.rs @@ -8,14 +8,6 @@ use crate::net::ErrorResponse; use super::*; use std::time::Duration; -#[derive(Debug, Clone)] -pub(crate) enum CutoverState { - WaitingForReplication { lag: u64 }, - WaitForCutover { action: CutoverAction }, - Abort { error: String }, - Complete, -} - #[derive(Debug, Clone)] pub(crate) enum OrchestratorState { SchemSyncPre, @@ -24,11 +16,6 @@ pub(crate) enum OrchestratorState { SchemaSyncPostCutover, DataSync, Replication, - Cutover(CutoverState), -} - -pub(crate) fn cutover_state(state: CutoverState) { - orchestrator_state(OrchestratorState::Cutover(state)); } pub(crate) fn orchestrator_state(state: OrchestratorState) {} diff --git a/pgdog/src/backend/replication/logical/error.rs b/pgdog/src/backend/replication/logical/error.rs index 6e0124e6e..ec953bdfb 100644 --- a/pgdog/src/backend/replication/logical/error.rs +++ b/pgdog/src/backend/replication/logical/error.rs @@ -131,12 +131,24 @@ pub(crate) enum Error { #[error("replication timeout")] ReplicationTimeout, + #[error("replication streams did not drain in time")] + DrainTimeout, + + #[error("replication slot \"{0}\" was not dropped in time")] + SlotDropTimeout(String), + + #[error("replication stream stopped before shutdown was requested")] + ReplicationStreamStopped, + #[error("publication \"{0}\" has no tables")] EmptyPublication(String), #[error("shard {0} has no replication slot")] NoReplicationSlot(usize), + #[error("shard {0} has no replication table entry")] + NoReplicationTables(usize), + #[error("parallel connection error")] ParallelConnection, @@ -161,6 +173,9 @@ pub(crate) enum Error { #[error("data sync has been aborted")] DataSyncAborted, + #[error("replication has been aborted")] + ReplicationAborted, + #[error("cutover abort timeout")] AbortTimeout, diff --git a/pgdog/src/backend/replication/logical/mod.rs b/pgdog/src/backend/replication/logical/mod.rs index b3eb47eae..b5a732423 100644 --- a/pgdog/src/backend/replication/logical/mod.rs +++ b/pgdog/src/backend/replication/logical/mod.rs @@ -12,8 +12,6 @@ pub(crate) mod tables_sync; pub(crate) use copy_statement::CopyStatement; pub(crate) use error::*; -use ee::*; -use orchestrator::*; -pub(crate) use publisher::publisher_impl::{Publisher, Waiter}; +pub(crate) use publisher::publisher_impl::Publisher; -use crate::{backend::databases::databases, config::config}; +use crate::backend::databases::databases; diff --git a/pgdog/src/backend/replication/logical/orchestrator.rs b/pgdog/src/backend/replication/logical/orchestrator.rs index ab556a00f..a5b2f54d5 100644 --- a/pgdog/src/backend/replication/logical/orchestrator.rs +++ b/pgdog/src/backend/replication/logical/orchestrator.rs @@ -1,28 +1,11 @@ -use crate::api::replication::ReplicationTask; -use crate::api::schema_sync::{SchemaSyncPhase, SchemaSyncTask}; -use crate::api::task::TaskContext; -use crate::{ - backend::{ - Cluster, - databases::{cancel_all, cutover}, - maintenance_mode, - }, - tasks, - util::{format_bytes, human_duration, random_string}, -}; -use pgdog_config::{ConfigAndUsers, CutoverTimeoutAction}; +use crate::tasks; +use crate::{backend::Cluster, util::random_string}; use pgdog_stats::Databases; -use std::{fmt::Display, sync::Arc, time::Duration}; -use tokio::{ - select, - sync::{Mutex, MutexGuard}, - time::Instant, -}; -use tokio_util::sync::CancellationToken; -use tracing::{error, info, warn}; +use std::{fmt::Display, sync::Arc}; +use tokio::sync::{Mutex, MutexGuard}; +use tracing::warn; use super::*; -use crate::util::safe_interval; #[derive(Debug, Clone)] pub(crate) struct Orchestrator { @@ -104,8 +87,8 @@ impl Orchestrator { Ok(()) } - /// Replace the publisher entirely (discards LSN state). Only valid - /// when starting a fresh replication phase, e.g. after cutover. + /// Replace the publisher entirely (discards LSN state). Only valid + /// before any replication slot exists, e.g. after the pre-data schema sync. pub(crate) fn refresh_publisher(&mut self) { let publisher = Publisher::new(&self.publication, self.replication_slot.clone()); self.publisher = Arc::new(Mutex::new(publisher)); @@ -126,38 +109,6 @@ impl Orchestrator { } } - /// Replicate forever. - /// - /// Useful for CLI interface only, since this will never stop. - /// - pub(crate) async fn replicate(&self) -> Result { - let mut publisher = self.publisher.lock().await; - let waiter = publisher.replicate(&self.source, &self.destination).await?; - - orchestrator_state(OrchestratorState::Replication); - - Ok(ReplicationWaiter { - orchestrator: self.clone(), - waiter, - config: config(), - }) - } - - /// Get the largest replication lag out of all the shards. - async fn replication_lag(&self) -> Option { - let shards_count = self.source.shards().len(); - let lag = self.publisher.lock().await.replication_lag(); - - if lag.len() != shards_count { - // if the len of lag map is not equal to source shards_count - // then some entries are not initialized yet and the lag value - // is not yet reported. - return None; - } - - lag.values().copied().max().map(|lag| lag as u64) - } - /// The two ends of the migration this orchestrator drives. pub(crate) fn databases(&self) -> Databases { Databases { @@ -177,712 +128,3 @@ impl Display for Orchestrator { ) } } - -#[derive(Debug, Display)] -#[display("{orchestrator}")] -pub(crate) struct ReplicationWaiter { - orchestrator: Orchestrator, - waiter: Waiter, - config: Arc, -} - -#[derive(Debug, Clone, PartialEq, Eq, Copy)] -pub(crate) enum CutoverReason { - Lag, - Timeout, - LastTransaction, -} - -#[derive(Debug, Clone, PartialEq, Eq, Copy)] -pub(crate) enum CutoverAction { - Go(CutoverReason), - NoGo(CutoverData), -} - -#[derive(Debug, Clone, PartialEq, Eq, Copy)] -pub(crate) struct CutoverData { - pub(crate) lag: u64, - pub(crate) last_transaction: Option, - pub(crate) elapsed: Duration, -} - -impl Display for CutoverReason { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - Self::Lag => write!(f, "lag"), - Self::Timeout => write!(f, "timeout"), - Self::LastTransaction => write!(f, "last_transaction"), - } - } -} - -impl ReplicationWaiter { - pub(crate) async fn wait(&mut self) -> Result<(), Error> { - self.waiter.wait().await - } - - /// The two ends of the migration this waiter replicates. - pub(crate) fn databases(&self) -> Databases { - self.orchestrator.databases() - } - - pub(crate) fn stop(&self) { - self.waiter.stop(); - } - - /// Wait for replication to catch up. - async fn wait_for_replication(&mut self) -> Result<(), Error> { - let traffic_stop = self.config.config.general.cutover_traffic_stop_threshold; - - info!( - "[cutover] started, waiting for traffic stop threshold={}", - format_bytes(traffic_stop) - ); - - // Check once a second how far we got. - let mut check = safe_interval(Duration::from_secs(1)); - - loop { - select! { - _ = check.tick() => {} - - // In case replication breaks now. - res = self.waiter.wait() => { - res?; - } - } - - let Some(lag) = self.orchestrator.replication_lag().await else { - info!("[cutover] replication lag is not calculated for all shards, yet"); - continue; - }; - - cutover_state(CutoverState::WaitingForReplication { lag }); - info!("[cutover] replication lag: {}", format_bytes(lag)); - - // Time to go. - if lag <= traffic_stop { - info!( - "[cutover] stopping traffic, lag={}, threshold={}", - format_bytes(lag), - format_bytes(traffic_stop), - ); - - // Pause traffic. - maintenance_mode::start(None); - - // Cancel any running queries. - ok_or_abort!(cancel_all(&self.orchestrator.source.identifier().database).await); - - break; - // TODO: wait for clients to all stop. - } - } - - Ok(()) - } - - async fn should_cutover(&self, elapsed: Duration) -> CutoverAction { - let cutover_timeout = Duration::from_millis(self.config.config.general.cutover_timeout); - let cutover_threshold = self.config.config.general.cutover_replication_lag_threshold; - let last_transaction_delay = - Duration::from_millis(self.config.config.general.cutover_last_transaction_delay); - - let lag = self.orchestrator.replication_lag().await; - let last_transaction = self.orchestrator.publisher.lock().await.last_transaction(); - let cutover_timeout_exceeded = elapsed >= cutover_timeout; - - if cutover_timeout_exceeded { - CutoverAction::Go(CutoverReason::Timeout) - } else if lag.is_some_and(|lag| lag <= cutover_threshold) { - CutoverAction::Go(CutoverReason::Lag) - } else if last_transaction.is_none_or(|t| t > last_transaction_delay) { - CutoverAction::Go(CutoverReason::LastTransaction) - } else { - CutoverAction::NoGo(CutoverData { - lag: lag.unwrap_or(u64::MAX), - last_transaction, - elapsed, - }) - } - } - - /// Wait for cutover. - async fn wait_for_cutover(&mut self) -> Result<(), Error> { - let cutover_threshold = self.config.config.general.cutover_replication_lag_threshold; - let last_transaction_delay = - Duration::from_millis(self.config.config.general.cutover_last_transaction_delay); - let cutover_timeout = Duration::from_millis(self.config.config.general.cutover_timeout); - let cutover_timeout_action = self.config.config.general.cutover_timeout_action; - - info!( - "[cutover] waiting for first cutover threshold: timeout={}, transaction={}, lag={}", - human_duration(cutover_timeout), - human_duration(last_transaction_delay), - format_bytes(cutover_threshold) - ); - - // Check more frequently. - let mut check = safe_interval(Duration::from_millis(50)); - let mut log = safe_interval(Duration::from_secs(1)); - // Abort clock starts now. - let start = Instant::now(); - - let mut cutover_data = None; - - loop { - select! { - _ = check.tick() => {} - - _ = log.tick() => { - if let Some(CutoverData { lag, last_transaction, elapsed }) = cutover_data { - info!("[cutover] lag={}, last_transaction={}, timeout={}", - format_bytes(lag), - if let Some(last_transaction) = last_transaction { - human_duration(last_transaction) - } else { - "none".into() - }, - human_duration(elapsed), - ); - } - - } - - // In case replication breaks now. - res = self.waiter.wait() => { - ok_or_abort!(res); - } - } - - let elapsed = start.elapsed(); - let cutover_reason = self.should_cutover(elapsed).await; - - cutover_state(CutoverState::WaitForCutover { - action: cutover_reason, - }); - - match cutover_reason { - CutoverAction::Go(CutoverReason::Timeout) => { - if cutover_timeout_action == CutoverTimeoutAction::Abort { - maintenance_mode::stop(None); - warn!("[cutover] abort timeout reached, resuming traffic"); - return Err(Error::AbortTimeout); - } else { - info!( - "[cutover] performing cutover now, reason: {}", - CutoverReason::Timeout - ); - break; - } - } - - CutoverAction::NoGo(data) => { - cutover_data = Some(data); - continue; - } - CutoverAction::Go(reason) => { - info!("[cutover] performing cutover now, reason: {}", reason); - break; - } - } - } - - Ok(()) - } - - /// Perform traffic cutover between source and destination. - /// - /// A `STOP_TASK` before the switch resumes traffic and stops the streams, - /// moving nothing. The switch itself is not cancellable. A `STOP_TASK` - /// during the cutover schema sync fails that subtask, so the cutover - /// aborts and traffic goes back to the source. - pub(crate) async fn cutover( - &mut self, - cancel: &CancellationToken, - ctx: &TaskContext, - schema_sync: SchemaSyncTask, - ) -> Result<(), Error> { - select! { - // Nothing has moved yet (`wait_for_replication` only pauses traffic - // at its very end). Resume traffic (no-op if never paused) and wind - // the streams down, so the aborted cutover leaves nothing running. - _ = cancel.cancelled() => { - maintenance_mode::stop(None); - self.waiter.stop(); - warn!("[cutover] stop requested before the traffic switch, aborting cutover"); - cutover_state(CutoverState::Abort { - error: "stopped before cutover".into(), - }); - return Ok(()); - } - res = async { - self.wait_for_replication().await?; - self.wait_for_cutover().await - } => { res?; } - } - - // We're going, point of no return. - self.waiter.stop(); - ok_or_abort!(self.waiter.wait().await); - ok_or_abort!(ctx.run(schema_sync).await); - // Traffic is about to go to the new cluster. - // If this fails, we'll resume traffic to the old cluster instead - // and the whole thing needs to be done from scratch. - ok_or_abort!( - cutover( - &self.orchestrator.source.identifier().database, - &self.orchestrator.destination.identifier().database, - ) - .await - ); - - // Source is now destination and vice versa; reload cluster refs and - // create a fresh publisher for reverse replication. - ok_or_abort!(self.orchestrator.refresh()); - self.orchestrator.refresh_publisher(); - - info!("[cutover] setting up reverse replication"); - - // Create reverse replication in case we need to rollback. - let waiter = ok_or_abort!(self.orchestrator.replicate().await); - - // Drive the running waiter as a background api task so it stays visible - // in SHOW TASKS and can be cut over (rollback) or stopped. - crate::api::run_task( - crate::api::replication::ReplicationTask::builder() - .waiter(waiter) - .direction(crate::api::replication::Direction::Reverse) - .schema_sync( - SchemaSyncTask::builder() - .databases(self.orchestrator.databases()) - .publication(self.orchestrator.publication.clone()) - .phase(SchemaSyncPhase::Cutover) - .ignore_errors(true) - .build(), - ) - .build(), - ); - - // Slot is established and capturing — now safe to resume traffic. - info!("[cutover] complete, resuming traffic"); - - // Point traffic to the other database and resume. - maintenance_mode::stop(None); - - cutover_state(CutoverState::Complete); - - Ok(()) - } -} - -macro_rules! ok_or_abort { - ($expr:expr_2021) => { - match $expr { - Ok(res) => res, - Err(err) => { - error!("Orchestrator failed: {err}"); - maintenance_mode::stop(None); - cutover_state(CutoverState::Abort { - error: err.to_string(), - }); - return Err(Error::from(err)); - } - } - }; -} - -use ok_or_abort; - -#[cfg(test)] -mod tests { - use super::*; - use crate::backend::pool::Cluster; - use crate::util::{safe_sleep, safe_timeout}; - use pgdog_config::ConfigAndUsers; - use std::assert_matches; - use std::sync::Arc; - use tokio::time::Instant; - - impl Orchestrator { - fn new_test(config: &ConfigAndUsers) -> Self { - let cluster = Cluster::new_test(config); - let publication = "test_pub".to_owned(); - let replication_slot = "test_slot".to_owned(); - let publisher = Publisher::new(&publication, replication_slot.clone()); - Self { - source: cluster.clone(), - destination: cluster, - publication, - publisher: Arc::new(Mutex::new(publisher)), - replication_slot, - } - } - } - - impl ReplicationWaiter { - fn new_test(orchestrator: Orchestrator, config: Arc) -> Self { - Self { - orchestrator, - waiter: Waiter::new_test(), - config, - } - } - } - - #[tokio::test] - async fn test_wait_for_replication_exits_when_lag_below_threshold() { - // Ensure maintenance mode is off at start - maintenance_mode::stop(None); - assert!(!maintenance_mode::is_on("")); // Will return true because all databases are paused. - - let mut config = ConfigAndUsers::default(); - config.config.general.cutover_traffic_stop_threshold = 1000; - - let orchestrator = Orchestrator::new_test(&config); - - // Set replication lag below threshold for every shard. - { - let publisher = orchestrator.publisher.lock().await; - publisher.set_replication_lag(0, 500); - publisher.set_replication_lag(1, 500); - } - - let config = Arc::new(config); - let mut waiter = ReplicationWaiter::new_test(orchestrator, config); - - // Should exit immediately since lag (500) <= threshold (1000) - let result = waiter.wait_for_replication().await; - assert!(result.is_ok()); - - // Maintenance mode should be on after wait_for_replication - assert!(maintenance_mode::is_on("")); - - // Clean up maintenance mode - maintenance_mode::stop(None); - assert!(!maintenance_mode::is_on("")); - } - - #[tokio::test] - async fn test_wait_for_cutover_exits_when_lag_below_threshold() { - let mut config = ConfigAndUsers::default(); - config.config.general.cutover_replication_lag_threshold = 100; - config.config.general.cutover_timeout = 10000; - - let orchestrator = Orchestrator::new_test(&config); - - // Set replication lag below cutover threshold for every shard. - { - let publisher = orchestrator.publisher.lock().await; - publisher.set_replication_lag(0, 50); - publisher.set_replication_lag(1, 50); - } - - let config = Arc::new(config); - let mut waiter = ReplicationWaiter::new_test(orchestrator, config); - - // should_cutover returns Lag when lag is below threshold - let result = waiter.should_cutover(Duration::from_millis(100)).await; - assert_eq!(result, CutoverAction::Go(CutoverReason::Lag)); - - // Should exit immediately since lag (50) <= threshold (100) - let result = waiter.wait_for_cutover().await; - assert!(result.is_ok()); - } - - #[tokio::test] - async fn test_wait_for_cutover_exits_when_last_transaction_old() { - let mut config = ConfigAndUsers::default(); - config.config.general.cutover_replication_lag_threshold = 10; - config.config.general.cutover_last_transaction_delay = 100; - config.config.general.cutover_timeout = 10000; - - let orchestrator = Orchestrator::new_test(&config); - - { - let publisher = orchestrator.publisher.lock().await; - // Set lag above threshold so we don't exit on that condition - publisher.set_replication_lag(0, 1000); - // Set last_transaction to a time in the past (> 100ms ago) - publisher.set_last_transaction(Some(Instant::now() - Duration::from_millis(200))); - } - - let config = Arc::new(config); - let mut waiter = ReplicationWaiter::new_test(orchestrator, config); - - // should_cutover returns LastTransaction when last transaction is old - let result = waiter.should_cutover(Duration::from_millis(100)).await; - assert_eq!(result, CutoverAction::Go(CutoverReason::LastTransaction)); - - // Should exit because last_transaction (200ms) > threshold (100ms) - let result = waiter.wait_for_cutover().await; - assert!(result.is_ok()); - } - - #[tokio::test] - async fn test_should_cutover_when_no_transaction() { - let mut config = ConfigAndUsers::default(); - config.config.general.cutover_replication_lag_threshold = 10; - config.config.general.cutover_last_transaction_delay = 100; - config.config.general.cutover_timeout = 10000; - - let orchestrator = Orchestrator::new_test(&config); - - { - let publisher = orchestrator.publisher.lock().await; - // Set lag above threshold so we don't exit on that condition - publisher.set_replication_lag(0, 1000); - // No transaction set (None) - publisher.set_last_transaction(None); - } - - let config = Arc::new(config); - let waiter = ReplicationWaiter::new_test(orchestrator, config); - - // should_cutover returns LastTransaction when there's no transaction - let result = waiter.should_cutover(Duration::from_millis(100)).await; - assert_eq!(result, CutoverAction::Go(CutoverReason::LastTransaction)); - } - - #[tokio::test] - async fn test_should_not_cutover_when_lag_above_threshold_and_recent_transaction() { - let mut config = ConfigAndUsers::default(); - config.config.general.cutover_timeout = 10000; - config.config.general.cutover_replication_lag_threshold = 100; - config.config.general.cutover_last_transaction_delay = 500; - - let orchestrator = Orchestrator::new_test(&config); - - { - let publisher = orchestrator.publisher.lock().await; - // Lag above threshold - publisher.set_replication_lag(0, 1000); - // Recent transaction (50ms ago, threshold is 500ms) - publisher.set_last_transaction(Some(Instant::now() - Duration::from_millis(50))); - } - - let config = Arc::new(config); - let waiter = ReplicationWaiter::new_test(orchestrator, config); - - // Not timed out (100ms elapsed, timeout is 10000ms) - let result = waiter.should_cutover(Duration::from_millis(100)).await; - assert!(matches!(result, CutoverAction::NoGo { .. })); - } - - #[tokio::test] - async fn test_should_not_cutover_when_timeout_not_reached() { - let mut config = ConfigAndUsers::default(); - config.config.general.cutover_timeout = 1000; - config.config.general.cutover_replication_lag_threshold = 10; - config.config.general.cutover_last_transaction_delay = 500; - - let orchestrator = Orchestrator::new_test(&config); - - { - let publisher = orchestrator.publisher.lock().await; - // Lag above threshold - publisher.set_replication_lag(0, 1000); - // Recent transaction - publisher.set_last_transaction(Some(Instant::now() - Duration::from_millis(100))); - } - - let config = Arc::new(config); - let waiter = ReplicationWaiter::new_test(orchestrator, config); - - // Elapsed is 999ms, timeout is 1000ms - should not trigger timeout - let result = waiter.should_cutover(Duration::from_millis(999)).await; - assert!(matches!(result, CutoverAction::NoGo { .. })); - } - - #[tokio::test] - async fn test_should_not_cutover_when_lag_just_above_threshold() { - let mut config = ConfigAndUsers::default(); - config.config.general.cutover_timeout = 10000; - config.config.general.cutover_replication_lag_threshold = 100; - config.config.general.cutover_last_transaction_delay = 500; - - let orchestrator = Orchestrator::new_test(&config); - - { - let publisher = orchestrator.publisher.lock().await; - // Lag just above threshold (101 > 100) - publisher.set_replication_lag(0, 101); - // Recent transaction - publisher.set_last_transaction(Some(Instant::now() - Duration::from_millis(50))); - } - - let config = Arc::new(config); - let waiter = ReplicationWaiter::new_test(orchestrator, config); - - let result = waiter.should_cutover(Duration::from_millis(100)).await; - assert!(matches!(result, CutoverAction::NoGo { .. })); - } - - /// Cutover holds off until every shard has reported a lag measurement; an - /// unreported shard reads as unknown (`None`), not zero. - #[tokio::test] - async fn should_not_cutover_before_every_shard_reports() { - let mut config = ConfigAndUsers::default(); - config.config.general.cutover_timeout = 10000; - config.config.general.cutover_replication_lag_threshold = 1000; - config.config.general.cutover_last_transaction_delay = 500; - - // No shard has reported a lag yet. - let orchestrator = Orchestrator::new_test(&config); - // Recent transaction so only the lag arm decides the outcome. - orchestrator - .publisher - .lock() - .await - .set_last_transaction(Some(Instant::now())); - - let config = Arc::new(config); - let waiter = ReplicationWaiter::new_test(orchestrator.clone(), config); - let elapsed = Duration::from_millis(100); - - // Empty map: lag is unknown -> None, no cutover. - assert_eq!(orchestrator.replication_lag().await, None); - assert_matches!( - waiter.should_cutover(elapsed).await, - CutoverAction::NoGo { .. } - ); - - // Only shard 0 reported; shard 1 is missing, so the lag stays unknown. - { - let publisher = orchestrator.publisher.lock().await; - publisher.set_replication_lag(0, 500); - } - assert_eq!(orchestrator.replication_lag().await, None); - assert_matches!( - waiter.should_cutover(elapsed).await, - CutoverAction::NoGo { .. } - ); - - // Once every shard has a real, below-threshold measurement, it cuts over. - orchestrator - .publisher - .lock() - .await - .set_replication_lag(1, 400); - assert_eq!( - waiter.should_cutover(elapsed).await, - CutoverAction::Go(CutoverReason::Lag) - ); - } - - /// Writes to a table outside the publication must not block cutover. - /// Unrelated WAL advances the instance's LSN but not the publication's - /// `confirmed_flush_lsn`; keepalives advance it to `wal_end` between - /// transactions, so the drained slot reports ~0 lag and - /// `wait_for_replication` completes. - /// - /// Runs against the live `pgdog` database (integration/setup.sh). - #[tokio::test] - async fn wait_for_replication_finishes_with_unrelated_writes() { - use crate::backend::server::test::test_server; - - crate::logger(); - maintenance_mode::stop(None); - - const TRAFFIC_STOP: u64 = 1_000; - - let mut config = ConfigAndUsers::default(); - config.config.general.cutover_traffic_stop_threshold = TRAFFIC_STOP; - config.config.general.cutover_timeout = 120_000; - - let orchestrator = Orchestrator::new_test(&config); - let publication = orchestrator.publication.clone(); - let slot = orchestrator.replication_slot().to_owned(); - let shards = orchestrator.source.shards().len(); - - // Publication covers only `issue1_main`; `issue1_noise` is never published. - let mut source = test_server().await; - let _ = source - .execute(format!("DROP PUBLICATION IF EXISTS {publication}")) - .await; - for shard in 0..shards { - let _ = source - .execute(format!("SELECT pg_drop_replication_slot('{slot}_{shard}')")) - .await; - } - source - .execute("DROP TABLE IF EXISTS issue1_main, issue1_noise") - .await - .unwrap(); - source - .execute("CREATE TABLE issue1_main (id BIGINT PRIMARY KEY)") - .await - .unwrap(); - source - .execute("CREATE TABLE issue1_noise (id BIGINT, payload TEXT)") - .await - .unwrap(); - source - .execute(format!( - "CREATE PUBLICATION {publication} FOR TABLE issue1_main" - )) - .await - .unwrap(); - - orchestrator.source.launch(); - - let stream = orchestrator - .publisher - .lock() - .await - .replicate(&orchestrator.source, &orchestrator.destination) - .await - .unwrap(); - - let config = Arc::new(config); - let mut waiter = ReplicationWaiter { - orchestrator: orchestrator.clone(), - waiter: stream, - config, - }; - - // ~10MB of incompressible WAL into the unpublished table: the slot - // decodes none of it, but the instance LSN advances, inflating lag. - source - .execute( - "INSERT INTO issue1_noise \ - SELECT g, (SELECT string_agg(md5(random()::text), '') FROM generate_series(1, 64)) \ - FROM generate_series(1, 5000) g", - ) - .await - .unwrap(); - - // Let one check_lag tick (1s) refresh the lag cache with the post-write - // value before sampling the gate. - safe_sleep(Duration::from_secs(1)).await; - - // 30s margin: keepalive cadence (wal_sender_timeout) is not bounded by - // the 1s sleep, so give confirmed_flush_lsn room to advance. - let result = safe_timeout(Duration::from_secs(20), waiter.wait_for_replication()).await; - let maintenance_on = maintenance_mode::is_on(""); - - // Clean up before asserting so a failure can't leak slots or maintenance mode. - waiter.stop(); - maintenance_mode::stop(None); - for shard in 0..shards { - let _ = source - .execute(format!("SELECT pg_drop_replication_slot('{slot}_{shard}')")) - .await; - } - let _ = source - .execute(format!("DROP PUBLICATION IF EXISTS {publication}")) - .await; - let _ = source - .execute("DROP TABLE IF EXISTS issue1_main, issue1_noise") - .await; - - let waited = result - .expect("wait_for_replication never finished: lag stays inflated by unrelated WAL"); - waited.expect("wait_for_replication returned an error"); - // Cutover fired: once the lag fell below the threshold, traffic stopped. - assert!( - maintenance_on, - "wait_for_replication returned without stopping traffic (cutover did not fire)" - ); - } -} diff --git a/pgdog/src/backend/replication/logical/publisher/cutover_policy.rs b/pgdog/src/backend/replication/logical/publisher/cutover_policy.rs new file mode 100644 index 000000000..62ccf3aa6 --- /dev/null +++ b/pgdog/src/backend/replication/logical/publisher/cutover_policy.rs @@ -0,0 +1,429 @@ +use std::time::Duration; + +use pgdog_config::{ConfigAndUsers, CutoverTimeoutAction}; +use tokio::{select, time::Instant}; +use tracing::{info, warn}; + +use super::super::Error; +use super::replication_progress::ReplicationProgress; +use crate::util::{format_bytes, human_duration, safe_interval}; +use pgdog_stats::ReplicationCutoverReason as CutoverReason; + +#[derive(Debug)] +pub(crate) struct CutoverConfig { + /// Check when replication_lag becomes less than this value + /// to stop the traffic on source. + pub(crate) traffic_stop_threshold: u64, + /// Start the cutover if replication_lag is less than this value + pub(crate) replication_lag_threshold: u64, + /// Start the cutover if last_transaction was more than this value time ago + pub(crate) last_transaction_delay: Duration, + /// Start/abort the wait for cutover after timeout + pub(crate) timeout: Duration, + pub(crate) timeout_action: CutoverTimeoutAction, +} + +impl From<&ConfigAndUsers> for CutoverConfig { + fn from(config: &ConfigAndUsers) -> Self { + let general = &config.config.general; + Self { + traffic_stop_threshold: general.cutover_traffic_stop_threshold, + replication_lag_threshold: general.cutover_replication_lag_threshold, + last_transaction_delay: Duration::from_millis(general.cutover_last_transaction_delay), + timeout: Duration::from_millis(general.cutover_timeout), + timeout_action: general.cutover_timeout_action, + } + } +} + +#[derive(Debug)] +pub(crate) struct CutoverPolicy { + config: CutoverConfig, + progress: ReplicationProgress, +} + +#[derive(Debug, Clone, PartialEq, Eq, Copy)] +pub(crate) enum CutoverAction { + Go(CutoverReason), + NoGo(CutoverData), +} + +#[derive(Debug, Clone, PartialEq, Eq, Copy)] +pub(crate) struct CutoverData { + pub(crate) lag: u64, + pub(crate) last_transaction: Option, + pub(crate) elapsed: Duration, +} + +impl CutoverPolicy { + pub(crate) fn new(config: CutoverConfig, progress: ReplicationProgress) -> Self { + Self { config, progress } + } + + /// Resolves when replication_lag reaches the value less than + /// configured [CutoverConfig::traffic_stop_threshold]. + /// After this source should stop any write activity and + /// [`CutoverPolicy::wait_for_catchup`] should be started. + pub(crate) async fn wait_for_stop_threshold(&self) { + let traffic_stop = self.config.traffic_stop_threshold; + + info!( + "[cutover] started, waiting for traffic stop threshold={}", + format_bytes(traffic_stop) + ); + + let mut check = safe_interval(Duration::from_secs(1)); + + loop { + check.tick().await; + + let Some(lag) = self.progress.snapshot().lag_bytes else { + info!("[cutover] replication lag is not calculated for all shards, yet"); + continue; + }; + + info!("[cutover] replication lag: {}", format_bytes(lag)); + + if lag <= traffic_stop { + info!( + "[cutover] stopping traffic, lag={}, threshold={}", + format_bytes(lag), + format_bytes(traffic_stop), + ); + break; + } + } + } + + fn should_cutover(&self, elapsed: Duration) -> CutoverAction { + let cutover_timeout = self.config.timeout; + let cutover_threshold = self.config.replication_lag_threshold; + let last_transaction_delay = self.config.last_transaction_delay; + + let progress = self.progress.snapshot(); + let lag = progress.lag_bytes; + let last_transaction = progress.last_transaction_ms.map(Duration::from_millis); + let cutover_timeout_exceeded = elapsed >= cutover_timeout; + + if cutover_timeout_exceeded { + CutoverAction::Go(CutoverReason::Timeout) + } else if lag.is_some_and(|lag| lag <= cutover_threshold) { + CutoverAction::Go(CutoverReason::Lag) + } else if last_transaction.is_some_and(|t| t > last_transaction_delay) { + CutoverAction::Go(CutoverReason::LastTransaction) + } else { + CutoverAction::NoGo(CutoverData { + lag: lag.unwrap_or(u64::MAX), + last_transaction, + elapsed, + }) + } + } + + /// Wait until cutover conditions are met depending on the + /// [`CutoverConfig`] settings + pub(crate) async fn wait_for_catchup(&self) -> Result { + let cutover_timeout_action = self.config.timeout_action; + + info!( + "[cutover] waiting for first cutover threshold: timeout={}, transaction={}, lag={}", + human_duration(self.config.timeout), + human_duration(self.config.last_transaction_delay), + format_bytes(self.config.replication_lag_threshold) + ); + + let mut check = safe_interval(Duration::from_millis(50)); + let mut log = safe_interval(Duration::from_secs(1)); + let start = Instant::now(); + + let mut cutover_data = None; + + loop { + select! { + _ = check.tick() => {} + _ = log.tick() => { + if let Some(CutoverData { lag, last_transaction, elapsed }) = cutover_data { + info!( + "[cutover] lag={}, last_transaction={}, timeout={}", + format_bytes(lag), + if let Some(last_transaction) = last_transaction { + human_duration(last_transaction) + } else { + "none".into() + }, + human_duration(elapsed), + ); + } + } + } + + let elapsed = start.elapsed(); + + match self.should_cutover(elapsed) { + CutoverAction::Go(CutoverReason::Timeout) => match cutover_timeout_action { + CutoverTimeoutAction::Abort => { + warn!("[cutover] abort timeout reached, resuming traffic"); + return Err(Error::AbortTimeout); + } + CutoverTimeoutAction::Cutover => { + info!( + "[cutover] performing cutover now, reason: {}", + CutoverReason::Timeout + ); + return Ok(CutoverReason::Timeout); + } + }, + CutoverAction::Go(reason) => { + info!("[cutover] performing cutover now, reason: {reason}"); + return Ok(reason); + } + CutoverAction::NoGo(data) => { + cutover_data = Some(data); + continue; + } + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::backend::replication::logical::publisher::replication_progress::ReplicationProgress; + use crate::util::safe_timeout; + use std::assert_matches; + use tokio::time::Instant; + + fn cutover_config() -> CutoverConfig { + CutoverConfig { + traffic_stop_threshold: 1000, + replication_lag_threshold: 100, + last_transaction_delay: Duration::from_millis(500), + timeout: Duration::from_secs(10), + timeout_action: CutoverTimeoutAction::Abort, + } + } + + #[tokio::test] + async fn test_wait_for_replication_exits_when_lag_below_threshold() { + let config = cutover_config(); + + let progress = ReplicationProgress::new(2); + progress + .updater_for_shard(0) + .update(|s| s.replication_lag = Some(500)); + progress + .updater_for_shard(1) + .update(|s| s.replication_lag = Some(500)); + + let waiter = CutoverPolicy::new(config, progress); + + safe_timeout(Duration::from_secs(5), waiter.wait_for_stop_threshold()) + .await + .expect("the wait must exit once every shard is below the threshold"); + } + + #[tokio::test] + async fn test_wait_for_cutover_exits_when_lag_below_threshold() { + let config = cutover_config(); + + let progress = ReplicationProgress::new(2); + progress + .updater_for_shard(0) + .update(|s| s.replication_lag = Some(50)); + progress + .updater_for_shard(1) + .update(|s| s.replication_lag = Some(50)); + + let waiter = CutoverPolicy::new(config, progress); + + assert_eq!( + waiter.should_cutover(Duration::from_millis(100)), + CutoverAction::Go(CutoverReason::Lag) + ); + + assert_eq!(waiter.wait_for_catchup().await.unwrap(), CutoverReason::Lag); + } + + #[tokio::test] + async fn test_wait_for_cutover_exits_when_last_transaction_old() { + let config = CutoverConfig { + replication_lag_threshold: 10, + last_transaction_delay: Duration::from_millis(100), + ..cutover_config() + }; + + let progress = ReplicationProgress::new(1); + progress.updater_for_shard(0).update(|s| { + s.replication_lag = Some(1000); + s.last_transaction = Some(Instant::now() - Duration::from_millis(200)); + }); + + let waiter = CutoverPolicy::new(config, progress); + + assert_eq!( + waiter.should_cutover(Duration::from_millis(100)), + CutoverAction::Go(CutoverReason::LastTransaction) + ); + + assert_eq!( + waiter.wait_for_catchup().await.unwrap(), + CutoverReason::LastTransaction + ); + } + + #[tokio::test] + async fn test_cutover_timeout_aborts_when_configured() { + let config = CutoverConfig { + timeout: Duration::ZERO, + timeout_action: CutoverTimeoutAction::Abort, + ..cutover_config() + }; + + let progress = ReplicationProgress::new(1); + progress + .updater_for_shard(0) + .update(|s| s.replication_lag = Some(5000)); + + let waiter = CutoverPolicy::new(config, progress); + + assert_eq!( + waiter.should_cutover(Duration::ZERO), + CutoverAction::Go(CutoverReason::Timeout) + ); + assert_matches!(waiter.wait_for_catchup().await, Err(Error::AbortTimeout)); + } + + #[tokio::test] + async fn test_cutover_timeout_cuts_over_when_configured() { + let config = CutoverConfig { + timeout: Duration::ZERO, + timeout_action: CutoverTimeoutAction::Cutover, + ..cutover_config() + }; + + let progress = ReplicationProgress::new(1); + progress + .updater_for_shard(0) + .update(|s| s.replication_lag = Some(5000)); + + let waiter = CutoverPolicy::new(config, progress); + + assert_eq!( + waiter.wait_for_catchup().await.unwrap(), + CutoverReason::Timeout + ); + } + + #[tokio::test] + async fn test_should_not_cutover_when_no_transaction_was_applied() { + let config = CutoverConfig { + replication_lag_threshold: 10, + last_transaction_delay: Duration::from_millis(100), + ..cutover_config() + }; + + let progress = ReplicationProgress::new(1); + progress + .updater_for_shard(0) + .update(|s| s.replication_lag = Some(1000)); + + let waiter = CutoverPolicy::new(config, progress); + + assert!(matches!( + waiter.should_cutover(Duration::from_millis(100)), + CutoverAction::NoGo(_) + )); + } + + #[tokio::test] + async fn test_should_not_cutover_when_lag_above_threshold_and_recent_transaction() { + let config = cutover_config(); + + let progress = ReplicationProgress::new(1); + progress.updater_for_shard(0).update(|s| { + s.replication_lag = Some(1000); + s.last_transaction = Some(Instant::now() - Duration::from_millis(50)); + }); + + let waiter = CutoverPolicy::new(config, progress); + + assert!(matches!( + waiter.should_cutover(Duration::from_millis(100)), + CutoverAction::NoGo { .. } + )); + } + + #[tokio::test] + async fn test_should_not_cutover_when_timeout_not_reached() { + let config = CutoverConfig { + timeout: Duration::from_secs(1), + replication_lag_threshold: 10, + ..cutover_config() + }; + + let progress = ReplicationProgress::new(1); + progress.updater_for_shard(0).update(|s| { + s.replication_lag = Some(1000); + s.last_transaction = Some(Instant::now() - Duration::from_millis(100)); + }); + + let waiter = CutoverPolicy::new(config, progress); + + assert!(matches!( + waiter.should_cutover(Duration::from_millis(999)), + CutoverAction::NoGo { .. } + )); + } + + #[tokio::test] + async fn test_should_not_cutover_when_lag_just_above_threshold() { + let config = cutover_config(); + + let progress = ReplicationProgress::new(1); + progress.updater_for_shard(0).update(|s| { + s.replication_lag = Some(101); + s.last_transaction = Some(Instant::now() - Duration::from_millis(50)); + }); + + let waiter = CutoverPolicy::new(config, progress); + + assert!(matches!( + waiter.should_cutover(Duration::from_millis(100)), + CutoverAction::NoGo { .. } + )); + } + + #[tokio::test] + async fn should_not_cutover_before_every_shard_reports() { + let config = CutoverConfig { + replication_lag_threshold: 1000, + ..cutover_config() + }; + + let progress = ReplicationProgress::new(2); + progress + .updater_for_shard(0) + .update(|s| s.last_transaction = Some(Instant::now())); + + let waiter = CutoverPolicy::new(config, progress.clone()); + let elapsed = Duration::from_millis(100); + + assert_eq!(progress.snapshot().lag_bytes, None); + assert_matches!(waiter.should_cutover(elapsed), CutoverAction::NoGo { .. }); + + progress + .updater_for_shard(0) + .update(|s| s.replication_lag = Some(500)); + assert_eq!(progress.snapshot().lag_bytes, None); + assert_matches!(waiter.should_cutover(elapsed), CutoverAction::NoGo { .. }); + + progress + .updater_for_shard(1) + .update(|s| s.replication_lag = Some(400)); + assert_eq!( + waiter.should_cutover(elapsed), + CutoverAction::Go(CutoverReason::Lag) + ); + } +} diff --git a/pgdog/src/backend/replication/logical/publisher/mod.rs b/pgdog/src/backend/replication/logical/publisher/mod.rs index 858c84331..d7141adb2 100644 --- a/pgdog/src/backend/replication/logical/publisher/mod.rs +++ b/pgdog/src/backend/replication/logical/publisher/mod.rs @@ -4,9 +4,11 @@ pub(crate) use non_identity_columns_presence::*; pub(crate) mod slot; pub(crate) use slot::*; pub(crate) mod copy; -pub(crate) mod progress; +pub(crate) mod cutover_policy; pub(crate) mod publisher_impl; pub(crate) mod queries; +pub(crate) mod replication_progress; +pub(crate) mod replication_stream; pub(crate) mod resharding_replicas; pub(crate) mod table; pub(crate) use copy::*; diff --git a/pgdog/src/backend/replication/logical/publisher/progress.rs b/pgdog/src/backend/replication/logical/publisher/progress.rs deleted file mode 100644 index 2255cd2f5..000000000 --- a/pgdog/src/backend/replication/logical/publisher/progress.rs +++ /dev/null @@ -1,79 +0,0 @@ -use std::sync::Arc; -use std::sync::atomic::{AtomicI64, AtomicUsize, Ordering}; -use std::time::Duration; - -use tokio::select; -use tokio::sync::Notify; -use tracing::info; - -use crate::backend::replication::publisher::Lsn; -use crate::tasks; -use crate::util::safe_sleep; - -#[derive(Debug)] -struct Inner { - bytes_sharded: AtomicUsize, - lsn: AtomicI64, - done: Notify, -} - -#[derive(Debug, Clone)] -pub(crate) struct Progress { - inner: Arc, -} - -impl Progress { - pub(crate) fn new_stream() -> Self { - let inner = Arc::new(Inner { - bytes_sharded: AtomicUsize::new(0), - lsn: AtomicI64::new(0), - done: Notify::new(), - }); - - let notify = inner.clone(); - - tasks::spawn("logical publisher progress", async move { - let mut prev = 0; - loop { - select! { - _ = safe_sleep(Duration::from_secs(5)) => { - let written = notify.bytes_sharded.load(Ordering::Relaxed); - let lsn = notify.lsn.load(Ordering::Relaxed); - - info!( - "replicated {:.3} MB position {} [{:.3} MB/sec]", - written as f64 / 1024.0 / 1024.0, - Lsn::from_i64(lsn), - (written - prev) as f64 / 5.0 / 1024.0 / 1024.0 - ); - - prev = written; - } - - _ = notify.done.notified() => { - break; - } - } - } - }); - - Progress { inner } - } - - pub(crate) fn update(&self, total_bytes: usize, lsn: i64) { - self.inner - .bytes_sharded - .store(total_bytes, Ordering::Relaxed); - self.inner.lsn.store(lsn, Ordering::Relaxed); - } - - pub(crate) fn done(&self) { - self.inner.done.notify_one(); - } -} - -impl Drop for Progress { - fn drop(&mut self) { - self.done() - } -} diff --git a/pgdog/src/backend/replication/logical/publisher/publisher_impl.rs b/pgdog/src/backend/replication/logical/publisher/publisher_impl.rs index 00e021a67..ab96fefa7 100644 --- a/pgdog/src/backend/replication/logical/publisher/publisher_impl.rs +++ b/pgdog/src/backend/replication/logical/publisher/publisher_impl.rs @@ -1,27 +1,12 @@ use std::collections::HashMap; -use std::sync::Arc; -use std::time::Duration; - -use parking_lot::Mutex; -use tokio::select; -use tokio::task::JoinHandle; -use tokio::time::Instant; -use tokio::try_join; + use tokio_util::sync::CancellationToken; -use tracing::{debug, warn}; +use tracing::warn; use super::super::{Error, publisher::Table}; use super::ReplicationSlot; - -use crate::backend::replication::logical::publisher::ReplicationData; -use crate::backend::replication::logical::subscriber::stream::StreamSubscriber; -use crate::backend::replication::publisher::Lsn; -use crate::backend::replication::publisher::progress::Progress; use crate::backend::replication::tables_sync::tables_sync; use crate::backend::{Cluster, pool::Request}; -use crate::net::replication::ReplicationMeta; -use crate::tasks; -use crate::util::{safe_interval, safe_sleep}; #[derive(Debug, Default)] pub(crate) struct Publisher { @@ -30,12 +15,7 @@ pub(crate) struct Publisher { /// Shard -> Tables mapping. pub(crate) tables: HashMap>, /// Replication slots. - slots: HashMap, - /// Replication lag. - replication_lag: Arc>>, - /// Last transaction. - last_transaction: Arc>>, - /// Slot name. + pub(crate) slots: HashMap, slot_name: String, } @@ -45,12 +25,22 @@ impl Publisher { publication: publication.to_string(), tables: HashMap::new(), slots: HashMap::new(), - replication_lag: Arc::new(Mutex::new(HashMap::new())), - last_transaction: Arc::new(Mutex::new(None)), slot_name, } } + pub(crate) fn pop_tables(&mut self, shard: usize) -> Result, Error> { + self.tables + .remove(&shard) + .ok_or(Error::NoReplicationTables(shard)) + } + + pub(crate) fn pop_slot(&mut self, shard: usize) -> Result { + self.slots + .remove(&shard) + .ok_or(Error::NoReplicationSlot(shard)) + } + /// Synchronize tables for all shards. pub(crate) async fn sync_tables( &mut self, @@ -91,210 +81,55 @@ impl Publisher { &mut self, source: &Cluster, cancel: &CancellationToken, + ) -> Result<(), Error> { + let result = self.create_every_slot(source, cancel).await; + if result.is_err() + && let Err(cleanup) = Box::pin(self.cleanup()).await + { + warn!("failed to drop partially created replication slots: {cleanup}"); + } + result + } + + async fn create_every_slot( + &mut self, + source: &Cluster, + cancel: &CancellationToken, ) -> Result<(), Error> { for (number, shard) in source.shards().iter().enumerate() { - // Cancel at slot boundaries so we never tear down an in-flight - // CREATE_REPLICATION_SLOT: the current slot completes, the next is - // not started. Slots already created are dropped by the caller. if cancel.is_cancelled() { - return Err(Error::DataSyncAborted); + return Err(Error::ReplicationAborted); } let addr = shard.primary(&Request::default()).await?.addr().clone(); - let mut slot = ReplicationSlot::replication( + let slot = ReplicationSlot::replication( &self.publication, &addr, Some(self.slot_name.clone()), number, ); + let slot = self.slots.entry(number).or_insert(slot); Box::pin(slot.create_slot()).await?; - - self.slots.insert(number, slot); } Ok(()) } - /// Replicate and fan-out data from a shard to N shards. - /// - /// This uses a dedicated replication slot which will survive crashes and reboots. - /// N.B.: The slot needs to be manually dropped! - pub(crate) async fn replicate( + pub(crate) async fn prepare_replication( &mut self, source: &Cluster, - dest: &Cluster, - ) -> Result { - // Replicate shards in parallel. - let mut streams = vec![]; - - let stop = CancellationToken::new(); - + cancel: &CancellationToken, + ) -> Result<(), Error> { // Synchronize tables from publication. self.sync_tables(false, source).await?; // Create replication slots if we haven't already. if self.slots.is_empty() { - Box::pin(self.create_slots(source, &stop)).await?; - } - - for (number, _) in source.shards().iter().enumerate() { - // Use table offsets from data sync - // or from loading them above. - let tables = self - .tables - .get(&number) - .map(Vec::as_slice) - .unwrap_or_default(); - - let mut stream = StreamSubscriber::new(dest, tables); - - // Take ownership of the slot for replication. - let mut slot = self - .slots - .remove(&number) - .ok_or(Error::NoReplicationSlot(number))?; - stream.set_current_lsn(slot.lsn().lsn); - - let mut check_lag = safe_interval(Duration::from_secs(1)); - let replication_lag = self.replication_lag.clone(); - let stop = stop.clone(); - let last_transaction = self.last_transaction.clone(); - - let source_cluster = source.clone(); - let dest = dest.clone(); - - // Replicate in parallel. - let handle = tasks::spawn("replication", async move { - slot.start_replication().await?; - let progress = Progress::new_stream(); - let max_attempts = dest.resharding_replication_retry_max_attempts(); - let delay = dest.resharding_replication_retry_min_delay(); - let mut attempt = 0usize; - // Latches on the first cancellation so the `cancelled()` arm fires - // once (it stays ready forever after `cancel()`); the drain below - // then runs to completion. - let mut stopping = false; - loop { - select! { - _ = stop.cancelled(), if !stopping => { - slot.stop_replication().await?; - stopping = true; - } - - // This is cancellation-safe. - replication_data = slot.replicate(Duration::MAX) => { - // Returns Ok(true) when the slot is drained and the loop - // should break; Ok(false) to continue. All errors bubble up - // to the single retry/abort site below. - let done: Result = async { - let Some(replication_data) = replication_data? else { - slot.drop_slot().await?; - return Ok(true); - }; - match replication_data { - ReplicationData::CopyData(data) => { - if let Some(ReplicationMeta::KeepAlive(ka)) = - data.replication_meta() - { - // Advance the lsn if we are not in the transaction currently - // (we don't use transactions actually without streaming on protocol version 4, - // but let it be as a safeguard). - // If we got the keep-alive message and not the update message - // then it's for the unrelated changes that advanced WAL. - // Since it's unrelated we can advance our progress and - // consider that lag replication - let advanced = !stream.in_transaction() - && stream.set_current_lsn(ka.wal_end); - - // Reply to walsender if it asked for reply or - // if we advanced due to the WAL progress but - // the update was not related - if advanced || ka.reply() { - slot.status_update(stream.status_update()).await?; - } - debug!( - "origin at lsn {} [{}]", - Lsn::from_i64(ka.wal_end), - slot.server()?.addr() - ); - progress.update(stream.bytes_sharded(), ka.wal_end); - } else { - if let Some(su) = stream.handle(data).await? { - slot.status_update(su).await?; - *last_transaction.lock() = Some(Instant::now()); - } - attempt = 0; - progress.update(stream.bytes_sharded(), stream.lsn()); - } - Ok(false) - } - ReplicationData::CopyDone => Ok(false), - } - } - .await; - - match done { - Ok(true) => break, - Ok(false) => {} - Err(err) - if err.is_retryable() - && (max_attempts == 0 || attempt < max_attempts) => - { - attempt += 1; - warn!( - "[replication] error ({attempt}/{max_attempts}): {err}, reconnecting in {}ms", - delay.as_millis() - ); - safe_sleep(delay).await; - if let Err(reconnect_err) = - try_join!(slot.reconnect(), stream.reconnect()) - { - if !reconnect_err.is_retryable() { - return Err(reconnect_err); - } - stream.reset_connections(); - warn!( - "[replication] reconnect error ({attempt}/{max_attempts}): {reconnect_err}, will retry" - ); - } - } - Err(err) => return Err(err), - } - } - - _ = check_lag.tick() => { - let lag = slot.replication_lag().await?; - - let mut guard = replication_lag.lock(); - guard.insert(number, lag); - - let missed = stream.missed_rows(); - if missed.non_zero() { - warn!("replication {} => {} has missing rows: {}", source_cluster.name(), dest.name(), missed); - } - - } - } - } - - Ok::<(), Error>(()) - }); - - streams.push(handle); + Box::pin(self.create_slots(source, cancel)).await?; } - Ok(Waiter { streams, stop }) - } - - /// Get current replication lag. - pub(crate) fn replication_lag(&self) -> HashMap { - self.replication_lag.lock().clone() - } - - /// Get how long ago last transaction was committed. - pub(crate) fn last_transaction(&self) -> Option { - (*self.last_transaction.lock()).map(|last| last.elapsed()) + Ok(()) } pub(crate) fn post_data_sync(&mut self, tables: HashMap>) { @@ -319,50 +154,10 @@ impl Publisher { } } -#[cfg(test)] -impl Publisher { - pub(crate) fn set_replication_lag(&self, shard: usize, lag: i64) { - self.replication_lag.lock().insert(shard, lag); - } - - pub(crate) fn set_last_transaction(&self, instant: Option) { - *self.last_transaction.lock() = instant; - } -} - -#[derive(Debug)] -pub(crate) struct Waiter { - streams: Vec>>, - stop: CancellationToken, -} - -impl Waiter { - pub(crate) fn stop(&self) { - self.stop.cancel(); - } - - pub(crate) async fn wait(&mut self) -> Result<(), Error> { - for stream in &mut self.streams { - stream.await??; - } - - Ok(()) - } -} - -#[cfg(test)] -impl Waiter { - pub(crate) fn new_test() -> Self { - Self { - streams: vec![], - stop: CancellationToken::new(), - } - } -} - #[cfg(test)] mod test { use super::*; + use crate::backend::replication::logical::subscriber::stream::StreamSubscriber; use crate::backend::server::test::test_replication_server; use crate::config::config; @@ -393,13 +188,9 @@ mod test { let result = publisher.create_slots(&source, &cancel).await; assert!( - matches!(result, Err(Error::DataSyncAborted)), + matches!(result, Err(Error::ReplicationAborted)), "slot creation must abort on a cancelled token; got: {result:?}" ); - assert!( - publisher.slots.is_empty(), - "the cancelled token aborts before any slot is created" - ); source.shutdown(); for ddl in &[ @@ -458,7 +249,7 @@ mod test { let cfg = config(); let cluster = Cluster::new_test(&cfg); cluster.launch(); - let mut stream = StreamSubscriber::new(&cluster, &[]); + let mut stream = StreamSubscriber::new(&cluster, vec![]); stream.connect().await.unwrap(); let result = stream.handle(begin_copy_data(1)).await; @@ -478,7 +269,7 @@ mod test { let cfg = config(); let cluster = Cluster::new_test(&cfg); cluster.launch(); - let mut stream = StreamSubscriber::new(&cluster, &[]); + let mut stream = StreamSubscriber::new(&cluster, vec![]); stream.connect().await.unwrap(); let result = stream.handle(commit_copy_data(1)).await; diff --git a/pgdog/src/backend/replication/logical/publisher/replication_progress.rs b/pgdog/src/backend/replication/logical/publisher/replication_progress.rs new file mode 100644 index 000000000..a1a20764a --- /dev/null +++ b/pgdog/src/backend/replication/logical/publisher/replication_progress.rs @@ -0,0 +1,226 @@ +use std::sync::Arc; + +use parking_lot::Mutex; +use tokio::time::Instant; + +use crate::backend::replication::publisher::Lsn; +use crate::util::stats::average_rate; +use pgdog_stats::MissedRows; + +/// Tracks the progress of replication for +/// a single source shard +#[derive(Debug, Default, Clone, Copy)] +pub(crate) struct ReplicationShardProgress { + pub(crate) replication_lag: Option, + pub(crate) last_transaction: Option, + pub(crate) applied_lsn: Option, + pub(crate) missed_rows: MissedRows, + pub(crate) bytes_sharded: usize, + pub(crate) rows_sharded: usize, + pub(crate) origin_lsn: Lsn, + pub(crate) started: Option, +} + +impl ReplicationShardProgress { + pub(crate) fn advance_applied_lsn(&mut self, applied: Lsn) { + self.applied_lsn = Some( + self.applied_lsn + .map_or(applied, |current| current.max(applied)), + ); + } + + fn rate(&self, count: u64) -> Option { + self.started + .and_then(|started| average_rate(count, started.into_std())) + } + + fn lag_bytes(&self) -> Option { + self.replication_lag.map(|lag| lag.max(0) as u64) + } + + pub(crate) fn snapshot(&self, fallback_lsn: Lsn) -> pgdog_stats::ReplicationShardStatus { + let rows = self.rows_sharded as u64; + let bytes = self.bytes_sharded as u64; + pgdog_stats::ReplicationShardStatus { + lsn: self.applied_lsn.unwrap_or(fallback_lsn), + lag_bytes: self.lag_bytes(), + missed_rows: self.missed_rows, + rows, + bytes, + rows_per_sec: self.rate(rows), + bytes_per_sec: self.rate(bytes), + } + } +} + +/// Tracks the progress for all of the source shards +#[derive(Clone, Debug)] +pub(crate) struct ReplicationProgress { + shards: Arc<[Mutex]>, +} + +impl ReplicationProgress { + pub(crate) fn new(shard_count: usize) -> Self { + let shards = (0..shard_count).map(|_| Mutex::default()).collect(); + Self { shards } + } + + /// Returns the entity to update the progress for a single source shard + pub(crate) fn updater_for_shard(&self, shard: usize) -> ReplicationProgressShardUpdater { + assert!( + shard < self.shards.len(), + "shard {shard} is out of range for {} shards", + self.shards.len() + ); + + ReplicationProgressShardUpdater { + shards: self.shards.clone(), + shard, + } + } + + /// The combined progress of every shard, as reported to `SHOW TASKS` and + /// read by the cutover policy. `lag_bytes` stays `None` until every shard + /// has reported one. + pub(crate) fn snapshot(&self) -> pgdog_stats::ReplicationProgress { + let mut lag: Option = None; + let mut every_shard_reported = true; + let mut last_transaction: Option = None; + let mut rows = 0; + let mut bytes = 0; + let mut rows_per_sec = None; + let mut bytes_per_sec = None; + + for shard in self.shards.iter() { + let shard = *shard.lock(); + match shard.lag_bytes() { + Some(shard_lag) => lag = Some(lag.map_or(shard_lag, |max| max.max(shard_lag))), + None => every_shard_reported = false, + } + if let Some(applied) = shard.last_transaction { + last_transaction = Some(last_transaction.map_or(applied, |max| max.max(applied))); + } + let shard_rows = shard.rows_sharded as u64; + let shard_bytes = shard.bytes_sharded as u64; + rows += shard_rows; + bytes += shard_bytes; + rows_per_sec = sum_rates(rows_per_sec, shard.rate(shard_rows)); + bytes_per_sec = sum_rates(bytes_per_sec, shard.rate(shard_bytes)); + } + + pgdog_stats::ReplicationProgress { + lag_bytes: every_shard_reported.then_some(lag).flatten(), + last_transaction_ms: last_transaction + .map(|applied| applied.elapsed().as_millis() as u64), + rows, + bytes, + rows_per_sec, + bytes_per_sec, + } + } +} + +fn sum_rates(total: Option, rate: Option) -> Option { + match (total, rate) { + (None, None) => None, + (total, rate) => Some(total.unwrap_or(0) + rate.unwrap_or(0)), + } +} + +/// Used to update the progress of a single source shard stream. +#[derive(Clone, Debug)] +pub(crate) struct ReplicationProgressShardUpdater { + shards: Arc<[Mutex]>, + shard: usize, +} + +impl ReplicationProgressShardUpdater { + pub(crate) fn update(&self, f: impl FnOnce(&mut ReplicationShardProgress)) { + f(&mut self.shards[self.shard].lock()); + } + + pub(crate) fn snapshot(&self) -> ReplicationShardProgress { + *self.shards[self.shard].lock() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::time::Duration; + + #[test] + fn lag_none_until_all_shards_report() { + let progress = ReplicationProgress::new(3); + + assert_eq!(progress.snapshot().lag_bytes, None); + + progress + .updater_for_shard(0) + .update(|p| p.replication_lag = Some(100)); + assert_eq!(progress.snapshot().lag_bytes, None); + + progress + .updater_for_shard(1) + .update(|p| p.replication_lag = Some(200)); + assert_eq!(progress.snapshot().lag_bytes, None); + + progress + .updater_for_shard(2) + .update(|p| p.replication_lag = Some(150)); + assert_eq!(progress.snapshot().lag_bytes, Some(200)); + } + + #[test] + fn cloned_updater_shares_shard_state() { + let progress = ReplicationProgress::new(2); + let a = progress.updater_for_shard(0); + let b = a.clone(); + + a.update(|p| p.replication_lag = Some(77)); + assert_eq!(b.snapshot().replication_lag, Some(77)); + + b.update(|p| p.replication_lag = Some(99)); + assert_eq!(a.snapshot().replication_lag, Some(99)); + } + + #[test] + fn updaters_for_different_shards_are_independent() { + let progress = ReplicationProgress::new(2); + progress + .updater_for_shard(0) + .update(|p| p.replication_lag = Some(10)); + progress + .updater_for_shard(1) + .update(|p| p.replication_lag = Some(20)); + + assert_eq!( + progress.updater_for_shard(0).snapshot().replication_lag, + Some(10) + ); + assert_eq!( + progress.updater_for_shard(1).snapshot().replication_lag, + Some(20) + ); + } + + #[tokio::test(start_paused = true)] + async fn last_transaction_returns_most_recent_across_shards() { + let progress = ReplicationProgress::new(2); + + assert_eq!(progress.snapshot().last_transaction_ms, None); + + let older = tokio::time::Instant::now() - Duration::from_millis(300); + progress + .updater_for_shard(0) + .update(|p| p.last_transaction = Some(older)); + + tokio::time::advance(Duration::from_millis(10)).await; + let recent = tokio::time::Instant::now(); + progress + .updater_for_shard(1) + .update(|p| p.last_transaction = Some(recent)); + + assert_eq!(progress.snapshot().last_transaction_ms, Some(0)); + } +} diff --git a/pgdog/src/backend/replication/logical/publisher/replication_stream.rs b/pgdog/src/backend/replication/logical/publisher/replication_stream.rs new file mode 100644 index 000000000..7741db5ae --- /dev/null +++ b/pgdog/src/backend/replication/logical/publisher/replication_stream.rs @@ -0,0 +1,560 @@ +use std::time::Duration; + +use tokio::select; +use tokio::time::{Instant, MissedTickBehavior}; +use tokio::try_join; +use tokio_util::sync::CancellationToken; +use tracing::{debug, warn}; + +use super::replication_progress::{ReplicationProgressShardUpdater, ReplicationShardProgress}; +use super::{Lsn, ReplicationData, ReplicationSlot, Table}; +use crate::backend::Cluster; +use crate::backend::replication::logical::Error; +use crate::backend::replication::logical::subscriber::stream::StreamSubscriber; +use crate::net::replication::ReplicationMeta; +use crate::util::{safe_interval, safe_sleep}; + +/// Runs the replication stream from a single shard (slot) +/// to the destination cluster. +#[derive(Debug)] +pub(crate) struct ReplicationStream { + source_name: String, + dest_cluster: Cluster, + updater: ReplicationProgressShardUpdater, +} + +impl ReplicationStream { + pub(crate) fn new( + source: &Cluster, + dest: &Cluster, + updater: ReplicationProgressShardUpdater, + ) -> Self { + Self { + source_name: source.name().to_owned(), + dest_cluster: dest.clone(), + updater, + } + } + + pub(crate) fn progress(&self) -> ReplicationShardProgress { + self.updater.snapshot() + } + + pub(crate) async fn run( + &self, + slot: &mut ReplicationSlot, + tables: Vec
, + stop: &CancellationToken, + ) -> Result<(), Error> { + let mut stream = StreamSubscriber::new(&self.dest_cluster, tables); + stream.set_current_lsn(slot.lsn().lsn); + self.updater.update(|p| { + p.advance_applied_lsn(slot.lsn()); + p.started = Some(Instant::now()); + }); + let result = self.replicate(slot, &mut stream, stop).await; + let final_lsn = Lsn::from_i64(stream.status_update().last_applied); + let missed = stream.missed_rows(); + self.updater.update(|p| { + p.advance_applied_lsn(final_lsn); + p.missed_rows.merge(missed); + }); + result + } + + async fn update_progress( + &self, + slot: &mut ReplicationSlot, + stream: &mut StreamSubscriber, + ) -> Result<(), Error> { + let missed = stream.missed_rows(); + let applied = Lsn::from_i64(stream.status_update().last_applied); + let bytes_sharded = stream.bytes_sharded(); + let rows_sharded = stream.rows_sharded(); + let origin_lsn = Lsn::from_i64(stream.lsn()); + self.updater.update(|p| { + p.advance_applied_lsn(applied); + p.missed_rows.merge(missed); + p.bytes_sharded = bytes_sharded; + p.rows_sharded = rows_sharded; + p.origin_lsn = origin_lsn; + }); + if missed.non_zero() { + warn!( + "replication {} => {} has missing rows: {}", + self.source_name, + self.dest_cluster.name(), + missed + ); + } + let lag = slot.replication_lag().await?; + self.updater.update(|p| p.replication_lag = Some(lag)); + Ok(()) + } + + async fn replicate( + &self, + slot: &mut ReplicationSlot, + stream: &mut StreamSubscriber, + stop: &CancellationToken, + ) -> Result<(), Error> { + let mut check_lag = safe_interval(Duration::from_secs(1)); + check_lag.set_missed_tick_behavior(MissedTickBehavior::Delay); + slot.start_replication().await?; + + let max_attempts = self + .dest_cluster + .resharding_replication_retry_max_attempts(); + let delay = self.dest_cluster.resharding_replication_retry_min_delay(); + + let mut attempt = 0usize; + + loop { + let stopping = slot.stopped(); + + select! { + biased; + + _ = stop.cancelled(), if !stopping => { + slot.stop_replication().await?; + } + + _ = check_lag.tick() => { + if let Err(err) = self.update_progress(slot, stream).await { + self.updater.update(|p| p.replication_lag = None); + warn!( + "[replication] progress update failed for slot \"{}\": {err}", + slot.name() + ); + } + } + + replication_data = slot.replicate(Duration::MAX) => { + // Returns Ok(true) when the slot is drained and the loop + // should break; Ok(false) to continue on next loop. All errors bubble up + // to the single retry/abort site below. + let done: Result = async { + let Some(replication_data) = replication_data? else { + return Ok(true); + }; + match replication_data { + ReplicationData::CopyData(data) => { + if let Some(ReplicationMeta::KeepAlive(ka)) = + data.replication_meta() + { + // Advance the lsn if we are not in the transaction currently + // (we don't use transactions actually without streaming on protocol version 4, + // but let it be as a safeguard). + // If we got the keep-alive message and not the update message + // then it's for the unrelated changes that advanced WAL. + // Since it's unrelated we can advance our progress and + // consider that lag replication + let advanced = !stream.in_transaction() + && ka.wal_end > stream.lsn() + && stream.set_current_lsn(ka.wal_end); + + // Reply to walsender if it asked for reply or + // if we advanced due to the WAL progress but + // the update was not related + if advanced || ka.reply() { + slot.status_update(stream.status_update()).await?; + } + debug!( + "origin at lsn {} [{}]", + Lsn::from_i64(ka.wal_end), + slot.addr() + ); + } else { + if let Some(su) = stream.handle(data).await? { + let applied = Lsn::from_i64(su.last_applied); + slot.status_update(su).await?; + self.updater.update(|p| { + p.last_transaction = Some(Instant::now()); + p.advance_applied_lsn(applied); + }); + } + attempt = 0; + } + Ok(false) + } + ReplicationData::CopyDone => Ok(false), + } + } + .await; + + match done { + Ok(true) => break, + Ok(false) => {} + Err(err) + if err.is_retryable() + && (max_attempts == 0 || attempt < max_attempts) => + { + attempt += 1; + warn!( + "[replication] error ({attempt}/{max_attempts}): {err}, reconnecting in {}ms", + delay.as_millis() + ); + safe_sleep(delay).await; + let missed = stream.missed_rows(); + self.updater.update(|p| p.missed_rows.merge(missed)); + if let Err(reconnect_err) = + try_join!(slot.reconnect(), stream.reconnect()) + { + if !reconnect_err.is_retryable() { + return Err(reconnect_err); + } + stream.reset_connections(); + warn!( + "[replication] reconnect error ({attempt}/{max_attempts}): {reconnect_err}, will retry" + ); + } + } + Err(err) => return Err(err), + } + } + } + } + + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use std::{error::Error as StdError, sync::Arc, time::Duration}; + + use tokio::{ + task::JoinHandle, + time::{sleep, timeout}, + }; + + use crate::{ + backend::replication::logical::publisher::replication_progress::{ + ReplicationProgress, ReplicationShardProgress, + }, + backend::{Server, server::test::test_server}, + config::config, + util::random_string, + }; + + use super::*; + + type TestResult = Result<(), Box>; + + struct Fixture { + source_table: String, + destination_table: String, + publication: String, + slot_base: String, + server: Server, + source: Cluster, + replication: Arc, + stop: CancellationToken, + worker: Option>>, + } + + impl Fixture { + async fn new() -> Self { + let suffix = random_string(12).to_lowercase(); + let source = Cluster::new_test_single_shard(&config()); + let progress = ReplicationProgress::new(1); + let updater = progress.updater_for_shard(0); + let replication = Arc::new(ReplicationStream::new(&source, &source, updater)); + Self { + source_table: format!("replication_source_{suffix}"), + destination_table: format!("replication_destination_{suffix}"), + publication: format!("replication_publication_{suffix}"), + slot_base: format!("replication_slot_{suffix}"), + server: test_server().await, + source, + replication, + stop: CancellationToken::new(), + worker: None, + } + } + + async fn start(&mut self) -> TestResult { + self.source.launch(); + for table in [&self.source_table, &self.destination_table] { + self.server + .execute_checked(format!( + "CREATE TABLE {table} (id BIGINT PRIMARY KEY, val TEXT NOT NULL)" + )) + .await?; + } + self.server + .execute_checked(format!( + "CREATE PUBLICATION {} FOR TABLE {}", + self.publication, self.source_table + )) + .await?; + let mut tables = Table::load(&self.publication, &mut self.server).await?; + let table = tables.first_mut().ok_or("publication has no table")?; + table.table.parent_name = self.destination_table.clone(); + table.table.parent_schema = table.table.schema.clone(); + let mut slot = ReplicationSlot::replication( + &self.publication, + self.server.addr(), + Some(self.slot_base.clone()), + 0, + ); + slot.create_slot().await?; + if self.replication.progress().replication_lag.is_some() { + return Err("lag was measured before replication started".into()); + } + let replication = Arc::clone(&self.replication); + let stop = self.stop.clone(); + self.worker = Some(tokio::spawn(async move { + let result = Box::pin(replication.run(&mut slot, tables, &stop)).await; + let dropped = slot.drop_slot().await; + result.and(dropped) + })); + Ok(()) + } + + async fn wait_for( + &mut self, + query: String, + ready: impl Fn(&[String], ReplicationShardProgress) -> bool, + ) -> TestResult { + timeout(Duration::from_secs(10), async { + loop { + let rows: Vec = self.server.fetch_all(query.clone()).await?; + if ready(&rows, self.replication.progress()) { + return Ok::<(), Box>(()); + } + if self.worker.as_ref().is_some_and(JoinHandle::is_finished) { + return Err("replication stopped before the expected result".into()); + } + sleep(Duration::from_millis(20)).await; + } + }) + .await??; + Ok(()) + } + + async fn stop(&mut self) -> TestResult { + self.stop.cancel(); + let result = if let Some(worker) = self.worker.as_mut() { + match timeout(Duration::from_secs(10), &mut *worker).await { + Ok(result) => result + .map_err(Box::::from) + .and_then(|result| result.map_err(Box::::from)), + Err(error) => { + worker.abort(); + let _ = worker.await; + Err(error.into()) + } + } + } else { + Ok(()) + }; + self.worker.take(); + result + } + + async fn cleanup(&mut self) -> TestResult { + let mut result = self.stop().await; + self.source.shutdown(); + let mut server = test_server().await; + for query in [ + format!( + "SELECT pg_drop_replication_slot(slot_name) FROM pg_replication_slots WHERE slot_name = '{}_0'", + self.slot_base + ), + format!("DROP PUBLICATION IF EXISTS {}", self.publication), + format!("DROP TABLE IF EXISTS {}", self.source_table), + format!("DROP TABLE IF EXISTS {}", self.destination_table), + ] { + let cleanup: TestResult = async { + timeout(Duration::from_secs(10), server.execute_checked(query)).await??; + Ok(()) + } + .await; + result = result.and(cleanup); + } + result + } + } + + async fn with_fixture(test: impl AsyncFnOnce(&mut Fixture) -> TestResult) -> TestResult { + crate::logger(); + let mut fixture = Fixture::new().await; + let result: TestResult = async { + timeout(Duration::from_secs(30), fixture.start()).await??; + timeout(Duration::from_secs(30), test(&mut fixture)).await??; + Ok(()) + } + .await; + let cleanup = fixture.cleanup().await; + result.and(cleanup) + } + + #[tokio::test] + async fn replication_applies_dml_and_reports_progress() -> TestResult { + with_fixture(async |fixture| { + fixture + .server + .execute_checked(format!( + "INSERT INTO {} VALUES (1, 'alpha')", + fixture.source_table + )) + .await?; + let query = format!("SELECT val FROM {} ORDER BY id", fixture.destination_table); + fixture + .wait_for(query.clone(), |rows, info| { + rows == ["alpha"] && info.last_transaction.is_some() + }) + .await?; + let first_transaction = fixture.replication.progress().last_transaction; + + fixture + .server + .execute_checked(format!( + "UPDATE {} SET val = 'beta' WHERE id = 1", + fixture.source_table + )) + .await?; + fixture + .wait_for(query.clone(), |rows, info| { + rows == ["beta"] && info.last_transaction > first_transaction + }) + .await?; + + fixture + .server + .execute_checked(format!("DELETE FROM {} WHERE id = 1", fixture.source_table)) + .await?; + fixture + .wait_for(query, |rows, info| { + rows.is_empty() && info.replication_lag.is_some_and(|lag| lag >= 0) + }) + .await + }) + .await + } + + #[tokio::test] + async fn replication_cancellation_drops_slot() -> TestResult { + with_fixture(async |fixture| { + let slot_query = format!( + "SELECT slot_name FROM pg_replication_slots WHERE slot_name = '{}_0'", + fixture.slot_base + ); + let before: Vec = fixture.server.fetch_all(slot_query.clone()).await?; + if before != [format!("{}_0", fixture.slot_base)] { + return Err("replication slot is missing before cancellation".into()); + } + fixture + .server + .execute_checked(format!( + "INSERT INTO {} VALUES (1, 'committed')", + fixture.source_table + )) + .await?; + let query = format!("SELECT val FROM {} ORDER BY id", fixture.destination_table); + fixture + .wait_for(query.clone(), |rows, info| { + rows == ["committed"] && info.last_transaction.is_some() + }) + .await?; + + fixture.stop().await?; + let after: Vec = fixture.server.fetch_all(slot_query).await?; + if !after.is_empty() { + return Err("replication slot remains after shutdown".into()); + } + let rows: Vec = fixture.server.fetch_all(query).await?; + if rows != ["committed"] { + return Err("committed destination row changed during shutdown".into()); + } + Ok(()) + }) + .await + } + + async fn missed_rows_kill_walsender(fixture: &mut Fixture) -> TestResult { + let killed: Vec = fixture + .server + .fetch_all(format!( + "SELECT pg_terminate_backend(active_pid)::text FROM pg_replication_slots \ + WHERE slot_name = '{}_0' AND active_pid IS NOT NULL", + fixture.slot_base + )) + .await?; + if killed != ["true"] { + return Err("replication connection was not terminated".into()); + } + Ok(()) + } + + async fn missed_rows_cause_missed_update( + fixture: &mut Fixture, + id: i64, + sentinel_id: i64, + ) -> TestResult { + fixture + .server + .execute_checked(format!( + "DELETE FROM {} WHERE id = {id}; \ + UPDATE {} SET val = 'miss' WHERE id = {id}; \ + INSERT INTO {} VALUES ({sentinel_id}, 'sentinel')", + fixture.destination_table, fixture.source_table, fixture.source_table + )) + .await?; + let dest_query = format!( + "SELECT id::text FROM {} ORDER BY id", + fixture.destination_table + ); + let sentinel = sentinel_id.to_string(); + fixture + .wait_for(dest_query, |rows, _| rows.iter().any(|r| r == &sentinel)) + .await + } + #[tokio::test] + async fn replication_missed_rows_survive_reconnect_and_accumulate() -> TestResult { + with_fixture(async |fixture| { + fixture + .server + .execute_checked(format!( + "INSERT INTO {} VALUES (1, 'row')", + fixture.source_table + )) + .await?; + let dest_query = format!( + "SELECT id::text FROM {} ORDER BY id", + fixture.destination_table + ); + fixture + .wait_for(dest_query.clone(), |rows, _| rows.iter().any(|r| r == "1")) + .await?; + + missed_rows_cause_missed_update(fixture, 1, 2).await?; + missed_rows_kill_walsender(fixture).await?; + + fixture + .server + .execute_checked(format!( + "INSERT INTO {} VALUES (3, 'post_reconnect')", + fixture.source_table + )) + .await?; + fixture + .wait_for(dest_query.clone(), |rows, _| rows.iter().any(|r| r == "3")) + .await?; + + if fixture.replication.progress().missed_rows.updates == 0 { + return Err("missed update count was lost during reconnect".into()); + } + + missed_rows_cause_missed_update(fixture, 3, 4).await?; + + fixture.stop().await?; + if fixture.replication.progress().missed_rows.updates < 2 { + return Err("missed updates did not accumulate across reconnect".into()); + } + Ok(()) + }) + .await + } +} diff --git a/pgdog/src/backend/replication/logical/publisher/slot.rs b/pgdog/src/backend/replication/logical/publisher/slot.rs index 1f7936a3d..f90b45522 100644 --- a/pgdog/src/backend/replication/logical/publisher/slot.rs +++ b/pgdog/src/backend/replication/logical/publisher/slot.rs @@ -17,6 +17,7 @@ use std::{fmt::Display, str::FromStr, time::Duration}; use tracing::{debug, info, trace, warn}; pub(crate) use pgdog_stats::Lsn; +use pgdog_stats::TaskId; #[derive(Debug, Clone, Copy)] pub(crate) enum Snapshot { @@ -101,6 +102,12 @@ impl ReplicationSlot { } } + pub(crate) fn set_task_id(&mut self, task_id: TaskId) { + if let Some(tracker) = &mut self.tracker { + tracker.set_task_id(task_id); + } + } + /// Connect to database using replication mode. pub(crate) async fn connect(&mut self) -> Result<(), Error> { self.server = Some( @@ -136,6 +143,16 @@ impl ReplicationSlot { /// Replication lag in bytes for this slot. pub(crate) async fn replication_lag(&mut self) -> Result { + let lag = self.query_replication_lag().await; + + if lag.is_err() { + self.server_meta = None; + } + + lag + } + + async fn query_replication_lag(&mut self) -> Result { let query = format!( "SELECT pg_current_wal_lsn() - confirmed_flush_lsn \ FROM pg_replication_slots \ @@ -148,7 +165,7 @@ impl ReplicationSlot { .pop() .ok_or(Error::MissingReplicationSlot(self.name.clone()))?; - if let Some(ref tracker) = self.tracker { + if let Some(tracker) = &self.tracker { tracker.update_lag(lag); } @@ -164,6 +181,7 @@ impl ReplicationSlot { if self.server.is_none() { self.connect().await?; } + drop(self.tracker.take()); debug!( "creating replication slot \"{}\" [{}]", @@ -276,6 +294,10 @@ impl ReplicationSlot { /// Drop the slot. pub(crate) async fn drop_slot(&mut self) -> Result<(), Error> { + if !self.server.as_ref().is_some_and(Server::in_sync) { + self.server = None; + self.connect().await?; + } let drop_slot = self.drop_slot_query(true); self.server()?.execute(&drop_slot).await?; @@ -396,14 +418,26 @@ impl ReplicationSlot { /// Drop the source connection and reconnect, restarting replication from the /// last confirmed position (`self.lsn`, kept in sync by `status_update`). + /// A stream that was already asked to stop is asked again. pub(crate) async fn reconnect(&mut self) -> Result<(), Error> { + let stopped = self.stopped; self.server = None; self.connect().await?; - self.start_replication().await + self.start_replication().await?; + + if stopped { + self.stop_replication().await?; + } + + Ok(()) } /// Ask remote to close stream. pub(crate) async fn stop_replication(&mut self) -> Result<(), Error> { + if self.stopped { + return Ok(()); + } + self.server()?.send_one(&CopyDone.into()).await?; self.server()?.flush().await?; self.stopped = true; @@ -411,10 +445,22 @@ impl ReplicationSlot { Ok(()) } + pub(crate) fn stopped(&self) -> bool { + self.stopped + } + /// Current slot LSN. pub(crate) fn lsn(&self) -> Lsn { self.lsn } + + pub(crate) fn name(&self) -> &str { + &self.name + } + + pub(crate) fn addr(&self) -> &Address { + &self.address + } } #[derive(Debug, Clone)] diff --git a/pgdog/src/backend/replication/logical/status.rs b/pgdog/src/backend/replication/logical/status.rs index f9ba0ecde..b9a82370e 100644 --- a/pgdog/src/backend/replication/logical/status.rs +++ b/pgdog/src/backend/replication/logical/status.rs @@ -3,7 +3,7 @@ use std::{ops::Deref, sync::Arc, time::SystemTime}; use dashmap::DashMap; use once_cell::sync::Lazy; -use pgdog_stats::{Lsn, SchemaStatementTask}; +use pgdog_stats::{Lsn, SchemaStatementTask, TaskId}; use crate::backend::pool::Address; use crate::backend::replication::ee::{ @@ -14,9 +14,10 @@ use crate::net::ErrorResponse; static REPLICATION_SLOTS: Lazy = Lazy::new(ReplicationSlots::default); /// Replication slot. -#[derive(Debug, Clone)] +#[derive(Debug)] pub(crate) struct ReplicationSlot { inner: pgdog_stats::ReplicationSlot, + key: String, } impl Deref for ReplicationSlot { @@ -43,10 +44,12 @@ impl ReplicationSlot { lag: 0, address: address.clone().into(), last_transaction: None, + task_id: None, }, + key: format!("{}@{}", name, address), }; - ReplicationSlots::get().insert(name.to_owned(), slot.clone()); + ReplicationSlots::get().insert(slot.key.clone(), slot.inner.clone()); replication_slot_create(&slot.inner); @@ -54,22 +57,30 @@ impl ReplicationSlot { } pub(crate) fn update_lsn(&self, lsn: &Lsn) { - if let Some(mut slot) = ReplicationSlots::get().get_mut(&self.name) { + if let Some(mut slot) = ReplicationSlots::get().get_mut(&self.key) { slot.lsn = *lsn; slot.last_transaction = Some(SystemTime::now()); - replication_slot_update(&slot.inner); + replication_slot_update(&slot); } } pub(crate) fn update_lag(&self, lag: i64) { - if let Some(mut slot) = ReplicationSlots::get().get_mut(&self.name) { + if let Some(mut slot) = ReplicationSlots::get().get_mut(&self.key) { slot.lag = lag; - replication_slot_update(&slot.inner); + replication_slot_update(&slot); + } + } + + pub(crate) fn set_task_id(&mut self, task_id: TaskId) { + self.inner.task_id = Some(task_id); + if let Some(mut slot) = ReplicationSlots::get().get_mut(&self.key) { + slot.task_id = Some(task_id); + replication_slot_update(&slot); } } pub(crate) fn dropped(&self) { - ReplicationSlots::get().remove(&self.name); + ReplicationSlots::get().remove(&self.key); replication_slot_drop(&self.inner); } @@ -80,18 +91,13 @@ impl ReplicationSlot { impl Drop for ReplicationSlot { fn drop(&mut self) { - // The slot is dropped automatically by the connection, - // and we don't call fn dropped manually, so we need to do that here - // to track the slot is gone. - if self.copy_data { - self.dropped(); - } + self.dropped(); } } #[derive(Default, Clone, Debug)] pub(crate) struct ReplicationSlots { - slots: Arc>, + slots: Arc>, } impl ReplicationSlots { @@ -101,7 +107,7 @@ impl ReplicationSlots { } impl Deref for ReplicationSlots { - type Target = Arc>; + type Target = Arc>; fn deref(&self) -> &Self::Target { &self.slots diff --git a/pgdog/src/backend/replication/logical/subscriber/pipeline.rs b/pgdog/src/backend/replication/logical/subscriber/pipeline.rs index 2a2de8612..e176b1239 100644 --- a/pgdog/src/backend/replication/logical/subscriber/pipeline.rs +++ b/pgdog/src/backend/replication/logical/subscriber/pipeline.rs @@ -10,13 +10,13 @@ use tokio::sync::{ }; use tracing::trace; -use super::stream::MissedRows; use crate::backend::Server; use crate::backend::pool::Address; use crate::net::{ Bind, CommandComplete, ErrorResponse, Execute, Flush, FromBytes, Message, Parse, Protocol, ProtocolMessage, Sync, ToBytes, }; +use pgdog_stats::MissedRows; use super::super::Error; @@ -350,7 +350,13 @@ impl Listener { && let Ok(complete) = CommandComplete::try_from(message) && matches!(complete.rows(), Ok(Some(0))) { - self.shared.lock().missed.record(complete.tag()); + let mut shared = self.shared.lock(); + match complete.tag() { + "INSERT" => shared.missed.inserts += 1, + "UPDATE" => shared.missed.updates += 1, + "DELETE" => shared.missed.deletes += 1, + _ => (), + } } if !self .queue @@ -702,6 +708,8 @@ mod test { // (insert, update, delete): one 0-row direct UPDATE and one 0-row direct // DELETE counted; the non-direct DELETE and the 1-row DELETE are not. // Insert never missed here, so it must stay 0 (no spurious counter). - assert_eq!(missed.counts(), (0, 1, 1)); + assert_eq!(missed.inserts, 0); + assert_eq!(missed.updates, 1); + assert_eq!(missed.deletes, 1); } } diff --git a/pgdog/src/backend/replication/logical/subscriber/stream.rs b/pgdog/src/backend/replication/logical/subscriber/stream.rs index 903382060..ade69c026 100644 --- a/pgdog/src/backend/replication/logical/subscriber/stream.rs +++ b/pgdog/src/backend/replication/logical/subscriber/stream.rs @@ -12,6 +12,7 @@ use std::{ use futures::future::try_join_all; use once_cell::sync::Lazy; use pgdog_postgres_types::Oid; +use pgdog_stats::MissedRows; use tracing::{debug, trace, warn}; use super::super::publisher::{NonIdentityColumnsPresence, tables_missing_unique_index}; @@ -128,10 +129,13 @@ pub(crate) struct StreamSubscriber { // Bytes sharded bytes_sharded: usize, + rows_sharded: usize, + + missed_rows: MissedRows, } impl StreamSubscriber { - pub(crate) fn new(cluster: &Cluster, tables: &[Table]) -> Self { + pub(crate) fn new(cluster: &Cluster, tables: Vec
) -> Self { let cluster = cluster.logical_stream(); Self { cluster, @@ -140,14 +144,14 @@ impl StreamSubscriber { table_lsns: HashMap::new(), changed_tables: HashSet::new(), tables: tables - .iter() + .into_iter() .map(|table| { ( Key { schema: table.table.schema.clone(), name: table.table.name.clone(), }, - table.clone(), + table, ) }) .collect(), @@ -155,9 +159,11 @@ impl StreamSubscriber { committed_lsn: 0, lsn: 0, // Unknown, bytes_sharded: 0, + rows_sharded: 0, lsn_changed: true, in_transaction: false, keys: HashMap::default(), + missed_rows: MissedRows::default(), } } @@ -804,6 +810,12 @@ impl StreamSubscriber { self.connections.clear(); } + fn capture_missed_rows(&mut self) { + for conn in &self.connections { + self.missed_rows.merge(conn.take_missed_rows()); + } + } + /// `docs/REPLICATION.md` → "Error rollback". pub(crate) async fn handle(&mut self, data: CopyData) -> Result, Error> { match self.handle_inner(data).await { @@ -835,11 +847,21 @@ impl StreamSubscriber { && let Some(payload) = xlog.payload() { match payload { - XLogPayload::Insert(insert) => self.insert(insert).await?, - XLogPayload::Update(update) => self.update(update).await?, - XLogPayload::Delete(delete) => self.delete(delete).await?, + XLogPayload::Insert(insert) => { + self.insert(insert).await?; + self.rows_sharded += 1; + } + XLogPayload::Update(update) => { + self.update(update).await?; + self.rows_sharded += 1; + } + XLogPayload::Delete(delete) => { + self.delete(delete).await?; + self.rows_sharded += 1; + } XLogPayload::Commit(commit) => { self.commit(commit).await?; + self.capture_missed_rows(); status_update = Some(self.status_update()); self.in_transaction = false; } @@ -873,6 +895,11 @@ impl StreamSubscriber { self.bytes_sharded } + /// Number of rows applied. + pub(crate) fn rows_sharded(&self) -> usize { + self.rows_sharded + } + /// Advance both LSN fields. Call after commit and on publisher init. pub(crate) fn set_current_lsn(&mut self, lsn: i64) -> bool { self.lsn_changed = lsn != self.lsn; @@ -897,13 +924,11 @@ impl StreamSubscriber { self.in_transaction } - /// Aggregate and reset the missed-row counters across all shard connections. + /// Missed rows of all transactions committed so far. Resets on read. + /// Rows of a transaction that failed are never counted, because the + /// source sends that transaction again after a reconnect. pub(crate) fn missed_rows(&mut self) -> MissedRows { - let mut total = MissedRows::default(); - for conn in &self.connections { - total.merge(conn.take_missed_rows()); - } - total + std::mem::take(&mut self.missed_rows) } /// Verify every destination shard has a qualifying unique index for all `tables`. @@ -936,71 +961,6 @@ impl StreamSubscriber { } } -#[derive(Debug, Default)] -pub(crate) struct MissedRows { - insert: usize, - delete: usize, - update: usize, -} - -impl MissedRows { - pub(crate) fn non_zero(&self) -> bool { - self.insert > 0 || self.delete > 0 || self.update > 0 - } - - /// Missed-row counts as `(insert, update, delete)`. - #[cfg(test)] - pub(crate) fn counts(&self) -> (usize, usize, usize) { - (self.insert, self.update, self.delete) - } - - /// Count a direct-to-shard DML that touched 0 rows, keyed by command tag. - pub(crate) fn record(&mut self, tag: &str) { - match tag { - "UPDATE" => self.update += 1, - "DELETE" => self.delete += 1, - "INSERT" => self.insert += 1, - _ => (), - } - } - - /// Fold another shard's counts into this one. - pub(crate) fn merge(&mut self, other: MissedRows) { - self.insert += other.insert; - self.update += other.update; - self.delete += other.delete; - } -} - -impl Display for MissedRows { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - let mut written = false; - if self.insert > 0 { - write!(f, "insert={}", self.insert)?; - written = true; - } - if self.update > 0 { - write!( - f, - "{}update={}", - if written { " " } else { "" }, - self.update - )?; - written = true; - } - if self.delete > 0 { - write!( - f, - "{}delete={}", - if written { " " } else { "" }, - self.delete - )?; - } - - Ok(()) - } -} - #[cfg(test)] mod tests { use super::super::tests::begin_copy_data; @@ -1009,7 +969,7 @@ mod tests { fn make_subscriber() -> StreamSubscriber { let cluster = Cluster::new_test(&config()); - StreamSubscriber::new(&cluster, &[]) + StreamSubscriber::new(&cluster, vec![]) } #[test] diff --git a/pgdog/src/backend/replication/logical/subscriber/tests.rs b/pgdog/src/backend/replication/logical/subscriber/tests.rs index 208931793..b962e58ae 100644 --- a/pgdog/src/backend/replication/logical/subscriber/tests.rs +++ b/pgdog/src/backend/replication/logical/subscriber/tests.rs @@ -343,18 +343,18 @@ fn x_update(u: XLogUpdate) -> CopyData { fn make_subscriber() -> StreamSubscriber { let cluster = Cluster::new_test(&config()); let tables = vec![make_sharded_table(), make_sharded_test_b_table()]; - StreamSubscriber::new(&cluster, &tables) + StreamSubscriber::new(&cluster, tables) } fn make_subscriber_with_tables(tables: Vec
) -> StreamSubscriber { let cluster = Cluster::new_test(&config()); - StreamSubscriber::new(&cluster, &tables) + StreamSubscriber::new(&cluster, tables) } fn make_subscriber_single_shard() -> StreamSubscriber { let cluster = Cluster::new_test_single_shard(&config()); let tables = vec![make_sharded_table(), make_sharded_test_b_table()]; - StreamSubscriber::new(&cluster, &tables) + StreamSubscriber::new(&cluster, tables) } /// Count rows matching the given `WHERE` predicate using a separate connection. @@ -664,7 +664,7 @@ async fn partition_leaves_share_destination() { leaf_b.table.parent_name = "sharded".to_string(); let cluster = Cluster::new_test_single_shard(&config()); - let mut sub = StreamSubscriber::new(&cluster, &[leaf_a, leaf_b]); + let mut sub = StreamSubscriber::new(&cluster, vec![leaf_a, leaf_b]); let mut verify = test_server().await; sub.connect().await.unwrap(); @@ -1595,7 +1595,7 @@ fn omni_insert_copy_data(oid: Oid, a: &str, b: &str) -> CopyData { #[tokio::test] async fn full_identity_nothing_rejected() { let cluster = Cluster::new_test_single_shard(&config()); - let mut sub = StreamSubscriber::new(&cluster, &[make_replica_identity_nothing_table()]); + let mut sub = StreamSubscriber::new(&cluster, vec![make_replica_identity_nothing_table()]); sub.connect().await.unwrap(); let oid = Oid(16390); @@ -1629,7 +1629,7 @@ async fn full_identity_nothing_rejected() { #[tokio::test] async fn full_identity_omni_no_unique_index_rejected() { let cluster = Cluster::new_test_single_shard(&config()); - let mut sub = StreamSubscriber::new(&cluster, &[make_full_identity_omni_table()]); + let mut sub = StreamSubscriber::new(&cluster, vec![make_full_identity_omni_table()]); // Enforce precondition: the table must exist but have no qualifying unique index. // A stale unique index from a prior run would make tables_missing_unique_index() return empty, @@ -1668,7 +1668,7 @@ async fn full_identity_omni_no_unique_index_rejected() { #[tokio::test] async fn full_identity_insert_sharded() { let cluster = Cluster::new_test_single_shard(&config()); - let mut sub = StreamSubscriber::new(&cluster, &[make_full_identity_sharded_table()]); + let mut sub = StreamSubscriber::new(&cluster, vec![make_full_identity_sharded_table()]); let mut verify = test_server().await; sub.connect().await.unwrap(); @@ -1694,7 +1694,7 @@ async fn full_identity_insert_sharded() { #[tokio::test] async fn full_identity_update_fast_path() { let cluster = Cluster::new_test_single_shard(&config()); - let mut sub = StreamSubscriber::new(&cluster, &[make_full_identity_sharded_table()]); + let mut sub = StreamSubscriber::new(&cluster, vec![make_full_identity_sharded_table()]); let mut verify = test_server().await; sub.connect().await.unwrap(); @@ -1749,7 +1749,7 @@ async fn full_identity_update_fast_path() { #[tokio::test] async fn full_identity_update_slow_path() { let cluster = Cluster::new_test_single_shard(&config()); - let mut sub = StreamSubscriber::new(&cluster, &[make_full_identity_sharded_table()]); + let mut sub = StreamSubscriber::new(&cluster, vec![make_full_identity_sharded_table()]); let mut verify = test_server().await; sub.connect().await.unwrap(); @@ -1809,7 +1809,7 @@ async fn full_identity_update_slow_path() { #[tokio::test] async fn full_identity_update_slow_path_realistic_old_tuple() { let cluster = Cluster::new_test_single_shard(&config()); - let mut sub = StreamSubscriber::new(&cluster, &[make_full_identity_sharded_table()]); + let mut sub = StreamSubscriber::new(&cluster, vec![make_full_identity_sharded_table()]); let mut verify = test_server().await; sub.connect().await.unwrap(); @@ -1866,7 +1866,7 @@ async fn full_identity_update_slow_path_realistic_old_tuple() { #[tokio::test] async fn full_identity_update_all_toasted_is_noop() { let cluster = Cluster::new_test_single_shard(&config()); - let mut sub = StreamSubscriber::new(&cluster, &[make_full_identity_sharded_table()]); + let mut sub = StreamSubscriber::new(&cluster, vec![make_full_identity_sharded_table()]); let mut verify = test_server().await; sub.connect().await.unwrap(); @@ -1908,7 +1908,7 @@ async fn full_identity_update_all_toasted_is_noop() { #[tokio::test] async fn full_identity_delete() { let cluster = Cluster::new_test_single_shard(&config()); - let mut sub = StreamSubscriber::new(&cluster, &[make_full_identity_sharded_table()]); + let mut sub = StreamSubscriber::new(&cluster, vec![make_full_identity_sharded_table()]); let mut verify = test_server().await; sub.connect().await.unwrap(); @@ -1946,7 +1946,7 @@ async fn full_identity_delete() { #[tokio::test] async fn full_identity_insert_omni_dedup() { let cluster = Cluster::new_test_single_shard(&config()); - let mut sub = StreamSubscriber::new(&cluster, &[make_full_identity_omni_dedup_table()]); + let mut sub = StreamSubscriber::new(&cluster, vec![make_full_identity_omni_dedup_table()]); let mut verify = test_server().await; // Ensure destination table exists with unique index before relation() runs. @@ -2005,7 +2005,7 @@ async fn full_identity_insert_omni_dedup() { #[tokio::test] async fn full_identity_update_duplicate_rows() { let cluster = Cluster::new_test_single_shard(&config()); - let mut sub = StreamSubscriber::new(&cluster, &[make_full_identity_dup_rows_table()]); + let mut sub = StreamSubscriber::new(&cluster, vec![make_full_identity_dup_rows_table()]); let mut verify = test_server().await; ensure_table(&mut verify, "public.full_dup_rows").await; @@ -2071,7 +2071,7 @@ async fn full_identity_update_duplicate_rows() { #[tokio::test] async fn full_identity_delete_duplicate_rows() { let cluster = Cluster::new_test_single_shard(&config()); - let mut sub = StreamSubscriber::new(&cluster, &[make_full_identity_dup_rows_table()]); + let mut sub = StreamSubscriber::new(&cluster, vec![make_full_identity_dup_rows_table()]); let mut verify = test_server().await; ensure_table(&mut verify, "public.full_dup_rows").await; @@ -2138,7 +2138,7 @@ async fn full_identity_delete_duplicate_rows() { #[tokio::test] async fn full_identity_update_matches_null_column() { let cluster = Cluster::new_test_single_shard(&config()); - let mut sub = StreamSubscriber::new(&cluster, &[make_full_identity_dup_rows_table()]); + let mut sub = StreamSubscriber::new(&cluster, vec![make_full_identity_dup_rows_table()]); let mut verify = test_server().await; // full_dup_rows has no NOT NULL on value — we can seed a NULL row. @@ -2199,7 +2199,7 @@ async fn full_identity_update_matches_null_column() { #[tokio::test] async fn full_identity_delete_matches_null_column() { let cluster = Cluster::new_test_single_shard(&config()); - let mut sub = StreamSubscriber::new(&cluster, &[make_full_identity_dup_rows_table()]); + let mut sub = StreamSubscriber::new(&cluster, vec![make_full_identity_dup_rows_table()]); let mut verify = test_server().await; ensure_table(&mut verify, "public.full_dup_rows").await; diff --git a/pgdog/src/backend/replication/logical/tables_sync.rs b/pgdog/src/backend/replication/logical/tables_sync.rs index 366252768..2425d3ee2 100644 --- a/pgdog/src/backend/replication/logical/tables_sync.rs +++ b/pgdog/src/backend/replication/logical/tables_sync.rs @@ -51,6 +51,10 @@ pub(crate) async fn tables_sync( return Err(Error::EmptyPublication(publication.to_owned())); } + for shard in source.shards() { + result.entry(shard.number()).or_default(); + } + Ok(result) } diff --git a/pgdog/src/backend/replication/tests.rs b/pgdog/src/backend/replication/tests.rs index efe487e5b..2edb39e37 100644 --- a/pgdog/src/backend/replication/tests.rs +++ b/pgdog/src/backend/replication/tests.rs @@ -1,14 +1,21 @@ +use std::num::NonZeroUsize; +use std::sync::Arc; use std::time::Duration; use pgdog_config::{ConfigAndUsers, Database, ShardedTableConfig, User}; use tokio_util::sync::CancellationToken; -use super::logical::{Error, data_sync::DataSync, publisher::publisher_impl::Publisher}; +use super::logical::Error; +use super::logical::orchestrator::Orchestrator; +use super::logical::publisher::Table; +use super::logical::publisher::replication_progress::ReplicationProgress; use crate::{ api::{ + copy_data::TableDataSyncTask, + replication::{ReplicationClusterStop, ReplicationClusterTask}, run_task, schema_sync::{SchemaSyncPhase, SchemaSyncTask}, - task::TaskError, + task::{TaskError, TaskWaiter}, }, backend::{ Cluster, ConnectReason, Error as BackendError, Server, ServerOptions, databases, @@ -17,7 +24,9 @@ use crate::{ server::test::test_server, }, config::{config, set}, + util::sync::WorkerPool, }; +use pgdog_stats::ReplicationDirection; async fn setup_replication_test( admin: &mut Server, @@ -54,53 +63,96 @@ async fn setup_replication_test( Ok(()) } -async fn replicate_until_caught_up( - publisher: &mut Publisher, - source: &Cluster, - destination: &Cluster, +fn start_replication( + orchestrator: &Orchestrator, +) -> (TaskWaiter<(), Error>, ReplicationClusterStop) { + let progress = ReplicationProgress::new(orchestrator.source.shards().len()); + let (cluster, stop) = ReplicationClusterTask::new( + orchestrator.clone(), + ReplicationDirection::Forward, + progress, + ); + + (run_task(cluster), stop) +} + +async fn drain_replication(task: TaskWaiter<(), Error>) -> Result<(), Box> { + tokio::time::timeout(Duration::from_secs(60), task).await??; + Ok(()) +} + +async fn wait_for_slot( + server: &mut Server, slot_name: &str, ) -> Result<(), Box> { - let mut server = source.primary(0, &Request::default()).await?; let target: Vec = server .fetch_all("SELECT pg_current_wal_lsn()::text") .await?; let target = target.first().ok_or(Error::MissingData)?; let query = format!( "SELECT 1::bigint FROM pg_replication_slots \ - WHERE slot_name = '{slot_name}_0' \ + WHERE slot_name = '{slot_name}' \ AND confirmed_flush_lsn >= '{target}'::pg_lsn" ); - let mut waiter = publisher.replicate(source, destination).await?; - let caught_up = tokio::select! { - result = waiter.wait() => { - result?; - return Err(Error::MissingData.into()); - } - result = tokio::time::timeout(Duration::from_secs(10), async { - loop { - let caught_up: Vec = server.fetch_all(&query).await?; - if caught_up == [1] { - break; - } - tokio::time::sleep(Duration::from_millis(10)).await; + + tokio::time::timeout(Duration::from_secs(20), async { + loop { + let rows: Vec = server.fetch_all(&query).await?; + if rows == [1] { + return Ok::<_, Box>(()); } - Ok::<_, Box>(()) - }) => result, - }; - waiter.stop(); - waiter.wait().await?; - caught_up??; + tokio::time::sleep(Duration::from_millis(50)).await; + } + }) + .await? +} + +async fn replicate_until_caught_up( + orchestrator: &Orchestrator, + slot_name: &str, +) -> Result<(), Box> { + let mut server = orchestrator.source.primary(0, &Request::default()).await?; + let (task, stop) = start_replication(orchestrator); + + let caught_up = wait_for_slot(&mut server, &format!("{slot_name}_0")).await; + + stop.stop(None); + let drained = drain_replication(task).await; + drained?; + caught_up?; Ok(()) } +async fn copy_table( + source: &Cluster, + dest: &Cluster, + table: &Table, + address: &Address, +) -> Result> { + let pool = Arc::new(WorkerPool::new( + vec![address.clone()], + NonZeroUsize::new(1).unwrap(), + )?); + let table = run_task( + TableDataSyncTask::builder() + .pool(pool) + .table(table.clone()) + .source(source.clone()) + .dest(dest.clone()) + .format(config().config.general.resharding_copy_format) + .source_shard(0) + .build(), + ) + .await?; + Ok(table) +} + async fn cleanup_replication_test( - publisher: &mut Publisher, admin: &mut Server, original_config: &ConfigAndUsers, test_databases: [&str; 2], ) -> Result<(), Box> { let cleanup = Box::pin(async { - publisher.cleanup().await?; for database in test_databases { let slots: Vec = admin .fetch_all(format!( @@ -146,6 +198,66 @@ async fn cleanup_replication_test( cleanup } +#[tokio::test] +async fn wait_for_replication_finishes_with_unrelated_writes() +-> Result<(), Box> { + let schema = "unrelated_writes_test"; + let destination = "unrelated_writes_test_dest"; + let original_config = config(); + let mut admin = test_server().await; + let result = async { + setup_replication_test(&mut admin, schema, destination).await?; + let source = databases::databases().schema_owner(schema)?; + let dest = databases::databases().schema_owner(destination)?; + let mut server = source.primary(0, &Request::default()).await?; + server + .execute_checked(format!( + "CREATE SCHEMA {schema}; \ + CREATE TABLE {schema}.main (id BIGINT PRIMARY KEY); \ + CREATE TABLE {schema}.noise (id BIGINT, payload TEXT); \ + CREATE PUBLICATION {schema} FOR TABLE {schema}.main" + )) + .await?; + let mut dest_server = dest.primary(0, &Request::default()).await?; + dest_server + .execute_checked(format!( + "CREATE SCHEMA {schema}; \ + CREATE TABLE {schema}.main (id BIGINT PRIMARY KEY)" + )) + .await?; + + let orchestrator = Orchestrator::new(schema, destination, schema, Some(schema.into()))?; + orchestrator + .publisher() + .await + .prepare_replication(&source, &CancellationToken::new()) + .await?; + let (task, stop) = start_replication(&orchestrator); + let result = async { + server + .execute_checked(format!( + "INSERT INTO {schema}.noise \ + SELECT g, (SELECT string_agg(md5(random()::text), '') FROM generate_series(1, 64)) \ + FROM generate_series(1, 5000) g" + )) + .await?; + wait_for_slot(&mut server, &format!("{schema}_0")).await?; + Ok::<_, Box>(()) + } + .await; + + stop.stop(None); + let drained = drain_replication(task).await; + drained?; + result?; + Ok::<_, Box>(()) + } + .await; + + cleanup_replication_test(&mut admin, &original_config, [schema, destination]).await?; + result +} + // Verify the case when the data related to fk update happened during tables // copy. We can hit the constraint violations when between the related // tables copies the updates to fk rows happened. Since we copy the table @@ -164,7 +276,6 @@ async fn test_replication_fk_conflicts_after_delete_during_copy() let child = format!("{schema}.children"); let original_config = config(); let mut admin = test_server().await; - let mut publisher = Publisher::new(&schema, schema.clone()); let result = async { setup_replication_test(&mut admin, &schema, &destination).await?; let source = databases::databases().schema_owner(&schema)?; @@ -196,26 +307,27 @@ async fn test_replication_fk_conflicts_after_delete_during_copy() let source = databases::databases().schema_owner(&schema)?; let dest = databases::databases().schema_owner(&destination)?; let cancel = CancellationToken::new(); - publisher.sync_tables(true, &source).await?; - publisher.create_slots(&source, &cancel).await?; - let tables = publisher.tables.get(&0).ok_or(Error::MissingData)?; - let child_table = tables - .iter() - .find(|table| table.table.name == "children") - .ok_or(Error::MissingData)?; - let parent_table = tables - .iter() - .find(|table| table.table.name == "parents") - .ok_or(Error::MissingData)?; - let sync = DataSync { - source: &source, - dest: &dest, - format: config().config.general.resharding_copy_format, + let orchestrator = Orchestrator::new(&schema, &destination, &schema, Some(schema.clone()))?; + let (child_table, parent_table) = { + let mut publisher = orchestrator.publisher().await; + publisher.sync_tables(true, &source).await?; + publisher.create_slots(&source, &cancel).await?; + let tables = publisher.tables.get(&0).ok_or(Error::MissingData)?; + let child_table = tables + .iter() + .find(|table| table.table.name == "children") + .ok_or(Error::MissingData)? + .clone(); + let parent_table = tables + .iter() + .find(|table| table.table.name == "parents") + .ok_or(Error::MissingData)? + .clone(); + + (child_table, parent_table) }; // copy the child table first, so it won't have updates we'll do during copy - let child_table = sync - .copy_table(child_table, source_server.addr(), &cancel, |_| {}) - .await?; + let child_table = copy_table(&source, &dest, &child_table, source_server.addr()).await?; // update the fk related data, so it would be present // only in parent table snapshot @@ -249,15 +361,15 @@ async fn test_replication_fk_conflicts_after_delete_during_copy() // and now start the copy of parent table. // that should have an updated snapshot already with the queries // executed above. - let parent_table = sync - .copy_table(parent_table, source_server.addr(), &cancel, |_| {}) - .await?; - publisher.post_data_sync([(0, vec![child_table, parent_table])].into()); + let parent_table = copy_table(&source, &dest, &parent_table, source_server.addr()).await?; + orchestrator + .publisher() + .await + .post_data_sync([(0, vec![child_table, parent_table])].into()); drop(source_server); run_task(schema_sync.clone().phase(SchemaSyncPhase::Post).build()).await?; - // run the replication and wait for all data to be copied - replicate_until_caught_up(&mut publisher, &source, &dest, &schema).await?; + replicate_until_caught_up(&orchestrator, &schema).await?; run_task(schema_sync.phase(SchemaSyncPhase::Cutover).build()).await?; Ok::<_, Box>(dest) } @@ -281,13 +393,7 @@ async fn test_replication_fk_conflicts_after_delete_during_copy() } .await; - cleanup_replication_test( - &mut publisher, - &mut admin, - &original_config, - [&schema, &destination], - ) - .await?; + cleanup_replication_test(&mut admin, &original_config, [&schema, &destination]).await?; let (parents, children, parent_ids) = validation?; assert_eq!(parents, [2, 3]); assert_eq!(children, [2, 3, 4]); @@ -305,7 +411,6 @@ async fn test_replication_fk_constraints_after_copy_child_before_parent() let destination = "fk_copy_test_dest"; let original_config = config(); let mut admin = test_server().await; - let mut publisher = Publisher::new(schema, schema.into()); let result = async { setup_replication_test(&mut admin, schema, destination).await?; let source = databases::databases().schema_owner(schema)?; @@ -335,33 +440,34 @@ async fn test_replication_fk_constraints_after_copy_child_before_parent() let source = databases::databases().schema_owner(schema)?; let dest = databases::databases().schema_owner(destination)?; let cancel = CancellationToken::new(); - publisher.sync_tables(true, &source).await?; - publisher.create_slots(&source, &cancel).await?; - let tables = publisher.tables.get(&0).ok_or(Error::MissingData)?; - let child = tables - .iter() - .find(|table| table.table.name == "children") - .ok_or(Error::MissingData)?; - let parent = tables - .iter() - .find(|table| table.table.name == "parents") - .ok_or(Error::MissingData)?; - let sync = DataSync { - source: &source, - dest: &dest, - format: config().config.general.resharding_copy_format, + let orchestrator = Orchestrator::new(schema, destination, schema, Some(schema.into()))?; + let (child, parent) = { + let mut publisher = orchestrator.publisher().await; + publisher.sync_tables(true, &source).await?; + publisher.create_slots(&source, &cancel).await?; + let tables = publisher.tables.get(&0).ok_or(Error::MissingData)?; + let child = tables + .iter() + .find(|table| table.table.name == "children") + .ok_or(Error::MissingData)? + .clone(); + let parent = tables + .iter() + .find(|table| table.table.name == "parents") + .ok_or(Error::MissingData)? + .clone(); + + (child, parent) }; - // copy the child table first, while the parent data is not yet present - let child = sync - .copy_table(child, server.addr(), &cancel, |_| {}) - .await?; + let child = copy_table(&source, &dest, &child, server.addr()).await?; // and now copy the parent table - let parent = sync - .copy_table(parent, server.addr(), &cancel, |_| {}) - .await?; - publisher.post_data_sync([(0, vec![child, parent])].into()); + let parent = copy_table(&source, &dest, &parent, server.addr()).await?; + orchestrator + .publisher() + .await + .post_data_sync([(0, vec![child, parent])].into()); // add rows after copy so replication must deliver them server @@ -375,7 +481,7 @@ async fn test_replication_fk_constraints_after_copy_child_before_parent() drop(server); run_task(schema_sync.clone().phase(SchemaSyncPhase::Post).build()).await?; - replicate_until_caught_up(&mut publisher, &source, &dest, schema).await?; + replicate_until_caught_up(&orchestrator, schema).await?; run_task(schema_sync.phase(SchemaSyncPhase::Cutover).build()).await?; Ok::<_, Box>(dest) } @@ -400,13 +506,7 @@ async fn test_replication_fk_constraints_after_copy_child_before_parent() } .await; - cleanup_replication_test( - &mut publisher, - &mut admin, - &original_config, - [schema, destination], - ) - .await?; + cleanup_replication_test(&mut admin, &original_config, [schema, destination]).await?; let (parents, children) = validation?; assert_eq!(parents, ["1:1", "2:1"]); assert_eq!(children, ["1:1:1", "2:1:2"]); @@ -422,7 +522,6 @@ async fn test_replication_copy_custom_parent_trigger() -> Result<(), Box Result<(), Box Result<(), Box>(dest) } @@ -519,13 +620,7 @@ async fn test_replication_copy_custom_parent_trigger() -> Result<(), Box