pgdog.git / summary / log / commit / refs
commit f29dacf80645c4bfa05dcc62b85647540085548d
Author: Kiryl Mialeshka <8974488+meskill@users.noreply.github.com>
Commit: GitHub <noreply@github.com>
Date: Tue Sep 22 10:23:59 2026 +0000
refactor(resharding): migrate replication to api/task (#1602)
Migrate the replication process to task:
- separate the logic into nested tasks: ReplicationTask ->
ReplicationClusterTask -> ReplicationShardTask with separate state and
tracking
- extend the definitions and statuses used to display progress and send
to EE
- rip apart most of the orchestrator and publisher
docs/RESHARDING.md | 108 +--
.../rust/tests/integration/admin/resharding/mod.rs | 36 +-
.../integration/admin/resharding/replication.rs | 300 ++++++--
.../admin/resharding/replication_slots.rs | 95 +++
pgdog-stats/src/resharding.rs | 60 +-
pgdog-stats/src/task.rs | 403 ++--------
pgdog-stats/src/task/copy_data.rs | 96 +++
pgdog-stats/src/task/replication.rs | 191 +++++
pgdog-stats/src/task/reshard.rs | 42 ++
pgdog-stats/src/task/schema_sync.rs | 132 ++++
pgdog/src/admin/replicate.rs | 3 +-
pgdog/src/admin/show_replication_slots.rs | 8 +-
pgdog/src/api/copy_data.rs | 36 +-
pgdog/src/api/replication.rs | 822 ++++++++++++++++++---
pgdog/src/api/resharding.rs | 10 +-
pgdog/src/api/task.rs | 27 +-
pgdog/src/backend/databases.rs | 39 +
pgdog/src/backend/maintenance_mode.rs | 5 -
pgdog/src/backend/replication/logical/data_sync.rs | 3 +
pgdog/src/backend/replication/logical/ee/mod.rs | 13 -
pgdog/src/backend/replication/logical/error.rs | 15 +
pgdog/src/backend/replication/logical/mod.rs | 6 +-
.../backend/replication/logical/orchestrator.rs | 772 +------------------
.../logical/publisher/cutover_policy.rs | 429 +++++++++++
.../backend/replication/logical/publisher/mod.rs | 4 +-
.../replication/logical/publisher/progress.rs | 79 --
.../logical/publisher/publisher_impl.rs | 291 +-------
.../logical/publisher/replication_progress.rs | 226 ++++++
.../logical/publisher/replication_stream.rs | 560 ++++++++++++++
.../backend/replication/logical/publisher/slot.rs | 50 +-
pgdog/src/backend/replication/logical/status.rs | 38 +-
.../replication/logical/subscriber/pipeline.rs | 14 +-
.../replication/logical/subscriber/stream.rs | 116 +--
.../replication/logical/subscriber/tests.rs | 34 +-
.../src/backend/replication/logical/tables_sync.rs | 4 +
pgdog/src/backend/replication/tests.rs | 400 ++++++----
36 files changed, 3494 insertions(+), 1973 deletions(-)
diff --git a/docs/RESHARDING.md b/docs/RESHARDING.md
index cd819326..a14b9e03 100644
--- a/docs/RESHARDING.md
+++ b/docs/RESHARDING.md
@@ -16,8 +16,9 @@ RESHARD <source> <destination> <publication>;
```
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<Mutex<Publisher>>` — manages replication slots, table list, and lag tracking
- `replication_slot: String` — auto-generated as `__pgdog_repl_<random19>` 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<br>pre-data to dest<br>reload schema cache"]
C["3. data_sync<br>ParallelSyncManager<br>binary COPY"]
D["4. schema_sync_post<br>secondary indexes"]
- E["5. replicate().cutover()<br>WAL drain<br>traffic swap"]
+ E["5. ReplicationTask<br>WAL drain<br>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 ec433034..d8f42d7a 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<Postgres>) {
.await;
}
+async fn test_slot_names(pool: &Pool<Postgres>) -> Vec<String> {
+ 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<Postgres>) {
+ 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<Postgres>, direct: &Pool<Postgres>) {
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<Postgres>, direct: &Pool<Postgres>) {
drop_test_schemas(direct).await;
}
+async fn with_cleanup(
+ admin: &Pool<Postgres>,
+ direct: &Pool<Postgres>,
+ body: impl Future<Output = ()>,
+) {
+ 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<Postgres>, 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 86b8dbd0..1e398c13 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<Postgres>, direct: &Pool<Postgres>) -> 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::<String, _>("task_id").parse().unwrap();
+pub(super) async fn prepare_replication(admin: &Pool<Postgres>, direct: &Pool<Postgres>) {
+ 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<Postgres>, 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<Postgres>,
+ 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::<String, _>("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::<String, _>("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::<String, _>("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<String> = 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<Postgres>, task_id: i64) {
+ let accepted = timeout(Duration::from_secs(10), async {
loop {
if let Ok(row) = admin.fetch_one("CUTOVER").await
&& row.get::<String, _>("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<Postgres>, 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 00000000..b1d4145d
--- /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<Postgres>) -> Option<PgRow> {
+ 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::<String, _>("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::<String, _>("database_name"), "pgdog");
+ assert!(!row.get::<bool, _>("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::<String, _>("lsn")
+ .parse()
+ .expect("displayed WAL position must be valid");
+ (lsn.lsn > before.lsn).then_some(row)
+ })
+ .await;
+ assert!(row.get::<Option<String>, _>("last_transaction").is_some());
+ assert!(row.get::<Option<i64>, _>("last_transaction_ms").is_some());
+ assert!(row.get::<Option<i64>, _>("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 aa5e0c38..da8494b1 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<SystemTime>,
+ pub task_id: Option<TaskId>,
}
/// 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 8c061d80..b9f87b44 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<Vec<u64>>,
-}
-
-impl From<CopyDataStage> 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<String>) -> 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<Vec<SchemaSyncStatement>>,
- },
- /// 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<SchemaStatementFailure>,
-}
-
-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<i64>,
-}
-
-/// 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<u64>,
- pub estimated_bytes: Option<u64>,
- pub rows_per_sec: Option<u64>,
- pub bytes_per_sec: Option<u64>,
- pub last_error: Option<String>,
-}
-
-/// 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 00000000..f5b738a0
--- /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<Vec<u64>>,
+}
+
+impl From<CopyDataStage> 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<u64>,
+ pub estimated_bytes: Option<u64>,
+ pub rows_per_sec: Option<u64>,
+ pub bytes_per_sec: Option<u64>,
+ pub last_error: Option<String>,
+}
diff --git a/pgdog-stats/src/task/replication.rs b/pgdog-stats/src/task/replication.rs
new file mode 100644
index 00000000..3376701e
--- /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<u64>,
+ pub last_transaction_ms: Option<u64>,
+ pub rows: u64,
+ pub bytes: u64,
+ pub rows_per_sec: Option<u64>,
+ pub bytes_per_sec: Option<u64>,
+}
+
+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<u64>,
+ pub missed_rows: MissedRows,
+ pub rows: u64,
+ pub bytes: u64,
+ pub rows_per_sec: Option<u64>,
+ pub bytes_per_sec: Option<u64>,
+}
+
+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 00000000..c83b2496
--- /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 00000000..6ca8c292
--- /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<String>) -> 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<Vec<SchemaSyncStatement>>,
+ },
+ /// 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<SchemaStatementFailure>,
+}
+
+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 a2749451..66060397 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 7c1df83a..368a5489 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 7cb6eb6a..75246c01 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<WorkerPool<Address>>,
- table: Table,
- source: Cluster,
- dest: Cluster,
- format: CopyFormat,
- source_shard: usize,
+#[derive(Debug, bon::Builder)]
+pub(crate) struct TableDataSyncTask {
+ pub(crate) pool: Arc<WorkerPool<Address>>,
+ 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 e8b28c6d..965c7d04 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<Mutex<HashMap<TaskId, CancellationToken>>> =
- 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<TaskDefinition> {
+ ReplicationDefinition {
+ databases: self.orchestrator.databases(),
+ auto_cutover: self.auto_cutover,
+ }
+ }
+
+ async fn run(self, ctx: TaskContext<Self>) -> 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<TaskId>) -> 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<ReplicationTask>,
+ orchestrator: Orchestrator,
+ direction: ReplicationDirection,
+ maintenance: MaintenanceMode,
+}
+
+impl<'a> Replication<'a> {
+ fn new(ctx: &'a TaskContext<ReplicationTask>, 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<ReplicationCutoverReason, Error> {
+ 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<Option<ReplicationCutoverReason>>,
+}
+
+impl ReplicationClusterStop {
+ pub(crate) fn stop(self, cutover_reason: Option<ReplicationCutoverReason>) {
+ 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<Option<ReplicationCutoverReason>>,
+}
+
+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<TaskDefinition> {
- 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<Self>) -> Result<(), Error> {
- let token = ctx.cancellation_token();
+ async fn run(self, ctx: TaskContext<Self>) -> 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<Self>,
+ 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<Table>,
+ 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<TaskDefinition> {
+ 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<Self>) -> 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<Self>,
- 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<JoinHandle<Result<(), Error>>>;
+
+/// 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<ReplicationSlot>,
+}
+
+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<TaskId>) -> 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<DashMap<TaskId, CancellationToken>> = 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<tokio::sync::Mutex<()>> =
+ 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 f30bfed0..a14ac1c3 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 a53feee8..a4a8070a 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<T: Task> TaskContext<T> {
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<T: Task> TaskContext<T> {
}
}
+ 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 200a9c9d..5d7800e4 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 e88e0d4f..639bf576 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<MaintenanceState>,
diff --git a/pgdog/src/backend/replication/logical/data_sync.rs b/pgdog/src/backend/replication/logical/data_sync.rs
index 4be1ccf3..97b8e91f 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 0efadd04..322accaf 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 6e0124e6..ec953bdf 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 b3eb47ea..b5a73242 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 ab556a00..a5b2f54d 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<ReplicationWaiter, Error> {
- 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<u64> {
- 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<ConfigAndUsers>,
-}
-
-#[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<Duration>,
- 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<ReplicationTask>,
- 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<ConfigAndUsers>) -> 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 00000000..62ccf3aa
--- /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<Duration>,
+ 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<CutoverReason, Error> {
+ 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 858c8433..d7141adb 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 2255cd2f..00000000
--- 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<Inner>,
-}
-
-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 00e021a6..ab96fefa 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<usize, Vec<Table>>,
/// Replication slots.
- slots: HashMap<usize, ReplicationSlot>,
- /// Replication lag.
- replication_lag: Arc<Mutex<HashMap<usize, i64>>>,
- /// Last transaction.
- last_transaction: Arc<Mutex<Option<Instant>>>,
- /// Slot name.
+ pub(crate) slots: HashMap<usize, ReplicationSlot>,
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<Vec<Table>, Error> {
+ self.tables
+ .remove(&shard)
+ .ok_or(Error::NoReplicationTables(shard))
+ }
+
+ pub(crate) fn pop_slot(&mut self, shard: usize) -> Result<ReplicationSlot, Error> {
+ 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<Waiter, Error> {
- // 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<bool, Error> = 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<usize, i64> {
- self.replication_lag.lock().clone()
- }
-
- /// Get how long ago last transaction was committed.
- pub(crate) fn last_transaction(&self) -> Option<Duration> {
- (*self.last_transaction.lock()).map(|last| last.elapsed())
+ Ok(())
}
pub(crate) fn post_data_sync(&mut self, tables: HashMap<usize, Vec<Table>>) {
@@ -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<Instant>) {
- *self.last_transaction.lock() = instant;
- }
-}
-
-#[derive(Debug)]
-pub(crate) struct Waiter {
- streams: Vec<JoinHandle<Result<(), Error>>>,
- 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 00000000..a1a20764
--- /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<i64>,
+ pub(crate) last_transaction: Option<Instant>,
+ pub(crate) applied_lsn: Option<Lsn>,
+ pub(crate) missed_rows: MissedRows,
+ pub(crate) bytes_sharded: usize,
+ pub(crate) rows_sharded: usize,
+ pub(crate) origin_lsn: Lsn,
+ pub(crate) started: Option<Instant>,
+}
+
+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<u64> {
+ self.started
+ .and_then(|started| average_rate(count, started.into_std()))
+ }
+
+ fn lag_bytes(&self) -> Option<u64> {
+ 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<ReplicationShardProgress>]>,
+}
+
+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<u64> = None;
+ let mut every_shard_reported = true;
+ let mut last_transaction: Option<Instant> = 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<u64>, rate: Option<u64>) -> Option<u64> {
+ 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<ReplicationShardProgress>]>,
+ 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 00000000..7741db5a
--- /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<Table>,
+ 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<bool, Error> = 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<dyn StdError>>;
+
+ struct Fixture {
+ source_table: String,
+ destination_table: String,
+ publication: String,
+ slot_base: String,
+ server: Server,
+ source: Cluster,
+ replication: Arc<ReplicationStream>,
+ stop: CancellationToken,
+ worker: Option<JoinHandle<Result<(), Error>>>,
+ }
+
+ 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<String> = self.server.fetch_all(query.clone()).await?;
+ if ready(&rows, self.replication.progress()) {
+ return Ok::<(), Box<dyn StdError>>(());
+ }
+ 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::<dyn StdError>::from)
+ .and_then(|result| result.map_err(Box::<dyn StdError>::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<String> = 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<String> = fixture.server.fetch_all(slot_query).await?;
+ if !after.is_empty() {
+ return Err("replication slot remains after shutdown".into());
+ }
+ let rows: Vec<String> = 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<String> = 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 1f7936a3..f90b4552 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<i64, Error> {
+ let lag = self.query_replication_lag().await;
+
+ if lag.is_err() {
+ self.server_meta = None;
+ }
+
+ lag
+ }
+
+ async fn query_replication_lag(&mut self) -> Result<i64, Error> {
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 f9ba0ecd..b9a82370 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<ReplicationSlots> = 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<DashMap<String, ReplicationSlot>>,
+ slots: Arc<DashMap<String, pgdog_stats::ReplicationSlot>>,
}
impl ReplicationSlots {
@@ -101,7 +107,7 @@ impl ReplicationSlots {
}
impl Deref for ReplicationSlots {
- type Target = Arc<DashMap<String, ReplicationSlot>>;
+ type Target = Arc<DashMap<String, pgdog_stats::ReplicationSlot>>;
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 2a2de861..e176b123 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 90338206..ade69c02 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<Table>) -> 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<Option<StatusUpdate>, 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 20893179..b962e58a 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<Table>) -> 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 36625276..2425d3ee 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 efe487e5..2edb39e3 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<dyn std::error::Error>> {
+ tokio::time::timeout(Duration::from_secs(60), task).await??;
+ Ok(())
+}
+
+async fn wait_for_slot(
+ server: &mut Server,
slot_name: &str,
) -> Result<(), Box<dyn std::error::Error>> {
- let mut server = source.primary(0, &Request::default()).await?;
let target: Vec<String> = 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<i64> = 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<i64> = server.fetch_all(&query).await?;
+ if rows == [1] {
+ return Ok::<_, Box<dyn std::error::Error>>(());
}
- Ok::<_, Box<dyn std::error::Error>>(())
- }) => 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<dyn std::error::Error>> {
+ 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<Table, Box<dyn std::error::Error>> {
+ 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<dyn std::error::Error>> {
let cleanup = Box::pin(async {
- publisher.cleanup().await?;
for database in test_databases {
let slots: Vec<String> = 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<dyn std::error::Error>> {
+ 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<dyn std::error::Error>>(())
+ }
+ .await;
+
+ stop.stop(None);
+ let drained = drain_replication(task).await;
+ drained?;
+ result?;
+ Ok::<_, Box<dyn std::error::Error>>(())
+ }
+ .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<dyn std::error::Error>>(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<dyn std::error::Error>>(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<dyn std
let destination = "trigger_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)?;
@@ -461,30 +560,32 @@ async fn test_replication_copy_custom_parent_trigger() -> Result<(), Box<dyn std
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 first, while its parent is still missing
- let child = sync
- .copy_table(child, server.addr(), &cancel, |_| {})
- .await?;
- let parent = sync
- .copy_table(parent, server.addr(), &cancel, |_| {})
- .await?;
- publisher.post_data_sync([(0, vec![child, parent])].into());
+ let child = copy_table(&source, &dest, &child, server.addr()).await?;
+ 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
@@ -497,7 +598,7 @@ async fn test_replication_copy_custom_parent_trigger() -> Result<(), Box<dyn std
.await?;
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<dyn std::error::Error>>(dest)
}
@@ -519,13 +620,7 @@ async fn test_replication_copy_custom_parent_trigger() -> Result<(), Box<dyn std
}
.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, 2]);
assert_eq!(children, [1, 2]);
@@ -543,7 +638,6 @@ async fn test_replication_fk_inconsistent_check_on_cutover()
let destination = "fk_post_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)?;
@@ -571,36 +665,38 @@ async fn test_replication_fk_inconsistent_check_on_cutover()
run_task(schema_sync.clone().phase(SchemaSyncPhase::Pre).build()).await?;
let source = databases::databases().schema_owner(schema)?;
let cancel = CancellationToken::new();
- publisher.sync_tables(true, &source).await?;
- publisher.create_slots(&source, &cancel).await?;
+ let orchestrator = Orchestrator::new(schema, destination, schema, Some(schema.into()))?;
let dest = databases::databases().schema_owner(destination)?;
- 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 (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)
};
- let child = sync
- .copy_table(child, server.addr(), &cancel, |_| {})
- .await?;
+ let child = copy_table(&source, &dest, &child, server.addr()).await?;
let mut destination_server = dest.primary(0, &Request::default()).await?;
// leave an orphan that replication cannot repair
destination_server
.execute_checked(format!("INSERT INTO {schema}.children VALUES (42, 1, 999)"))
.await?;
drop(destination_server);
- 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 valid rows that must arrive through replication
server
@@ -616,7 +712,7 @@ async fn test_replication_fk_inconsistent_check_on_cutover()
run_task(schema_sync.clone().phase(SchemaSyncPhase::Post).build()).await?;
// wait for the valid source rows to arrive without repairing the orphan
- replicate_until_caught_up(&mut publisher, &source, &dest, schema).await?;
+ replicate_until_caught_up(&orchestrator, schema).await?;
// cutover should reject the orphan left on the destination
let cutover = run_task(schema_sync.phase(SchemaSyncPhase::Cutover).build()).await;
@@ -639,13 +735,7 @@ async fn test_replication_fk_inconsistent_check_on_cutover()
}
.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, cutover) = validation?;
assert_eq!(parents, [1, 2]);
assert_eq!(children, [1, 2, 999]);