refactor(resharding): extract tables sync to separate function

docs/RESHARDING.md

@@ -143,6 +143,17 @@ Two behaviours are specific to the resharding context:
   COPY. Messages at or below that LSN are skipped; the row is already on the destination.
 - **Omnisharded tables** (`statements.omni = true`): upsert is broadcast to all shards
   simultaneously rather than routed to a single shard.
+- **Table ownership** ([`tables_sync()`](../pgdog/src/backend/replication/logical/tables_sync.rs)):
+  a table that is *sharded on the source* is copied and replayed from every source shard.
+  A table that is *omnisharded on the source* is copied and replayed from one source shard
+  only, chosen by publication order, because every source shard holds the same rows.
+- **Destination row contention**: a table that is sharded on the source and omnisharded on
+  the destination is replayed by every subscriber, and every subscriber writes to every
+  destination shard. Two subscribers therefore write the same destination row whenever one
+  key reaches two source shards, for example after a sharding-key update. Two subscribers
+  can then lock the same rows on two destinations in opposite order. No Postgres instance
+  sees the whole cycle, so no instance reports a deadlock. Set `lock_timeout` on the
+  destination user so a blocked apply is cancelled and retried by `Publisher::replicate()`.
 ---
 
 ### Cutover phases

integration/resharding/dev.sh

@@ -97,6 +97,8 @@ replace_copy_with_replicate projects name
 replace_copy_with_replicate tasks title
 replace_copy_with_replicate task_comments body
 replace_copy_with_replicate settings name
+replace_copy_with_replicate sharded_to_omni name
+replace_copy_with_replicate omni_to_sharded name
 
 # REPLICATION SENTINEL — must be the last DML issued against the source.
 # pgbench uses random(1, 1_000_000_000), so id=0 is reserved for this purpose.
@@ -150,6 +152,8 @@ wait_for_no_copy_rows projects name
 wait_for_no_copy_rows tasks title
 wait_for_no_copy_rows task_comments body
 wait_for_no_copy_rows settings name
+wait_for_no_copy_rows sharded_to_omni name
+wait_for_no_copy_rows omni_to_sharded name
 
 
 # pg_count PORT TABLE — row count via a direct postgres connection (bypasses pgdog).
@@ -193,11 +197,53 @@ check_omni_each_shard() {
     echo "OK omni ${table}: ${source_count} rows on each shard"
 }
 
+check_sharded_source_to_omni_destination() {
+    local table="$1"
+    local source_count dest0_count dest1_count
+
+    source_count=$(psql -d source -tAc "SELECT COUNT(*) FROM ${table}")
+    dest0_count=$(pg_count 15434 "${table}")
+    dest1_count=$(pg_count 15435 "${table}")
+
+    if [ "${source_count}" -ne "${dest0_count}" ] || [ "${source_count}" -ne "${dest1_count}" ]; then
+        echo "MISMATCH sharded->omni ${table}: source=${source_count} dest-0(15434)=${dest0_count} dest-1(15435)=${dest1_count} (expected ${source_count} on each)"
+        exit 1
+    fi
+
+    echo "OK sharded->omni ${table}: ${source_count} rows on each destination shard"
+}
+
+check_omni_source_to_sharded_destination() {
+    local table="$1"
+    local source0_count source1_count dest0_count dest1_count dest_total
+
+    source0_count=$(pg_count 15432 "${table}")
+    source1_count=$(pg_count 15433 "${table}")
+    if [ "${source0_count}" -ne "${source1_count}" ]; then
+        echo "MISMATCH omni->sharded ${table}: source shards disagree source-0(15432)=${source0_count} source-1(15433)=${source1_count}"
+        exit 1
+    fi
+
+    dest0_count=$(pg_count 15434 "${table}")
+    dest1_count=$(pg_count 15435 "${table}")
+    dest_total=$((dest0_count + dest1_count))
+
+    if [ "${source0_count}" -ne "${dest_total}" ]; then
+        echo "MISMATCH omni->sharded ${table}: source=${source0_count} dest total=${dest_total} (dest-0=${dest0_count} dest-1=${dest1_count})"
+        exit 1
+    fi
+
+    echo "OK omni->sharded ${table}: ${source0_count} rows split ${dest0_count}/${dest1_count}"
+}
+
+
 check_row_count_matches tenants
 check_row_count_matches accounts
 check_row_count_matches projects
 check_row_count_matches tasks
 check_row_count_matches task_comments
 check_omni_each_shard settings
