fix: table LSN not preserved during copy -> replication handover (#874)

- fix: regression introduced in #869 erasing table LSNs during copy to
replication handover, causing rows written between slot creation and
copy start replayed on the destination
- fix: off-by-one in replication to copy LSN check
- fix: advancing LSN prior to transaction commit, potentially skipping
rows in the same transaction

Author: levkk

integration/copy_data/dev.sh

@@ -11,13 +11,36 @@ export PGPORT=5432
 export PGPASSWORD=pgdog
 
 cleanup() {
+    if [ -n "${BENCH_PID}" ]; then
+        kill ${BENCH_PID} 2>/dev/null || true
+        wait ${BENCH_PID} 2>/dev/null || true
+    fi
     if [ -n "${REPL_PID}" ]; then
         kill ${REPL_PID} 2>/dev/null || true
         wait ${REPL_PID} 2>/dev/null || true
     fi
 }
 trap cleanup EXIT
 
+start_pgbench() {
+    (
+
+        pgbench -h 127.0.0.1 -p 5432 -U pgdog pgdog \
+            -t 100000000 -c 3 --protocol extended \
+            -f "${SCRIPT_DIR}/pgbench.sql" -P 1
+
+    ) &
+    BENCH_PID=$!
+}
+
+stop_pgbench() {
+    if [ -n "${BENCH_PID}" ]; then
+        kill ${BENCH_PID} 2>/dev/null || true
+        wait ${BENCH_PID} 2>/dev/null || true
+        BENCH_PID=""
+    fi
+}
+
 pushd ${SCRIPT_DIR}
 
 psql -f init.sql
@@ -26,42 +49,46 @@ psql -f init.sql
 # 0 -> 2
 #
 ${PGDOG_BIN} schema-sync --from-database source --to-database destination --publication pgdog
-${PGDOG_BIN} data-sync --sync-only --from-database source --to-database destination --publication pgdog --replication-slot copy_data
-${PGDOG_BIN} schema-sync --from-database source --to-database destination --publication pgdog --cutover
+start_pgbench
+${PGDOG_BIN} data-sync --from-database source --to-database destination --publication pgdog &
+REPL_PID=$!
 
-# Check row counts: source (pgdog) vs destination (pgdog1 + pgdog2)
-echo "Checking row counts: source -> destination..."
-SHARDED_TABLES="copy_data.users copy_data.orders copy_data.order_items copy_data.log_actions copy_data.with_identity"
-OMNI_TABLES="copy_data.countries copy_data.currencies copy_data.categories"
+# Give replication a moment to connect.
+sleep 2
 
-for TABLE in ${SHARDED_TABLES}; do
-    SRC=$(psql -d pgdog -tAc "SELECT COUNT(*) FROM ${TABLE}")
-    DST1=$(psql -d pgdog1 -tAc "SELECT COUNT(*) FROM ${TABLE}")
-    DST2=$(psql -d pgdog2 -tAc "SELECT COUNT(*) FROM ${TABLE}")
-    DST=$((DST1 + DST2))
-    if [ "${SRC}" -ne "${DST}" ]; then
-        echo "MISMATCH ${TABLE}: source=${SRC} destination=${DST} (shard0=${DST1} shard1=${DST2})"
-        exit 1
-    fi
-    echo "OK ${TABLE}: ${SRC} rows"
-done
+# Check that the replication process is still alive.
+if ! kill -0 ${REPL_PID} 2>/dev/null; then
+    echo "ERROR: replication process exited early"
+    wait ${REPL_PID}
+    exit $?
+fi
 
-for TABLE in ${OMNI_TABLES}; do
-    SRC=$(psql -d pgdog -tAc "SELECT COUNT(*) FROM ${TABLE}")
-    DST1=$(psql -d pgdog1 -tAc "SELECT COUNT(*) FROM ${TABLE}")
-    DST2=$(psql -d pgdog2 -tAc "SELECT COUNT(*) FROM ${TABLE}")
-    if [ "${SRC}" -ne "${DST1}" ] || [ "${SRC}" -ne "${DST2}" ]; then
-        echo "MISMATCH ${TABLE}: source=${SRC} shard0=${DST1} shard1=${DST2} (expected ${SRC} on each shard)"
-        exit 1
-    fi
-    echo "OK ${TABLE}: ${SRC} rows on each shard"
-done
+# Let data sync and replication catch up
+echo "Letting replication run for 20 seconds..."
+sleep 20
+
+# Stop replication and capture its exit code.
+kill ${REPL_PID} 2>/dev/null || true
+set +e
+wait ${REPL_PID}
+REPL_EXIT=$?
+set -e
+REPL_PID=""
+
+# 0, 130 (SIGINT), 143 (SIGTERM) are all normal shutdown codes.
+if [ ${REPL_EXIT} -ne 0 ] && [ ${REPL_EXIT} -ne 130 ] && [ ${REPL_EXIT} -ne 143 ]; then
+    echo "ERROR: replication process exited with code ${REPL_EXIT}"
+    exit ${REPL_EXIT}
+fi
+
+stop_pgbench
+${PGDOG_BIN} schema-sync --from-database source --to-database destination --publication pgdog --cutover
 
 #
-# 2 -> 2
+# 2 --> 2
 #
 ${PGDOG_BIN} schema-sync --from-database destination --to-database destination2 --publication pgdog
-${PGDOG_BIN} data-sync --sync-only --from-database destination --to-database destination2 --publication pgdog --replication-slot copy_data
+${PGDOG_BIN} data-sync --sync-only --from-database destination --to-database destination2 --publication pgdog --replication-slot copy_data_2
 ${PGDOG_BIN} schema-sync --from-database destination --to-database destination2 --publication pgdog --cutover
 
 # Check row counts: destination (pgdog1 + pgdog2) vs destination2 (shard_0 + shard_1)
@@ -91,43 +118,6 @@ for TABLE in ${OMNI_TABLES}; do
     echo "OK ${TABLE}: ${SRC} rows on each shard"
 done
 
-# Start replication in the background.
-${PGDOG_BIN} data-sync --replicate-only --from-database source --to-database destination --publication pgdog &
-REPL_PID=$!
-
-# Give replication a moment to connect.
-sleep 2
-
-# Check that the replication process is still alive.
-if ! kill -0 ${REPL_PID} 2>/dev/null; then
-    echo "ERROR: replication process exited early"
-    wait ${REPL_PID}
-    exit $?
-fi
-
-# Run pgbench against the source database — writes land on the source and
-# get replicated to the destination shards via logical replication.
-pgbench -h 127.0.0.1 -p 5432 -U pgdog pgdog \
-    -t 1000 -c 3 --protocol extended \
-    -f "${SCRIPT_DIR}/pgbench.sql" -P 1
-
-# Let replication catch up.
-sleep 3
-
-# Stop replication and capture its exit code.
-kill ${REPL_PID} 2>/dev/null || true
-set +e
-wait ${REPL_PID}
-REPL_EXIT=$?
-set -e
-REPL_PID=""
-
-# 0, 130 (SIGINT), 143 (SIGTERM) are all normal shutdown codes.
-if [ ${REPL_EXIT} -ne 0 ] && [ ${REPL_EXIT} -ne 130 ] && [ ${REPL_EXIT} -ne 143 ]; then
-    echo "ERROR: replication process exited with code ${REPL_EXIT}"
-    exit ${REPL_EXIT}
-fi
-
 psql -f init.sql
 
 popd

integration/copy_data/pgbench.sql

@@ -1,21 +1,19 @@
 \set tenant_id random(1, 20)
-\set user_id (1021 * random(10000, 100000))
-\set order_id random(100001, 999999999)
-\set log_id random(100001, 999999999)
 \set order_amount random(1, 50000) / 100.0
 
--- Upsert a user for this tenant.
+-- Upsert a persistent user for this transaction. txid_current() is stable
+-- within the transaction and unique enough across the benchmark run.
 INSERT INTO copy_data.users (id, tenant_id, email, settings)
-VALUES (:user_id, :tenant_id, 'bench_' || :user_id || '@example.com', '{"theme":"dark"}')
+VALUES (txid_current()::bigint, :tenant_id, 'bench_' || txid_current() || '@example.com', '{"theme":"dark"}')
 ON CONFLICT (id, tenant_id) DO UPDATE SET settings = EXCLUDED.settings;
 
 -- Read the user back.
 SELECT id, tenant_id, email, created_at FROM copy_data.users
-WHERE id = :user_id AND tenant_id = :tenant_id;
+WHERE id = txid_current()::bigint AND tenant_id = :tenant_id;
 
--- Insert an order with an explicit large id.
+-- Insert a persistent order with a transaction-unique id.
 INSERT INTO copy_data.orders (id, user_id, tenant_id, amount)
-VALUES (:order_id, :user_id, :tenant_id, :order_amount)
+VALUES (txid_current()::bigint, txid_current()::bigint, :tenant_id, :order_amount)
 ON CONFLICT (id) DO NOTHING;
 
 -- Read recent orders for this tenant.
@@ -24,11 +22,11 @@ WHERE tenant_id = :tenant_id ORDER BY created_at DESC LIMIT 5;
 
 -- Update the user settings.
 UPDATE copy_data.users SET settings = jsonb_set(settings, '{last_bench}', to_jsonb(now()::text))
-WHERE id = :user_id AND tenant_id = :tenant_id;
+WHERE id = txid_current()::bigint AND tenant_id = :tenant_id;
 
--- Log an action with an explicit large id.
+-- Insert a persistent log row with a transaction-unique id.
 INSERT INTO copy_data.log_actions (id, tenant_id, action)
-VALUES (:log_id, :tenant_id, 'bench')
+VALUES (txid_current()::bigint, :tenant_id, 'bench')
 ON CONFLICT (id) DO NOTHING;
 
 \set country_id random(1, 10)
@@ -56,15 +54,3 @@ ON CONFLICT (id) DO UPDATE SET name = EXCLUDED.name;
 
 -- Read a category.
 SELECT id, name, parent_id FROM copy_data.categories WHERE id = :category_id;
-
--- Clean up everything we created.
-DELETE FROM copy_data.orders WHERE id = :order_id;
-
-DELETE FROM copy_data.log_actions WHERE id = :log_id;
-
-DELETE FROM copy_data.users
-WHERE id = :user_id AND tenant_id = :tenant_id AND email = 'bench_' || :user_id || '@example.com';
-
-DELETE FROM copy_data.countries WHERE id = :country_id AND code = 'X' || :country_id;
-
-DELETE FROM copy_data.categories WHERE id = :category_id AND name = 'Bench Category ' || :category_id;

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

@@ -23,6 +23,21 @@ use crate::backend::{pool::Request, Cluster};
 use crate::config::Role;
 use crate::net::replication::ReplicationMeta;
 
+fn merge_table_lsns(
+    tables: Vec<Table>,
+    existing_lsns: Option<&HashMap<(String, String), Lsn>>,
+) -> Vec<Table> {
+    tables
+        .into_iter()
+        .map(|mut table| {
+            if let Some(lsn) = existing_lsns.and_then(|tables| tables.get(&table.key())) {
+                table.lsn = *lsn;
+            }
+            table
+        })
+        .collect()
+}
+
 #[derive(Debug, Default)]
 pub struct Publisher {
     /// Destination cluster.
@@ -72,6 +87,19 @@ impl Publisher {
     /// Synchronize tables for all shards.
     pub async fn sync_tables(&mut self, data_sync: bool, dest: &Cluster) -> Result<(), Error> {
         let sharding_tables = dest.sharding_schema().tables;
+        let existing_lsns: HashMap<usize, HashMap<(String, String), Lsn>> = self
+            .tables
+            .iter()
+            .map(|(shard, tables)| {
+                (
+                    *shard,
+                    tables
+                        .iter()
+                        .map(|table| (table.key(), table.lsn))
+                        .collect::<HashMap<_, _>>(),
+                )
+            })
+            .collect();
 
         // Omnisharded tables are split evenly between shards
         // during copy to avoid duplicate key errors.
@@ -96,6 +124,7 @@ impl Publisher {
                 }
             } else {
                 // For replication, process changes from all shards.
+                let tables = merge_table_lsns(tables, existing_lsns.get(&number));
                 self.tables.insert(number, tables);
             }
         }
@@ -389,3 +418,61 @@ impl Waiter {
         }
     }
 }
