Recover sync batches containing large transactions

crates/cloudsync/src/lib.rs

@@ -374,6 +374,198 @@ mod tests {
         pool.close().await;
     }
 
+    #[tokio::test]
+    async fn single_large_version_retries_after_restart_without_skipping_later_changes() {
+        let directory = tempfile::tempdir().unwrap();
+        let options = SqliteConnectOptions::new()
+            .filename(directory.path().join("sender.db"))
+            .create_if_missing(true);
+        let (options, _) = apply(options).unwrap();
+        let sender = SqlitePoolOptions::new()
+            .max_connections(1)
+            .connect_with(options.clone())
+            .await
+            .unwrap();
+        let (receiver_options, _) =
+            apply(SqliteConnectOptions::from_str("sqlite::memory:").unwrap()).unwrap();
+        let receiver = SqlitePoolOptions::new()
+            .max_connections(1)
+            .connect_with(receiver_options)
+            .await
+            .unwrap();
+        for pool in [&sender, &receiver] {
+            sqlx::query(
+                "CREATE TABLE items (id TEXT PRIMARY KEY NOT NULL, value TEXT NOT NULL DEFAULT '')",
+            )
+            .execute(pool)
+            .await
+            .unwrap();
+            init(pool, "items", None, None).await.unwrap();
+        }
+        // One statement gives every change the same indivisible database version.
+        sqlx::query(
+            "WITH RECURSIVE ids(id) AS (SELECT 1 UNION ALL SELECT id + 1 FROM ids WHERE id < 6564)
+             INSERT INTO items SELECT CAST(id AS TEXT), printf('%02048d', id) FROM ids",
+        )
+        .execute(&sender)
+        .await
+        .unwrap();
+        let first_version = db_version(&sender).await.unwrap();
+        let mut connection = sender.acquire().await.unwrap();
+        let only_version = pending_payload_batch(&mut connection, 8, 4096, 32 * 1024 * 1024)
+            .await
+            .unwrap();
+        assert!(only_version.fits && only_version.complete && !only_version.remaining);
+        assert_eq!(only_version.rows, 6564);
+        drop(connection);
+        sqlx::query("INSERT INTO items VALUES ('later', 'keep pending')")
+            .execute(&sender)
+            .await
+            .unwrap();
+        let mut connection = sender.acquire().await.unwrap();
+        let (watermark, final_chunk): (i64, bool) = sqlx::query_as(
+            "SELECT watermark_db_version, is_final FROM cloudsync_payload_chunks WHERE until_db_version = 0 LIMIT 1",
+        )
+        .fetch_one(&mut *connection)
+        .await
+        .unwrap();
+        assert!(!final_chunk);
+        assert!(watermark > first_version);
+        let truncated_scan = pending_payload_batch(&mut connection, 8, 1, 32 * 1024 * 1024)
+            .await
+            .unwrap();
+        assert!(truncated_scan.fits && truncated_scan.complete && truncated_scan.remaining);
+        assert_eq!(truncated_scan.watermark_db_version, Some(first_version));
+        let batch = pending_payload_batch(&mut connection, 8, 4096, 32 * 1024 * 1024)
+            .await
+            .unwrap();
+        assert!(batch.fits && batch.complete && batch.remaining, "{batch:?}");
+        assert_eq!(batch.rows, 6564);
+        assert!(batch.chunks > 1 && batch.chunks <= 8);
+        assert_eq!(batch.watermark_db_version, Some(first_version));
+
+        for (max_chunks, max_bytes) in [(1, 32 * 1024 * 1024), (8, batch.bytes - 1)] {
+            let rejected = pending_payload_batch(&mut connection, max_chunks, 4096, max_bytes)
+                .await
+                .unwrap();
+            assert!(
+                !rejected.fits,
+                "hard limits must still reject: {rejected:?}"
+            );
+        }
+        let payloads: Vec<Vec<u8>> = sqlx::query_scalar(
+            "SELECT payload FROM cloudsync_payload_chunks WHERE until_db_version = ?",
+        )
+        .bind(first_version)
+        .fetch_all(&mut *connection)
+        .await
+        .unwrap();
+        // A chunk can be applied before its acknowledgement is lost.
+        sqlx::query("SELECT cloudsync_payload_apply(?)")
+            .bind(&payloads[0])
+            .fetch_optional(&receiver)
+            .await
+            .unwrap();
+        let mut confirmed = NetworkStatus {
+            last_optimistic_version: first_version,
+            last_confirmed_version: batch.start_db_version,
+            gaps: Vec::new(),
+            failures: NetworkStatusFailures::default(),
+        };
+        assert!(
+            !reconcile_confirmed_pending_payload(&mut connection, batch, &confirmed)
+                .await
+                .unwrap()
+        );
+        drop(connection);
+        sender.close().await;
+        let sender = SqlitePoolOptions::new()
+            .max_connections(1)
+            .connect_with(options)
+            .await
+            .unwrap();
+        init(&sender, "items", None, None).await.unwrap();
+        let mut connection = sender.acquire().await.unwrap();
+        assert_eq!(
+            pending_payload_batch(&mut connection, 8, 4096, 32 * 1024 * 1024)
+                .await
+                .unwrap(),
+            batch
+        );
+        for payload in payloads {
+            sqlx::query("SELECT cloudsync_payload_apply(?)")
+                .bind(payload)
+                .fetch_optional(&receiver)
+                .await
+                .unwrap();
+        }
+        confirmed.last_confirmed_version = first_version;
+        assert!(
+            reconcile_confirmed_pending_payload(&mut connection, batch, &confirmed)
+                .await
+                .unwrap()
+        );
+        let tail = pending_payload_batch(&mut connection, 8, 4096, 32 * 1024 * 1024)
+            .await
+            .unwrap();
+        assert!(tail.fits && tail.complete && !tail.remaining);
+        assert_eq!(tail.rows, 1);
+        assert_eq!(tail.start_db_version, first_version);
+        let payload: Vec<u8> = sqlx::query_scalar("SELECT payload FROM cloudsync_payload_chunks")
+            .fetch_one(&mut *connection)
+            .await
+            .unwrap();
+        sqlx::query("SELECT cloudsync_payload_apply(?)")
+            .bind(payload)
+            .fetch_optional(&receiver)
+            .await
+            .unwrap();
+        let actual: Vec<(String, String)> =
+            sqlx::query_as("SELECT id, value FROM items ORDER BY id")
+                .fetch_all(&receiver)
+                .await
+                .unwrap();
+        let expected: Vec<(String, String)> =
+            sqlx::query_as("SELECT id, value FROM items ORDER BY id")
+                .fetch_all(&mut *connection)
+                .await
+                .unwrap();
+        assert_eq!(actual.len(), 6565);
+        assert_eq!(actual, expected);
+        drop(connection);
+        sender.close().await;
+        receiver.close().await;
+    }
+
+    #[tokio::test]
+    async fn single_large_version_with_incomplete_chunks_is_rejected() {
+        let pool = SqlitePoolOptions::new()
+            .max_connections(1)
+            .connect("sqlite::memory:")
+            .await
+            .unwrap();
+        sqlx::query(
+            "CREATE TABLE cloudsync_settings (key TEXT, value TEXT);
+             CREATE TABLE cloudsync_payload_chunks (
+                 payload_size INTEGER, rows INTEGER, watermark_db_version INTEGER,
+                 is_final BOOLEAN, db_version_min INTEGER, until_db_version INTEGER
+             );
+             INSERT INTO cloudsync_payload_chunks VALUES
+                 (1024, 6564, 1, FALSE, 1, 0),
+                 (1024, 6564, 1, FALSE, 1, 1)",
+        )
+        .execute(&pool)
+        .await
+        .unwrap();
+        let mut connection = pool.acquire().await.unwrap();
+        let batch = pending_payload_batch(&mut connection, 8, 4096, 32 * 1024 * 1024)
+            .await
+            .unwrap();
+        assert!(!batch.fits && !batch.complete);
+        drop(connection);
+        pool.close().await;
+    }
+
     #[tokio::test]
     async fn bounded_native_send_keeps_unsent_versions_and_retries_failed_windows() {
         use std::io::{BufRead, Read, Write};

crates/cloudsync/src/network.rs

@@ -237,15 +237,15 @@ where
 pub async fn pending_payload_batch(
     connection: &mut SqliteConnection,
     max_chunks: u32,
-    max_rows: u64,
+    target_rows: u64,
     max_bytes: u64,
 ) -> Result<PendingPayloadBatch, Error> {
-    if max_chunks == 0 || max_rows == 0 || max_bytes == 0 {
+    if max_chunks == 0 || target_rows == 0 || max_bytes == 0 {
         return Err(Error::InvalidPendingPayloadLimits);
     }
 
     let (batch, first_version) =
-        scan_pending_payload_batch(connection, max_chunks, max_rows, max_bytes, None).await?;
+        scan_pending_payload_batch(connection, max_chunks, target_rows, max_bytes, None).await?;
     if batch.fits {
         return Ok(batch);
     }
@@ -257,13 +257,28 @@ pub async fn pending_payload_batch(
     while until > first_version {
         until = first_version + (until - first_version) / 2;
         let (mut prefix, _) =
-            scan_pending_payload_batch(connection, max_chunks, max_rows, max_bytes, Some(until))
+            scan_pending_payload_batch(connection, max_chunks, target_rows, max_bytes, Some(until))
                 .await?;
         if prefix.fits && prefix.chunks > 0 {
             prefix.remaining = true;
             return Ok(prefix);
         }
     }
+    // One database version cannot be split. Let it exceed the row target only
+    // when its complete chunk stream still fits both hard resource limits.
+    let (mut first, _) = scan_pending_payload_batch(
+        connection,
+        max_chunks,
+        u64::MAX,
+        max_bytes,
+        Some(first_version),
+    )
+    .await?;
+    if first.fits && first.complete && first.chunks > 0 {
+        // Every chunk carries the watermark for the entire scan, even if we stop early.
+        first.remaining = first.watermark_db_version < batch.watermark_db_version;
+        return Ok(first);
+    }
     Ok(batch)
 }
 

crates/db-core/src/cloudsync/ops/tests.rs

@@ -607,7 +607,7 @@ fn cancelled_send_never_starts_status_reconciliation() {
 }
 
 #[tokio::test]
-async fn confirmed_send_recovery_needs_no_additional_network_request() {
+async fn confirmed_large_version_recovery_needs_no_additional_network_request() {
     use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader};
 
     let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
@@ -672,10 +672,13 @@ async fn confirmed_send_recovery_needs_no_additional_network_request() {
         .execute(&mut *connection)
         .await
         .unwrap();
-    sqlx::query("INSERT INTO items VALUES ('first', 'pending')")
-        .execute(&mut *connection)
-        .await
-        .unwrap();
+    sqlx::query(
+        "WITH RECURSIVE ids(id) AS (SELECT 1 UNION ALL SELECT id + 1 FROM ids WHERE id < 6564)
+         INSERT INTO items SELECT CAST(id AS TEXT), 'pending' FROM ids",
+    )
+    .execute(&mut *connection)
+    .await
+    .unwrap();
     let result = guarded_interruptible_network_send_changes(
         &mut connection,
         &db.cloudsync_interrupt,