+check_sharded_source_to_omni_destination sharded_to_omni
+check_omni_source_to_sharded_destination omni_to_sharded
 
 cleanup

integration/resharding/pgbench.sql

@@ -56,6 +56,24 @@ VALUES (
 )
 ON CONFLICT (id) DO NOTHING;
 
+INSERT INTO sharded_to_omni (id, org_id, name, value)
+VALUES (
+    :id_seed,
+    :id_seed,
+    'sharded-to-omni-' || :id_seed || '-copy',
+    'value-' || :id_seed || '-copy'
+)
+ON CONFLICT (id) DO NOTHING;
+
+INSERT INTO omni_to_sharded (id, org_id, name, value)
+VALUES (
+    :id_seed,
+    :id_seed,
+    'omni-to-sharded-' || :id_seed || '-copy',
+    'value-' || :id_seed || '-copy'
+)
+ON CONFLICT (id) DO NOTHING;
+
 UPDATE tenants
 SET name = 'tenant-' || :id_seed || '-replicate'
 WHERE id = :id_seed;
@@ -79,3 +97,11 @@ WHERE id = :id_seed;
 UPDATE settings
 SET name = 'setting-' || :id_seed || '-replicate'
 WHERE id = :id_seed;
+
+UPDATE sharded_to_omni
+SET name = 'sharded-to-omni-' || :id_seed || '-replicate'
+WHERE id = :id_seed;
+
+UPDATE omni_to_sharded
+SET name = 'omni-to-sharded-' || :id_seed || '-replicate'
+WHERE id = :id_seed;

integration/resharding/pgdog.toml

@@ -51,6 +51,18 @@ database = "destination"
 column = "tenant_id"
 data_type = "bigint"
 
+[[sharded_tables]]
+database = "source"
+name = "sharded_to_omni"
+column = "org_id"
+data_type = "bigint"
+
+[[sharded_tables]]
+database = "destination"
+name = "omni_to_sharded"
+column = "org_id"
+data_type = "bigint"
+
 [admin]
 password = "pgdog"
 user = "pgdog"

integration/resharding/schema.sql

@@ -62,6 +62,22 @@ CREATE TABLE settings (
     created_at timestamptz NOT NULL DEFAULT NOW()
 );
 
+CREATE TABLE sharded_to_omni (
+    id BIGINT PRIMARY KEY,
+    org_id BIGINT NOT NULL,
+    name VARCHAR NOT NULL,
+    value VARCHAR NOT NULL,
+    created_at timestamptz NOT NULL DEFAULT NOW()
+);
+
+CREATE TABLE omni_to_sharded (
+    id BIGINT PRIMARY KEY,
+    org_id BIGINT NOT NULL,
+    name VARCHAR NOT NULL,
+    value VARCHAR NOT NULL,
+    created_at timestamptz NOT NULL DEFAULT NOW()
+);
+
 CREATE INDEX idx_accounts_tenant_id ON accounts (tenant_id);
 CREATE INDEX idx_projects_tenant_id ON projects (tenant_id);
 CREATE INDEX idx_projects_owner ON projects (tenant_id, owner_account_id);

integration/resharding/users.toml

@@ -9,4 +9,3 @@ database = "destination"
 name = "pgdog"
 password = "pgdog"
 schema_admin = true
-lock_timeout = 100

pgdog/src/backend/connect_reason.rs

