[rust] Stale out-of-order responses can incorrectly reset the idempotent writer

fluss-rust/crates/fluss/src/client/write/accumulator.rs

@@ -924,11 +924,11 @@ impl RecordAccumulator {
             return;
         }
 
-        // Find the correct position sorted by batch_sequence
+        // Keep retries ordered ahead of batches that have never been sent.
         let batch_seq = ready_write_batch.write_batch.batch_sequence();
         let mut insert_pos = dq.len();
         for (i, existing) in dq.iter().enumerate() {
-            if existing.has_batch_sequence() && existing.batch_sequence() > batch_seq {
+            if !existing.has_batch_sequence() || existing.batch_sequence() > batch_seq {
                 insert_pos = i;
                 break;
             }
@@ -1352,6 +1352,58 @@ mod tests {
         Arc::new(IdempotenceManager::new(true, 5))
     }
 
+    #[tokio::test]
+    async fn test_retries_drain_before_fresh_batches() -> Result<()> {
+        let idempotence = Arc::new(IdempotenceManager::new(true, 2));
+        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, 2));
+        let cluster = Arc::new(build_cluster(&table_path, 1, 2));
+        let first = append_and_drain(&accumulator, &cluster, &table_path, 0)?;
+        let second = append_and_drain(&accumulator, &cluster, &table_path, 0)?;
+        let second_id = second.write_batch.batch_id();
+        let bucket = first.table_bucket.clone();
+        let row = GenericRow {
+            values: vec![Datum::Int32(1)],
+        };
+        let record = WriteRecord::for_append(table_info, Arc::clone(&physical_path), 1, &row);
+        let mut fresh_ids = Vec::new();
+        for _ in 0..2 {
+            accumulator.append(&record, 0, &cluster, false)?;
+            let entry = accumulator.write_batches.get(&physical_path).unwrap();
+            let mut queue = entry.batches.get(&0).unwrap().lock();
+            let batch = queue.back_mut().unwrap();
+            fresh_ids.push(batch.batch_id());
+            // Keep two distinct fresh batches queued while both slots are occupied.
+            batch.close()?;
+        }
+        accumulator.re_enqueue(second);
+        let nodes = HashSet::from([cluster.get_tablet_server(1).unwrap().clone()]);
+        assert!(
+            accumulator
+                .drain(cluster.clone(), &nodes, 1024 * 1024)?
+                .is_empty()
+        );
+        idempotence.handle_completed_batch(&bucket, first.write_batch.batch_id(), 42);
+
+        // The retry must precede both fresh batches, even with a free in-flight slot.
+        for (expected_seq, expected_id) in [second_id, fresh_ids[0], fresh_ids[1]]
+            .into_iter()
+            .enumerate()
+        {
+            let mut batches = accumulator.drain(cluster.clone(), &nodes, 1024 * 1024)?;
+            let batch = batches.remove(&1).unwrap().pop().unwrap();
+            assert_eq!(batch.write_batch.batch_id(), expected_id);
+            assert_eq!(batch.write_batch.batch_sequence(), expected_seq as i32 + 1);
+            idempotence.handle_completed_batch(&bucket, expected_id, 42);
+        }
+        assert_eq!(idempotence.in_flight_count(&bucket), 0);
+        assert!(accumulator.drain(cluster, &nodes, 1024 * 1024)?.is_empty());
+        Ok(())
+    }
+
     #[tokio::test]
     async fn re_enqueue_increments_attempts() -> Result<()> {
         let config = Config::default();

fluss-rust/crates/fluss/src/client/write/batch.rs

@@ -39,6 +39,7 @@ pub struct InnerWriteBatch {
     drained_ms: i64,
     batch_sequence: i32,
     writer_id: i64,
+    last_acked_sequence_at_send: i32,
 }
 
 impl InnerWriteBatch {
@@ -59,6 +60,7 @@ impl InnerWriteBatch {
             drained_ms: -1,
             batch_sequence: NO_BATCH_SEQUENCE,
             writer_id: NO_WRITER_ID,
+            last_acked_sequence_at_send: -1,
         }
     }
 
@@ -243,6 +245,16 @@ impl WriteBatch {
         self.inner_batch().has_batch_sequence()
     }
 
+    /// Last acknowledged sequence for this bucket when the current attempt was sent.
+    pub(crate) fn last_acked_sequence_at_send(&self) -> i32 {
+        self.inner_batch().last_acked_sequence_at_send
+    }
+
+    /// Refreshes the acknowledged-sequence snapshot on every send, including retries.
+    pub(crate) fn set_last_acked_sequence_at_send(&mut self, sequence: i32) {
+        self.inner_batch_mut().last_acked_sequence_at_send = sequence;
+    }
+
     pub fn set_writer_state(&mut self, writer_id: i64, batch_base_sequence: i32) {
         match self {
             WriteBatch::ArrowLog(batch) => batch.set_writer_state(writer_id, batch_base_sequence),

fluss-rust/crates/fluss/src/client/write/idempotence.rs

@@ -314,6 +314,14 @@ impl IdempotenceManager {
             .map(|b| b.batch_sequence)
     }
 
+    /// Returns the last acknowledged sequence, or -1 if no batch has been acknowledged.
+    pub(crate) fn last_acked_sequence(&self, bucket: &TableBucket) -> i32 {
+        self.bucket_entries
+            .lock()
+            .get(bucket)
+            .map_or(-1, |entry| entry.last_acked_sequence)
+    }
+
     pub fn is_next_sequence(&self, bucket: &TableBucket, batch_sequence: i32) -> bool {
         let entries = self.bucket_entries.lock();
         if let Some(entry) = entries.get(bucket) {
@@ -342,6 +350,7 @@ impl IdempotenceManager {
         bucket: &TableBucket,
         batch_sequence: i32,
         batch_id: i64,
+        last_acked_sequence_at_send: i32,
         error: FlussError,
     ) -> bool {
         if !self.has_writer_id() {
@@ -353,10 +362,11 @@ impl IdempotenceManager {
 
         if error == FlussError::OutOfOrderSequenceException {
             // Inline is_next_sequence logic to avoid double-locking
-            let is_next = entry.map_or(batch_sequence == 0, |e| {
-                e.last_acked_sequence + 1 == batch_sequence
-            });
-            return is_reset || !is_next;
+            let last_acked_sequence = entry.map_or(-1, |e| e.last_acked_sequence);
+            let is_next = last_acked_sequence + 1 == batch_sequence;
+            // A predecessor acknowledged since this attempt was sent makes the error stale.
+            // Without that progress, an unadjusted next batch must still fail and reset.
+            return is_reset || !is_next || last_acked_sequence > last_acked_sequence_at_send;
         }
         if error == FlussError::UnknownWriterIdException {
             return is_reset;
@@ -457,37 +467,40 @@ mod tests {
 
     #[test]
     fn test_can_retry_out_of_order() {
+        let error = FlussError::OutOfOrderSequenceException;
         let mgr = IdempotenceManager::new(true, 5);
         let b0 = test_bucket(0);
 
         // No writer_id → never retriable
-        assert!(!mgr.can_retry_for_error(&b0, 0, 100, FlussError::OutOfOrderSequenceException));
+        assert!(!mgr.can_retry_for_error(&b0, 0, 100, -1, error));
 
         mgr.set_writer_id(42);
         mgr.add_in_flight_batch(&b0, 0, 100);
         mgr.add_in_flight_batch(&b0, 1, 101);
 
         // seq=0 IS next expected (last_acked=-1+1=0) → genuine violation, NOT retriable
-        assert!(!mgr.can_retry_for_error(&b0, 0, 100, FlussError::OutOfOrderSequenceException));
+        assert!(!mgr.can_retry_for_error(&b0, 0, 100, -1, error));
         // seq=1 is NOT next expected → retriable
-        assert!(mgr.can_retry_for_error(&b0, 1, 101, FlussError::OutOfOrderSequenceException));
+        assert!(mgr.can_retry_for_error(&b0, 1, 101, -1, error));
     }
 
     #[test]
     fn test_can_retry_after_sequence_reset() {
+        let error = FlussError::OutOfOrderSequenceException;
         // OOS: batch whose seq was adjusted to match last_acked+1 is still retriable
         let (mgr, b0) = setup_three_in_flight();
         mgr.handle_completed_batch(&b0, 100, 42); // last_acked=0
         mgr.handle_failed_batch(&b0, 101, 42, None, true); // batch_id=102 adjusted to seq=1
 
         // seq=1 == last_acked(0)+1, but batch was reset → retriable
-        assert!(mgr.can_retry_for_error(&b0, 1, 102, FlussError::OutOfOrderSequenceException));
+        assert!(mgr.can_retry_for_error(&b0, 1, 102, 0, error));
 
         // UnknownWriterId: non-reset → NOT retriable, reset → retriable
+        let error = FlussError::UnknownWriterIdException;
         let (mgr, b0) = setup_three_in_flight();
-        assert!(!mgr.can_retry_for_error(&b0, 0, 100, FlussError::UnknownWriterIdException));
+        assert!(!mgr.can_retry_for_error(&b0, 0, 100, -1, error));
         mgr.handle_failed_batch(&b0, 101, 42, None, true); // batch_id=102 is reset
-        assert!(mgr.can_retry_for_error(&b0, 1, 102, FlussError::UnknownWriterIdException));
+        assert!(mgr.can_retry_for_error(&b0, 1, 102, -1, error));
     }
 
     #[test]
@@ -644,16 +657,17 @@ mod tests {
 
     #[test]
     fn scenario_multiple_inflight_retried_in_order() {
+        let error = FlussError::OutOfOrderSequenceException;
         // Java: testIdempotenceWithMultipleInflightBatchesRetriedInOrder
         // 3 batches in-flight, batch 0 times out, batches 1+2 get OOS.
         // All are retriable and must be retried one-at-a-time in sequence order.
         let (mgr, b0) = setup_three_in_flight();
 
         // Batch 0 (seq=0) times out → retriable, stays in in-flight
         // Batch 1 (seq=1) OOS → retriable (not next expected seq)
-        assert!(mgr.can_retry_for_error(&b0, 1, 101, FlussError::OutOfOrderSequenceException));
+        assert!(mgr.can_retry_for_error(&b0, 1, 101, -1, error));
         // Batch 2 (seq=2) OOS → retriable
-        assert!(mgr.can_retry_for_error(&b0, 2, 102, FlussError::OutOfOrderSequenceException));
+        assert!(mgr.can_retry_for_error(&b0, 2, 102, -1, error));
 
         // Retry phase: only first-in-flight batch should be drained
         assert!(mgr.is_first_in_flight_batch(&b0, 100));
@@ -676,6 +690,7 @@ mod tests {
 
     #[test]
     fn scenario_out_of_order_responses() {
+        let error = FlussError::OutOfOrderSequenceException;
         // Java: testCorrectHandlingOfOutOfOrderResponses
         // Server responds to batch 1 (OOS) before batch 0 (timeout).
         // Both re-enqueued, retried in order.
@@ -688,7 +703,7 @@ mod tests {
         mgr.add_in_flight_batch(&b0, 1, 101);
 
         // Batch 1 response arrives first: OOS → retriable (seq 1 ≠ next expected 0)
-        assert!(mgr.can_retry_for_error(&b0, 1, 101, FlussError::OutOfOrderSequenceException));
+        assert!(mgr.can_retry_for_error(&b0, 1, 101, -1, error));
         // Batch 0 response: timeout → retriable (no IdempotenceManager call)
 
         // Retry: batch 0 must go first
@@ -736,6 +751,7 @@ mod tests {
 
     #[test]
     fn scenario_unknown_writer_id_resets_and_restarts() {
+        let error = FlussError::UnknownWriterIdException;
         // Java: testRetryAfterResettingInFlightBatchSequence
         // Batch 0 times out (retriable), batch 1 gets UnknownWriterId (non-retriable).
         // UnknownWriterId resets all state. After new writer ID, sequences restart at 0.
@@ -749,7 +765,7 @@ mod tests {
 
         // Batch 0 times out → retriable (stays in in-flight)
         // Batch 1 UnknownWriterId → NOT retriable (non-reset batch)
-        assert!(!mgr.can_retry_for_error(&b0, 1, 101, FlussError::UnknownWriterIdException));
+        assert!(!mgr.can_retry_for_error(&b0, 1, 101, -1, error));
 
         // Sender calls fail_batch → handle_failed_batch with error → full reset
         mgr.handle_failed_batch(