diff --git a/fluss-rust/crates/fluss/src/client/write/accumulator.rs b/fluss-rust/crates/fluss/src/client/write/accumulator.rs index c3e9e330a1f..dba47209716 100644 --- a/fluss-rust/crates/fluss/src/client/write/accumulator.rs +++ b/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(); diff --git a/fluss-rust/crates/fluss/src/client/write/batch.rs b/fluss-rust/crates/fluss/src/client/write/batch.rs index f8cfbf994a4..03b5dcaa0fa 100644 --- a/fluss-rust/crates/fluss/src/client/write/batch.rs +++ b/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), diff --git a/fluss-rust/crates/fluss/src/client/write/idempotence.rs b/fluss-rust/crates/fluss/src/client/write/idempotence.rs index b59787ddb84..dfad1083ba7 100644 --- a/fluss-rust/crates/fluss/src/client/write/idempotence.rs +++ b/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,6 +657,7 @@ 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. @@ -651,9 +665,9 @@ mod tests { // 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( diff --git a/fluss-rust/crates/fluss/src/client/write/sender.rs b/fluss-rust/crates/fluss/src/client/write/sender.rs index 25233df7088..6cf5e8613da 100644 --- a/fluss-rust/crates/fluss/src/client/write/sender.rs +++ b/fluss-rust/crates/fluss/src/client/write/sender.rs @@ -385,9 +385,21 @@ impl Sender { } }; - // Put batches back into records_by_bucket since response handling - // will use them. - for request_batch in request_batches { + // Snapshot after connection setup and request construction, immediately before + // dispatch. Refresh on every attempt, including retries, so an old out-of-order + // response can be distinguished from one with no subsequent ACK progress. + // Put batches back into records_by_bucket for response handling. + for mut request_batch in request_batches { + if self.idempotence_manager.is_enabled() + && request_batch.write_batch.has_batch_sequence() + { + let last_acked = self + .idempotence_manager + .last_acked_sequence(&request_batch.table_bucket); + request_batch + .write_batch + .set_last_acked_sequence_at_send(last_acked); + } records_by_bucket.insert(request_batch.table_bucket.clone(), request_batch); } @@ -822,6 +834,7 @@ impl Sender { &ready_write_batch.table_bucket, seq, ready_write_batch.write_batch.batch_id(), + ready_write_batch.write_batch.last_acked_sequence_at_send(), error, ); } @@ -1219,11 +1232,12 @@ mod tests { use crate::row::{Datum, GenericRow}; use crate::rpc::FlussError; use crate::test_utils::{build_cluster_arc, build_cluster_arc_with_port, build_table_info}; + use futures::FutureExt; use prost::Message; use std::collections::{HashMap, HashSet}; use std::sync::atomic::AtomicUsize; use tokio::io::{AsyncReadExt, AsyncWriteExt}; - use tokio::net::TcpListener; + use tokio::net::{TcpListener, TcpStream}; fn disabled_idempotence() -> Arc { Arc::new(IdempotenceManager::new(false, 5)) @@ -1894,6 +1908,201 @@ mod tests { Ok(()) } + /// Handle the handshake, then let the test choose when to respond to each produce request. + async fn read_produce_request(stream: &mut TcpStream) -> i32 { + loop { + let len = stream.read_i32().await.expect("request length"); + let mut payload = vec![0u8; len as usize]; + stream.read_exact(&mut payload).await.expect("request"); + let api_key = i16::from_be_bytes(payload[..2].try_into().unwrap()); + let request_id = i32::from_be_bytes(payload[4..8].try_into().unwrap()); + if api_key == 1014 { + return request_id; + } + assert_eq!(api_key, 1000, "expected ApiVersions or ProduceLog"); + let response = ApiVersionsResponse { + api_versions: vec![ + PbApiVersion { + api_key: 1000, // ApiVersions + min_version: 0, + max_version: 0, + }, + PbApiVersion { + api_key: 1014, // ProduceLog + min_version: 0, + max_version: 0, + }, + ], + server_type: Some(ServerType::TabletServer.to_type_id()), + }; + write_controlled_response(stream, request_id, response).await; + } + } + + async fn write_controlled_response( + stream: &mut TcpStream, + request_id: i32, + response: impl Message, + ) { + let body = response.encode_to_vec(); + stream + .write_i32((5 + body.len()) as i32) + .await + .expect("response length"); + stream.write_u8(0).await.expect("success response type"); + stream.write_i32(request_id).await.expect("request id"); + stream.write_all(&body).await.expect("response body"); + } + + async fn respond_produce(stream: &mut TcpStream, request_id: i32, error: FlussError) { + let response = ProduceLogResponse { + buckets_resp: vec![PbProduceLogRespForBucket { + bucket_id: 0, + error_code: Some(error.code()), + ..Default::default() + }], + }; + write_controlled_response(stream, request_id, response).await; + } + + fn send_controlled_batch( + sender: &Arc, + batch: ReadyWriteBatch, + ) -> tokio::task::JoinHandle> { + let mut batches = HashMap::from([(1, vec![batch])]); + sender.add_to_inflight_batches(&batches); + let batches = batches.remove(&1).unwrap(); + let sender = Arc::clone(sender); + tokio::spawn(async move { sender.send_write_request(1, -1, batches).await }) + } + + #[derive(Clone, Copy)] + enum OutOfOrderScenario { + Stale, + Genuine, + Repeated, + } + + async fn check_out_of_order_response(scenario: OutOfOrderScenario) -> Result<()> { + let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind"); + let port = listener.local_addr().expect("address").port(); + let table_path = Arc::new(TablePath::new("db".to_string(), "tbl".to_string())); + let cluster = build_cluster_arc_with_port(table_path.as_ref(), 1, 1, port as u32); + let idempotence = enabled_idempotence(); + idempotence.set_writer_id(42); + let accumulator = Arc::new(RecordAccumulator::new( + Config::default(), + Arc::clone(&idempotence), + )); + let sender = Arc::new(Sender::new( + Arc::new(Metadata::new_for_test(cluster.clone())), + accumulator.clone(), + 1024 * 1024, + 10_000, + -1, + i32::MAX, + Arc::clone(&idempotence), + Arc::new(crate::metrics::WriterMetrics::new()), + )); + let (batch0, handle0) = + build_ready_batch(&accumulator, cluster.clone(), table_path.clone())?; + let bucket = batch0.table_bucket.clone(); + assert_eq!(batch0.write_batch.batch_sequence(), 0); + let send0 = send_controlled_batch(&sender, batch0); + let (mut stream, _) = listener.accept().await.expect("accept"); + let request0 = read_produce_request(&mut stream).await; + let mut pending_first_send = Some(send0); + + if matches!(scenario, OutOfOrderScenario::Genuine) { + // No predecessor is pending when seq1 is sent: an OOO must still reset. + let send = pending_first_send.take().unwrap(); + respond_produce(&mut stream, request0, FlussError::None).await; + send.await.expect("send seq0")?; + } + + let (batch1, handle1) = build_ready_batch(&accumulator, cluster.clone(), table_path)?; + assert_eq!(batch1.write_batch.batch_sequence(), 1); + let send1 = send_controlled_batch(&sender, batch1); + let request1 = read_produce_request(&mut stream).await; + + if let Some(send) = pending_first_send { + // Both sends saw last_acked=-1. Process seq0's ACK before seq1's old error. + respond_produce(&mut stream, request0, FlussError::None).await; + send.await.expect("send seq0")?; + } + assert!(handle0.wait().await?.is_ok()); + assert!(idempotence.is_next_sequence(&bucket, 1)); + respond_produce( + &mut stream, + request1, + FlussError::OutOfOrderSequenceException, + ) + .await; + send1.await.expect("send seq1")?; + + if !matches!(scenario, OutOfOrderScenario::Genuine) { + assert_eq!( + idempotence.writer_id(), + 42, + "stale error must not reset writer" + ); + assert!( + handle1.wait().now_or_never().is_none(), + "batch must remain pending" + ); + assert_eq!(idempotence.in_flight_count(&bucket), 1); + + let node = cluster.get_tablet_server(1).expect("server").clone(); + let mut batches = + accumulator.drain(cluster.clone(), &HashSet::from([node]), 1024 * 1024)?; + let retry = batches.remove(&1).expect("retry queued").pop().unwrap(); + assert_eq!(retry.write_batch.batch_sequence(), 1); + assert_eq!(retry.write_batch.attempts(), 1); + let send_retry = send_controlled_batch(&sender, retry); + let retry_request = read_produce_request(&mut stream).await; + let error = if matches!(scenario, OutOfOrderScenario::Repeated) { + // No new ACK since the resend: the refreshed snapshot must prevent retry. + FlussError::OutOfOrderSequenceException + } else { + FlussError::None + }; + respond_produce(&mut stream, retry_request, error).await; + send_retry.await.expect("send retry")?; + } + + let result = handle1.wait().now_or_never().expect("batch completed")?; + if matches!(scenario, OutOfOrderScenario::Stale) { + assert!(result.is_ok()); + assert_eq!(idempotence.writer_id(), 42); + assert!(idempotence.is_next_sequence(&bucket, 2)); + } else { + assert!(!idempotence.has_writer_id()); + assert!(matches!( + result, + Err(broadcast::Error::WriteFailed { code, .. }) + if code == FlussError::OutOfOrderSequenceException.code() + )); + } + assert_eq!(idempotence.in_flight_count(&bucket), 0); + assert!(sender.in_flight_batches.lock().is_empty()); + Ok(()) + } + + #[tokio::test] + async fn test_stale_out_of_order_response_retries_without_reset() -> Result<()> { + check_out_of_order_response(OutOfOrderScenario::Stale).await + } + + #[tokio::test] + async fn test_genuine_out_of_order_response_resets_writer() -> Result<()> { + check_out_of_order_response(OutOfOrderScenario::Genuine).await + } + + #[tokio::test] + async fn test_repeated_out_of_order_response_refreshes_snapshot() -> Result<()> { + check_out_of_order_response(OutOfOrderScenario::Repeated).await + } + #[tokio::test] async fn test_stale_writer_id_prevents_retry() -> Result<()> { let table_path = Arc::new(TablePath::new("db".to_string(), "tbl".to_string()));