@@ -1,31 +1,13 @@
-use std::fmt::Display;
-
-#[derive(Debug, Clone, Copy, Default, PartialEq)]
+#[derive(Debug, Display, Clone, Copy, Default, PartialEq)]
+#[display(rename_all = "snake_case")]
 pub(crate) enum ConnectReason {
     LsnCheck,
     BelowMin,
     ClientWaiting,
-    Replication,
+    Resharding,
     PubSub,
     Probe,
     Healthcheck,
     #[default]
     Other,
 }
-
-impl Display for ConnectReason {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        let reason = match self {
-            Self::LsnCheck => "lsn check",
-            Self::BelowMin => "min",
-            Self::ClientWaiting => "client",
-            Self::Replication => "replication",
-            Self::PubSub => "pub/sub",
-            Self::Probe => "probe",
-            Self::Healthcheck => "healthcheck",
-            Self::Other => "other",
-        };
-
-        write!(f, "{}", reason)
-    }
-}

pgdog/src/backend/pool/cluster.rs

@@ -10,7 +10,6 @@ use std::{sync::Arc, time::Duration};
 
 use crate::backend::schema::SchemaCache;
 use crate::backend::server::ServerRequest;
-use crate::frontend::router::sharding::ShardedTable;
 use crate::{
     backend::{
         Schema, ShardedTables, databases::User as DatabaseUser, replication::ShardedSchemas,
@@ -509,8 +508,8 @@ impl Cluster {
     }
 
     // Get sharded tables if any.
-    pub(crate) fn sharded_tables(&self) -> &[ShardedTable] {
-        self.sharded_tables.tables()
+    pub(crate) fn sharded_tables(&self) -> &ShardedTables {
+        &self.sharded_tables
     }
 
     /// Get query rewrite config.
@@ -916,35 +915,6 @@ mod test {
             cluster
         }
 
-        /// Two shards targeting different databases on the same server.
-        /// Gives separate lock namespaces without needing two Postgres instances.
-        pub(crate) fn new_test_two_databases(config: &ConfigAndUsers) -> Cluster {
-            let mut cluster = Self::new_test(config);
-            let shard1 = cluster.shards.last_mut().unwrap();
-            *shard1 = Shard::new(ShardConfig {
-                number: 1,
-                primary: Some(&PoolConfig {
-                    address: Address {
-                        database_name: "pgdog1".into(),
-                        ..Address::new_test()
-                    },
-                    config: Config::default(),
-                }),
-                replicas: &[PoolConfig {
-                    address: Address {
-                        database_name: "pgdog1".into(),
-                        configured_role: Role::Replica,
-                        ..Address::new_test()
-                    },
-                    config: Config::default(),
-                }],
-                identifier: cluster.identifier.clone(),
-                lsn_check_interval: Duration::MAX,
-                ..Default::default()
-            });
-            cluster
-        }
-
         pub(crate) fn new_test_single_primary(config: &ConfigAndUsers) -> Cluster {
             let identifier = Arc::new(DatabaseUser {
                 user: "pgdog".into(),

pgdog/src/backend/pool/monitor.rs

@@ -434,7 +434,7 @@ impl Monitor {
         let connect_timeout = pool.config().connect_timeout;
         let connect_attempts = pool.config().connect_attempts;
         let connect_attempt_delay = pool.config().connect_attempt_delay;
-        let options = pool.server_options();
+        let options = pool.server_options(reason);
 
         let mut error = Error::ServerError;
         let now = Instant::now();

pgdog/src/backend/pool/pool_impl.rs

@@ -431,7 +431,7 @@ impl Pool {
     }
 
     /// Get startup parameters for new server connections.
-    pub(super) fn server_options(&self) -> ServerOptions {
+    pub(super) fn server_options(&self, reason: ConnectReason) -> ServerOptions {
         let mut params = vec![
             Parameter {
                 name: "application_name".into(),
@@ -445,14 +445,22 @@ impl Pool {
 
         let config = self.inner.config;
 
+        let lock_timeout = config
+            .lock_timeout
+            // Enforce some lock_timeout during resharding to prevent possible deadlocks.
+            // This should be mostly avoided by pgdog, but in case some invariants are not met,
+            // the resharding could deadlock and with timeout we'll probably retry the update
+            // and either succeed or fail explicitly.
+            .or(matches!(reason, ConnectReason::Resharding).then_some(Duration::from_secs(5)));
+
         if let Some(statement_timeout) = config.statement_timeout {
             params.push(Parameter {
                 name: "statement_timeout".into(),
                 value: statement_timeout.as_millis().to_string().into(),
             });
         }
 
-        if let Some(lock_timeout) = config.lock_timeout {
+        if let Some(lock_timeout) = lock_timeout {
             params.push(Parameter {
                 name: "lock_timeout".into(),
                 value: lock_timeout.as_millis().to_string().into(),

pgdog/src/backend/replication/logical/error.rs

@@ -118,8 +118,8 @@ pub(crate) enum Error {
     #[error("replication timeout")]
     ReplicationTimeout,
 
-    #[error("shard {0} has no replication tables")]
-    NoReplicationTables(usize),
+    #[error("publication \"{0}\" has no tables")]
+    EmptyPublication(String),
 
     #[error("shard {0} has no replication slot")]
     NoReplicationSlot(usize),

pgdog/src/backend/replication/logical/mod.rs

@@ -6,6 +6,7 @@ pub(crate) mod publisher;
 pub(crate) mod schema_sync;
 pub(crate) mod status;
 pub(crate) mod subscriber;
+pub(crate) mod tables_sync;
 
 pub(crate) use copy_statement::CopyStatement;
 pub(crate) use error::*;

pgdog/src/backend/replication/logical/publisher/mod.rs

@@ -21,22 +21,19 @@ pub(crate) mod test {
 
     pub(crate) struct PublicationTest {
         pub(crate) server: Server,
+        pub(crate) publication: String,
+        pub(crate) tables: Vec<String>,
     }
 
     impl PublicationTest {
         pub(crate) async fn cleanup(&mut self) {
-            self.server
-                .execute("DROP PUBLICATION IF EXISTS publication_test")
-                .await
-                .unwrap();
-            self.server
-                .execute("DROP TABLE IF EXISTS publication_test_two")
-                .await
-                .unwrap();
-            self.server
-                .execute("DROP TABLE IF EXISTS publication_test_one")
-                .await
-                .unwrap();
+            let drop_publication = format!("DROP PUBLICATION IF EXISTS {}", self.publication);
+            self.server.execute(drop_publication).await.unwrap();
+
+            for table in self.tables.iter().rev() {
+                let drop_table = format!("DROP TABLE IF EXISTS {}", table);
+                self.server.execute(drop_table).await.unwrap();
+            }
         }
     }
 
@@ -69,6 +66,38 @@ pub(crate) mod test {
             .unwrap();
         server.execute("CREATE PUBLICATION publication_test FOR TABLE publication_test_one, publication_test_two").await.unwrap();
 
-        PublicationTest { server }
+        PublicationTest {
+            server,
+            publication: "publication_test".into(),
+            tables: vec!["publication_test_one".into(), "publication_test_two".into()],
+        }
+    }
+
+    pub(crate) async fn setup_publication_tables(
+        publication: &str,
+        tables: &[&str],
+    ) -> PublicationTest {
+        let mut test = PublicationTest {
+            server: test_replication_server().await,
+            publication: publication.to_owned(),
+            tables: tables.iter().map(|table| table.to_string()).collect(),
+        };
+
+        test.cleanup().await;
+
+        for table in &test.tables {
+            let create_table =
+                format!("CREATE TABLE {} (id BIGINT PRIMARY KEY, value TEXT)", table);
+            test.server.execute(create_table).await.unwrap();
+        }
+
+        let create_publication = format!(
+            "CREATE PUBLICATION {} FOR TABLE {}",
+            test.publication,
+            test.tables.join(", ")
+        );
+        test.server.execute(create_publication).await.unwrap();
+
+        test
     }
 }