+
+#[cfg(test)]
+mod test {
+    use super::*;
+    use crate::backend::replication::logical::publisher::{
+        PublicationTable, PublicationTableColumn, ReplicaIdentity,
+    };
+
+    fn make_table(schema: &str, name: &str, lsn: i64) -> Table {
+        Table {
+            publication: "test".to_string(),
+            table: PublicationTable {
+                schema: schema.to_string(),
+                name: name.to_string(),
+                attributes: String::new(),
+                parent_schema: String::new(),
+                parent_name: String::new(),
+            },
+            identity: ReplicaIdentity {
+                oid: pgdog_postgres_types::Oid(1),
+                identity: String::new(),
+                kind: String::new(),
+            },
+            columns: vec![PublicationTableColumn {
+                oid: 1,
+                name: "tenant_id".to_string(),
+                type_oid: pgdog_postgres_types::Oid(20),
+                identity: true,
+            }],
+            lsn: Lsn::from_i64(lsn),
+            query_parser_engine: QueryParserEngine::default(),
+        }
+    }
+
+    #[test]
+    fn merge_table_lsns_preserves_existing_offsets() {
+        let existing = HashMap::from([(
+            ("copy_data".to_string(), "users".to_string()),
+            Lsn::from_i64(123),
+        )]);
+
+        let merged = merge_table_lsns(vec![make_table("copy_data", "users", 0)], Some(&existing));
+
+        assert_eq!(merged[0].lsn, Lsn::from_i64(123));
+    }
+
+    #[test]
+    fn merge_table_lsns_leaves_unknown_tables_unset() {
+        let existing = HashMap::from([(
+            ("copy_data".to_string(), "users".to_string()),
+            Lsn::from_i64(123),
+        )]);
+
+        let merged = merge_table_lsns(vec![make_table("copy_data", "orders", 0)], Some(&existing));
+
+        assert_eq!(merged[0].lsn, Lsn::default());
+    }
+}