(fluss) branch main updated: [rust] Close drained batches before sending (#4258)

[ho...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

loserwang1024 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/fluss.git


The following commit(s) were added to refs/heads/main by this push:
     new b9a1653fa6 [rust] Close drained batches before sending (#4258)
b9a1653fa6 is described below

commit b9a1653fa6ff6711c506b150dc1b54e7f36ce125
Author: naivedogger <ya...@gmail.com>
AuthorDate: Wed Sep 9 14:02:12 2026 +0800

    [rust] Close drained batches before sending (#4258)
---
 .../crates/fluss/src/client/write/accumulator.rs   | 75 ++++++++++++++++++++++
 1 file changed, 75 insertions(+)

diff --git a/fluss-rust/crates/fluss/src/client/write/accumulator.rs b/fluss-rust/crates/fluss/src/client/write/accumulator.rs
index dba4720971..7ce7e01c7c 100644
--- a/fluss-rust/crates/fluss/src/client/write/accumulator.rs
+++ b/fluss-rust/crates/fluss/src/client/write/accumulator.rs
@@ -792,6 +792,8 @@ impl RecordAccumulator {
                 }
 
                 if let Some(mut batch) = maybe_batch {
+                    // A drained batch must not accept new records when re-enqueued for retry.
+                    batch.close()?;
                     let current_batch_size = batch.estimated_size_in_bytes();
                     size += current_batch_size;
 
@@ -1352,6 +1354,79 @@ mod tests {
         Arc::new(IdempotenceManager::new(true, 5))
     }
 
+    #[tokio::test]
+    async fn test_retry_keeps_new_records_in_a_separate_batch() -> Result<()> {
+        use futures::FutureExt;
+
+        for idempotent in [false, true] {
+            for kv in [false, true] {
+                let idempotence = Arc::new(IdempotenceManager::new(idempotent, 5));
+                idempotence.set_writer_id(42);
+                let accumulator =
+                    RecordAccumulator::new(Config::default(), Arc::clone(&idempotence));
+                let table_path = TablePath::new("db".to_string(), "tbl".to_string());
+                let physical_path = Arc::new(PhysicalTablePath::of(Arc::new(table_path.clone())));
+                let table_info = Arc::new(build_table_info(table_path.clone(), 1, 1));
+                let cluster = Arc::new(build_cluster(&table_path, 1, 1));
+                let nodes = HashSet::from([cluster.get_tablet_server(1).unwrap().clone()]);
+                let row = GenericRow {
+                    values: vec![Datum::Int32(1)],
+                };
+                let record = if kv {
+                    WriteRecord::for_upsert(
+                        table_info,
+                        physical_path,
+                        1,
+                        Bytes::from_static(b"key"),
+                        None,
+                        WriteFormat::CompactedKv,
+                        None,
+                        Some(RowBytes::Owned(Bytes::from_static(b"value"))),
+                    )
+                } else {
+                    WriteRecord::for_append(table_info, physical_path, 1, &row)
+                };
+                let first = accumulator
+                    .append(&record, 0, &cluster, false)?
+                    .result_handle
+                    .unwrap();
+                let mut batches = accumulator.drain(cluster.clone(), &nodes, 1024 * 1024)?;
+                let mut batch = batches.remove(&1).unwrap().pop().unwrap();
+                let batch_id = batch.write_batch.batch_id();
+                let original = batch.write_batch.build()?;
+                accumulator.re_enqueue(batch);
+
+                let appended = accumulator.append(&record, 0, &cluster, false)?;
+                assert!(
+                    appended.new_batch_created,
+                    "a retry must not accept additional records (kv={kv}, idempotent={idempotent})"
+                );
+                let second = appended.result_handle.unwrap();
+                let mut batches = accumulator.drain(cluster.clone(), &nodes, 1024 * 1024)?;
+                let mut retry = batches.remove(&1).unwrap().pop().unwrap();
+                assert_eq!(retry.write_batch.batch_id(), batch_id);
+                assert!(retry.write_batch.is_closed());
+                assert_eq!(retry.write_batch.record_count(), 1);
+                assert_eq!(retry.write_batch.build()?, original);
+                assert!(retry.write_batch.complete(Ok(())));
+                idempotence.handle_completed_batch(&retry.table_bucket, batch_id, 42);
+                assert!(first.wait().await?.is_ok());
+                assert!(
+                    second.wait().now_or_never().is_none(),
+                    "the old batch ACK must not complete the new record"
+                );
+
+                let mut batches = accumulator.drain(cluster, &nodes, 1024 * 1024)?;
+                let next = batches.remove(&1).unwrap().pop().unwrap();
+                assert_ne!(next.write_batch.batch_id(), batch_id);
+                assert_eq!(next.write_batch.record_count(), 1);
+                assert!(next.write_batch.complete(Ok(())));
+                assert!(second.wait().await?.is_ok());
+            }
+        }
+        Ok(())
+    }
+
     #[tokio::test]
     async fn test_retries_drain_before_fresh_batches() -> Result<()> {
         let idempotence = Arc::new(IdempotenceManager::new(true, 2));